Configuration¶

Let’s look at how we can configure DataFusion. When creating a SessionContext, you can pass in a SessionConfig and RuntimeEnvBuilder object. These two cover a wide range of options.

from datafusion import RuntimeEnvBuilder, SessionConfig, SessionContext

# create a session context with default settings
ctx = SessionContext()
print(ctx)

# create a session context with explicit runtime and config settings
runtime = RuntimeEnvBuilder().with_disk_manager_os().with_fair_spill_pool(10000000)
config = (
    SessionConfig()
    .with_create_default_catalog_and_schema(True)
    .with_default_catalog_and_schema("foo", "bar")
    .with_target_partitions(8)
    .with_information_schema(True)
    .with_repartition_joins(False)
    .with_repartition_aggregations(False)
    .with_repartition_windows(False)
    .with_parquet_pruning(False)
    .set("datafusion.execution.parquet.pushdown_filters", "true")
)
ctx = SessionContext(config, runtime)
print(ctx)

Maximizing CPU Usage¶

DataFusion uses partitions to parallelize work. For small queries the default configuration (number of CPU cores) is often sufficient, but to fully utilize available hardware you can tune how many partitions are created and when DataFusion will repartition data automatically.

Configure a SessionContext with a higher partition count:

from datafusion import SessionConfig, SessionContext

# allow up to 16 concurrent partitions
config = SessionConfig().with_target_partitions(16)
ctx = SessionContext(config)

Automatic repartitioning for joins, aggregations, window functions and other operations can be enabled to increase parallelism:

config = (
    SessionConfig()
    .with_target_partitions(16)
    .with_repartition_joins(True)
    .with_repartition_aggregations(True)
    .with_repartition_windows(True)
)

Manual repartitioning is available on DataFrames when you need precise control:

from datafusion import col

df = ctx.read_parquet("data.parquet")

# Evenly divide into 16 partitions
df = df.repartition(16)

# Or partition by the hash of a column
df = df.repartition_by_hash(col("a"), num=16)

result = df.collect()

Benchmark Example¶

The repository includes a benchmark script that demonstrates how to maximize CPU usage with DataFusion. The benchmarks/max_cpu_usage.py script shows a practical example of configuring DataFusion for optimal parallelism.

You can run the benchmark script to see the impact of different configuration settings:

# Run with default settings (uses all CPU cores)
python benchmarks/max_cpu_usage.py

# Run with specific number of rows and partitions
python benchmarks/max_cpu_usage.py --rows 5000000 --partitions 16

# See all available options
python benchmarks/max_cpu_usage.py --help

Here’s an example showing the performance difference between single and multiple partitions:

# Single partition - slower processing
$ python benchmarks/max_cpu_usage.py --rows=10000000 --partitions 1
Processed 10000000 rows using 1 partitions in 0.107s

# Multiple partitions - faster processing
$ python benchmarks/max_cpu_usage.py --rows=10000000 --partitions 10
Processed 10000000 rows using 10 partitions in 0.038s

This example demonstrates nearly 3x performance improvement (0.107s vs 0.038s) when using 10 partitions instead of 1, showcasing how proper partitioning can significantly improve CPU utilization and query performance.

The script demonstrates several key optimization techniques:

Higher target partition count: Uses with_target_partitions() to set the number of concurrent partitions
Automatic repartitioning: Enables repartitioning for joins, aggregations, and window functions
Manual repartitioning: Uses repartition() to ensure all partitions are utilized
CPU-intensive operations: Performs aggregations that can benefit from parallelization

The benchmark creates synthetic data and measures the time taken to perform a sum aggregation across the specified number of partitions. This helps you understand how partition configuration affects performance on your specific hardware.

Important Considerations¶

The provided benchmark script demonstrates partitioning concepts using synthetic in-memory data and simple aggregation operations. While useful for understanding basic configuration principles, actual performance in production environments may vary significantly based on numerous factors:

Data Sources and I/O Characteristics:

Table providers: Performance differs greatly between Parquet files, CSV files, databases, and cloud storage
Storage type: Local SSD, network-attached storage, and cloud storage have vastly different characteristics
Network latency: Remote data sources introduce additional latency considerations
File sizes and distribution: Large files may benefit differently from partitioning than many small files

Query and Workload Characteristics:

Operation complexity: Simple aggregations versus complex joins, window functions, or nested queries
Data distribution: Skewed data may not partition evenly, affecting parallel efficiency
Memory usage: Large datasets may require different memory management strategies
Concurrent workloads: Multiple queries running simultaneously affect resource allocation

Hardware and Environment Factors:

CPU architecture: Different processors have varying parallel processing capabilities
Available memory: Limited RAM may require different optimization strategies
System load: Other applications competing for resources affect DataFusion performance

Recommendations for Production Use:

To optimize DataFusion for your specific use case, it is strongly recommended to:

Create custom benchmarks using your actual data sources, formats, and query patterns
Test with representative data volumes that match your production workloads
Measure end-to-end performance including data loading, processing, and result handling
Evaluate different configuration combinations for your specific hardware and workload
Monitor resource utilization (CPU, memory, I/O) to identify bottlenecks in your environment

This approach will provide more accurate insights into how DataFusion configuration options will impact your particular applications and infrastructure.

For more information about available SessionConfig options, see the rust DataFusion Configuration guide, and about RuntimeEnvBuilder options in the rust online API documentation.

Custom Table Provider

SQL