How to Write DataFrames to Parquet Files in PySpark

Apache Parquet is a columnar storage file format more efficient as compared to traditional row-based files like CSV. Writing dataframes to Parquet files in PySpark is, therefore, an efficient way to store and retrieve large datasets. This tutorial will teach you how to write PySpark dataframes to parquet files.

Before You Begin

You should have PySpark installed:

$ pip3 install pyspark

Here’s the Google Colab notebook for this tutorial.

1. Start a PySpark Session

Let’s start a PySpark session:

from pyspark.sql import SparkSession

# Initialize a Spark session
spark = SparkSession.builder \
	.appName("WriteToParquet") \
	.getOrCreate()

2. Generate a Sample Dataset

We’ll create a simple dataset representing customer transactions. Which includes columns like ‘CustomerID’, ‘TransactionID’, ‘TransactionDate’, ‘Amount’, and ‘ProductCategory’.

import random
import pandas as pd

# Function to generate random transaction data
def generate_data(n):
    customer_ids = [f'C{str(i).zfill(5)}' for i in range(1, 101)]
    product_categories = ['Electronics', 'Books', 'Clothing', 'Groceries', 'Furniture']

    data = []
    for _ in range(n):
        customer_id = random.choice(customer_ids)
        transaction_id = f'T{str(random.randint(10000, 99999))}'
        transaction_date = pd.Timestamp('2023-01-01') + pd.to_timedelta(random.randint(0, 180), unit='d')
        amount = round(random.uniform(5, 500), 2)
        product_category = random.choice(product_categories)
        data.append((customer_id, transaction_id, transaction_date, amount, product_category))

    return data

We’ll create a sizeable dataset with about 100K records:

# Generate 100,000 rows of transaction data
data = generate_data(100_000)

# Convert to a Pandas DataFrame
columns = ['CustomerID', 'TransactionID', 'TransactionDate', 'Amount', 'ProductCategory']
df = pd.DataFrame(data, columns=columns)

# Convert to a PySpark DataFrame
spark_df = spark.createDataFrame(df)
spark_df.show(5)

Output:

+----------+-------------+-------------------+------+---------------+
|CustomerID|TransactionID|    TransactionDate|Amount|ProductCategory|
+----------+-------------+-------------------+------+---------------+
|    C00012|       T36462|2023-05-05 00:00:00| 90.91|      Furniture|
|    C00037|       T81031|2023-03-19 00:00:00|465.54|    Electronics|
|    C00092|       T98628|2023-02-25 00:00:00| 180.9|       Clothing|
|    C00050|       T46850|2023-04-16 00:00:00|494.67|      Furniture|
|    C00097|       T79766|2023-04-11 00:00:00|179.65|      Groceries|
+----------+-------------+-------------------+------+---------------+
only showing top 5 row

Here we create a Spark dataframe from a pandas dataframe. But you can use any other dataframe to follow along.

3. Write DataFrames to Parquet Files

We can now write the dataframe to a Parquet file. PySpark’s <dataframe>.write.parquet allows you to write the dataframe to Parquet format by specifying the output path like so:

output_path = "transactions.parquet"

# Write the DataFrame to Parquet format
spark_df.write.parquet(output_path)

This will save the dataframe spark_df in Parquet format at the specified location.

4. Write Partitioned Parquet Files

Partitioning Parquet files allows you to store data in sub-directories based on the values of specific columns. This can improve query performance when you need to filter the data by those columns.

# Write the dataframe to Parquet format, partitioned by 'ProductCategory'
partitioned_output_path = "transactions_partitioned.parquet"
spark_df.write.partitionBy("ProductCategory").parquet(partitioned_output_path)

This will create a directory structure under transactions_partitioned.parquet/ with sub-directories for each ‘ProductCategory’.

5. Read Parquet Files

Reading Parquet files in PySpark is as simple as writing them. Use the read.parquet method to load a Parquet file back into a dataframe.

# Read in the Parquet file
df_read = spark.read.parquet(output_path)

# Show the content of the DataFrame
df_read.show(5)

Output:

+----------+-------------+-------------------+------+---------------+
|CustomerID|TransactionID|    TransactionDate|Amount|ProductCategory|
+----------+-------------+-------------------+------+---------------+
|    C00012|       T36462|2023-05-05 00:00:00| 90.91|      Furniture|
|    C00037|       T81031|2023-03-19 00:00:00|465.54|    Electronics|
|    C00092|       T98628|2023-02-25 00:00:00| 180.9|       Clothing|
|    C00050|       T46850|2023-04-16 00:00:00|494.67|      Furniture|
|    C00097|       T79766|2023-04-11 00:00:00|179.65|      Groceries|
+----------+-------------+-------------------+------+---------------+
only showing top 5 rows

You can also read partitioned Parquet files the same way with read.parquet.

In this tutorial, we learned how to write PySpark dataframes to Parquet files, including partitioning options.

Additional Resources

Here are a few resources you can explore:

Happy learning!