How to Write a Polars DataFrame to a CSV File Using write_csv()

How to Write a Polars DataFrame to a CSV File Using write_csv()

Polars is a fast and efficient data manipulation library in Python, ideal for working with large datasets. Among its key functionalities is the ability to save DataFrames as CSV files using the write_csv() method, which simplifies data storage and sharing. This method also provides several customization options, such as defining delimiters and managing null values. In this post, we’ll explore how to use write_csv() with practical examples to demonstrate its flexibility and ease of use.

Example 1: Basic Usage

First, let’s create a simple DataFrame and write it to a CSV file using the default settings.

import polars as pl

# Create a sample DataFrame
df = pl.DataFrame({
    'Name': ['Alice', 'Bob', 'Charlie'],
    'Age': [25, 30, 35],
    'City': ['New York', 'San Francisco', 'Los Angeles']
})

# Write the DataFrame to a CSV file
df.write_csv("output.csv")

print("CSV file has been created successfully.")

This example creates a basic DataFrame with three columns and writes it to a file named “output.csv” in the current directory.

Example 2: Handling Null Values

When dealing with missing data, you can specify how null values should be represented in the CSV file.

# Create a DataFrame with null values
df_with_nulls = pl.DataFrame({
    'Name': ['Alice', None, 'Charlie'],
    'Age': [25, 30, None],
    'City': ['New York', 'San Francisco', 'Los Angeles']
})

# Write CSV with custom null value representation
df_with_nulls.write_csv("output_nulls.csv", null_value="N/A")

print("CSV file with custom null values has been created successfully.")

This example replaces null values with “N/A” in the output CSV file.

Example 3: Writing a Subset of Columns

You might want to export only specific columns from your DataFrame. Here’s how you can do that:

# Select specific columns and write to CSV
df.select(['Name', 'City']).write_csv("output_subset.csv")

print("CSV file with selected columns has been created successfully.")

This example writes only the ‘Name’ and ‘City’ columns to the CSV file.

Viewing the Results

After creating our CSV files, it’s important to verify the results. Let’s read and display the contents of each CSV file we created to see how they look.

import polars as pl

def display_csv_contents(file_path):
    print(f"\nContents of {file_path}:")
    df = pl.read_csv(file_path)
    print(df)

# Display contents of each CSV file we created
display_csv_contents("output.csv")
display_csv_contents("output_nulls.csv")
display_csv_contents("output_subset.csv")

Let’s go through the output for each file:

1. output.csv (Basic Usage)

Contents of output.csv:
shape: (3, 3)
┌─────────┬─────┬───────────────┐
│ Name    ┆ Age ┆ City          │
│ ---     ┆ --- ┆ ---           │
│ str     ┆ i64 ┆ str           │
╞═════════╪═════╪═══════════════╡
│ Alice   ┆ 25  ┆ New York      │
│ Bob     ┆ 30  ┆ San Francisco │
│ Charlie ┆ 35  ┆ Los Angeles   │
└─────────┴─────┴───────────────┘

This shows our original DataFrame with all columns and rows, including the header. Note that the ‘Age’ column is of type ‘i64’ (64-bit integer).

2. output_nulls.csv (Custom Null Value Representation)

Contents of output_nulls.csv:
shape: (3, 3)
┌─────────┬─────┬───────────────┐
│ Name    ┆ Age ┆ City          │
│ ---     ┆ --- ┆ ---           │
│ str     ┆ str ┆ str           │
╞═════════╪═════╪═══════════════╡
│ Alice   ┆ 25  ┆ New York      │
│ N/A     ┆ 30  ┆ San Francisco │
│ Charlie ┆ N/A ┆ Los Angeles   │
└─────────┴─────┴───────────────┘

Here we can see that null values have been replaced with “N/A” as specified. Note that the ‘Age’ column is now of type ‘str’ (string) because it contains the “N/A” value.

3. output_subset.csv (Subset of Columns)

Contents of output_subset.csv:
shape: (3, 2)
┌─────────┬───────────────┐
│ Name    ┆ City          │
│ ---     ┆ ---           │
│ str     ┆ str           │
╞═════════╪═══════════════╡
│ Alice   ┆ New York      │
│ Bob     ┆ San Francisco │
│ Charlie ┆ Los Angeles   │
└─────────┴───────────────┘

This output shows only the ‘Name’ and ‘City’ columns, as we selected in our example.

These outputs demonstrate that our write_csv() operations produced the expected results in each case. Let’s break down what we’re seeing:

  1. In the basic usage example, all data types are preserved, with ‘Age’ being an integer (i64) and the other columns being strings.
  2. In the null values example, replacing null values with “N/A” has changed the ‘Age’ column to a string type, as it now contains non-numeric data.
  3. In the subset example, we’ve successfully written only the ‘Name’ and ‘City’ columns to the CSV file.

This step of verifying the output is helpful in confirming that our CSV files are formatted correctly for their intended use. 

Best Practices and Tips

When using the write_csv() method in Polars, consider the following tips based on the available parameters:

  1. File Path: The file parameter accepts various types including string paths and file-like objects. Always provide a full file path if you want to save the CSV in a specific location other than the current directory.
  2. Header Control: Use include_header=False if you don’t want to include column names in your CSV file. By default, headers are included.
  3. Custom Separators: While commas are the default, you can use the separator parameter to specify a different character for separating values (e.g., separator=’\t’ for tab-separated values).
  4. Date and Time Formatting: When working with date/time columns, use the datetime_format, date_format, and time_format parameters to specify how these values should be formatted in the CSV.
  5. Float Precision: Use the float_precision parameter to control the number of decimal places for float values.

By applying these parameters, you can customize the CSV writing process to fit your specific needs. It’s important to review the output file, especially when working with large or complex datasets, to ensure it meets your expectations. Regular practice with the write_csv() method in Polars will help you efficiently export your data for analysis, reporting, or sharing, streamlining your data management tasks.

Leave a Reply

Your email address will not be published. Required fields are marked *