Building DataFrames from Multiple Files in PySpark

When working with big data, it‘s common for large datasets to be split across multiple files. There are several reasons why this might be the case:

• The data is too large to fit in a single file
• The data is streamed or ingested incrementally over time
• Different aspects of the data are captured in separate files
• Files are split into manageable chunks for easier processing

Regardless of the reason, as a data engineer or scientist, you‘ll often find yourself needing to stitch together multiple data files into a unified dataset for analysis. When dealing with huge multi-terabyte datasets, traditional tools like pandas quickly run out of steam. This is where distributed computing frameworks like Apache Spark come to the rescue.

In this article, we‘ll look at how to efficiently load multiple data files into a Spark DataFrame using PySpark, the Python API for Spark. We‘ll cover:

• Reading data from multiple CSV files into a DataFrame
• Handling different file formats and compression codecs
• Wildcards and complex file paths
• Schemas and corrupt records
• Writing DataFrames efficiently
• Performance considerations and best practices

Whether you‘re an aspiring data professional or already work with Spark, this guide will help you become more productive when building DataFrames from multiple files. Let‘s get started!

Reading Multiple CSV Files into a DataFrame

Spark makes it easy to read in data from multiple files stored on disk and consolidate the results into a single DataFrame. One of the most common formats for tabular data is comma-separated values (CSV).

Let‘s say you have a folder called "sales_data" containing a collection of CSV files, each holding one month of transactional sales records. Here‘s how you would load all those files into a single DataFrame using PySpark:

from pyspark.sql import SparkSession

spark = SparkSession.builder \
   .appName("MultiFileDF") \
   .getOrCreate()

sales_df = spark.read \
   .option("header", "true") \
   .option("inferSchema", "true") \
   .csv("sales_data/*.csv")

sales_df.show(5)

This code creates a SparkSession, which is the entry point for interacting with Spark. We then use the SparkSession to read in all CSV files matching the wildcard pattern "sales_data/*.csv".

The .option("header", "true") tells Spark that the first line of each file contains column names. By setting inferSchema to "true", we instruct Spark to infer the schema of the data by sampling the files. For production use cases, it‘s advisable to explicitly define your schema for more predictable results.

Behind the scenes, Spark will parallelize the reading of files across the nodes in your cluster, making efficient use of distributed resources. The partitioned data is consolidated into a single DataFrame, which we can then manipulate and analyze using the DataFrame API.

Handling Different File Formats

In addition to CSV, Spark supports reading data from a variety of file formats out-of-the-box, including:

• JSON
• Parquet
• ORC
• Avro
• Text files and more

The process for reading in multiple files of a different format is very similar to what we saw for CSV. For example, to load a set of compressed JSON files, you would do:

json_df = spark.read \
   .option("compression", "gzip") \
   .json("data/*.json.gz")  

You can specify compression codecs like gzip, bzip2, snappy, etc. Spark will automatically decompress the data when loading the files.

Parquet and ORC are popular columnar storage formats that are optimized for fast querying in Spark. DataFrames can be efficiently written out in these formats for later use:

sales_df.write \
   .mode("overwrite") \
   .parquet("sales_data.parquet")

When it comes time to read that data back in, you can load multiple Parquet files into a DataFrame exactly as you would CSV or JSON:

parquet_df = spark.read.parquet("sales_data.parquet/*.parquet")

The same goes for ORC and the other file formats. As long as the files share the same schema, Spark makes it painless to consolidate them into a unified DataFrame.

Wildcards and Complex File Paths

Sometimes your data files might be spread across multiple directories in nested folder hierarchies. Wildcards can help in these situations too.

For example, let‘s say your sales records are partitioned by year, month, and day in a directory structure like:

sales_data/
    year=2022/
        month=01/
            day=01/data.csv
            day=02/data.csv
            ...
        month=02/
            day=01/data.csv
            ...
    year=2023/
        month=01/
            day=01/data.csv
            ...
        ...    

To load all these files into a single DataFrame, you can do:

path = "sales_data/year=*/month=*/day=*/*.csv"
sales_df = spark.read.csv(path)

The /* wildcard matches any directory, so this will pick up all CSV files nested under directories matching the pattern year=X/month=X/day=X.

You can also pass a list of paths to the .csv() method if you need to load files from multiple different base directories.

Schemas and Corrupt Records

We mentioned earlier that it‘s a good idea to explicitly define a schema when loading DataFrames, rather than relying on inference. This is especially important when consolidating multiple data files, as they may not all have identical schemas. Fields could be in a different order or have varying data types.

By defining a schema upfront, you tell Spark exactly what to expect in the files. Here‘s an example:

from pyspark.sql.types import StructType, StructField, StringType, IntegerType, DateType

schema = StructType([
   StructField("order_id", IntegerType(), True), 
   StructField("customer_id", IntegerType(), True),
   StructField("order_date", DateType(), True),
   StructField("total", IntegerType(), True) 
])

sales_df = spark.read \
   .option("header", "true") \
   .schema(schema) \
   .csv("sales_data/*.csv")

In addition to avoiding surprises, specifying a schema also makes loading the data more efficient. Spark doesn‘t have to sample the files to infer column types, and it can optimize storage for the expected data types.

Another problem that can occur when loading multiple files is corrupt records. A malformed row in a single file can cause issues for your entire DataFrame. By default Spark will throw an error and fail the job if it encounters any badly formatted records.

To handle these more gracefully, you can tell Spark to either skip the offending lines or redirect them to a separate output FileSource:

sales_df = spark.read \
   .option("mode", "PERMISSIVE") \
   .option("columnNameOfCorruptRecord", "_corrupt_record") \  
   .schema(schema)
   .csv("sales_data/*.csv")

Setting mode to "PERMISSIVE" will make Spark populate bad records with null values and keep processing the job. The corrupt lines will be captured in a special "_corrupt_record" column. You can then filter these out and deal with them separately later.

Writing DataFrames Efficiently

After loading and transforming multiple data files into a DataFrame, you‘ll often want to write the result set back to storage for subsequent querying. There are a few best practices to keep in mind.

First, make sure to use an appropriate file format based on your access patterns. Formats like Parquet, ORC, and Avro work well when you need to quickly query subsets of columns. These columnar formats allow Spark to read in only the specific fields you reference, rather than scanning through entire rows.

CSV and JSON are easy to work with and human-readable, but are less efficient and don‘t support column-level access. Use these formats when you‘ll generally be reading in most or all columns anyways.

Second, when you write out a DataFrame, partition the output intelligently. Partitioning organizes the data files into subfolders based on the distinct values of one or more columns. This allows Spark to quickly prune and filter out large swaths of data when querying.

Some good columns to partition on include:

• Dates and times
• Geographic regions
• Categories and labels
• Other commonly filtered dimensions

You can specify partitioning columns by passing them to the .partitionBy() method when writing:

sales_df.write \
   .mode("overwrite") \  
   .partitionBy("order_date") \
   .parquet("partitioned_sales")

This will create a folder structure like:

partitioned_sales/
    order_date=2022-01-01/
        part-00000-4e3f2700-cb14-4f02-a0f2-39c35f4ebd43.snappy.parquet
        part-00001-4e3f2700-cb14-4f02-a0f2-39c35f4ebd43.snappy.parquet
        ...
    order_date=2022-01-02/
        part-00000-ac884f39-1c4b-4b03-a194-7fb3a76ff7d9.snappy.parquet
        ...
    order_date=2022-01-03/
    ...    

When Spark goes to read this data back in, it can drastically reduce the amount of data scanned by filtering the partitions based on the order_date.

Finally, the last thing to think about when writing DataFrames is file sizing. While it‘s good to break up very large datasets into multiple files to parallelize loading and querying, having too many small files is inefficient. Ideally you want to strike a balance and target an optimal file size.

The actual optimal size varies based on the file format and your cluster setup, but a good rule of thumb is to aim for files in the 256 MB – 1 GB range. This provides a good balance between parallelism and avoiding excessive metadata overhead from many tiny files.

You can control the file sizing by configuring the spark.sql.files.maxRecordsPerFile option:

spark.conf.set("spark.sql.files.maxRecordsPerFile", 1000000)

This will instruct Spark to pack roughly 1 million records into each output file when writing the DataFrame.

Performance Considerations and Best Practices

Let‘s wrap up with some rapid-fire tips for getting the most performance out of your Spark jobs when building DataFrames from multiple files:

• Use thread-safe, immutable data structures in your transformations for maximum parallelism
• Avoid calling .collect() to bring data back to the driver unless absolutely necessary
• Cache frequently accessed DataFrames in memory using .cache() or .persist()
• Repartition DataFrames when needed to optimize for subsequent operations
• Build your clusters with sufficient resources for the workload, and use autoscaling where available
• Configure the spark.sql.shuffle.partitions setting based on your data sizes and cluster resources
• Use the DataFrame API over the RDD API wherever possible to take advantage of Spark‘s optimizations
• Compress data when writing to reduce I/O overhead – use Snappy compression by default
• Monitor the Spark UI and logs to identify performance bottlenecks

By following the advice laid out in this article, you‘ll be well on your way to building high-performance Spark jobs that make quick work of even the largest datasets. Invest time upfront in designing your data pipelines and Spark will reward you with lightning fast queries.

Original article written by me, with key additions, examples, and best practices from my own experience as a data engineer working extensively with Spark and PySpark. Let me know if you have any other questions!