How to Read CSV Files into PySpark DataFrames

Reading CSV files into PySpark dataframes is often the first task in data processing workflows. This tutorial will guide you through the steps to generate a sample CSV file, read it into a PySpark DataFrame, and perform basic data exploration.

Before You Begin

You should have PySpark installed:

$ pip3 install pyspark

You can use PySpark in a cloud notebook environment of your choice. Or you can install PySpark virtual environment on your local machine.

Here’s the Google Colab notebook for this tutorial.

1. Start a PySpark Session

First, start a PySpark session:

from pyspark.sql import SparkSession

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

2. Generate a Sample CSV File

Let’s first generate a sample CSV file that we’ll use in this tutorial. We’ll create a dataset of customer transactions with columns such as ‘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 use the generate_data function to create a CSV file of 10000 transaction records:

# Generate 10000 rows of transaction data
data = generate_data(10_000)

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

# Create the CSV file
csv_path = "sample_transactions.csv"
df.to_csv(csv_path, index=False)

print(f"Sample CSV file '{csv_path}' generated.")

This should generate the CSV file ‘sample_transactions.csv’:

Sample CSV file 'sample_transactions.csv' generated.

Note: You can read in any other CSV file of your choice as well.

3. Read the CSV File into a PySpark DataFrame

Now that we have the CSV file, we can read it into a PySpark DataFrame using the read.csv method like so:

spark_df = spark.read.csv(csv_path, header=True, inferSchema=True)

# Show the first 5 rows
spark_df.show(5)

Output:

+----------+-------------+---------------+------+---------------+
|CustomerID|TransactionID|TransactionDate|Amount|ProductCategory|
+----------+-------------+---------------+------+---------------+
|    C00006|       T58996|     2023-01-09| 17.02|      Furniture|
|    C00076|       T30519|     2023-02-28|459.67|          Books|
|    C00076|       T89246|     2023-06-10|404.95|       Clothing|
|    C00049|       T11436|     2023-06-05| 103.9|          Books|
|    C00049|       T18176|     2023-04-03|406.55|      Furniture|
+----------+-------------+---------------+------+---------------+
only showing top 5 rows

Here, we use header=True to specify that the first row contains column names and inferSchema=True to automatically infer data types for each column.

4. Explore the DataFrame

Once the CSV file is loaded into a PySpark DataFrame, you can check the schema, filter and select specific columns and more.

# Print the schema of the DataFrame
spark_df.printSchema()

Output:

root
 |-- CustomerID: string (nullable = true)
 |-- TransactionID: string (nullable = true)
 |-- TransactionDate: date (nullable = true)
 |-- Amount: double (nullable = true)
 |-- ProductCategory: string (nullable = true)

You can also perform operations like filtering rows or selecting specific columns.

from pyspark.sql.functions import col

# Filter transactions with an Amount greater than 100
filtered_df = spark_df.filter(col("Amount") > 100)

# Select specific columns
selected_df = filtered_df.select("CustomerID", "TransactionID", "Amount")

# Show the results
selected_df.show(5)

Output:

+----------+-------------+------+
|CustomerID|TransactionID|Amount|
+----------+-------------+------+
|    C00076|       T30519|459.67|
|    C00076|       T89246|404.95|
|    C00049|       T11436| 103.9|
|    C00049|       T18176|406.55|
|    C00096|       T31087|349.47|
+----------+-------------+------+
only showing top 5 rows

In this tutorial, we walked through the process of generating a sample CSV file, reading it into a PySpark DataFrame, and performing basic data exploration and processing tasks.

Additional Resources

You can refer to the following resources for further learning:

Stay tuned for other PySpark tutorials!