From 506dabefda59c4812e867fb11ddd2ff92d365ae0 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Mon, 8 Jun 2026 23:54:41 -0700 Subject: [PATCH 1/2] feat(hash-join): re-port pluggable left-join accumulator seam onto DF53.1 Restore the extension seam dropped during the DF53.1 reconciliation so downstream code can plug a custom build-side (left) accumulator. Re-publishes the public API of the 52.x PR #116/#117 seam, adapted by hand to the substantially restructured DF53.1 hash join (HashJoinExecBuilder, PushdownStrategy/SharedBuildAccumulator, and the new ArrayMap perfect-hash path) since the original commits no longer cherry-pick cleanly: - pub trait CollectLeftAccumulator (try_new/update_batch/evaluate/name/static_name) - pub trait ColumnBounds (dyn-able, returns a bounds physical_expr); native MinMaxColumnBounds implements it - HashJoinExec + matching generic HashJoinExecBuilder (default preserves native min/max + capped InList behavior exactly) - HashJoinExec::recreate_with_accumulator::() - ConfigOptions::optimizer.exact_join_filter_max_bytes (default usize::MAX, no behavior change) The native perfect-hash-join path (try_create_array_map) recovers concrete min/max by downcasting dyn ColumnBounds to MinMaxColumnBounds; custom bounds types simply skip the ArrayMap optimization. cargo check -p datafusion is green and all 371 hash_join unit tests pass. --- datafusion/common/src/config.rs | 15 ++ .../physical-plan/src/joins/hash_join/exec.rs | 215 +++++++++++++----- .../physical-plan/src/joins/hash_join/mod.rs | 5 +- .../src/joins/hash_join/shared_bounds.rs | 113 ++++++--- datafusion/physical-plan/src/joins/mod.rs | 3 +- docs/source/user-guide/configs.md | 1 + 6 files changed, 267 insertions(+), 85 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 2c4959a769f4d..b8861a0d1410a 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1115,6 +1115,21 @@ config_namespace! { /// See: pub hash_join_inlist_pushdown_max_distinct_values: usize, default = 150 + /// Maximum size in bytes of the exact join filter (e.g. an exact `InList` of + /// build-side keys) that a custom build-side + /// [`CollectLeftAccumulator`](https://docs.rs/datafusion-physical-plan/latest/datafusion_physical_plan/joins/trait.CollectLeftAccumulator.html) + /// may materialize for dynamic filter pushdown. + /// + /// This is an additional budget consulted by pluggable exact-membership + /// accumulators (the accumulator seam); build sides whose exact filter would + /// exceed this size are expected to fall back to a coarser strategy. It does + /// not affect the native min/max accumulator, whose InList budget is governed + /// by `hash_join_inlist_pushdown_max_size`. + /// + /// The default (`usize::MAX`) imposes no additional limit, preserving existing + /// behavior. + pub exact_join_filter_max_bytes: usize, default = usize::MAX + /// The default filter selectivity used by Filter Statistics /// when an exact selectivity cannot be determined. Valid values are /// between 0 (no selectivity) and 100 (all rows are selected). diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 25b320f985507..d450dadd0b31b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -17,6 +17,7 @@ use std::collections::HashSet; use std::fmt; +use std::marker::PhantomData; use std::mem::size_of; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; @@ -35,7 +36,8 @@ use crate::joins::Map; use crate::joins::array_map::ArrayMap; use crate::joins::hash_join::inlist_builder::build_struct_inlist_values; use crate::joins::hash_join::shared_bounds::{ - ColumnBounds, PartitionBounds, PushdownStrategy, SharedBuildAccumulator, + ColumnBounds, MinMaxColumnBounds, PartitionBounds, PushdownStrategy, + SharedBuildAccumulator, }; use crate::joins::hash_join::stream::{ BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState, @@ -127,7 +129,14 @@ fn try_create_array_map( } let (min_val, max_val) = if let Some(bounds) = bounds { - let (min_val, max_val) = if let Some(cb) = bounds.get_column_bounds(0) { + // The perfect-hash-join optimization requires concrete min/max values. + // Recover them by downcasting to the native `MinMaxColumnBounds`. A custom + // (non-min/max) `ColumnBounds` implementation will not match, in which case + // we simply skip the ArrayMap optimization (returning `None`). + let (min_val, max_val) = if let Some(cb) = bounds + .get_column_bounds(0) + .and_then(|cb| cb.as_any().downcast_ref::()) + { (cb.min.clone(), cb.max.clone()) } else { return Ok(None); @@ -263,12 +272,13 @@ impl JoinLeftData { /// flag is set to false if modifying the field requires a recomputation of the plan's /// properties. /// -pub struct HashJoinExecBuilder { - exec: HashJoinExec, +pub struct HashJoinExecBuilder +{ + exec: HashJoinExec, preserve_properties: bool, } -impl HashJoinExecBuilder { +impl HashJoinExecBuilder { /// Make a new [`HashJoinExecBuilder`]. pub fn new( left: Arc, @@ -296,6 +306,7 @@ impl HashJoinExecBuilder { // Will be computed at when plan will be built. cache: stub_properties(), join_schema: Arc::new(Schema::empty()), + _phantom_accumulator: PhantomData, }, // As `exec` is initialized with stub properties, // they will be properly computed when plan will be built. @@ -395,7 +406,7 @@ impl HashJoinExecBuilder { } /// Build resulting execution plan. - pub fn build(self) -> Result { + pub fn build(self) -> Result> { let Self { exec, preserve_properties, @@ -437,6 +448,7 @@ impl HashJoinExecBuilder { null_aware, dynamic_filter, fetch, + _phantom_accumulator: phantom_accumulator, // Recomputed. join_schema: _, column_indices: _, @@ -458,7 +470,7 @@ impl HashJoinExecBuilder { // Check if the projection is valid. can_project(&join_schema, projection.as_deref())?; - let cache = HashJoinExec::compute_properties( + let cache = HashJoinExec::::compute_properties( &left, &right, &join_schema, @@ -486,6 +498,7 @@ impl HashJoinExecBuilder { cache: Arc::new(cache), dynamic_filter, fetch, + _phantom_accumulator: phantom_accumulator, }) } @@ -495,8 +508,10 @@ impl HashJoinExecBuilder { } } -impl From<&HashJoinExec> for HashJoinExecBuilder { - fn from(exec: &HashJoinExec) -> Self { +impl From<&HashJoinExec> + for HashJoinExecBuilder +{ + fn from(exec: &HashJoinExec) -> Self { Self { exec: HashJoinExec { left: Arc::clone(exec.left()), @@ -516,6 +531,7 @@ impl From<&HashJoinExec> for HashJoinExecBuilder { cache: Arc::clone(&exec.cache), dynamic_filter: exec.dynamic_filter.clone(), fetch: exec.fetch, + _phantom_accumulator: PhantomData, }, preserve_properties: true, } @@ -714,7 +730,7 @@ impl From<&HashJoinExec> for HashJoinExecBuilder { /// Note this structure includes a [`OnceAsync`] that is used to coordinate the /// loading of the left side with the processing in each output stream. /// Therefore it can not be [`Clone`] -pub struct HashJoinExec { +pub struct HashJoinExec { /// left (build) side which gets hashed pub left: Arc, /// right (probe) side which are filtered by the hash table @@ -757,6 +773,9 @@ pub struct HashJoinExec { dynamic_filter: Option, /// Maximum number of rows to return fetch: Option, + /// Marker for the pluggable build-side (left) accumulator type. + /// Defaults to [`MinMaxLeftAccumulator`] (native min/max behavior). + _phantom_accumulator: PhantomData, } #[derive(Clone)] @@ -768,7 +787,7 @@ struct HashJoinExecDynamicFilter { build_accumulator: OnceLock>, } -impl fmt::Debug for HashJoinExec { +impl fmt::Debug for HashJoinExec { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("HashJoinExec") .field("left", &self.left) @@ -785,19 +804,24 @@ impl fmt::Debug for HashJoinExec { .field("column_indices", &self.column_indices) .field("null_equality", &self.null_equality) .field("cache", &self.cache) + .field("accumulator", &A::static_name()) // Explicitly exclude dynamic_filter to avoid runtime state differences in tests .finish() } } -impl EmbeddedProjection for HashJoinExec { +impl EmbeddedProjection for HashJoinExec { fn with_projection(&self, projection: Option>) -> Result { self.with_projection(projection) } } -impl HashJoinExec { - /// Tries to create a new [`HashJoinExec`]. +impl HashJoinExec { + /// Tries to create a new [`HashJoinExec`] with the default (native min/max) + /// build-side accumulator. + /// + /// To use a custom [`CollectLeftAccumulator`], create the join with this + /// constructor and then call [`HashJoinExec::recreate_with_accumulator`]. /// /// # Error /// This function errors when it is not possible to join the left and right sides on keys `on`. @@ -821,16 +845,51 @@ impl HashJoinExec { .with_null_aware(null_aware) .build() } +} +impl HashJoinExec { /// Create a builder based on the existing [`HashJoinExec`]. /// /// Returned builder preserves all existing fields. If a field requiring properties /// recomputation is modified, this will be done automatically during the node build. /// - pub fn builder(&self) -> HashJoinExecBuilder { + pub fn builder(&self) -> HashJoinExecBuilder { self.into() } + /// Recreate this [`HashJoinExec`] with a different build-side (left) + /// accumulator type `B`. + /// + /// This is the extension seam that lets downstream code swap the native + /// [`MinMaxLeftAccumulator`] for a custom [`CollectLeftAccumulator`] + /// implementation while preserving all other configuration. Runtime state + /// (the `dynamic_filter`/build accumulator) is carried over as-is; the build + /// accumulator is reset during execution to reflect actual partition counts. + pub fn recreate_with_accumulator( + &self, + ) -> HashJoinExec { + HashJoinExec { + left: Arc::clone(&self.left), + right: Arc::clone(&self.right), + on: self.on.clone(), + filter: self.filter.clone(), + join_type: self.join_type, + join_schema: Arc::clone(&self.join_schema), + left_fut: Arc::clone(&self.left_fut), + random_state: self.random_state.clone(), + mode: self.mode, + metrics: self.metrics.clone(), + projection: self.projection.clone(), + column_indices: self.column_indices.clone(), + null_equality: self.null_equality, + null_aware: self.null_aware, + cache: Arc::clone(&self.cache), + dynamic_filter: self.dynamic_filter.clone(), + fetch: self.fetch, + _phantom_accumulator: PhantomData, + } + } + fn create_dynamic_filter(on: &JoinOn) -> Arc { // Extract the right-side keys (probe side keys) from the `on` clauses // Dynamic filter will be created from build side values (left side) and applied to probe side (right side) @@ -1087,7 +1146,7 @@ impl HashJoinExec { } } -impl DisplayAs for HashJoinExec { +impl DisplayAs for HashJoinExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { @@ -1174,7 +1233,7 @@ impl DisplayAs for HashJoinExec { } } -impl ExecutionPlan for HashJoinExec { +impl ExecutionPlan for HashJoinExec { fn name(&self) -> &'static str { "HashJoinExec" } @@ -1302,7 +1361,7 @@ impl ExecutionPlan for HashJoinExec { let reservation = MemoryConsumer::new("HashJoinInput").register(context.memory_pool()); - Ok(collect_left_input( + Ok(collect_left_input::( self.random_state.random_state().clone(), left_stream, on_left.clone(), @@ -1323,7 +1382,7 @@ impl ExecutionPlan for HashJoinExec { MemoryConsumer::new(format!("HashJoinInput[{partition}]")) .register(context.memory_pool()); - OnceFut::new(collect_left_input( + OnceFut::new(collect_left_input::( self.random_state.random_state().clone(), left_stream, on_left.clone(), @@ -1676,6 +1735,60 @@ fn lr_is_preserved(join_type: JoinType) -> (bool, bool) { } } +/// Trait defining an accumulator for collecting build-side data during hash joins. +/// +/// The accumulator processes batches of data from the build side and computes the +/// per-column bounds needed for dynamic filtering. This is the extension seam that lets +/// downstream code plug a custom left-side accumulator (for example, an exact +/// membership accumulator) in place of the native min/max one. +/// +/// [`HashJoinExec`] is generic over this trait and defaults to [`MinMaxLeftAccumulator`], +/// which collects minimum and maximum values for join key expressions across all +/// build-side batches and preserves DataFusion's native behavior. +pub trait CollectLeftAccumulator: Send + Sync { + /// Returns the name of the accumulator. + fn name(&self) -> &'static str; + + /// Returns the static name of the accumulator type. + fn static_name() -> &'static str + where + Self: Sized, + { + std::any::type_name::() + } + + /// Creates a new accumulator for the given expression and schema. + /// + /// # Arguments + /// * `expr` - The physical expression to track bounds for + /// * `schema` - The schema of the input data + /// + /// # Returns + /// A new accumulator instance configured for the expression's data type + fn try_new(expr: Arc, schema: &SchemaRef) -> Result + where + Self: Sized; + + /// Updates the accumulator with values from a new batch. + /// + /// # Arguments + /// * `batch` - The record batch to process + /// + /// # Returns + /// Ok(()) if the update succeeds, or an error if updating fails. + fn update_batch(&mut self, batch: &RecordBatch) -> Result<()>; + + /// Finalizes the accumulation and returns the computed bounds. + /// + /// Consumes self to extract the final bounds from the accumulator. + /// + /// # Returns + /// The [`ColumnBounds`] computed from the observed build-side data. + fn evaluate(self) -> Result> + where + Self: Sized; +} + /// Accumulator for collecting min/max bounds from build-side data during hash join. /// /// This struct encapsulates the logic for progressively computing column bounds @@ -1684,8 +1797,9 @@ fn lr_is_preserved(join_type: JoinType) -> (bool, bool) { /// /// The bounds are used for dynamic filter pushdown optimization, where filters /// based on the actual data ranges can be pushed down to the probe side to -/// eliminate unnecessary data early. -struct CollectLeftAccumulator { +/// eliminate unnecessary data early. This is the native default +/// [`CollectLeftAccumulator`] used by [`HashJoinExec`]. +pub struct MinMaxLeftAccumulator { /// The physical expression to evaluate for each batch expr: Arc, /// Accumulator for tracking the minimum value across all batches @@ -1694,15 +1808,19 @@ struct CollectLeftAccumulator { max: MaxAccumulator, } -impl CollectLeftAccumulator { +impl CollectLeftAccumulator for MinMaxLeftAccumulator { + fn name(&self) -> &'static str { + "MinMaxLeftAccumulator" + } + + fn static_name() -> &'static str + where + Self: Sized, + { + "MinMaxLeftAccumulator" + } + /// Creates a new accumulator for tracking bounds of a join key expression. - /// - /// # Arguments - /// * `expr` - The physical expression to track bounds for - /// * `schema` - The schema of the input data - /// - /// # Returns - /// A new `CollectLeftAccumulator` instance configured for the expression's data type fn try_new(expr: Arc, schema: &SchemaRef) -> Result { /// Recursively unwraps dictionary types to get the underlying value type. fn dictionary_value_type(data_type: &DataType) -> DataType { @@ -1725,16 +1843,8 @@ impl CollectLeftAccumulator { }) } - /// Updates the accumulators with values from a new batch. - /// /// Evaluates the expression on the batch and updates both min and max /// accumulators with the resulting values. - /// - /// # Arguments - /// * `batch` - The record batch to process - /// - /// # Returns - /// Ok(()) if the update succeeds, or an error if expression evaluation fails fn update_batch(&mut self, batch: &RecordBatch) -> Result<()> { let array = self.expr.evaluate(batch)?.into_array(batch.num_rows())?; self.min.update_batch(std::slice::from_ref(&array))?; @@ -1742,30 +1852,25 @@ impl CollectLeftAccumulator { Ok(()) } - /// Finalizes the accumulation and returns the computed bounds. - /// /// Consumes self to extract the final min and max values from the accumulators. - /// - /// # Returns - /// The `ColumnBounds` containing the minimum and maximum values observed - fn evaluate(mut self) -> Result { - Ok(ColumnBounds::new( + fn evaluate(mut self) -> Result> { + Ok(Arc::new(MinMaxColumnBounds::new( self.min.evaluate()?, self.max.evaluate()?, - )) + ))) } } /// State for collecting the build-side data during hash join -struct BuildSideState { +struct BuildSideState { batches: Vec, num_rows: usize, metrics: BuildProbeJoinMetrics, reservation: MemoryReservation, - bounds_accumulators: Option>, + bounds_accumulators: Option>, } -impl BuildSideState { +impl BuildSideState { /// Create a new BuildSideState with optional accumulators for bounds computation fn try_new( metrics: BuildProbeJoinMetrics, @@ -1783,7 +1888,7 @@ impl BuildSideState { .then(|| { on_left .into_iter() - .map(|expr| CollectLeftAccumulator::try_new(expr, schema)) + .map(|expr| A::try_new(expr, schema)) .collect::>>() }) .transpose()?, @@ -1833,7 +1938,7 @@ fn should_collect_min_max_for_perfect_hash( /// `JoinLeftData` containing the hash map, consolidated batch, join key values, /// visited indices bitmap, and computed bounds (if requested). #[expect(clippy::too_many_arguments)] -async fn collect_left_input( +async fn collect_left_input( random_state: RandomState, left_stream: SendableRecordBatchStream, on_left: Vec, @@ -1851,7 +1956,7 @@ async fn collect_left_input( let should_collect_min_max_for_phj = should_collect_min_max_for_perfect_hash(&on_left, &schema)?; - let initial = BuildSideState::try_new( + let initial = BuildSideState::::try_new( metrics, reservation, on_left.clone(), @@ -1898,7 +2003,7 @@ async fn collect_left_input( Some(accumulators) if num_rows > 0 => { let bounds = accumulators .into_iter() - .map(CollectLeftAccumulator::evaluate) + .map(A::evaluate) .collect::>>()?; Some(PartitionBounds::new(bounds)) } @@ -5367,7 +5472,8 @@ mod tests { )]; // Create a dynamic filter manually - let dynamic_filter = HashJoinExec::create_dynamic_filter(&on); + let dynamic_filter = + HashJoinExec::::create_dynamic_filter(&on); let dynamic_filter_clone = Arc::clone(&dynamic_filter); // Create HashJoinExec with the dynamic filter @@ -5416,7 +5522,8 @@ mod tests { )]; // Create a dynamic filter manually - let dynamic_filter = HashJoinExec::create_dynamic_filter(&on); + let dynamic_filter = + HashJoinExec::::create_dynamic_filter(&on); let dynamic_filter_clone = Arc::clone(&dynamic_filter); // Create HashJoinExec with the dynamic filter diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index b915802ea4015..b998ced889249 100644 --- a/datafusion/physical-plan/src/joins/hash_join/mod.rs +++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs @@ -17,8 +17,11 @@ //! [`HashJoinExec`] Partitioned Hash Join Operator -pub use exec::{HashJoinExec, HashJoinExecBuilder}; +pub use exec::{ + CollectLeftAccumulator, HashJoinExec, HashJoinExecBuilder, MinMaxLeftAccumulator, +}; pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState}; +pub use shared_bounds::{ColumnBounds, MinMaxColumnBounds}; mod exec; mod inlist_builder; diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index f32dc7fa80268..4a5d741fde12a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -18,7 +18,8 @@ //! Utilities for shared build-side information. Used in dynamic filter pushdown in Hash Joins. // TODO: include the link to the Dynamic Filter blog post. -use std::fmt; +use std::any::Any; +use std::fmt::{self, Debug}; use std::sync::Arc; use crate::ExecutionPlan; @@ -44,37 +45,101 @@ use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef, ScalarFunctionExpr use parking_lot::Mutex; use tokio::sync::Barrier; +/// Trait representing some set of bounds for a column used in join dynamic filtering. +/// +/// Bounds could be min/max values, or some other type of custom bounds that can be +/// represented by a physical expression. This is the dyn-able bounds type returned by +/// [`CollectLeftAccumulator::evaluate`](super::exec::CollectLeftAccumulator::evaluate), +/// allowing downstream code to plug a custom left-side accumulator. +/// +/// Refer to the [`MinMaxColumnBounds`] implementation for an example of min/max bounds, +/// which is the native default behavior of [`HashJoinExec`](super::exec::HashJoinExec). +pub trait ColumnBounds: Send + Sync + Debug { + /// Upcast to [`Any`] so callers can downcast to a concrete bounds type. + /// + /// The native perfect-hash-join path downcasts to [`MinMaxColumnBounds`] to recover + /// the concrete min/max values; custom bounds types simply will not match. + fn as_any(&self) -> &dyn Any; + + /// Creates a physical expression representing the bounds for this column (the + /// probe-side / right expression the dynamic filter is applied to). + /// + /// # Arguments + /// + /// * `expr` - The physical expression to constrain with these bounds. + /// + /// # Returns + /// `Ok(Arc)` if creating the bounds expression succeeds, or an error otherwise. + fn physical_expr( + &self, + expr: Arc, + ) -> Result>; +} + /// Represents the minimum and maximum values for a specific column. /// Used in dynamic filter pushdown to establish value boundaries. +/// +/// This is the native [`ColumnBounds`] implementation backing the default +/// `MinMaxLeftAccumulator`; it preserves the min/max range-predicate behavior and is the +/// concrete type the perfect-hash-join path downcasts to. #[derive(Debug, Clone, PartialEq)] -pub(crate) struct ColumnBounds { +pub struct MinMaxColumnBounds { /// The minimum value observed for this column pub(crate) min: ScalarValue, - /// The maximum value observed for this column + /// The maximum value observed for this column pub(crate) max: ScalarValue, } -impl ColumnBounds { +impl MinMaxColumnBounds { pub(crate) fn new(min: ScalarValue, max: ScalarValue) -> Self { Self { min, max } } } +impl ColumnBounds for MinMaxColumnBounds { + fn as_any(&self) -> &dyn Any { + self + } + + fn physical_expr( + &self, + expr: Arc, + ) -> Result> { + // Create predicate: col >= min AND col <= max + let min_expr = Arc::new(BinaryExpr::new( + Arc::clone(&expr), + Operator::GtEq, + lit(self.min.clone()), + )) as Arc; + let max_expr = Arc::new(BinaryExpr::new( + Arc::clone(&expr), + Operator::LtEq, + lit(self.max.clone()), + )) as Arc; + let range_expr = Arc::new(BinaryExpr::new(min_expr, Operator::And, max_expr)) + as Arc; + Ok(range_expr) + } +} + /// Represents the bounds for all join key columns from a single partition. -/// This contains the min/max values computed from one partition's build-side data. +/// This contains the bounds computed from one partition's build-side data. #[derive(Debug, Clone)] pub(crate) struct PartitionBounds { - /// Min/max bounds for each join key column in this partition. + /// Bounds for each join key column in this partition. /// Index corresponds to the join key expression index. - column_bounds: Vec, + column_bounds: Vec>, } impl PartitionBounds { - pub(crate) fn new(column_bounds: Vec) -> Self { + pub(crate) fn new(column_bounds: Vec>) -> Self { Self { column_bounds } } - pub(crate) fn get_column_bounds(&self, index: usize) -> Option<&ColumnBounds> { + pub(crate) fn get_column_bounds( + &self, + index: usize, + ) -> Option<&Arc> { self.column_bounds.get(index) } } @@ -146,32 +211,22 @@ fn create_membership_predicate( fn create_bounds_predicate( on_right: &[PhysicalExprRef], bounds: &PartitionBounds, -) -> Option> { +) -> Result>> { let mut column_predicates = Vec::new(); for (col_idx, right_expr) in on_right.iter().enumerate() { if let Some(column_bounds) = bounds.get_column_bounds(col_idx) { - // Create predicate: col >= min AND col <= max - let min_expr = Arc::new(BinaryExpr::new( - Arc::clone(right_expr), - Operator::GtEq, - lit(column_bounds.min.clone()), - )) as Arc; - let max_expr = Arc::new(BinaryExpr::new( - Arc::clone(right_expr), - Operator::LtEq, - lit(column_bounds.max.clone()), - )) as Arc; - let range_expr = Arc::new(BinaryExpr::new(min_expr, Operator::And, max_expr)) - as Arc; + // Delegate to the bounds implementation. For the native + // `MinMaxColumnBounds` this produces `col >= min AND col <= max`. + let range_expr = column_bounds.physical_expr(Arc::clone(right_expr))?; column_predicates.push(range_expr); } } if column_predicates.is_empty() { - None + Ok(None) } else { - Some( + Ok(Some( column_predicates .into_iter() .reduce(|acc, pred| { @@ -179,7 +234,7 @@ fn create_bounds_predicate( as Arc }) .unwrap(), - ) + )) } } @@ -414,7 +469,7 @@ impl SharedBuildAccumulator { let bounds_expr = create_bounds_predicate( &self.on_right, &partition_data.bounds, - ); + )?; // Combine membership and bounds expressions for multi-layer optimization: // - Bounds (min/max): Enable statistics-based pruning (Parquet row group/file skipping) @@ -524,7 +579,7 @@ impl SharedBuildAccumulator { let bounds_expr = create_bounds_predicate( &self.on_right, &partition.bounds, - ); + )?; // 3. Combine membership and bounds expressions let then_expr = match (membership_expr, bounds_expr) { @@ -587,7 +642,7 @@ impl SharedBuildAccumulator { } } -impl fmt::Debug for SharedBuildAccumulator { +impl Debug for SharedBuildAccumulator { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "SharedBuildAccumulator") } diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 2cdfa1e6ac020..5e51bc2bbfbca 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -21,7 +21,8 @@ use arrow::array::BooleanBufferBuilder; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; pub use hash_join::{ - HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr, SeededRandomState, + CollectLeftAccumulator, ColumnBounds, HashExpr, HashJoinExec, HashJoinExecBuilder, + HashTableLookupExpr, MinMaxColumnBounds, MinMaxLeftAccumulator, SeededRandomState, }; pub use nested_loop_join::{NestedLoopJoinExec, NestedLoopJoinExecBuilder}; use parking_lot::Mutex; diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 1245e59d477a7..e322d099a8940 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -161,6 +161,7 @@ The following configuration settings are available: | datafusion.optimizer.hash_join_single_partition_threshold_rows | 131072 | The maximum estimated size in rows for one input side of a HashJoin will be collected into a single partition | | datafusion.optimizer.hash_join_inlist_pushdown_max_size | 131072 | Maximum size in bytes for the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides larger than this will use hash table lookups instead. Set to 0 to always use hash table lookups. InList pushdown can be more efficient for small build sides because it can result in better statistics pruning as well as use any bloom filters present on the scan side. InList expressions are also more transparent and easier to serialize over the network in distributed uses of DataFusion. On the other hand InList pushdown requires making a copy of the data and thus adds some overhead to the build side and uses more memory. This setting is per-partition, so we may end up using `hash_join_inlist_pushdown_max_size` \* `target_partitions` memory. The default is 128kB per partition. This should allow point lookup joins (e.g. joining on a unique primary key) to use InList pushdown in most cases but avoids excessive memory usage or overhead for larger joins. | | datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values | 150 | Maximum number of distinct values (rows) in the build side of a hash join to be pushed down as an InList expression for dynamic filtering. Build sides with more rows than this will use hash table lookups instead. Set to 0 to always use hash table lookups. This provides an additional limit beyond `hash_join_inlist_pushdown_max_size` to prevent very large IN lists that might not provide much benefit over hash table lookups. This uses the deduplicated row count once the build side has been evaluated. The default is 150 values per partition. This is inspired by Trino's `max-filter-keys-per-column` setting. See: | +| datafusion.optimizer.exact_join_filter_max_bytes | 18446744073709551615 | Maximum size in bytes of the exact join filter (e.g. an exact `InList` of build-side keys) that a custom build-side `CollectLeftAccumulator` may materialize for dynamic filter pushdown. This is an additional budget consulted by pluggable exact-membership accumulators (the accumulator seam); build sides whose exact filter would exceed this size are expected to fall back to a coarser strategy. It does not affect the native min/max accumulator, whose InList budget is governed by `hash_join_inlist_pushdown_max_size`. The default (`usize::MAX`) imposes no additional limit, preserving existing behavior. | | datafusion.optimizer.default_filter_selectivity | 20 | The default filter selectivity used by Filter Statistics when an exact selectivity cannot be determined. Valid values are between 0 (no selectivity) and 100 (all rows are selected). | | datafusion.optimizer.prefer_existing_union | false | When set to true, the optimizer will not attempt to convert Union to Interleave | | datafusion.optimizer.expand_views_at_output | false | When set to true, if the returned type is a view type then the output will be coerced to a non-view. Coerces `Utf8View` to `LargeUtf8`, and `BinaryView` to `LargeBinary`. | From 148ae5b752f7824b328a7a73482b3b37143d831a Mon Sep 17 00:00:00 2001 From: Sergei Grebnov Date: Tue, 9 Jun 2026 15:08:58 +0300 Subject: [PATCH 2/2] fix: bounds-guard project_statistics for metadata/partition columns beyond input stats --- datafusion/datasource/src/file_scan_config.rs | 468 +++++++++++++++++- datafusion/datasource/src/test_util.rs | 15 +- datafusion/physical-expr/src/projection.rs | 58 ++- 3 files changed, 537 insertions(+), 4 deletions(-) diff --git a/datafusion/datasource/src/file_scan_config.rs b/datafusion/datasource/src/file_scan_config.rs index 16269e1df1ef1..4fa1a4649a43c 100644 --- a/datafusion/datasource/src/file_scan_config.rs +++ b/datafusion/datasource/src/file_scan_config.rs @@ -1579,6 +1579,7 @@ mod tests { use super::*; use crate::TableSchema; + use crate::metadata::MetadataColumn; use crate::test_util::col; use crate::{ generate_test_files, test_util::MockSource, tests::aggr_test_schema, @@ -1586,13 +1587,15 @@ mod tests { }; use arrow::datatypes::Field; + use datafusion_common::ScalarValue; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, internal_err}; use datafusion_expr::{Operator, SortExpr}; use datafusion_physical_expr::create_physical_sort_expr; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; - use datafusion_physical_expr::projection::ProjectionExpr; + use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use datafusion_physical_plan::filter_pushdown::PushedDown; #[derive(Clone)] struct InexactSortPushdownSource { @@ -2647,4 +2650,467 @@ mod tests { Ok(()) } + + #[test] + fn test_partition_statistics_with_metadata_columns() { + use crate::source::DataSourceExec; + use datafusion_physical_plan::ExecutionPlan; + + // File schema: [compression] + let file_schema = Arc::new(Schema::new(vec![Field::new( + "compression", + DataType::Utf8, + true, + )])); + + // Partition column: [day] + let partition_cols = vec![Arc::new(Field::new("day", DataType::Utf8, true))]; + + let table_schema = TableSchema::new(Arc::clone(&file_schema), partition_cols); + + // Metadata column: location + let metadata_cols = vec![MetadataColumn::Location(None)]; + + let file_group_stats = Statistics { + num_rows: Precision::Exact(1), + total_byte_size: Precision::Exact(100), + column_statistics: vec![ + // compression + ColumnStatistics { + null_count: Precision::Exact(0), + ..ColumnStatistics::new_unknown() + }, + // day (partition) + ColumnStatistics { + null_count: Precision::Exact(0), + ..ColumnStatistics::new_unknown() + }, + ], + }; + + let file_group = FileGroup::new(vec![PartitionedFile::new("data.parquet", 100)]) + .with_statistics(Arc::new(file_group_stats)); + + // table_schema columns: [compression(0), day(1)] + // metadata columns: [location(2)] + // SELECT location, day, compression -> projection = [2, 1, 0] + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(MockSource::new(table_schema)), + ) + .with_metadata_cols(metadata_cols) + .expect("metadata cols valid") + .with_projection_indices(Some(vec![2, 1, 0])) + .expect("projection indices should be valid") + .with_file_groups(vec![file_group]) + .build(); + + let exec = DataSourceExec::from_data_source(config); + + // The projected schema has 3 columns: [location, day, compression] + assert_eq!(exec.schema().fields().len(), 3); + + // partition_statistics must return 3 column_statistics entries — + // one for each column in the projected schema — so that + // ProjectionExec doesn't panic on index access. + let stats = exec + .partition_statistics(Some(0)) + .expect("partition_statistics should succeed"); + assert_eq!( + stats.column_statistics.len(), + 3, + "Expected 3 column statistics (2 file/partition + 1 metadata), got {}", + stats.column_statistics.len() + ); + + // Also verify aggregate statistics (partition = None) + let agg_stats = exec + .partition_statistics(None) + .expect("aggregate partition_statistics should succeed"); + assert_eq!( + agg_stats.column_statistics.len(), + 3, + "Expected 3 aggregate column statistics, got {}", + agg_stats.column_statistics.len() + ); + } + + /// Test that `try_pushdown_filters` works when a filter references + /// a metadata column (e.g. `location`) that is appended after the projection. + /// Metadata column indices are beyond the projection length, so they must be + /// separated out rather than remapped through `update_expr`. + #[test] + fn test_try_pushdown_filters_with_metadata_column_filter() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let object_store_url = ObjectStoreUrl::parse("test:///").expect("valid url"); + let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let file_source: Arc = + Arc::new(MockSource::new(table_schema.clone())); + + let metadata_cols = vec![MetadataColumn::Location(None)]; + + // Build config with metadata columns and a projection that includes metadata. + let config = + FileScanConfigBuilder::new(object_store_url, Arc::clone(&file_source)) + .with_metadata_cols(metadata_cols) + .expect("metadata cols valid") + .with_projection_indices(Some(vec![0, 1, 2])) + .expect("valid projection") + .build(); + + // Push a projection so that `file_source.projection()` is Some. + let initial_schema = config.projected_schema().expect("projected schema"); + let proj_exprs = ProjectionExprs::new(vec![ + ProjectionExpr::new(col("id", &initial_schema).expect("col"), "id"), + ProjectionExpr::new(col("value", &initial_schema).expect("col"), "value"), + ProjectionExpr::new( + col("_location", &initial_schema).expect("col"), + "_location", + ), + ]); + let data_source = config + .try_swapping_with_projection(&proj_exprs) + .expect("swap ok") + .expect("should produce new source"); + + // projected schema: [id, value, location] + let projected = data_source + .as_any() + .downcast_ref::() + .expect("is FileScanConfig") + .projected_schema() + .expect("projected schema"); + assert_eq!(projected.fields().len(), 3); + assert_eq!(projected.field(2).name(), "_location"); + + // Create a filter on the metadata column: location@2 = 's3://bucket' + // (index 2 because projected output is [id@0, value@1, location@2]) + let location_filter: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("_location", 2)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "s3://bucket".to_string(), + )))), + )); + + // This should work without crashing, even though the filter references a metadata column + let config_options = ConfigOptions::default(); + let result = + data_source.try_pushdown_filters(vec![location_filter], &config_options); + let propagation = result.expect("to pushdown filters"); + + // The metadata filter cannot be pushed down to the file source. + assert_eq!(propagation.filters.len(), 1); + assert!( + matches!(propagation.filters[0], PushedDown::No), + "metadata column filter should not be pushed down" + ); + } + + /// Test that `try_pushdown_filters` correctly handles a mix of regular + /// filters and metadata column filters — regular filters are remapped + /// through the projection while metadata filters are returned as not pushed. + #[test] + fn test_try_pushdown_filters_mixed_regular_and_metadata_filters() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let object_store_url = ObjectStoreUrl::parse("test:///").expect("valid url"); + let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let file_source: Arc = + Arc::new(MockSource::new(table_schema.clone())); + + let metadata_cols = vec![MetadataColumn::Location(None)]; + + let config = + FileScanConfigBuilder::new(object_store_url, Arc::clone(&file_source)) + .with_metadata_cols(metadata_cols) + .expect("metadata cols valid") + .with_projection_indices(Some(vec![0, 1, 2])) + .expect("valid projection") + .build(); + + let initial_schema = config.projected_schema().expect("projected schema"); + let proj_exprs = ProjectionExprs::new(vec![ + ProjectionExpr::new(col("id", &initial_schema).expect("col"), "id"), + ProjectionExpr::new(col("value", &initial_schema).expect("col"), "value"), + ProjectionExpr::new( + col("_location", &initial_schema).expect("col"), + "_location", + ), + ]); + let data_source = config + .try_swapping_with_projection(&proj_exprs) + .expect("swap ok") + .expect("should produce new source"); + + // Filter 0: regular filter on id@0 (within projection) + let id_filter: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("id", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + // Filter 1: metadata column filter on location@2 (beyond projection) + let location_filter: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("_location", 2)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "s3://bucket".to_string(), + )))), + )); + + let config_options = ConfigOptions::default(); + let result = data_source + .try_pushdown_filters(vec![id_filter, location_filter], &config_options); + let propagation = result.expect("to pushdown filters"); + + // Both filters should be present in results, in the original order. + assert_eq!(propagation.filters.len(), 2); + // Regular filter: MockSource returns PushedDown::No (default impl). + assert!( + matches!(propagation.filters[0], PushedDown::No), + "regular filter: MockSource default returns No" + ); + // Metadata filter: separated out, returned as PushedDown::No. + assert!( + matches!(propagation.filters[1], PushedDown::No), + "metadata column filter should not be pushed down" + ); + } + + /// Test that `eq_properties` includes metadata columns in its schema. + /// `DataSourceExec::schema()` is derived from `eq_properties().schema()`, + /// so metadata columns must be appended to match the actual batches + /// produced by the file source. + #[test] + fn test_eq_properties_schema_includes_metadata_columns() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let object_store_url = ObjectStoreUrl::parse("test:///").expect("valid url"); + let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let file_source: Arc = + Arc::new(MockSource::new(table_schema.clone())); + + let metadata_cols = vec![MetadataColumn::Location(None), MetadataColumn::Size]; + + let config = FileScanConfigBuilder::new(object_store_url, file_source) + .with_metadata_cols(metadata_cols) + .expect("metadata cols valid") + .build(); + + let eq_props = config.eq_properties(); + let schema = eq_props.schema(); + + // Schema should include file columns + metadata columns + assert_eq!( + schema.fields().len(), + 4, + "Expected 4 fields (2 file + 2 metadata), got {}", + schema.fields().len() + ); + assert_eq!(schema.field(0).name(), "id"); + assert_eq!(schema.field(1).name(), "value"); + assert_eq!(schema.field(2).name(), "_location"); + assert_eq!(schema.field(3).name(), "_size"); + } + + /// Same as above but with a projection applied — metadata columns must + /// still appear in the schema after projection. + #[test] + fn test_eq_properties_schema_includes_metadata_columns_with_projection() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let object_store_url = ObjectStoreUrl::parse("test:///").expect("valid url"); + let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let file_source: Arc = + Arc::new(MockSource::new(table_schema.clone())); + + let metadata_cols = vec![MetadataColumn::Location(None)]; + + let config = + FileScanConfigBuilder::new(object_store_url, Arc::clone(&file_source)) + .with_metadata_cols(metadata_cols) + .expect("metadata cols valid") + .with_projection_indices(Some(vec![0, 1, 2])) + .expect("valid projection") + .build(); + + // Push a projection so that file_source.projection() is Some + let initial_schema = config.projected_schema().expect("projected schema"); + let proj_exprs = ProjectionExprs::new(vec![ + ProjectionExpr::new(col("id", &initial_schema).expect("col"), "id"), + ProjectionExpr::new(col("value", &initial_schema).expect("col"), "value"), + ProjectionExpr::new( + col("_location", &initial_schema).expect("col"), + "_location", + ), + ]); + let data_source = config + .try_swapping_with_projection(&proj_exprs) + .expect("swap ok") + .expect("should produce new source"); + + let eq_props = data_source.eq_properties(); + let schema = eq_props.schema(); + + // Projected schema: [id, value] + metadata: [location] + assert_eq!( + schema.fields().len(), + 3, + "Expected 3 fields (2 projected + 1 metadata), got {}", + schema.fields().len() + ); + assert_eq!(schema.field(0).name(), "id"); + assert_eq!(schema.field(1).name(), "value"); + assert_eq!(schema.field(2).name(), "_location"); + + // Verify it matches projected_schema() + let projected = data_source + .as_any() + .downcast_ref::() + .expect("is FileScanConfig") + .projected_schema() + .expect("projected schema"); + assert_eq!(schema.fields().len(), projected.fields().len()); + for (eq_field, proj_field) in + schema.fields().iter().zip(projected.fields().iter()) + { + assert_eq!(eq_field.name(), proj_field.name()); + } + } + + /// Test that eq_properties schema matches projected_schema when only + /// metadata columns are projected (no file/partition columns). + #[test] + fn test_eq_properties_metadata_only_projection() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let object_store_url = ObjectStoreUrl::parse("test:///").expect("valid url"); + + let partition_cols = vec![ + Arc::new(Field::new( + "year", + wrap_partition_type_in_dict(DataType::Utf8), + false, + )), + Arc::new(Field::new( + "month", + wrap_partition_type_in_dict(DataType::Utf8), + false, + )), + Arc::new(Field::new( + "day", + wrap_partition_type_in_dict(DataType::Utf8), + false, + )), + ]; + + let table_schema = + TableSchema::new(Arc::clone(&file_schema), partition_cols.clone()); + let file_source: Arc = + Arc::new(MockSource::new(table_schema.clone())); + + let metadata_cols = vec![ + MetadataColumn::LastModified, + MetadataColumn::Location(None), + MetadataColumn::Size, + ]; + + // table_schema = [id(0), value(1), year(2), month(3), day(4)] = 5 cols + // metadata = [last_modified(5), location(6), size(7)] + // SELECT location -> projection = [6] + let config = FileScanConfigBuilder::new(object_store_url, file_source) + .with_metadata_cols(metadata_cols) + .expect("metadata cols valid") + .with_projection_indices(Some(vec![6])) + .expect("valid projection") + .build(); + + let eq_props = config.eq_properties(); + let eq_schema = eq_props.schema(); + + let projected = config.projected_schema().expect("projected schema"); + + // Both schemas must match: [location] + assert_eq!( + eq_schema.fields().len(), + projected.fields().len(), + "eq_properties schema has {} fields but projected_schema has {}", + eq_schema.fields().len(), + projected.fields().len(), + ); + for (eq_field, proj_field) in + eq_schema.fields().iter().zip(projected.fields().iter()) + { + assert_eq!( + eq_field.name(), + proj_field.name(), + "Field name mismatch: eq_properties has '{}' but projected_schema has '{}'", + eq_field.name(), + proj_field.name(), + ); + } + } + + /// Test that eq_properties schema matches projected_schema when a mix of + /// file columns and metadata columns are projected, with metadata appearing + /// at the beginning of the output. + #[test] + fn test_eq_properties_metadata_before_file_columns() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let object_store_url = ObjectStoreUrl::parse("test:///").expect("valid url"); + let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let file_source: Arc = + Arc::new(MockSource::new(table_schema.clone())); + + let metadata_cols = vec![MetadataColumn::Location(None), MetadataColumn::Size]; + + // table_schema = [id(0), value(1)] = 2 cols + // metadata = [location(2), size(3)] + // SELECT location, id -> projection = [2, 0] + // Metadata "_location" should be at output position 0, "id" after it. + let config = FileScanConfigBuilder::new(object_store_url, file_source) + .with_metadata_cols(metadata_cols) + .expect("metadata cols valid") + .with_projection_indices(Some(vec![2, 0])) + .expect("valid projection") + .build(); + + let eq_props = config.eq_properties(); + let eq_schema = eq_props.schema(); + + let projected = config.projected_schema().expect("projected schema"); + + assert_eq!( + eq_schema.fields().len(), + projected.fields().len(), + "eq_properties schema has {} fields but projected_schema has {}", + eq_schema.fields().len(), + projected.fields().len(), + ); + for (eq_field, proj_field) in + eq_schema.fields().iter().zip(projected.fields().iter()) + { + assert_eq!( + eq_field.name(), + proj_field.name(), + "Field name mismatch: eq_properties has '{}' but projected_schema has '{}'", + eq_field.name(), + proj_field.name(), + ); + } + } } diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs index c8d5dd54cb8a2..160f5ee89fd7c 100644 --- a/datafusion/datasource/src/test_util.rs +++ b/datafusion/datasource/src/test_util.rs @@ -88,6 +88,17 @@ impl FileSource for MockSource { Arc::new(Self { ..self.clone() }) } + fn with_metadata_cols( + &self, + metadata_cols: Vec, + ) -> Option> { + let mut source = self.clone(); + source.table_schema = source.table_schema.with_metadata_cols(metadata_cols); + source.projection = + crate::projection::SplitProjection::unprojected(&source.table_schema); + Some(Arc::new(source)) + } + fn metrics(&self) -> &ExecutionPlanMetricsSet { &self.metrics } @@ -106,8 +117,8 @@ impl FileSource for MockSource { ) -> Result>> { let mut source = self.clone(); let new_projection = self.projection.source.try_merge(projection)?; - let split_projection = crate::projection::SplitProjection::new( - self.table_schema.file_schema(), + let split_projection = crate::projection::SplitProjection::new_with_table_schema( + &self.table_schema, &new_projection, ); source.projection = split_projection; diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index dbbd289415277..3e492d705dd14 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -661,7 +661,15 @@ impl ProjectionExprs { for proj_expr in self.exprs.iter() { let expr = &proj_expr.expr; let col_stats = if let Some(col) = expr.as_any().downcast_ref::() { - std::mem::take(&mut stats.column_statistics[col.index()]) + // Spice metadata columns (and projected partition columns) reference + // table-schema indices that fall beyond the upstream file-level + // statistics, which only cover the file (and partition) columns. Those + // injected columns have no known statistics, so fall back to unknown + if col.index() < stats.column_statistics.len() { + std::mem::take(&mut stats.column_statistics[col.index()]) + } else { + ColumnStatistics::new_unknown() + } } else if let Some(literal) = expr.as_any().downcast_ref::() { // Handle literal expressions (constants) by calculating proper statistics let data_type = expr.data_type(output_schema)?; @@ -2737,6 +2745,54 @@ pub(crate) mod tests { Ok(()) } + #[test] + fn test_project_statistics_column_index_beyond_input_stats() -> Result<()> { + // File-level statistics only cover the three file columns. + let input_stats = get_stats(); + assert_eq!(input_stats.column_statistics.len(), 3); + + // The table schema additionally exposes a metadata column at index 3. + let table_schema = Schema::new(vec![ + Field::new("col0", DataType::Int64, false), + Field::new("col1", DataType::Utf8, false), + Field::new("col2", DataType::Float32, false), + Field::new("_location", DataType::Utf8, true), + ]); + + // SELECT col0, _location — the metadata column (index 3) is beyond the + // input statistics width (3). + let projection = ProjectionExprs::new(vec![ + ProjectionExpr { + expr: Arc::new(Column::new("col0", 0)), + alias: "col0".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("_location", 3)), + alias: "_location".to_string(), + }, + ]); + + let output_stats = projection + .project_statistics(input_stats, &projection.project_schema(&table_schema)?)?; + + assert_eq!(output_stats.num_rows, Precision::Exact(5)); + assert_eq!(output_stats.column_statistics.len(), 2); + + // col0 keeps its known statistics. + assert_eq!( + output_stats.column_statistics[0].max_value, + Precision::Exact(ScalarValue::Int64(Some(21))) + ); + + // The metadata column has no upstream statistics -> unknown, not a panic. + assert_eq!( + output_stats.column_statistics[1], + ColumnStatistics::new_unknown() + ); + + Ok(()) + } + #[test] fn test_project_statistics_primitive_width_only() -> Result<()> { let input_stats = get_stats();