<!-- Canonical: https://lucioduran.com/blog/arrow-datafusion-query-engine-rust
Language: en
Content-SHA256: df16ecda767f0c5b223571c553396b798f69574fafcd88f04d9470c471da2d91
-->

Data Engineering

# Apache Arrow DataFusion: Building Custom Query Engines in Rust

2026-01-057 min read

Lucio Durán
Engineering Manager & AI Solutions Architect

Also available in: [Español](<https://lucioduran.com/es/blog/arrow-datafusion-query-engine-rust>)

## DataFusion Overview and Architecture

[Apache Arrow DataFusion](<https://datafusion.apache.org/>) is a query engine, not a database. It provides:

- A SQL parser and planner
- A cost-based query optimizer (with 60+ optimization rules)
- A parallel execution engine using Arrow RecordBatches
- An extension API for custom everything

What it does NOT give you: storage management, transactions, replication, user management. It's a library, not a server.

```rust
use datafusion::prelude::*;

#[tokio::main]
async fn main() -> datafusion::error::Result<()> {
 let ctx = SessionContext::new();

 // Register a Parquet file as a table
 ctx.register_parquet(
 "events",
 "s3://mybucket/events/*.parquet",
 ParquetReadOptions::default(),
 ).await?;

 // SQL just works
 let df = ctx.sql("
 SELECT date_trunc('hour', timestamp) as hour,
 count(*) as event_count,
 approx_percentile_cont(latency_ms, 0.99) as p99
 FROM events
 WHERE timestamp > '2026-02-01'
 GROUP BY 1
 ORDER BY 1
 ").await?;

 df.show().await?;
 Ok(())
}
```

That's a fully parallel query with predicate pushdown to Parquet row groups, columnar execution on Arrow arrays, and approximate percentile computation. In about 15 lines.

## The Execution Model: How Queries Actually Run

DataFusion uses a pull-based streaming model. Here's the flow:

```
SQL String
 ↓ (parse)
Logical Plan (AST of relational operators)
 ↓ (optimize, 60+ rules)
Optimized Logical Plan
 ↓ (physical planning)
Physical Plan (concrete execution operators)
 ↓ (execute)
Stream of RecordBatches
```

Each physical operator implements ` ExecutionPlan `:

```rust
pub trait ExecutionPlan: Send + Sync {
 /// Output schema
 fn schema(&self) -> SchemaRef;

 /// How the output is partitioned
 fn output_partitioning(&self) -> Partitioning;

 /// Required input ordering
 fn required_input_ordering(&self) -> Vec<Option<Vec<PhysicalSortExpr>>>;

 /// Create a stream of RecordBatches for partition `partition`
 fn execute(
 &self,
 partition: usize,
 context: Arc<TaskContext>,
 ) -> Result<SendableRecordBatchStream>;

 /// Children of this node
 fn children(&self) -> Vec<Arc<dyn ExecutionPlan>>;
}
```

The key insight: ` execute() ` returns a ` Stream `, not a ` Vec `. Data flows in batches (default 8,192 rows), and each partition runs on its own tokio task. This means a query over 16 Parquet files naturally parallelizes to 16 tasks, each streaming batches through the operator pipeline.

## Custom Table Providers: Your Data, Your Rules

This is where DataFusion gets really interesting. The ` TableProvider ` trait lets you wrap any data source:

```rust
use arrow::datatypes::SchemaRef;
use datafusion::datasource::TableProvider;
use datafusion::logical_expr::TableType;
use datafusion::execution::context::SessionState;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::logical_expr::Expr;

pub struct MetricsTableProvider {
 schema: SchemaRef,
 metrics_client: Arc<MetricsClient>,
}

#[async_trait]
impl TableProvider for MetricsTableProvider {
 fn as_any(&self) -> &dyn Any { self }

 fn schema(&self) -> SchemaRef {
 self.schema.clone()
 }

 fn table_type(&self) -> TableType {
 TableType::Base
 }

 async fn scan(
 &self,
 state: &SessionState,
 projection: Option<&Vec<usize>>,
 filters: &[Expr],
 limit: Option<usize>,
 ) -> Result<Arc<dyn ExecutionPlan>> {
 // Extract time range from filters for pushdown
 let time_range = extract_time_range(filters);

 // Only project requested columns
 let projected_schema = project_schema(&self.schema, projection);

 Ok(Arc::new(MetricsScanExec {
 schema: projected_schema,
 client: self.metrics_client.clone(),
 time_range,
 limit,
 }))
 }

 fn supports_filters_pushdown(
 &self,
 filters: &[&Expr],
 ) -> Result<Vec<TableProviderFilterPushDown>> {
 filters.iter().map(|f| {
 if is_time_filter(f) {
 Ok(TableProviderFilterPushDown::Exact)
 } else {
 Ok(TableProviderFilterPushDown::Unsupported)
 }
 }).collect()
 }
}
```

The ` supports_filters_pushdown ` method tells the optimizer which predicates your provider can handle natively. Returning ` Exact ` means the optimizer won't apply that filter again after the scan, your provider handled it completely. This eliminates redundant filtering and can cut metric queries from ~800ms to ~200ms.

## User-Defined Functions: Extending the SQL Surface

DataFusion supports scalar UDFs, aggregate UDFs, and window UDFs:

```rust
use datafusion::logical_expr::{create_udf, Volatility};
use datafusion::arrow::datatypes::DataType;
use arrow::array::{ArrayRef, StringArray, Float64Array};

// Scalar UDF: custom haversine distance
fn haversine_distance(args: &[ArrayRef]) -> Result<ArrayRef> {
 let lat1 = args[0].as_any().downcast_ref::<Float64Array>().unwrap();
 let lon1 = args[1].as_any().downcast_ref::<Float64Array>().unwrap();
 let lat2 = args[2].as_any().downcast_ref::<Float64Array>().unwrap();
 let lon2 = args[3].as_any().downcast_ref::<Float64Array>().unwrap();

 let result: Float64Array = lat1.iter().zip(lon1.iter())
 .zip(lat2.iter().zip(lon2.iter()))
 .map(|((la1, lo1), (la2, lo2))| {
 match (la1, lo1, la2, lo2) {
 (Some(la1), Some(lo1), Some(la2), Some(lo2)) => {
 Some(haversine(la1, lo1, la2, lo2))
 }
 _ => None,
 }
 })
 .collect();

 Ok(Arc::new(result))
}

// Register it
let udf = create_udf(
 "haversine",
 vec![DataType::Float64, DataType::Float64,
 DataType::Float64, DataType::Float64],
 Arc::new(DataType::Float64),
 Volatility::Immutable,
 Arc::new(haversine_distance),
);

ctx.register_udf(udf);

// Now use it in SQL
let df = ctx.sql("
 SELECT city_name,
 haversine(lat, lon, -34.6037, -58.3816) as distance_to_bsas
 FROM cities
 ORDER BY distance_to_bsas
 LIMIT 10
").await?;
```

For aggregate UDFs, you implement the ` Accumulator ` trait which handles partial aggregation across partitions automatically:

```rust
use datafusion::logical_expr::Accumulator;
use datafusion::scalar::ScalarValue;

#[derive(Debug)]
struct ApproxCountDistinctAccumulator {
 hll: HyperLogLog,
}

impl Accumulator for ApproxCountDistinctAccumulator {
 fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
 let array = values[0].as_any().downcast_ref::<StringArray>().unwrap();
 for value in array.iter().flatten() {
 self.hll.insert(value);
 }
 Ok(())
 }

 fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
 // Merge HLL sketches from other partitions
 let state_array = states[0].as_any()
 .downcast_ref::<BinaryArray>().unwrap();
 for state in state_array.iter().flatten() {
 let other: HyperLogLog = deserialize(state);
 self.hll.merge(&other);
 }
 Ok(())
 }

 fn state(&mut self) -> Result<Vec<ScalarValue>> {
 Ok(vec![ScalarValue::Binary(Some(serialize(&self.hll)))])
 }

 fn evaluate(&mut self) -> Result<ScalarValue> {
 Ok(ScalarValue::UInt64(Some(self.hll.count())))
 }
}
```

## Push-Based vs Pull-Based: The Architecture Decision

DataFusion uses pull-based execution (the Volcano model): each operator calls ` next() ` on its child to get the next batch. This has natural backpressure, a slow consumer automatically slows down the producer.

Push-based systems (like Velox or some parts of ClickHouse) invert this: a source pushes data downstream through a chain of operators. This can be faster because it avoids virtual function call overhead per batch and enables better SIMD vectorization.

```
Pull-based (DataFusion):
 FilterExec.next()
 → calls ProjectionExec.next()
 → calls ParquetExec.next()
 → reads batch from Parquet file
 → returns RecordBatch upstream
 → applies projection
 → returns RecordBatch upstream
 → applies filter
 → returns RecordBatch to caller

Push-based (Velox-style):
 ParquetSource.produce()
 → reads batch from Parquet file
 → pushes to ProjectionOp.process(batch)
 → applies projection
 → pushes to FilterOp.process(batch)
 → applies filter
 → pushes to OutputSink.consume(batch)
```

In practice, the performance difference on analytical queries is minimal (5-15%) because the batch size (8K rows) amortizes the per-call overhead. Where pull-based shines is composability: you can easily build complex plans with joins, subqueries, and correlated expressions without managing explicit data flow.

## Performance Tuning: Real Numbers

Here are the optimizations that make the biggest difference in production deployments:

### 1. Partition pruning

```rust
// Register partitioned data
ctx.register_parquet(
 "events",
 "s3://bucket/events/year=*/month=*/day=*/*.parquet",
 ParquetReadOptions::default()
 .table_partition_cols(vec![
 ("year".to_string(), DataType::Int32),
 ("month".to_string(), DataType::Int32),
 ("day".to_string(), DataType::Int32),
 ]),
).await?;

// Query with partition filter, only reads relevant files
// Before: scans 365 files (~50 GB)
// After: scans 7 files (~960 MB)
let df = ctx.sql("
 SELECT count(*)
 FROM events
 WHERE year = 2026 AND month = 2 AND day BETWEEN 18 AND 24
").await?;
```

### 2. Memory management

```rust
let config = SessionConfig::new()
 .with_batch_size(16384) // Larger batches = less overhead
 .with_target_partitions(16) // Match CPU cores
 .with_repartition_joins(true) // Repartition for hash joins
 .with_repartition_aggregations(true);

let runtime = RuntimeConfig::new()
 .with_memory_limit(4 * 1024 * 1024 * 1024, 0.8); // 4 GB, spill at 80%

let ctx = SessionContext::new_with_config_rt(
 config,
 Arc::new(RuntimeEnv::new(runtime)?),
);
```

### 3. Benchmark results

On a 32-core machine, querying 50 GB of Parquet data (1.2 billion rows):

```
Query: SELECT date_trunc('hour', ts), count(*), avg(value)
 FROM events WHERE ts > '2026-02-01' GROUP BY 1

DuckDB (in-process): 1.8s
DataFusion (default): 2.4s
DataFusion (tuned): 1.1s
Spark (local[32]): 14.2s

Tuning applied:
 - batch_size: 16384 (up from 8192)
 - target_partitions: 32
 - Predicate pushdown to row groups
 - Custom Parquet reader with async prefetch
```

The tuned DataFusion was actually faster than DuckDB on this specific workload because the custom Parquet reader overlapped I/O with computation, while DuckDB blocked on reads, DataFusion was processing the previous batch.

## A Production Architecture

A mature DataFusion deployment architecture:

```
┌─────────────────────────────────────────────┐
│ SQL / DataFrame API │
├─────────────────────────────────────────────┤
│ DataFusion Query Engine │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ Parquet │ │ Redis │ │ Metrics │ │
│ │ Provider │ │ Provider │ │ Provider │ │
│ └──────────┘ └──────────┘ └──────────┘ │
├─────────────────────────────────────────────┤
│ Custom UDFs │
│ - haversine(), geo_contains() │
│ - percentile_tdigest() aggregate │
│ - session_window() window function │
├─────────────────────────────────────────────┤
│ Object Store Layer (S3, local, cache) │
└─────────────────────────────────────────────┘
```

Custom providers for each data source, a shared UDF library, and an object store abstraction that handles S3 with local caching. The whole thing compiles to a single binary that starts in under 100ms.

For analytics on heterogeneous data without a managed database server, DataFusion is one of the strongest options available. The Rust ecosystem around Arrow is maturing fast, and the community is very active. The extension APIs are powerful but not trivial, solid Rust knowledge is a prerequisite.

## Frequently asked questions

### What is Apache Arrow DataFusion?

DataFusion is an extensible query engine written in Rust, built on top of Apache Arrow's columnar memory format. It provides a SQL and DataFrame API, a cost-based query optimizer, and a parallel execution engine that you can embed in your own applications or extend with custom data sources, functions, and execution strategies.

### When should I use DataFusion instead of an existing database?

DataFusion excels when you need embedded analytics without a separate database server, custom data sources that don't fit traditional databases, or when you want to build domain-specific query engines. If you need transactions, multi-user concurrency, or full ACID guarantees, a traditional database is better.

### What's the difference between push-based and pull-based execution?

In pull-based (Volcano) execution, each operator pulls batches from its children by calling next(). In push-based execution, data flows downstream without operators needing to request it. DataFusion uses a pull-based streaming model with RecordBatch as the unit of work, which naturally supports backpressure and integrates well with Rust's async model.

## Tips

1. **Tip:** DataFusion's ` TableProvider ` trait is deceptively powerful. You can wrap literally any data source, S3 Parquet files, REST APIs, even gRPC streams, behind it, and the query optimizer just works. A custom provider for an internal metrics store can be built in about 200 lines of Rust.
2. **Tip:** If your queries filter on a specific column 90% of the time, implement ` supports_filters_pushdown ` in your ` TableProvider `. The optimizer will push predicates down to your scan, and you can skip reading entire row groups. This alone can cut query latency by 4x.
3. **Tip:** Don't write custom aggregation functions from scratch. Use DataFusion's ` Accumulator ` trait, it handles partitioning and merging across partitions automatically. Rolling a custom parallel HyperLogLog from scratch is unnecessary when the trait handles all the hard parts automatically.

- arrow
- datafusion
- rust
- query-engine
- columnar
- data-engineering
- analytics
- performance

[Back to Blog](<https://lucioduran.com/blog>)