Re: Auto Loader for copying files on s3

lingareddy_Alva · ‎05-08-2025

You're facing a classic small files problem in S3, which is challenging to solve efficiently. Your current Auto Loader approach has performance limitations when processing millions of small files individually. Let me suggest several optimizations and alternative approaches to speed up this process.
Key Performance Issues

Your copy_files function processes files one by one with dbutils.fs.cp, creating a lot of overhead
Using batch_df.collect() brings all file paths to the driver, creating memory pressure
Individual S3 operations have high latency, especially when done sequentially.

Try below code:

from pyspark.sql.functions import col, lit, input_file_name, regexp_replace, regexp_extract

from pyspark.sql.types import BooleanType, StringType

from pyspark import SparkFiles

import os

import time

# S3 paths configuration

source_path = "s3://source-bucket/source-folder/"

destination_path = "s3://destination-bucket/destination-folder/"

checkpoint_path = "/tmp/autoloader_checkpoints/file_copy_job"

schema_path = "/tmp/autoloader_schemas/file_copy_job"

# Performance tuning configuration

spark.conf.set("spark.sql.files.maxPartitionBytes", "128m") # Smaller partitions for more parallelism

spark.conf.set("spark.sql.adaptive.enabled", "true") # Enable adaptive query execution

spark.conf.set("spark.default.parallelism", 100) # Adjust based on your cluster size

spark.conf.set("spark.sql.shuffle.partitions", 100) # Adjust based on your cluster size

spark.conf.set("spark.databricks.io.cache.enabled", "true") # Enable IO cache if on Databricks

# Define a UDF for copying files that preserves structure and handles errors

def copy_single_file(src_path, dest_path😞

try:

# Ensure parent directory exists

parent_dir = "/".join(dest_path.split("/")[:-1]) + "/"

dbutils.fs.mkdirs(parent_dir)

# Copy the file

dbutils.fs.cp(src_path, dest_path)

return "success"

except Exception as e:

return f"error: {str(e)}"

# Register the UDF

copy_file_udf = udf(copy_single_file, StringType())

# Set up the Auto Loader stream

file_stream = (spark.readStream

.format("cloudFiles")

.option("cloudFiles.format", "binaryFile") # Read files in binary format

.option("cloudFiles.schemaLocation", schema_path)

.option("cloudFiles.includeExistingFiles", "false")

.option("recursiveFileLookup", "true") # Search all subdirectories

.option("pathGlobFilter", "*") # Process all file types

.option("cloudFiles.useNotifications", "true") # Use S3 notifications if available

.option("cloudFiles.fetchParallelism", 64) # Increase parallelism for listing files

.option("cloudFiles.maxFilesPerTrigger", 10000) # Process more files per batch

.option("cloudFiles.region", "eu-central-1")

.load(source_path)

.select("path", "length", "modificationTime")) # Only select needed columns to reduce memory

# Process each batch efficiently

def process_batch(batch_df, batch_id😞

start_time = time.time()

# Skip empty batches

if batch_df.count() == 0:

print(f"Batch {batch_id}: No files to process")

return

# Get timestamp for logging

timestamp = time.strftime("%Y-%m-%d %H:%M:%S")

# Calculate the destination path for each file by preserving directory structure

processed_df = (batch_df

# Create destination path by replacing source path with destination path

.withColumn("relative_path",

regexp_replace("path", source_path, ""))

.withColumn("destination_path",

concat(lit(destination_path), col("relative_path")))

# Apply the copy operation to each file in parallel

.withColumn("copy_result",

copy_file_udf(col("path"), col("destination_path")))

)

# Repartition for better parallelism based on file size distribution

# More partitions = more parallel operations

file_count = batch_df.count()

optimal_partitions = min(max(file_count // 1000, 8), 128) # Between 8-128 partitions

# Force execution and collect metrics

result_df = processed_df.repartition(optimal_partitions).cache()

# Trigger execution and collect stats

success_count = result_df.filter(col("copy_result").startswith("success")).count()

error_count = result_df.filter(col("copy_result").startswith("error")).count()

# Log errors for investigation

if error_count > 0:

errors_df = result_df.filter(col("copy_result").startswith("error"))

print(f"Batch {batch_id}: Found {error_count} errors. Sample errors:")

errors_df.select("path", "copy_result").show(10, truncate=False)

# Optionally write errors to a log location

errors_df.write.mode("append").parquet(f"{destination_path}/_error_logs/{batch_id}")

# Calculate performance metrics

duration = time.time() - start_time

files_per_second = file_count / duration if duration > 0 else 0

# Log summary

print(f"""

Batch {batch_id} completed at {timestamp}:

- Files processed: {file_count}

- Success: {success_count}

- Errors: {error_count}

- Duration: {duration:.2f} seconds

- Performance: {files_per_second:.2f} files/second

""")

# Unpersist to free memory

result_df.unpersist()

# Execute the streaming job

(file_stream.writeStream

.foreachBatch(process_batch)

.option("checkpointLocation", checkpoint_path)

.trigger(availableNow=True) # Process available files and terminate

.start()

.awaitTermination())

LR