vendor/iceberg-datafusion-arrow59/src/physical_plan/scan.rs
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
// LOCAL PATCH (nornir-catalog bench; re-applied 2026-10-02 onto
// apache/datafusion-iceberg a2bc942): parallel, per-file partitioned scan.
// Upstream still exposes the whole table as `UnknownPartitioning(1)` and reads
// every file into one flattened stream (`to_arrow()`), so the Parquet decode is
// single-threaded — on a 32-core box ~1 core works and ~31 idle. Here we plan
// the file scan tasks once (async, at construction), expose ONE DataFusion
// partition per file, and `execute(partition)` reads just that file via
// `ArrowReaderBuilder`. DataFusion then decodes the files concurrently across all
// cores. (Drop-in once iceberg-datafusion ships output partitioning upstream.)
use std::pin::Pin;
use std::sync::Arc;
use std::vec;
use datafusion::arrow::array::RecordBatch;
use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
use datafusion::error::Result as DFResult;
use datafusion::execution::{SendableRecordBatchStream, TaskContext};
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties};
use datafusion::prelude::Expr;
use futures::{Stream, TryStreamExt};
use iceberg::expr::Predicate;
use iceberg::scan::{FileScanTask, FileScanTaskStream};
use iceberg::table::Table;
use super::expr_to_predicate::convert_filters_to_predicate;
use crate::to_datafusion_error;
/// Manages the scanning process of an Iceberg [`Table`], encapsulating the
/// necessary details and computed properties required for execution planning.
#[derive(Debug)]
pub struct IcebergTableScan {
/// A table in the catalog.
table: Table,
/// Snapshot of the table to scan.
snapshot_id: Option<i64>,
/// The planned file scan tasks — one DataFusion partition reads each.
tasks: Vec<FileScanTask>,
/// Stores certain, often expensive to compute,
/// plan properties used in query optimization.
plan_properties: Arc<PlanProperties>,
/// Projection column names, None means all columns
projection: Option<Vec<String>>,
/// Filters to apply to the table scan
predicates: Option<Predicate>,
/// Optional limit on the number of rows to return
limit: Option<usize>,
}
impl IcebergTableScan {
/// Creates a new [`IcebergTableScan`] object. Plans the file scan tasks up
/// front (async) so the scan can expose one partition per file.
pub(crate) async fn new(
table: Table,
snapshot_id: Option<i64>,
schema: ArrowSchemaRef,
projection: Option<&[usize]>,
filters: &[Expr],
limit: Option<usize>,
) -> DFResult<Self> {
let output_schema = match projection {
None => schema.clone(),
Some(projection) => Arc::new(schema.project(projection).unwrap()),
};
let projection = get_column_names(schema.clone(), projection);
let predicates = convert_filters_to_predicate(filters);
// Plan the file scan tasks once (same build steps as the upstream
// single-stream path, just collected so we can split them per partition).
let scan_builder = match snapshot_id {
Some(snapshot_id) => table.scan().snapshot_id(snapshot_id),
None => table.scan(),
};
let mut scan_builder = match &projection {
Some(column_names) => scan_builder.select(column_names.clone()),
None => scan_builder.select_all(),
};
if let Some(pred) = &predicates {
scan_builder = scan_builder.with_filter(pred.clone());
}
let table_scan = scan_builder.build().map_err(to_datafusion_error)?;
let tasks: Vec<FileScanTask> = table_scan
.plan_files()
.await
.map_err(to_datafusion_error)?
.try_collect()
.await
.map_err(to_datafusion_error)?;
let partitions = tasks.len().max(1);
let plan_properties = Arc::new(Self::compute_properties(output_schema, partitions));
Ok(Self {
table,
snapshot_id,
tasks,
plan_properties,
projection,
predicates,
limit,
})
}
pub fn table(&self) -> &Table {
&self.table
}
pub fn snapshot_id(&self) -> Option<i64> {
self.snapshot_id
}
pub fn projection(&self) -> Option<&[String]> {
self.projection.as_deref()
}
pub fn predicates(&self) -> Option<&Predicate> {
self.predicates.as_ref()
}
pub fn limit(&self) -> Option<usize> {
self.limit
}
/// Computes [`PlanProperties`]. `partitions` = number of file scan tasks, so
/// DataFusion runs one decode track per file across all cores.
fn compute_properties(schema: ArrowSchemaRef, partitions: usize) -> PlanProperties {
PlanProperties::new(
EquivalenceProperties::new(schema),
Partitioning::UnknownPartitioning(partitions),
EmissionType::Incremental,
Boundedness::Bounded,
)
}
}
/// The Iceberg field id an Arrow field carries (`PARQUET:field_id`), as the
/// schema conversion writes it.
fn field_id_of(field: &datafusion::arrow::datatypes::Field) -> Option<i32> {
field.metadata().get(parquet::arrow::PARQUET_FIELD_ID_META_KEY)?.parse().ok()
}
/// Fold one column's per-file manifest metrics into DataFusion column
/// statistics. `exact` says whether the scan reads the files whole (no
/// residual predicate, no limit): bounds and null counts are then exact for
/// the output; otherwise they bound it.
fn column_stats_over(
tasks: &[FileScanTask],
field_id: i32,
exact: bool,
arrow_type: &datafusion::arrow::datatypes::DataType,
) -> datafusion::common::ColumnStatistics {
use datafusion::common::ColumnStatistics;
use datafusion::common::stats::Precision;
use iceberg::spec::PrimitiveLiteral;
fn precise<T: std::fmt::Debug + Clone + PartialEq + Eq + PartialOrd>(exact: bool, v: T) -> Precision<T> {
if exact { Precision::Exact(v) } else { Precision::Inexact(v) }
}
let mut out = ColumnStatistics::new_unknown();
// Null count: every file must report it, or the total is unknown.
let nulls: Option<u64> = tasks.iter().map(|t| t.metrics().null_value_counts.get(&field_id).copied()).sum();
if let Some(n) = nulls {
out.null_count = precise(exact, n as usize);
}
// Bounds: the minimum lower bound and the maximum upper bound over the
// files, every file contributing or the bound is unknown. Datum orders
// by value within one type, which the field id guarantees.
let lower = tasks.iter().map(|t| t.metrics().lower_bounds.get(&field_id)).collect::<Option<Vec<_>>>();
let upper = tasks.iter().map(|t| t.metrics().upper_bounds.get(&field_id)).collect::<Option<Vec<_>>>();
let (Some(lower), Some(upper)) = (lower, upper) else {
return out;
};
let (Some(min), Some(max)) = (
lower.iter().copied().min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)),
upper.iter().copied().max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)),
) else {
return out;
};
if let (Some(lo), Some(hi)) = (datum_to_scalar(min, arrow_type), datum_to_scalar(max, arrow_type)) {
out.min_value = precise(exact, lo);
out.max_value = precise(exact, hi);
}
// An integer column cannot hold more distinct values than its range — the
// ceiling the join estimator needs. Always Inexact: a range is not a count.
let range: Option<u128> = match (min.literal(), max.literal()) {
(PrimitiveLiteral::Int(a), PrimitiveLiteral::Int(b)) if b >= a => Some((*b as i64 - *a as i64) as u128 + 1),
(PrimitiveLiteral::Long(a), PrimitiveLiteral::Long(b)) if b >= a => Some((*b as i128 - *a as i128) as u128 + 1),
_ => None,
};
if let Some(r) = range {
let rows_ceiling = tasks.iter().map(|t| t.record_count().unwrap_or(u64::MAX)).fold(0u128, |a, b| a.saturating_add(b as u128));
out.distinct_count = Precision::Inexact(r.min(rows_ceiling).min(usize::MAX as u128) as usize);
}
out
}
/// A manifest bound as the DataFusion scalar of the column's Arrow type —
/// `None` for a type this does not map (the bound is then left unknown,
/// never guessed).
fn datum_to_scalar(d: &iceberg::spec::Datum, arrow_type: &datafusion::arrow::datatypes::DataType) -> Option<datafusion::common::ScalarValue> {
use datafusion::arrow::datatypes::{DataType, TimeUnit};
use datafusion::common::ScalarValue;
use iceberg::spec::PrimitiveLiteral as L;
Some(match (d.literal(), arrow_type) {
(L::Boolean(v), DataType::Boolean) => ScalarValue::Boolean(Some(*v)),
(L::Int(v), DataType::Int32) => ScalarValue::Int32(Some(*v)),
(L::Int(v), DataType::Date32) => ScalarValue::Date32(Some(*v)),
(L::Long(v), DataType::Int64) => ScalarValue::Int64(Some(*v)),
(L::Long(v), DataType::Time64(TimeUnit::Microsecond)) => ScalarValue::Time64Microsecond(Some(*v)),
(L::Long(v), DataType::Timestamp(TimeUnit::Microsecond, tz)) => ScalarValue::TimestampMicrosecond(Some(*v), tz.clone()),
(L::Long(v), DataType::Timestamp(TimeUnit::Nanosecond, tz)) => ScalarValue::TimestampNanosecond(Some(*v), tz.clone()),
(L::Float(v), DataType::Float32) => ScalarValue::Float32(Some(v.0)),
(L::Double(v), DataType::Float64) => ScalarValue::Float64(Some(v.0)),
(L::String(v), DataType::Utf8) => ScalarValue::Utf8(Some(v.clone())),
(L::String(v), DataType::LargeUtf8) => ScalarValue::LargeUtf8(Some(v.clone())),
(L::Int128(v), DataType::Decimal128(p, s)) => ScalarValue::Decimal128(Some(*v), *p, *s),
_ => return None,
})
}
impl ExecutionPlan for IcebergTableScan {
fn name(&self) -> &str {
"IcebergTableScan"
}
fn children(&self) -> Vec<&Arc<dyn ExecutionPlan + 'static>> {
vec![]
}
fn with_new_children(
self: Arc<Self>,
_children: Vec<Arc<dyn ExecutionPlan>>,
) -> DFResult<Arc<dyn ExecutionPlan>> {
Ok(self)
}
fn properties(&self) -> &Arc<PlanProperties> {
&self.plan_properties
}
fn apply_expressions(
&self,
_f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> DFResult<TreeNodeRecursion>,
) -> DFResult<TreeNodeRecursion> {
Ok(TreeNodeRecursion::Continue)
}
/// **What the planned files say about size** — the row count and byte
/// length summed over the scan's tasks (one task per partition), so
/// DataFusion's physical optimiser can weigh this scan against another.
///
/// This node returned `Statistics::new_unknown` for everything, and every
/// decision that compares two inputs read that: hash-join build/probe
/// selection above all. MEASURED (skade-flight, TPC-H SF=1, 2026-09-10):
/// Q17 built its hash table on the 6 001 215-row lineitem and probed with
/// the ~2 000-row filtered part, because with both sides `Absent` the
/// `JoinSelection` rule has nothing to swap on.
///
/// Precision: `Exact` when the scan carries no residual predicate and no
/// limit — the manifest's per-file record counts ARE the row count then.
/// `Inexact` otherwise: Iceberg pruned whole files by their metrics, the
/// filter DataFusion keeps above (`Inexact` pushdown) still drops rows
/// inside the files that survived, so the sum is a ceiling.
///
/// Byte size is deliberately left `Absent`. The only number at hand is the
/// files' length on disk — Parquet, compressed — and reporting it as the
/// Arrow output size was MEASURED to mislead: DataFusion's hash-join mode
/// selection compares that figure with `hash_join_single_partition_threshold`
/// (1 MiB), took 6M-row inputs for small ones, and chose single-partition
/// `CollectLeft` builds — Q7 +39 %, Q21 +16 % at SF=1 (min of 10) against
/// the unknown-statistics baseline. A wrong number is worse than none.
///
/// Column statistics are folded from the manifest's per-file metrics
/// ([`FileScanTask::metrics`]): min of the lower bounds, max of the upper
/// bounds, the null counts summed — and for integer columns a
/// `distinct_count` CEILING of `max − min + 1`, `Inexact`. That last one is
/// the number DataFusion's join estimator actually reads. With rows alone
/// it assumes the key is unique on the BIGGER side and estimated
/// `supplier ⋈ lineitem` at 10 000 rows where 1.8 M arrive, then chose
/// `CollectLeft` for every join above it (Q7 +19 %, min of 10, both
/// runs). `max − min + 1` on `s_suppkey` and `l_suppkey` is 10 000 on both
/// sides, and the estimate comes out at lineitem's size, as it should.
fn statistics_from_inputs(
&self,
_input_stats: &[Arc<datafusion::common::Statistics>],
args: &datafusion::physical_plan::statistics::StatisticsArgs,
) -> DFResult<Arc<datafusion::common::Statistics>> {
use datafusion::common::stats::Precision;
use datafusion::common::Statistics;
let tasks: &[FileScanTask] = match args.partition() {
Some(i) => {
let n = self.plan_properties.partitioning.partition_count();
if i >= n {
return Err(datafusion::error::DataFusionError::Internal(format!(
"IcebergTableScan: invalid partition index {i}, the partition count is {n}"
)));
}
// An empty table exposes one partition and no task.
self.tasks.get(i..i + 1).unwrap_or(&[])
}
None => &self.tasks,
};
let schema = self.schema();
let mut stats = Statistics::new_unknown(&schema);
let rows: Option<u64> = tasks.iter().map(|t| t.record_count()).sum();
let exact = self.predicates.is_none() && self.limit.is_none();
stats.num_rows = match rows {
Some(n) if exact => Precision::Exact(n as usize),
Some(n) => Precision::Inexact(n as usize),
None => Precision::Absent,
};
for (i, field) in schema.fields().iter().enumerate() {
if let Some(id) = field_id_of(field) {
stats.column_statistics[i] = column_stats_over(tasks, id, exact, field.data_type());
}
}
Ok(Arc::new(stats))
}
fn execute(
&self,
partition: usize,
_context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
// Each partition decodes exactly one file scan task (independent work,
// so DataFusion fans these out across cores). A partition with no task
// (only when the table is empty and we reported 1) yields nothing.
let stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> =
match self.tasks.get(partition).cloned() {
Some(task) => {
let tasks: FileScanTaskStream = Box::pin(futures::stream::iter(vec![Ok(task)]));
// The table's reader: its FileIO and runtime, with the same
// defaults `TableScan::to_arrow` builds from.
let batches = self
.table
.reader_builder()
.build()
.read(tasks)
.map_err(to_datafusion_error)?
.stream()
.map_err(to_datafusion_error);
Box::pin(batches)
}
None => Box::pin(futures::stream::empty()),
};
// Per-partition limit: a soft cap (a GlobalLimitExec above enforces the
// exact total), kept as the upstream optimization.
let limited_stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> =
if let Some(limit) = self.limit {
let mut remaining = limit;
Box::pin(stream.try_filter_map(move |batch| {
futures::future::ready(if remaining == 0 {
Ok(None)
} else if batch.num_rows() <= remaining {
remaining -= batch.num_rows();
Ok(Some(batch))
} else {
let limited_batch = batch.slice(0, remaining);
remaining = 0;
Ok(Some(limited_batch))
})
}))
} else {
stream
};
Ok(Box::pin(RecordBatchStreamAdapter::new(
self.schema(),
limited_stream,
)))
}
}
impl DisplayAs for IcebergTableScan {
fn fmt_as(
&self,
_t: datafusion::physical_plan::DisplayFormatType,
f: &mut std::fmt::Formatter,
) -> std::fmt::Result {
write!(
f,
"IcebergTableScan projection:[{}] predicate:[{}] partitions:[{}]",
self.projection
.clone()
.map_or(String::new(), |v| v.join(",")),
self.predicates
.clone()
.map_or(String::from(""), |p| format!("{p}")),
self.tasks.len()
)?;
if let Some(limit) = self.limit {
write!(f, " limit:[{limit}]")?;
}
Ok(())
}
}
fn get_column_names(
schema: ArrowSchemaRef,
projection: Option<&[usize]>,
) -> Option<Vec<String>> {
projection.map(|v| {
v.iter()
.map(|p| schema.field(*p).name().clone())
.collect::<Vec<String>>()
})
}