5 Best Ways to Maintain Selected Columns in Python CSV Files

by

💡 Problem Formulation: When working with CSV files in Python, there are frequent scenarios where you only need to keep certain columns while excluding others. For instance, given a CSV file with columns “ID”, “Name”, “Age”, and “Email”, you might want the output to only include “Name” and “Email”. This article explores different methods for selecting and keeping only the specified columns from a CSV file using Python.

Method 1: Using the CSV module with column names

This method leverages Python’s built-in CSV module to read the entire CSV file and then write out only the columns of interest by specifying their names. It’s excellent for when you know the column names ahead of time and when the CSV structure is consistent.

Here’s an example:

import csv

# specify the columns to keep
columns_to_keep = ['Name', 'Email']

with open('input.csv', 'r') as infile, open('output.csv', 'w', newline='') as outfile:
    reader = csv.DictReader(infile)
    writer = csv.DictWriter(outfile, fieldnames=columns_to_keep)
    
    writer.writeheader()  # write the header
    
    for row in reader:
        writer.writerow({col: row[col] for col in columns_to_keep})

Here’s the output:

Name,Email
Jane Doe,jane@example.com
John Smith,john@example.com

This code snippet is relatively straightforward. It reads each row as a dictionary and writes a new row to the output CSV file containing only the keys and values for the specified columns. This method is row-agnostic and will process all rows in the file.

Method 2: Using the CSV module with column indices

If the column names are unknown or if it’s easier to use indices, this method uses Python’s CSV module to keep columns by their index positions. It’s a good choice when dealing with a consistent CSV format where the relative positions of columns are known.

Here’s an example:

import csv

# specify the column indices to keep
indices_to_keep = [1, 3]

with open('input.csv', 'r') as infile, open('output.csv', 'w', newline='') as outfile:
    reader = csv.reader(infile)
    writer = csv.writer(outfile)
    
    for row in reader:
        selected_columns = [row[index] for index in indices_to_keep]
        writer.writerow(selected_columns)

The output would be akin to Method 1 but generated through column indices rather than names.

In this snippet, the csv.reader object allows iteration over rows directly as lists. Selected columns are written to the new CSV by indexing the row directly with the list of column indices.

Method 3: Using pandas

The pandas library provides a more powerful and flexible approach for manipulating CSV data, enabling you to specify column names to keep and handle larger datasets efficiently. It is best for those who are familiar with pandas and require additional data processing capabilities.

Here’s an example:

import pandas as pd

# specify the columns to keep
columns_to_keep = ['Name', 'Email']

# read the CSV file and keep only the specified columns
df = pd.read_csv('input.csv', usecols=columns_to_keep)

# write the filtered dataframe to a new CSV file
df.to_csv('output.csv', index=False)

Just like above, the output conveys the same data structure, focusing only on the “Name” and “Email” columns.

This code uses pandas to directly select the columns to keep when reading the CSV file. The resulting dataframe is then written out to a new CSV file with the to_csv method, excluding the index column.

Method 4: Using list comprehension with the file object

When you need a solution without external libraries, you can use Python’s file object and list comprehension to keep specific columns. This method is particularly useful when working with smaller files or when CSV module and pandas are not available.

Here’s an example:

columns_to_keep = [1, 3]  # assuming the header is in the first line

with open('input.csv', 'r') as infile, open('output.csv', 'w') as outfile:
    # Read the header and keep specified columns
    header = infile.readline().strip().split(',')
    selected_header = ','.join([header[i] for i in columns_to_keep])
    outfile.write(selected_header + '\n')

    # Read and write the selected data columns
    for line in infile:
        data = line.strip().split(',')
        selected_data = ','.join([data[i] for i in columns_to_keep])
        outfile.write(selected_data + '\n')

The output again would resemble the ones from the previous methods.

This code snippet directly manipulates the CSV file as a text file. It uses file reading and list comprehension to select and join columns, then writes them to the new file.

Bonus One-Liner Method 5: Filter with generator expression

For a compact, memory-efficient solution, you can filter specific columns in a CSV file using a one-liner generator expression in combination with the file object. It’s a quick and elegant solution when you need to manipulate files with less code.

Here’s an example:

columns_to_keep = [1, 3]  # Header expected in the first line

with open('input.csv', 'r') as infile, open('output.csv', 'w') as outfile:
    [outfile.write(','.join(item.strip().split(',')[i] for i in columns_to_keep) + '\n') for item in infile]

The output will be identical to the previous examples, providing the selected data columns.

This one-liner uses a list comprehension to iterate through each line of the input file, selects the specified columns, and writes them to the output file in one go. It is worth noting that although this is a compact solution, using a list comprehension for side effects like file writing is not generally recommended for clarity and maintainability.

Summary/Discussion

  • Method 1: CSV module with column names. Ideal for when column names are known. Strong in readability and ease of use, weak in requiring known column names.
  • Method 2: CSV module with column indices. Useful when column positions are fixed and known. Its strengths are similar to Method 1 but does not require column names.
  • Method 3: pandas. Powerful for large datasets and additional data processing. Strengths include flexibility and additional features, weakness is dependency on an external library.
  • Method 4: List comprehension with file object. Good when no external libraries are available, it’s strong in simplicity but may lack readability for larger blocks of code.
  • Bonus Method 5: Generator expression one-liner. Compact and memory-efficient, but is less readable and can confuse beginners or obscure intent.