PySpark S3 Small Files Compaction

Streaming jobs, frequent micro-batches, and per-event writers leave object stores full of tiny files. The script below reads a prefix of small Parquet files on S3 and rewrites it as a smaller number of files near a target size (128 MB by default), optionally partitioned by date columns. The sections after the code explain why this matters, how each function works, and what to watch for in production.

from pyspark.sql import SparkSession
from pyspark.sql.functions import *
import boto3
from datetime import datetime
import logging

def create_optimized_spark_session(app_name="S3SmallFilesProcessor"):
    """
    Creates a Spark session optimized for handling many small files in S3
    """
    spark = (SparkSession.builder
        .appName(app_name)
        # Optimize for small files
        .config("spark.sql.files.maxPartitionBytes", "128MB")
        .config("spark.sql.files.openCostInBytes",   "1048576") # 1MB
        .config("spark.hadoop.mapreduce.fileoutputcommitter.algorithm.version", "2")

        # S3 specific configurations
        .config("spark.hadoop.fs.s3a.fast.upload",        "true")
        .config("spark.hadoop.fs.s3a.multipart.size",     "104857600") # 100MB
        .config("spark.hadoop.fs.s3a.block.size",         "33554432")  # 32MB
        .config("spark.hadoop.fs.s3a.connection.maximum", "100")

        # Enable dynamic partition pruning
        .config("spark.sql.adaptive.enabled",                    "true")
        .config("spark.sql.adaptive.coalescePartitions.enabled", "true")
        .getOrCreate())
    
    return spark

def list_s3_files(bucket, prefix):
    """
    Lists all files in an S3 bucket with given prefix
    """
    s3_client = boto3.client('s3')
    files = []
    
    paginator = s3_client.get_paginator('list_objects_v2')
    pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
    
    for page in pages:
        if 'Contents' in page:
            for obj in page['Contents']:
                files.append(f"s3a://{bucket}/{obj['Key']}")
    
    return files

def coalesce_small_files(spark, input_path, output_path, partition_cols=None, target_size_mb=128):
    """
    Reads small files and writes them back as larger, optimized files
    
    Args:
        spark: SparkSession object
        input_path: S3 path to input files
        output_path: S3 path for output files
        partition_cols: List of columns to partition by
        target_size_mb: Target size for output files in MB
    """
    try:
        # Read the input data
        df = spark.read.option("mergeSchema", "true").parquet(input_path)
        
        # Calculate number of partitions based on data size
        total_size_bytes = df.rdd.map(lambda x: len(str(x))).sum()
        target_size_bytes = target_size_mb * 1024 * 1024
        num_partitions = max(1, int(total_size_bytes / target_size_bytes))
        
        # Repartition and write optimized files
        if partition_cols:
            df = df.repartition(num_partitions, *partition_cols)
            df.write.mode("overwrite").partitionBy(*partition_cols).parquet(output_path)
        else:
            df = df.repartition(num_partitions)
            df.write.mode("overwrite").parquet(output_path)
            
        return True
        
    except Exception as e:
        logging.error(f"Error processing files: {str(e)}")
        return False

def process_small_files(bucket, input_prefix, output_prefix, partition_cols=None):
    """
    Main function to process small files in S3
    """
    # Initialize Spark with optimized settings
    spark = create_optimized_spark_session()
    
    # Process timestamp for output path
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = f"s3a://{bucket}/{output_prefix}/processed_{timestamp}"
    
    # List input files
    input_files = list_s3_files(bucket, input_prefix)
    if not input_files:
        logging.warning(f"No files found in s3://{bucket}/{input_prefix}")
        return
    
    logging.info(f"Found {len(input_files)} files to process")
    
    # Process the files
    success = coalesce_small_files(
        spark=spark,
        input_path=f"s3a://{bucket}/{input_prefix}",
        output_path=output_path,
        partition_cols=partition_cols
    )
    
    if success:
        logging.info(f"Successfully processed files. Output written to: {output_path}")
    else:
        logging.error("Failed to process files")
    
    spark.stop()

if __name__ == "__main__":
    # Configuration
    BUCKET         = "bucket-name"
    INPUT_PREFIX   = "raw/data"
    OUTPUT_PREFIX  = "processed/data"
    PARTITION_COLS = ["year", "month", "day"]
    
    # Set up logging
    logging.basicConfig(level=logging.INFO)
    
    # Process the files
    process_small_files(
        bucket=BUCKET,
        input_prefix=INPUT_PREFIX,
        output_prefix=OUTPUT_PREFIX,
        partition_cols=PARTITION_COLS
    )

1. What Compaction Solves

Spark plans one or more tasks per input file split, and on S3 every file costs a LIST entry, an HTTP GET, and for Parquet a separate read of the footer to discover row groups and schema. With tens of thousands of 50 KB files, the job spends most of its time in request latency and driver-side planning rather than reading bytes. Query engines such as Athena, Trino, and Spark SQL all suffer the same way, and S3 request charges grow with the object count.

Compaction trades one periodic rewrite for faster reads forever after. The target size is a balance: files around 128 MB to 1 GB keep parallelism high while making per-file overhead negligible. Files much larger than that reduce parallelism and make retries expensive.

2. How the Script Works

create_optimized_spark_session sets spark.sql.files.maxPartitionBytes to 128 MB and openCostInBytes to 1 MB. The open cost tells Spark to treat each file as if it were at least 1 MB when packing files into read partitions, so many tiny files are bundled into one task instead of one task each. File output committer algorithm version 2 moves task output directly to the destination, which avoids a slow serial rename phase on S3. The S3A settings enable multipart fast upload and raise the connection pool, and adaptive query execution lets Spark coalesce shuffle partitions at runtime.

list_s3_files uses the boto3 paginator because list_objects_v2 returns at most 1,000 keys per call. The script uses the list only to check that input exists and to log a count; Spark performs its own listing when reading the prefix.

coalesce_small_files reads the whole prefix with mergeSchema so files written by different producer versions are unified, estimates the data size, computes the number of output partitions, and calls repartition. When partition columns are given, rows are hashed by those columns so each date directory receives a bounded number of files. process_small_files wires this together and writes into a timestamped output prefix.

3. Worked Example and Better Size Estimation

The size estimate in the original code, df.rdd.map(lambda x: len(str(x))).sum(), serializes every row to a Python string: it is a full extra pass over the data and measures text length, not compressed Parquet size. A cheaper and more accurate approach sums the input object sizes from the S3 listing, since compacted Parquet will be roughly the same size as its input.

import boto3, math

def prefix_bytes(bucket, prefix):
    s3 = boto3.client("s3")
    total = 0
    for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix):
        total += sum(o["Size"] for o in page.get("Contents", []))
    return total

def target_partitions(bucket, prefix, target_mb=128):
    return max(1, math.ceil(prefix_bytes(bucket, prefix) / (target_mb * 1024 * 1024)))

# 40 GB of input at 128 MB per file  ->  320 output files
n = target_partitions("bucket-name", "raw/data")
df = spark.read.parquet("s3a://bucket-name/raw/data")
(df.repartition(n, "year", "month", "day")
   .write.mode("overwrite")
   .option("maxRecordsPerFile", 5_000_000)
   .partitionBy("year", "month", "day")
   .parquet("s3a://bucket-name/processed/data/run_20260101"))

Note that repartition(n, cols) hashes rows by the partition columns, so one day can end up in only a few of the n tasks; with heavily skewed days, add a salt column or use maxRecordsPerFile to cap file size.

4. Trade-offs and Alternatives

If the data lives in a table format, prefer the format's own compaction: it commits atomically and concurrent readers never observe a half-written state. The raw-Parquet approach here is appropriate for landing zones and legacy datasets.

5. Pitfalls

• Overwriting the input in place. Reading and writing the same path with mode("overwrite") deletes the source before it is fully read. Always write to a new prefix and switch readers afterwards.
• Concurrent writers. Files that arrive during compaction are missed or double-counted. Compact closed partitions only (for example, days older than today).
• Schema drift. mergeSchema unifies compatible changes but fails on type conflicts; validate before scheduling.
• Committer behaviour. Algorithm version 2 can leave partial output if a job fails mid-write. On S3, the S3A magic or staging committers are safer choices.
• Bare except returning False. The script logs and returns; in a scheduler, re-raise so the task is marked failed.

6. Related Topics

• Databricks Delta Lake for transactional compaction with OPTIMIZE.
• Large-Scale Data Ingestion Tools for the producers that create small files.

Featured Videos