From 4da6ebf6fe9b52f0e082931207860ddd16db027b Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 18:58:22 +0800 Subject: [PATCH 01/14] fix(planner): derive SMJ ordering from left input Port the sort-merge join behavior carried by fork commit 0f7361ac7606cb7421b3c9ac6b4ef70aa62d9cb0 onto DataFusion 55. Preserve descending sort options already provided by the left input and default only unmatched join keys. --- datafusion/core/src/physical_planner.rs | 151 +++++++++++++++++++++++- 1 file changed, 145 insertions(+), 6 deletions(-) diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 3c1e7b50780a5..6256b0f650484 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -57,7 +57,6 @@ use crate::physical_plan::{ use crate::schema_equivalence::schema_satisfied_by; use arrow::array::{RecordBatch, builder::StringBuilder}; -use arrow::compute::SortOptions; use arrow::datatypes::Schema; use arrow_schema::Field; use datafusion_catalog::ScanArgs; @@ -1708,15 +1707,27 @@ impl DefaultPhysicalPlanner { // implement null-aware anti-join semantics and would return wrong // results when the right side contains a null join key. { - // Use SortMergeJoin if hash join is not preferred - let join_on_len = join_on.len(); + let sort_options = join_on + .iter() + .map(|(left_col, _)| { + physical_left + .output_ordering() + .and_then(|ordering| { + ordering + .iter() + .find(|sort_expr| sort_expr.expr.eq(left_col)) + .map(|sort_expr| sort_expr.options) + }) + .unwrap_or_default() + }) + .collect(); Arc::new(SortMergeJoinExec::try_new( physical_left, physical_right, join_on, join_filter, *join_type, - vec![SortOptions::default(); join_on_len], + sort_options, *null_equality, )?) } else if session_state.config().target_partitions() > 1 @@ -3316,8 +3327,8 @@ mod tests { use datafusion_expr::{ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, HigherOrderUDF, LogicalPlanBuilder, Partitioning as LogicalPartitioning, RangePartitioning, - ScalarUDF, Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, - WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery, + ScalarUDF, Signature, SortExpr, TableSource, UserDefinedLogicalNodeCore, + Volatility, WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; @@ -3820,6 +3831,134 @@ mod tests { ctx.sql(query).await?.collect().await } + fn smj_test_context( + left_ordering: Option>, + right_ordering: Option>, + ) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![3, 2, 1])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + )?; + + let table = |batch, ordering| -> Result { + let table = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?; + Ok(match ordering { + Some(ordering) => table.with_sort_order(vec![ordering]), + None => table, + }) + }; + let left = table(batch.clone(), left_ordering)?; + let right = table(batch, right_ordering)?; + + let config = SessionConfig::new() + .with_target_partitions(4) + .set_bool("datafusion.optimizer.prefer_hash_join", false) + .set_bool("datafusion.optimizer.skip_failed_rules", false); + let state = SessionStateBuilder::new() + .with_config(config) + .with_runtime_env(Arc::new(RuntimeEnv::default())) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + ctx.register_table("left_table", Arc::new(left))?; + ctx.register_table("right_table", Arc::new(right))?; + Ok(ctx) + } + + async fn smj_plan(ctx: &SessionContext, sql: &str) -> Result { + let state = ctx.state(); + let logical_plan = state.create_logical_plan(sql).await?; + let logical_plan = state.optimize(&logical_plan)?; + let plan = DefaultPhysicalPlanner::default() + .create_physical_plan(&logical_plan, &state) + .await?; + Ok(format!("{}", displayable(plan.as_ref()).indent(false))) + } + + #[tokio::test] + async fn smj_uses_left_descending_sort_options() -> Result<()> { + let descending = vec![col("a").sort(false, false)]; + let ctx = smj_test_context(Some(descending), None)?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)] + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + Ok(()) + } + + #[tokio::test] + async fn smj_defaults_to_ascending_when_left_is_unsorted() -> Result<()> { + let ctx = smj_test_context(None, None)?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)] + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + Ok(()) + } + + #[tokio::test] + async fn smj_defaults_only_unmatched_join_columns_to_ascending() -> Result<()> { + let descending = vec![col("a").sort(false, false)]; + let ctx = smj_test_context(Some(descending), None)?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table \ + ON left_table.a = right_table.a AND left_table.b = right_table.b", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0), (b@1, b@1)] + SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + Ok(()) + } + + #[tokio::test] + async fn smj_reuses_matching_descending_order_on_both_sides() -> Result<()> { + let descending = vec![col("a").sort(false, false)]; + let ctx = smj_test_context(Some(descending.clone()), Some(descending))?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)] + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + "); + Ok(()) + } + #[tokio::test] async fn test_all_operators() -> Result<()> { let logical_plan = test_csv_scan() From 7e2f2179abeb1264fb56a541fef4d9db9717a3b2 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 19:01:08 +0800 Subject: [PATCH 02/14] fix(physical-plan): fall back when interleave rewrites diverge Reimplement the fallback carried by fork commit 0f7361ac7606cb7421b3c9ac6b4ef70aa62d9cb0 for DataFusion 55's replace_children API. Compatible children retain InterleaveExec; incompatible partitioning safely degrades to UnionExec. --- datafusion/physical-plan/src/union.rs | 48 +++++++++++++++++++++++---- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index c1cc5da31abaf..de85e5c20bc14 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -746,12 +746,14 @@ impl ExecutionPlan for InterleaveExec { ..Self::clone(&*self) })), ChildrenPropertiesMode::Recompute => { - // New children are no longer interleavable, which might be a bug of optimization rewrite. - assert_or_internal_err!( - can_interleave(children.iter()), - "Can not create InterleaveExec: new children can not be interleaved" - ); - Ok(Arc::new(InterleaveExec::try_new(children)?)) + if can_interleave(children.iter()) { + Ok(Arc::new(InterleaveExec::try_new(children)?)) + } else { + // An optimizer can legitimately change child partitioning after + // introducing InterleaveExec. UnionExec preserves correctness when + // the rewritten children can no longer share output partitions. + UnionExec::try_new(children) + } } } } @@ -1764,6 +1766,40 @@ mod tests { Ok(()) } + #[test] + fn interleave_child_rewrite_falls_back_to_union() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let interleave: Arc = + Arc::new(InterleaveExec::try_new(vec![ + make_hash_exec(&schema, vec!["a"], 3)?, + make_hash_exec(&schema, vec!["a"], 3)?, + ])?); + + let compatible = Arc::clone(&interleave).replace_children( + vec![ + make_hash_exec(&schema, vec!["b"], 3)?, + make_hash_exec(&schema, vec!["b"], 3)?, + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + assert!(compatible.downcast_ref::().is_some()); + assert_eq!(compatible.output_partitioning().partition_count(), 3); + + let incompatible = interleave.replace_children( + vec![ + make_hash_exec(&schema, vec!["a"], 3)?, + make_hash_exec(&schema, vec!["b"], 3)?, + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + assert!(incompatible.downcast_ref::().is_some()); + assert_eq!(incompatible.output_partitioning().partition_count(), 6); + Ok(()) + } + #[test] fn test_union_cardinality_effect() -> Result<()> { let schema = create_test_schema()?; From 148904432fe3185af5118fedb0a5082b97faacd6 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 19:07:26 +0800 Subject: [PATCH 03/14] perf: reuse projected equivalence groups Port fork commit b8f161807ad2af15e58cfa14d4ed1fa37e65092a using the final Apache implementation from PR #24445 (f95102074b1e7ebe8b398767de5064fc74d962f2), which landed after the 55.0.0 tag. --- .../physical-expr/src/equivalence/class.rs | 117 +++++- .../src/equivalence/properties/mod.rs | 191 ++++++++++ datafusion/physical-plan/src/projection.rs | 349 +++++++++++++++++- 3 files changed, 649 insertions(+), 8 deletions(-) diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index 1f9a6a583cc44..a9593d3dfd6ce 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -310,6 +310,33 @@ pub struct EquivalenceGroup { } impl EquivalenceGroup { + /// A cheap, deliberately conservative check that two groups hold the same + /// equivalence classes. + /// + /// This is not `PartialEq`, and the distinction is the point. `classes` is a + /// `Vec` whose order carries no meaning -- `remove_class_at_idx` uses + /// `swap_remove` -- so two groups describing exactly the same equalities can + /// hold their classes in different orders and this returns `false` for them. + /// Naming it `PartialEq` would invite callers to read it as semantic equality, + /// which it is not. + /// + /// The comparison is positional because it runs on a hot path: + /// `ProjectionExec` consults it every time a rule replaces its child. Set + /// semantics would mean scanning the other group once per class, and that + /// quadratic term costs more than the recomputation the caller is trying to + /// skip, by a margin that widens with the number of classes. + /// + /// Only the false direction is reachable: a group can be reported different + /// when it is not, never the same when it is not. Callers using this to skip + /// work must be built so that a `false` merely costs them that work, which is + /// exactly how the projection fast path uses it. + /// + /// `map` is an index into `classes` and carries no information the classes do + /// not already have, so it takes no part in the comparison. + pub fn has_same_classes(&self, other: &Self) -> bool { + self.classes == other.classes + } + /// Creates an equivalence group from the given equivalence classes. pub fn new(classes: impl IntoIterator) -> Self { classes.into_iter().collect::>().into() @@ -924,7 +951,7 @@ mod tests { use super::*; use crate::equivalence::tests::create_test_params; use crate::expressions::{BinaryExpr, Column, binary, col, lit}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_expr::Operator; @@ -1245,4 +1272,92 @@ mod tests { Ok(()) } + + /// Builds a group from a list of equated column pairs. + fn group_of(schema: &SchemaRef, pairs: &[(&str, &str)]) -> Result { + let mut group = EquivalenceGroup::default(); + for (lhs, rhs) in pairs { + group.add_equal_conditions(col(lhs, schema)?, col(rhs, schema)?); + } + Ok(group) + } + + fn abc_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + ])) + } + + #[test] + fn test_has_same_classes_is_conservative_about_class_order() -> Result<()> { + // `classes` is a `Vec` and `remove_class_at_idx` uses `swap_remove`, so + // the order two groups hold their classes in depends on how they were + // built. This check is positional and reports such a pair as different. + // + // That is deliberate, and it is why this is not `PartialEq`: set + // semantics would mean scanning the other group per class, and the + // quadratic cost dwarfs the recomputation the caller is trying to skip. + // The error can only go this way -- different when they match, never the + // reverse -- so a caller only forfeits an optimization. + let schema = abc_schema(); + let ab_then_cd = group_of(&schema, &[("a", "b"), ("c", "d")])?; + let cd_then_ab = group_of(&schema, &[("c", "d"), ("a", "b")])?; + + assert_eq!(ab_then_cd.len(), 2, "expected two disjoint classes"); + assert!(!ab_then_cd.has_same_classes(&cd_then_ab)); + + Ok(()) + } + + #[test] + fn test_has_same_classes_compares_classes() -> Result<()> { + let schema = abc_schema(); + + // Two empty groups agree. + assert!(group_of(&schema, &[])?.has_same_classes(&group_of(&schema, &[])?)); + // The same class, built the same way. + assert!( + group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "b")])?) + ); + // A class is a set, so the order within a pair is immaterial. + assert!( + group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("b", "a")])?) + ); + // A populated group is not an empty one. + assert!( + !group_of(&schema, &[("a", "b")])?.has_same_classes(&group_of(&schema, &[])?) + ); + // Equating a different pair yields a different group. + assert!( + !group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "c")])?) + ); + // Widening a class yields a different group. + assert!( + !group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "b"), ("b", "c")])?) + ); + + Ok(()) + } + + #[test] + fn test_has_same_classes_ignores_the_map() -> Result<()> { + // `map` indexes into `classes`, so equal classes must imply equal + // groups no matter how the classes were arrived at. Bridging `a = b` + // and `b = c` into one class must match stating `a = c` and `a = b`. + let schema = abc_schema(); + let bridged = group_of(&schema, &[("a", "b"), ("b", "c")])?; + let direct = group_of(&schema, &[("a", "c"), ("a", "b")])?; + + assert_eq!(bridged.len(), 1, "expected a single bridged class"); + assert!(bridged.has_same_classes(&direct)); + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 22b3382f50638..6a0fc76afe778 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -194,6 +194,31 @@ impl OrderingEquivalenceCache { } } +/// Counts how often [`EquivalenceProperties::project_reusing`] reused a cached +/// equivalence group instead of reprojecting it. +/// +/// Thread local rather than a global counter: the test binary runs tests in +/// parallel, and a shared count would let one test observe another's hits. +#[cfg(test)] +mod eq_group_reuse_probe { + use std::cell::Cell; + + thread_local! { + static HITS: Cell = const { Cell::new(0) }; + } + + pub(super) fn record_hit() { + HITS.with(|h| h.set(h.get() + 1)); + } + + /// Runs `f` and reports how many reuses happened while it did. + pub(super) fn count(f: impl FnOnce() -> T) -> (T, usize) { + let before = HITS.with(Cell::get); + let value = f(); + (value, HITS.with(Cell::get) - before) + } +} + impl EquivalenceProperties { /// Helper used by the ordering equivalence rule when considering whether /// an expression can replace an existing sort key without invalidating @@ -1169,6 +1194,48 @@ impl EquivalenceProperties { /// `output_schema`. pub fn project(&self, mapping: &ProjectionMapping, output_schema: SchemaRef) -> Self { let eq_group = self.eq_group.project(mapping); + // Built here, so it satisfies the precondition by construction; going + // through the checked entry point would reproject it under + // `debug_assertions` for nothing. + self.project_with_eq_group_unchecked(mapping, output_schema, eq_group) + } + + /// Projects `self`, reusing `cached`'s already-projected equivalence group + /// when `self`'s group is unchanged from `previous`'s. + /// + /// [`EquivalenceGroup::project`] is a pure function of the group and the + /// mapping, so when the group is unchanged the previous result can be handed + /// back rather than recomputed. Orderings are still derived here: they are + /// precisely what changes when a sort is introduced below this node. + /// + /// Falls back to a full projection when the groups differ, so there is no + /// precondition to violate. The caller does still have to pass a `cached` + /// that came from projecting `previous` through this same `mapping`; that + /// part cannot be checked here, and in practice it holds because the caller + /// carries its projection over untouched. + pub fn project_reusing( + &self, + mapping: &ProjectionMapping, + output_schema: SchemaRef, + previous: &EquivalenceProperties, + cached: &EquivalenceProperties, + ) -> Self { + let eq_group = if self.eq_group.has_same_classes(&previous.eq_group) { + #[cfg(test)] + eq_group_reuse_probe::record_hit(); + cached.eq_group.clone() + } else { + self.eq_group.project(mapping) + }; + self.project_with_eq_group_unchecked(mapping, output_schema, eq_group) + } + + fn project_with_eq_group_unchecked( + &self, + mapping: &ProjectionMapping, + output_schema: SchemaRef, + eq_group: EquivalenceGroup, + ) -> Self { let orderings = self.projected_orderings(mapping, self.oeq_cache.normal_cls.clone()); let normal_orderings = orderings @@ -1540,3 +1607,127 @@ fn get_expr_properties( expr.get_properties(&child_states) } } + +#[cfg(test)] +mod tests { + use super::*; + + use crate::equivalence::tests::create_test_params; + use crate::expressions::col; + + use arrow::datatypes::{DataType, Field, Schema}; + + /// Renames `a`, `b`, `c` and `d` so the projection is non-trivial while + /// still carrying every ordering and the `a = c` class across. + fn renaming_mapping( + schema: &SchemaRef, + output_schema: &SchemaRef, + ) -> Result { + [ + ("a", "a1", 0), + ("b", "b1", 1), + ("c", "c1", 2), + ("d", "d1", 3), + ] + .into_iter() + .map(|(source, target, index)| { + Ok(( + col(source, schema)?, + vec![(col(target, output_schema)?, index)].into(), + )) + }) + .collect::>>() + .map(|entries| entries.into_iter().collect()) + } + + fn renamed_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, true), + Field::new("b1", DataType::Int32, true), + Field::new("c1", DataType::Int32, true), + Field::new("d1", DataType::Int32, true), + ])) + } + + #[test] + fn test_project_reusing_matches_project() -> Result<()> { + let (schema, eq_properties) = create_test_params()?; + let output_schema = renamed_schema(); + let mapping = renaming_mapping(&schema, &output_schema)?; + + let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema)); + + // Reusing the projection of an identical group must reproduce it exactly. + let (reused, hits) = eq_group_reuse_probe::count(|| { + eq_properties.project_reusing( + &mapping, + Arc::clone(&output_schema), + &eq_properties, + &baseline, + ) + }); + assert_eq!( + hits, 1, + "the reuse path was not taken, so this compares nothing" + ); + + // Guard against a vacuous comparison. + assert!( + !baseline.eq_group().is_empty(), + "the projection dropped the equivalence class, nothing is being tested" + ); + assert!( + !baseline.oeq_class().is_empty(), + "the projection dropped every ordering, nothing is being tested" + ); + + assert!( + baseline.eq_group().has_same_classes(reused.eq_group()), + "equivalence group" + ); + assert_eq!(baseline.oeq_class(), reused.oeq_class(), "orderings"); + assert_eq!(baseline.constraints(), reused.constraints(), "constraints"); + assert_eq!(baseline.schema(), reused.schema(), "schema"); + + Ok(()) + } + + #[test] + fn test_project_reusing_falls_back_when_the_group_moved() -> Result<()> { + // A `previous` whose group differs must not have its projection carried + // over. There is no precondition to violate here: the fallback is what + // keeps a mismatched pair from producing properties that disagree with + // the projection. + let (schema, eq_properties) = create_test_params()?; + let output_schema = renamed_schema(); + let mapping = renaming_mapping(&schema, &output_schema)?; + + let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema)); + let unrelated = EquivalenceProperties::new(Arc::clone(&schema)); + assert!( + !eq_properties + .eq_group() + .has_same_classes(unrelated.eq_group()), + "the two groups were meant to differ" + ); + + let (recomputed, hits) = eq_group_reuse_probe::count(|| { + eq_properties.project_reusing( + &mapping, + Arc::clone(&output_schema), + &unrelated, + &baseline, + ) + }); + assert_eq!(hits, 0, "a mismatched group was carried over anyway"); + + // Falling back must land on exactly what `project` would have produced. + assert!( + baseline.eq_group().has_same_classes(recomputed.eq_group()), + "equivalence group" + ); + assert_eq!(baseline.oeq_class(), recomputed.oeq_class(), "orderings"); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index cf362cdee55d3..1ec278eb9c377 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -52,6 +52,7 @@ use datafusion_common::tree_node::{ use datafusion_common::{DataFusionError, JoinSide, Result, internal_err, plan_err}; use datafusion_execution::TaskContext; use datafusion_expr::ExpressionPlacement; +use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr::equivalence::ProjectionMapping; use datafusion_physical_expr::projection::Projector; use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; @@ -177,6 +178,32 @@ impl ProjectionExec { fn try_from_projector( projector: Projector, input: Arc, + ) -> Result { + Self::try_from_projector_with_eq_group(projector, input, None) + } + + /// As [`Self::try_from_projector`], but `reuse_from` may carry the previous + /// child's equivalence properties together with the projection they produced, + /// letting [`EquivalenceProperties::project_reusing`] skip reprojecting a + /// group that has not changed. + /// + /// Projecting an equivalence group is a pure function of the group and the + /// mapping, so reuse is sound exactly when both are unchanged. + /// + /// The caller establishes the first by comparing the old and new child + /// groups. The second holds because the mapping comes from + /// `projector.projection()`, carried over untouched, and from the child's + /// schema, which `ProjectionMapping::try_new` consults only for field names + /// and indices -- never for types or nullability. So a child differing only + /// in nullability keeps the same mapping. A child that renamed or reordered + /// those fields would change the group too, since its members are `Column`s + /// carrying those names, and the comparison above would reject it; were one + /// to slip through anyway, `try_new`'s name assertion errors out rather than + /// letting a stale group into the plan. + fn try_from_projector_with_eq_group( + projector: Projector, + input: Arc, + reuse_from: Option<(&EquivalenceProperties, &EquivalenceProperties)>, ) -> Result { // Construct a map from the input expressions to the output expression of the Projection let projection_mapping = @@ -185,6 +212,7 @@ impl ProjectionExec { &input, &projection_mapping, Arc::clone(projector.output_schema()), + reuse_from, )?; Ok(Self { projector, @@ -214,10 +242,21 @@ impl ProjectionExec { input: &Arc, projection_mapping: &ProjectionMapping, schema: SchemaRef, + reuse_from: Option<(&EquivalenceProperties, &EquivalenceProperties)>, ) -> Result { - // Calculate equivalence properties: + // Calculate equivalence properties. Whether the group is reprojected or + // handed back is the only thing reuse changes; everything below is + // common, so the two paths cannot drift apart. let input_eq_properties = input.equivalence_properties(); - let eq_properties = input_eq_properties.project(projection_mapping, schema); + let eq_properties = match reuse_from { + Some((previous, cached)) => input_eq_properties.project_reusing( + projection_mapping, + schema, + previous, + cached, + ), + None => input_eq_properties.project(projection_mapping, schema), + }; // Calculate output partitioning, which needs to respect aliases: let output_partitioning = input .output_partitioning() @@ -352,11 +391,28 @@ impl ExecutionPlan for ProjectionExec { metrics: ExecutionPlanMetricsSet::new(), ..Self::clone(&*self) })), - ChildrenPropertiesMode::Recompute => ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), - ) - .map(|p| Arc::new(p) as _), + ChildrenPropertiesMode::Recompute => { + // `Keep` above requires the child's properties to be unchanged + // outright. A rule that introduces a sort below this projection + // does not qualify, yet the child's *equivalence group* is still + // identical: sorting changes which orderings hold, not which + // expressions are equal to one another. Projecting that group + // again would reproduce the group already cached here, so reuse + // it and derive only the orderings. + // Hand over what this projection was built from and what that + // produced; `project_reusing` decides whether the group can be + // carried over and falls back to a full projection otherwise. + let reuse_from = Some(( + self.input.equivalence_properties(), + self.cache.equivalence_properties(), + )); + ProjectionExec::try_from_projector_with_eq_group( + self.projector.clone(), + children.swap_remove(0), + reuse_from, + ) + .map(|p| Arc::new(p) as _) + } } } @@ -1436,6 +1492,8 @@ mod tests { use crate::common::collect; use crate::empty::EmptyExec; + use crate::filter::FilterExec; + use crate::sorts::sort::SortExec; use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; @@ -2215,4 +2273,281 @@ mod tests { Ok(()) } + + /// `EmptyExec(a, b, c)` under a filter that equates `lhs` and `rhs`, so the + /// child carries a non-trivial equivalence group. + fn filtered_source(lhs: &str, rhs: &str) -> Result> { + filtered_source_with_nullability(lhs, rhs, false) + } + + /// As [`filtered_source`], but `nullable` varies the schema's nullability + /// while leaving field names and order alone. + fn filtered_source_with_nullability( + lhs: &str, + rhs: &str, + nullable: bool, + ) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, nullable), + Field::new("b", DataType::Int32, nullable), + Field::new("c", DataType::Int32, nullable), + ])); + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let predicate = binary( + col(lhs, &schema)?, + Operator::Eq, + col(rhs, &schema)?, + &schema, + )?; + Ok(Arc::new(FilterExec::try_new(predicate, input)?)) + } + + /// `[a AS x, b AS y, c AS z]` against `filtered_source`'s schema. + fn renaming_exprs(schema: &SchemaRef) -> Result> { + [("a", "x"), ("b", "y"), ("c", "z")] + .into_iter() + .map(|(source, alias)| { + Ok(ProjectionExpr { + expr: col(source, schema)?, + alias: alias.to_string(), + }) + }) + .collect() + } + + fn assert_same_properties(actual: &dyn ExecutionPlan, expected: &ProjectionExec) { + let actual_props = actual.properties().equivalence_properties(); + let expected_props = expected.properties().equivalence_properties(); + assert!( + actual_props + .eq_group() + .has_same_classes(expected_props.eq_group()), + "equivalence group: {:?} vs {:?}", + actual_props.eq_group(), + expected_props.eq_group() + ); + assert_eq!( + actual_props.oeq_class(), + expected_props.oeq_class(), + "orderings" + ); + assert_eq!( + actual_props.constraints(), + expected_props.constraints(), + "constraints" + ); + assert_eq!(actual_props.schema(), expected_props.schema(), "schema"); + // `Partitioning` has no `PartialEq`, so compare the partition count and + // the explicit `Display` form. Derived `Debug` would change with any + // field addition, making this brittle for no gain. + let actual_partitioning = actual.properties().output_partitioning(); + let expected_partitioning = expected.properties().output_partitioning(); + assert_eq!( + actual_partitioning.partition_count(), + expected_partitioning.partition_count(), + "partition count" + ); + assert_eq!( + actual_partitioning.to_string(), + expected_partitioning.to_string(), + "partitioning" + ); + } + + #[test] + fn test_sort_below_changes_orderings_but_not_the_equivalence_group() -> Result<()> { + // The premise the fast path rests on. If a sort ever starts altering + // the equivalence group, reusing the cached group becomes unsound and + // this test is the one that should fail first. + let child = filtered_source("a", "b")?; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "c", + &child.schema(), + )?)]) + .expect("non-empty ordering"); + let sorted = SortExec::new(ordering, Arc::clone(&child)); + + let child_props = child.properties().equivalence_properties(); + let sorted_props = sorted.properties().equivalence_properties(); + + assert!( + !child_props.eq_group().is_empty(), + "the filter did not produce an equivalence class" + ); + assert!( + child_props + .eq_group() + .has_same_classes(sorted_props.eq_group()), + "sorting altered the equivalence group" + ); + assert_ne!( + child_props.oeq_class(), + sorted_props.oeq_class(), + "sorting did not alter the orderings" + ); + + Ok(()) + } + + #[test] + fn test_replace_children_reuses_eq_group_when_only_orderings_change() -> Result<()> { + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = + Arc::new(ProjectionExec::try_new(exprs.clone(), Arc::clone(&child))?); + + // Sorting below the projection changes which orderings hold but leaves + // the equivalence group untouched -- the case the fast path targets. + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "c", + &child.schema(), + )?)]) + .expect("non-empty ordering"); + let sorted: Arc = + Arc::new(SortExec::new(ordering, Arc::clone(&child))); + + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&sorted)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + // Guard against vacuity: the group must be worth reusing, and the sort + // must genuinely have added an ordering the original did not have. + assert!( + !projection + .properties() + .equivalence_properties() + .eq_group() + .is_empty(), + "the projection carries no equivalence class, nothing is being reused" + ); + assert!( + projection + .properties() + .equivalence_properties() + .oeq_class() + .is_empty(), + "the unsorted projection was already ordered" + ); + assert!( + !replaced + .properties() + .equivalence_properties() + .oeq_class() + .is_empty(), + "the sort did not introduce an ordering" + ); + + // The fast path must agree with building the projection from scratch. + let expected = ProjectionExec::try_new(exprs, sorted)?; + assert_same_properties(replaced.as_ref(), &expected); + + Ok(()) + } + + #[test] + fn test_replace_children_reuses_eq_group_across_a_nullability_change() -> Result<()> { + // Reuse is sound only if the projection mapping is unchanged as well. + // `ProjectionMapping::try_new` reads the child schema for field names + // and indices alone, so a child differing only in nullability keeps the + // same mapping and must still take the fast path. + // + // The swap tightens nullability rather than loosening it, matching what + // `is_allowed_field_change` permits of a physical optimizer rule. + // + // The comparison is against `try_from_projector`, the path this one + // replaces, rather than a freshly built projection: `replace_children` + // carries the existing `Projector` over, so the output schema stays as + // it was, while `try_new` would derive a new one from the new child. + // That difference is inherent to `replace_children` and not something + // this fast path introduces, so the meaningful contract is that the two + // `replace_children` paths agree. + let child = filtered_source_with_nullability("a", "b", true)?; + let exprs = renaming_exprs(&child.schema())?; + let projection = Arc::new(ProjectionExec::try_new(exprs, Arc::clone(&child))?); + + let tightened_child = filtered_source_with_nullability("a", "b", false)?; + assert_ne!( + child.schema(), + tightened_child.schema(), + "the two children were meant to differ in nullability" + ); + assert!( + child + .properties() + .equivalence_properties() + .eq_group() + .has_same_classes( + tightened_child + .properties() + .equivalence_properties() + .eq_group() + ), + "nullability moved the equivalence group, so the fast path is no longer under test" + ); + + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&tightened_child)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let recomputed = ProjectionExec::try_from_projector( + projection.projector.clone(), + tightened_child, + )?; + + assert_same_properties(replaced.as_ref(), &recomputed); + + Ok(()) + } + + #[test] + fn test_replace_children_recomputes_when_eq_group_changes() -> Result<()> { + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = + Arc::new(ProjectionExec::try_new(exprs.clone(), Arc::clone(&child))?); + + // This child equates a different pair, so the cached group is stale and + // reusing it would be unsound: the guard has to fall through. + let other = filtered_source("a", "c")?; + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&other)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + assert!( + !replaced + .properties() + .equivalence_properties() + .eq_group() + .has_same_classes( + projection.properties().equivalence_properties().eq_group() + ), + "the projection kept the previous child's equivalence group" + ); + + let expected = ProjectionExec::try_new(exprs, other)?; + assert_same_properties(replaced.as_ref(), &expected); + + Ok(()) + } + + #[test] + fn test_replace_children_keep_mode_carries_properties_over() -> Result<()> { + // `Keep` is the caller's promise that the new child's properties match + // the old one's, so the cached properties must survive verbatim rather + // than being derived again. + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = Arc::new(ProjectionExec::try_new(exprs, Arc::clone(&child))?); + + let replaced = Arc::clone(&projection).replace_children( + vec![filtered_source("a", "b")?], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + )?; + + assert_same_properties(replaced.as_ref(), &projection); + + Ok(()) + } } From 8445c36e981bff27d45ffee0137e99df7e79db64 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 19:11:42 +0800 Subject: [PATCH 04/14] test(parquet): preserve fork schema adaptation behavior Carry forward the end-to-end regressions from fork commits 0f7361ac7606cb7421b3c9ac6b4ef70aa62d9cb0, 5df21f6a0e2e73c0030dd30933ab4e4a3e27e2b1, and 87fa8980744bcf27c3f890a283e4191f88894f60. DataFusion 55's new reader architecture passes all three without implementation changes. --- .../datasource-parquet/src/opener/mod.rs | 140 +++++++++++++++++- 1 file changed, 135 insertions(+), 5 deletions(-) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 693e9bd2cbf31..8e29a528d5b37 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1749,7 +1749,11 @@ mod test { CachedParquetFileReaderFactory, DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetRowSelection, RowGroupAccess, }; - use arrow::array::{RecordBatch, record_batch}; + use arrow::array::{ + Array, ArrayRef, Date32Array, Int32Array, ListArray, RecordBatch, StringArray, + record_batch, + }; + use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use bytes::{BufMut, BytesMut}; use datafusion_common::{ @@ -1774,8 +1778,8 @@ mod test { }; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_pruning::MAX_IN_LIST_SIZE; - use futures::StreamExt; use futures::stream::BoxStream; + use futures::{StreamExt, TryStreamExt}; use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path}; use parquet::arrow::{ArrowSchemaConverter, ArrowWriter}; use parquet::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData}; @@ -2340,13 +2344,12 @@ mod test { async fn collect_int32_values( mut stream: BoxStream<'static, Result>, ) -> Vec { - use arrow::array::Array; let mut values = vec![]; while let Some(Ok(batch)) = stream.next().await { let array = batch .column(0) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); for i in 0..array.len() { if !array.is_null(i) { @@ -2365,6 +2368,133 @@ mod test { write_parquet_batches(store, filename, vec![batch], None).await } + async fn read_batches_with_schema( + store: Arc, + filename: &str, + data_size: usize, + schema: SchemaRef, + ) -> Result> { + let projection_indices = (0..schema.fields().len()).collect::>(); + let file = + PartitionedFile::new(filename.to_string(), u64::try_from(data_size).unwrap()); + let morselizer = ParquetMorselizerBuilder::new() + .with_store(store) + .with_schema(schema) + .with_projection_indices(&projection_indices) + .build(); + open_file(&morselizer, file).await?.try_collect().await + } + + #[tokio::test] + async fn schema_adaptation_casts_utf8_to_date32() -> Result<()> { + let store = Arc::new(InMemory::new()) as Arc; + let file_schema = + Arc::new(Schema::new(vec![Field::new("date", DataType::Utf8, true)])); + let batch = RecordBatch::try_new( + file_schema, + vec![Arc::new(StringArray::from(vec![ + "2026-01-01", + "2026-02-01", + ]))], + )?; + let data_size = + write_parquet(Arc::clone(&store), "schema_adapt_dates.parquet", batch).await; + + let logical_schema = Arc::new(Schema::new(vec![Field::new( + "date", + DataType::Date32, + true, + )])); + let batches = read_batches_with_schema( + store, + "schema_adapt_dates.parquet", + data_size, + logical_schema, + ) + .await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 2); + let dates = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(dates.values(), &[20_454, 20_485]); + Ok(()) + } + + #[tokio::test] + async fn schema_adaptation_normalizes_list_fields() -> Result<()> { + let store = Arc::new(InMemory::new()) as Arc; + let inner_field = Arc::new(Field::new("conditions", DataType::Int32, false)); + let file_schema = Arc::new(Schema::new(vec![Field::new( + "conditions", + DataType::List(Arc::clone(&inner_field)), + true, + )])); + let values = Int32Array::from(vec![1, 2, 3, 4]); + let offsets = OffsetBuffer::from_lengths([2, 2]); + let list = ListArray::new(inner_field, offsets, Arc::new(values), None); + let batch = RecordBatch::try_new(file_schema, vec![Arc::new(list) as ArrayRef])?; + let data_size = + write_parquet(Arc::clone(&store), "schema_adapt_list.parquet", batch).await; + + let logical_inner = Arc::new(Field::new("element", DataType::Int32, true)); + let logical_schema = Arc::new(Schema::new(vec![Field::new( + "conditions", + DataType::List(logical_inner), + true, + )])); + let batches = read_batches_with_schema( + store, + "schema_adapt_list.parquet", + data_size, + Arc::clone(&logical_schema), + ) + .await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 2); + assert_eq!(batches[0].schema(), logical_schema); + Ok(()) + } + + #[tokio::test] + async fn schema_adaptation_preserves_logical_nullability() -> Result<()> { + let store = Arc::new(InMemory::new()) as Arc; + let file_schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + file_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + let data_size = + write_parquet(Arc::clone(&store), "schema_adapt_nullable.parquet", batch) + .await; + + let logical_schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); + let batches = read_batches_with_schema( + store, + "schema_adapt_nullable.parquet", + data_size, + logical_schema, + ) + .await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + assert!(batches[0].schema().field(0).is_nullable()); + let ids = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values(), &[1, 2, 3]); + Ok(()) + } + /// Write multiple batches to a parquet file with optional writer properties async fn write_parquet_batches( store: Arc, @@ -3742,7 +3872,7 @@ mod test { let part = batch .column(1) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); assert!(part.iter().all(|v| v == Some(5))); From 785c2fe46ea1c30ec3b29ee6745ffcb511f7c515 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 19:21:15 +0800 Subject: [PATCH 05/14] feat(parquet): expose reverse row-group controls Port the downstream-facing ParquetSource API from fork commit 0f7361ac7606cb7421b3c9ac6b4ef70aa62d9cb0. DataFusion 55 retained the behavior internally but restricted the setter and getter to crate tests. --- datafusion/datasource-parquet/src/source.rs | 23 +++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 097b4563af5df..3d4ec2cb25081 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -507,13 +507,28 @@ impl ParquetSource { } } - #[cfg(test)] - pub(crate) fn with_reverse_row_groups(mut self, reverse_row_groups: bool) -> Self { + /// Sets whether row groups are read in reverse order. + /// + /// This changes row-group traversal only; rows within each row group retain + /// their file order. + /// + /// ``` + /// use std::sync::Arc; + /// + /// use arrow_schema::Schema; + /// use datafusion_datasource_parquet::source::ParquetSource; + /// + /// let source = ParquetSource::new(Arc::new(Schema::empty())) + /// .with_reverse_row_groups(true); + /// assert!(source.reverse_row_groups()); + /// ``` + pub fn with_reverse_row_groups(mut self, reverse_row_groups: bool) -> Self { self.reverse_row_groups = reverse_row_groups; self } - #[cfg(test)] - pub(crate) fn reverse_row_groups(&self) -> bool { + + /// Returns whether row groups are read in reverse order. + pub fn reverse_row_groups(&self) -> bool { self.reverse_row_groups } } From 3631413a1b2dd64902a43ee51bd24777d5e3db79 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 19:45:36 +0800 Subject: [PATCH 06/14] fix(proto): preserve max IN-list presence Reimplements fork commit 32b07786c for DataFusion 55. Keep max_row_group_bytes on upstream field 37 while using an optional oneof on field 38 so omitted values retain the documented default and explicit zero still disables IN-list pruning. The old fork's field-37 encoding collides with DataFusion 55's max_row_group_bytes and cannot be decoded unambiguously; the upgrade compatibility guide documents the required rolling-upgrade boundary. --- .../datasource-parquet/src/file_format.rs | 6 ++- .../proto/datafusion_common.proto | 8 +++- datafusion/proto-common/src/from_proto/mod.rs | 37 ++++++++++++++- .../proto-common/src/generated/pbjson.rs | 46 ++++++++++--------- .../proto-common/src/generated/prost.rs | 13 +++++- datafusion/proto-common/src/to_proto/mod.rs | 2 +- datafusion/proto-models/src/from_proto.rs | 12 ++++- .../src/generated/datafusion_proto_common.rs | 13 +++++- .../proto/tests/cases/public_conversions.rs | 26 +++++++++++ 9 files changed, 131 insertions(+), 32 deletions(-) diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 6358201c06fa5..057f77815320c 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -732,7 +732,11 @@ impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions { parquet_options::StatisticsEnabledOpt::StatisticsEnabled(enabled) }), max_row_group_size: global_options.global.max_row_group_size as u64, - max_in_list_size: global_options.global.max_in_list_size as u64, + max_in_list_size_opt: Some( + parquet_options::MaxInListSizeOpt::MaxInListSize( + global_options.global.max_in_list_size as u64, + ), + ), created_by: global_options.global.created_by.clone(), column_index_truncate_length_opt: global_options.global.column_index_truncate_length.map(|length| { parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length as u64) diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 27d1101036d9b..71fca2d214100 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -617,7 +617,11 @@ message ParquetOptions { uint64 max_row_group_size = 15; - uint64 max_in_list_size = 38; + // Presence distinguishes an omitted value (default 20) from an explicit + // zero, which disables IN-list pruning. + oneof max_in_list_size_opt { + uint64 max_in_list_size = 38; + } string created_by = 16; @@ -717,4 +721,4 @@ enum MetricCategory { message ExplainAnalyzeCategoriesNode { bool all = 1; repeated MetricCategory only = 2; -} \ No newline at end of file +} diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 169ff7f3d9ff2..21a1614c969c3 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1081,7 +1081,12 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { }) .unwrap_or(None), max_row_group_size: value.max_row_group_size as usize, - max_in_list_size: value.max_in_list_size as usize, + // Plans written before this field was introduced do not carry a value. + // Preserve the documented default instead of treating absence as zero, + // which disables IN-list pruning. + max_in_list_size: value.max_in_list_size_opt.map(|opt| match opt { + protobuf::parquet_options::MaxInListSizeOpt::MaxInListSize(v) => v as usize, + }).unwrap_or(20), created_by: value.created_by.clone(), column_index_truncate_length: value .column_index_truncate_length_opt.as_ref() @@ -1393,6 +1398,36 @@ mod tests { ); } + #[test] + fn test_parquet_options_max_in_list_size_round_trip() { + let opts = ParquetOptions { + max_in_list_size: 64, + ..ParquetOptions::default() + }; + let recovered = parquet_options_proto_round_trip(opts); + assert_eq!(recovered.max_in_list_size, 64); + } + + #[test] + fn test_parquet_options_max_in_list_size_zero_round_trip() { + let opts = ParquetOptions { + max_in_list_size: 0, + ..ParquetOptions::default() + }; + let recovered = parquet_options_proto_round_trip(opts); + assert_eq!(recovered.max_in_list_size, 0); + } + + #[test] + fn test_parquet_options_max_in_list_size_absent_uses_default() { + let opts = ParquetOptions::default(); + let mut proto: crate::protobuf_common::ParquetOptions = + (&opts).try_into().expect("to_proto"); + proto.max_in_list_size_opt = None; + let recovered = ParquetOptions::try_from(&proto).expect("from_proto"); + assert_eq!(recovered.max_in_list_size, 20); + } + #[test] fn test_table_parquet_options_coerce_int96_tz_round_trip() { let mut opts = TableParquetOptions::default(); diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index c222cd1cb8687..6a29ad3400ff3 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -6409,9 +6409,6 @@ impl serde::Serialize for ParquetOptions { if self.max_row_group_size != 0 { len += 1; } - if self.max_in_list_size != 0 { - len += 1; - } if !self.created_by.is_empty() { len += 1; } @@ -6445,6 +6442,9 @@ impl serde::Serialize for ParquetOptions { if self.bloom_filter_ndv_opt.is_some() { len += 1; } + if self.max_in_list_size_opt.is_some() { + len += 1; + } if self.coerce_int96_opt.is_some() { len += 1; } @@ -6532,11 +6532,6 @@ impl serde::Serialize for ParquetOptions { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("maxRowGroupSize", ToString::to_string(&self.max_row_group_size).as_str())?; } - if self.max_in_list_size != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("maxInListSize", ToString::to_string(&self.max_in_list_size).as_str())?; - } if !self.created_by.is_empty() { struct_ser.serialize_field("createdBy", &self.created_by)?; } @@ -6614,6 +6609,15 @@ impl serde::Serialize for ParquetOptions { } } } + if let Some(v) = self.max_in_list_size_opt.as_ref() { + match v { + parquet_options::MaxInListSizeOpt::MaxInListSize(v) => { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("maxInListSize", ToString::to_string(&v).as_str())?; + } + } + } if let Some(v) = self.coerce_int96_opt.as_ref() { match v { parquet_options::CoerceInt96Opt::CoerceInt96(v) => { @@ -6695,8 +6699,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dataPageRowCountLimit", "max_row_group_size", "maxRowGroupSize", - "max_in_list_size", - "maxInListSize", "created_by", "createdBy", "content_defined_chunking", @@ -6717,6 +6719,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "bloomFilterFpp", "bloom_filter_ndv", "bloomFilterNdv", + "max_in_list_size", + "maxInListSize", "coerce_int96", "coerceInt96", "max_predicate_cache_size", @@ -6749,7 +6753,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { DictionaryPageSizeLimit, DataPageRowCountLimit, MaxRowGroupSize, - MaxInListSize, CreatedBy, ContentDefinedChunking, MetadataSizeHint, @@ -6761,6 +6764,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Encoding, BloomFilterFpp, BloomFilterNdv, + MaxInListSize, CoerceInt96, MaxPredicateCacheSize, MaxRowGroupBytes, @@ -6806,7 +6810,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dictionaryPageSizeLimit" | "dictionary_page_size_limit" => Ok(GeneratedField::DictionaryPageSizeLimit), "dataPageRowCountLimit" | "data_page_row_count_limit" => Ok(GeneratedField::DataPageRowCountLimit), "maxRowGroupSize" | "max_row_group_size" => Ok(GeneratedField::MaxRowGroupSize), - "maxInListSize" | "max_in_list_size" => Ok(GeneratedField::MaxInListSize), "createdBy" | "created_by" => Ok(GeneratedField::CreatedBy), "contentDefinedChunking" | "content_defined_chunking" => Ok(GeneratedField::ContentDefinedChunking), "metadataSizeHint" | "metadata_size_hint" => Ok(GeneratedField::MetadataSizeHint), @@ -6818,6 +6821,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "encoding" => Ok(GeneratedField::Encoding), "bloomFilterFpp" | "bloom_filter_fpp" => Ok(GeneratedField::BloomFilterFpp), "bloomFilterNdv" | "bloom_filter_ndv" => Ok(GeneratedField::BloomFilterNdv), + "maxInListSize" | "max_in_list_size" => Ok(GeneratedField::MaxInListSize), "coerceInt96" | "coerce_int96" => Ok(GeneratedField::CoerceInt96), "maxPredicateCacheSize" | "max_predicate_cache_size" => Ok(GeneratedField::MaxPredicateCacheSize), "maxRowGroupBytes" | "max_row_group_bytes" => Ok(GeneratedField::MaxRowGroupBytes), @@ -6861,7 +6865,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut dictionary_page_size_limit__ = None; let mut data_page_row_count_limit__ = None; let mut max_row_group_size__ = None; - let mut max_in_list_size__ = None; let mut created_by__ = None; let mut content_defined_chunking__ = None; let mut metadata_size_hint_opt__ = None; @@ -6873,6 +6876,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut encoding_opt__ = None; let mut bloom_filter_fpp_opt__ = None; let mut bloom_filter_ndv_opt__ = None; + let mut max_in_list_size_opt__ = None; let mut coerce_int96_opt__ = None; let mut max_predicate_cache_size_opt__ = None; let mut max_row_group_bytes_opt__ = None; @@ -7013,14 +7017,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } - GeneratedField::MaxInListSize => { - if max_in_list_size__.is_some() { - return Err(serde::de::Error::duplicate_field("maxInListSize")); - } - max_in_list_size__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; - } GeneratedField::CreatedBy => { if created_by__.is_some() { return Err(serde::de::Error::duplicate_field("createdBy")); @@ -7087,6 +7083,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { } bloom_filter_ndv_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::BloomFilterNdvOpt::BloomFilterNdv(x.0)); } + GeneratedField::MaxInListSize => { + if max_in_list_size_opt__.is_some() { + return Err(serde::de::Error::duplicate_field("maxInListSize")); + } + max_in_list_size_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::MaxInListSizeOpt::MaxInListSize(x.0)); + } GeneratedField::CoerceInt96 => { if coerce_int96_opt__.is_some() { return Err(serde::de::Error::duplicate_field("coerceInt96")); @@ -7134,7 +7136,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { dictionary_page_size_limit: dictionary_page_size_limit__.unwrap_or_default(), data_page_row_count_limit: data_page_row_count_limit__.unwrap_or_default(), max_row_group_size: max_row_group_size__.unwrap_or_default(), - max_in_list_size: max_in_list_size__.unwrap_or_default(), created_by: created_by__.unwrap_or_default(), content_defined_chunking: content_defined_chunking__, metadata_size_hint_opt: metadata_size_hint_opt__, @@ -7146,6 +7147,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { encoding_opt: encoding_opt__, bloom_filter_fpp_opt: bloom_filter_fpp_opt__, bloom_filter_ndv_opt: bloom_filter_ndv_opt__, + max_in_list_size_opt: max_in_list_size_opt__, coerce_int96_opt: coerce_int96_opt__, max_predicate_cache_size_opt: max_predicate_cache_size_opt__, max_row_group_bytes_opt: max_row_group_bytes_opt__, diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index bdbe38538e1d7..2c588aec829d6 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -862,8 +862,6 @@ pub struct ParquetOptions { pub data_page_row_count_limit: u64, #[prost(uint64, tag = "15")] pub max_row_group_size: u64, - #[prost(uint64, tag = "38")] - pub max_in_list_size: u64, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] @@ -896,6 +894,10 @@ pub struct ParquetOptions { pub bloom_filter_fpp_opt: ::core::option::Option, #[prost(oneof = "parquet_options::BloomFilterNdvOpt", tags = "22")] pub bloom_filter_ndv_opt: ::core::option::Option, + /// Presence distinguishes an omitted value (default 20) from an explicit + /// zero, which disables IN-list pruning. + #[prost(oneof = "parquet_options::MaxInListSizeOpt", tags = "38")] + pub max_in_list_size_opt: ::core::option::Option, #[prost(oneof = "parquet_options::CoerceInt96Opt", tags = "32")] pub coerce_int96_opt: ::core::option::Option, #[prost(oneof = "parquet_options::MaxPredicateCacheSizeOpt", tags = "33")] @@ -960,6 +962,13 @@ pub mod parquet_options { #[prost(uint64, tag = "22")] BloomFilterNdv(u64), } + /// Presence distinguishes an omitted value (default 20) from an explicit + /// zero, which disables IN-list pruning. + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum MaxInListSizeOpt { + #[prost(uint64, tag = "38")] + MaxInListSize(u64), + } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] pub enum CoerceInt96Opt { #[prost(string, tag = "32")] diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 360981746585b..75a2ded4c9979 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -912,7 +912,7 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { dictionary_page_size_limit: value.dictionary_page_size_limit as u64, statistics_enabled_opt: value.statistics_enabled.clone().map(protobuf::parquet_options::StatisticsEnabledOpt::StatisticsEnabled), max_row_group_size: value.max_row_group_size as u64, - max_in_list_size: value.max_in_list_size as u64, + max_in_list_size_opt: Some(protobuf::parquet_options::MaxInListSizeOpt::MaxInListSize(value.max_in_list_size as u64)), created_by: value.created_by.clone(), column_index_truncate_length_opt: value.column_index_truncate_length.map(|v| protobuf::parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(v as u64)), statistics_truncate_length_opt: value.statistics_truncate_length.map(|v| protobuf::parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(v as u64)), diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index 74ead8c52049b..2936143bb6dda 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -374,7 +374,17 @@ impl TryFrom<&ParquetOptionsProto> for ParquetOptions { }, ), max_row_group_size: proto.max_row_group_size as usize, - max_in_list_size: proto.max_in_list_size as usize, + // Preserve the documented default when an older plan omits the + // field; primitive zero disables IN-list pruning. + max_in_list_size: proto + .max_in_list_size_opt + .as_ref() + .map(|opt| match opt { + parquet_options::MaxInListSizeOpt::MaxInListSize(size) => { + *size as usize + } + }) + .unwrap_or(20), created_by: proto.created_by.clone(), column_index_truncate_length: proto .column_index_truncate_length_opt diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index bdbe38538e1d7..2c588aec829d6 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -862,8 +862,6 @@ pub struct ParquetOptions { pub data_page_row_count_limit: u64, #[prost(uint64, tag = "15")] pub max_row_group_size: u64, - #[prost(uint64, tag = "38")] - pub max_in_list_size: u64, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] @@ -896,6 +894,10 @@ pub struct ParquetOptions { pub bloom_filter_fpp_opt: ::core::option::Option, #[prost(oneof = "parquet_options::BloomFilterNdvOpt", tags = "22")] pub bloom_filter_ndv_opt: ::core::option::Option, + /// Presence distinguishes an omitted value (default 20) from an explicit + /// zero, which disables IN-list pruning. + #[prost(oneof = "parquet_options::MaxInListSizeOpt", tags = "38")] + pub max_in_list_size_opt: ::core::option::Option, #[prost(oneof = "parquet_options::CoerceInt96Opt", tags = "32")] pub coerce_int96_opt: ::core::option::Option, #[prost(oneof = "parquet_options::MaxPredicateCacheSizeOpt", tags = "33")] @@ -960,6 +962,13 @@ pub mod parquet_options { #[prost(uint64, tag = "22")] BloomFilterNdv(u64), } + /// Presence distinguishes an omitted value (default 20) from an explicit + /// zero, which disables IN-list pruning. + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum MaxInListSizeOpt { + #[prost(uint64, tag = "38")] + MaxInListSize(u64), + } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] pub enum CoerceInt96Opt { #[prost(string, tag = "32")] diff --git a/datafusion/proto/tests/cases/public_conversions.rs b/datafusion/proto/tests/cases/public_conversions.rs index 1cda8e01a0765..f486ad06214ff 100644 --- a/datafusion/proto/tests/cases/public_conversions.rs +++ b/datafusion/proto/tests/cases/public_conversions.rs @@ -126,3 +126,29 @@ fn file_format_option_conversions_are_std_traits() { assert_from::<&JsonFormatFactory, protobuf::JsonOptions>(); assert_from::<&ParquetFormatFactory, protobuf::TableParquetOptions>(); } + +#[test] +fn parquet_max_in_list_size_preserves_zero_and_absent_default() { + let mut options = TableParquetOptions::default(); + options.global.max_in_list_size = 0; + let factory = ParquetFormatFactory::new_with_options(options); + let mut proto = protobuf::TableParquetOptions::from(&factory); + + let global = proto.global.as_ref().expect("global parquet options"); + assert!(matches!( + global.max_in_list_size_opt, + Some(protobuf::parquet_options::MaxInListSizeOpt::MaxInListSize( + 0 + )) + )); + let decoded = TableParquetOptions::try_from(&proto).expect("from_proto"); + assert_eq!(decoded.global.max_in_list_size, 0); + + proto + .global + .as_mut() + .expect("global parquet options") + .max_in_list_size_opt = None; + let decoded = TableParquetOptions::try_from(&proto).expect("from_proto"); + assert_eq!(decoded.global.max_in_list_size, 20); +} From c8801b38ae2e218b242a2a1eee4b832be0265c22 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 19:58:08 +0800 Subject: [PATCH 07/14] test(expr): preserve extension expression fast path Adds focused regression coverage for the expressionless extension-node fast path carried by fork commit 0f7361ac. Mapping expressions must keep the original extension Arc and avoid with_exprs_and_inputs reconstruction. --- datafusion/expr/src/logical_plan/tree_node.rs | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c4c1d743b58b6..7e30568c450d2 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -1034,3 +1034,73 @@ impl LogicalPlan { }) } } + +#[cfg(test)] +mod tests { + use std::fmt::Formatter; + use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + use std::sync::{Arc, LazyLock}; + + use datafusion_common::{DFSchema, DFSchemaRef, Result}; + + use super::*; + use crate::UserDefinedLogicalNodeCore; + + static REBUILD_COUNT: AtomicUsize = AtomicUsize::new(0); + static EMPTY_SCHEMA: LazyLock = + LazyLock::new(|| Arc::new(DFSchema::empty())); + + #[derive(Debug, Eq, Hash, PartialEq, PartialOrd)] + struct EmptyExpressionNode; + + impl UserDefinedLogicalNodeCore for EmptyExpressionNode { + fn name(&self) -> &str { + "EmptyExpressionNode" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &EMPTY_SCHEMA + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "EmptyExpressionNode") + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec, + inputs: Vec, + ) -> Result { + assert!(exprs.is_empty()); + assert!(inputs.is_empty()); + REBUILD_COUNT.fetch_add(1, AtomicOrdering::Relaxed); + Ok(Self) + } + } + + #[test] + fn map_expressions_does_not_rebuild_expressionless_extension() -> Result<()> { + REBUILD_COUNT.store(0, AtomicOrdering::Relaxed); + let original: Arc = Arc::new(EmptyExpressionNode); + let plan = LogicalPlan::Extension(Extension { + node: Arc::clone(&original), + }); + + let transformed = plan.map_expressions(|expr| Ok(Transformed::no(expr)))?; + let LogicalPlan::Extension(Extension { node }) = transformed.data else { + panic!("expected extension plan") + }; + + assert!(Arc::ptr_eq(&node, &original)); + assert_eq!(REBUILD_COUNT.load(AtomicOrdering::Relaxed), 0); + Ok(()) + } +} From 026199f56931ea607978ea73cca086133c3fd3f1 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Sun, 23 Aug 2026 20:06:07 +0800 Subject: [PATCH 08/14] chore(deps): update h2 for security advisory Ports Apache DataFusion commit f6ad7810d (PR #24467) onto the 55.0.0 fork baseline. h2 0.4.16 fixes RUSTSEC-2026-0258 and restores the maintained fork's security-audit requirement without adding an ignore. --- Cargo.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 62c04d332b98e..f045e7697f576 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3284,9 +3284,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.13" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" dependencies = [ "atomic-waker", "bytes", From c997b36d592ed83ade72fd7fb1ff8b8c1ed01119 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 24 Aug 2026 11:04:44 +0800 Subject: [PATCH 09/14] perf: restore concrete grouped Top-K storage Port fork commit 2465b2cd14de (X-3345) to DataFusion 55 while preserving the new all-NULL group bookkeeping. Store concrete Arrow arrays in primitive heaps and hash tables, borrow string keys for comparisons, and swap heap slots in place. --- .../src/aggregates/topk/hash_table.rs | 177 +++++++++--------- .../physical-plan/src/aggregates/topk/heap.rs | 30 ++- 2 files changed, 102 insertions(+), 105 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index adc8f8c315b32..771bf69182acf 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -106,6 +106,54 @@ pub trait ArrowHashTable { fn null_map_idxs(&self) -> Vec; } +enum StringArrayType { + Utf8(StringArray), + Utf8View(StringViewArray), + LargeUtf8(LargeStringArray), +} + +impl StringArrayType { + fn value(&self, row_idx: usize) -> Option<&str> { + match self { + Self::Utf8(array) if !array.is_null(row_idx) => Some(array.value(row_idx)), + Self::Utf8View(array) if !array.is_null(row_idx) => { + Some(array.value(row_idx)) + } + Self::LargeUtf8(array) if !array.is_null(row_idx) => { + Some(array.value(row_idx)) + } + _ => None, + } + } +} + +impl TryFrom<&DataType> for StringArrayType { + type Error = DataType; + + fn try_from(data_type: &DataType) -> std::result::Result { + let values = Vec::<&str>::new(); + match data_type { + DataType::Utf8 => Ok(Self::Utf8(values.into())), + DataType::Utf8View => Ok(Self::Utf8View(values.into())), + DataType::LargeUtf8 => Ok(Self::LargeUtf8(values.into())), + data_type => Err(data_type.clone()), + } + } +} + +impl TryFrom for StringArrayType { + type Error = DataType; + + fn try_from(array: ArrayRef) -> std::result::Result { + match array.data_type() { + DataType::Utf8 => Ok(Self::Utf8(array.as_string::().clone())), + DataType::Utf8View => Ok(Self::Utf8View(array.as_string_view().clone())), + DataType::LargeUtf8 => Ok(Self::LargeUtf8(array.as_string::().clone())), + data_type => Err(data_type.clone()), + } + } +} + /// Returns true if the given data type can be used as a top-K aggregation hash key. /// /// Supported types include Arrow primitives (integers, floats, decimals, intervals) @@ -121,10 +169,9 @@ pub fn is_supported_hash_key_type(kt: &DataType) -> bool { // An implementation of ArrowHashTable for String keys pub struct StringHashTable { - owned: ArrayRef, + owned: StringArrayType, map: TopKHashTable>, rnd: RandomState, - data_type: DataType, } // An implementation of ArrowHashTable for any `ArrowPrimitiveType` key @@ -132,69 +179,30 @@ struct PrimitiveHashTable where Option<::Native>: Comparable, { - owned: ArrayRef, + owned: PrimitiveArray, map: TopKHashTable>, rnd: RandomState, - kt: DataType, } impl StringHashTable { - pub fn new(limit: usize, data_type: DataType) -> Self { - let vals: Vec<&str> = Vec::new(); - let owned: ArrayRef = match data_type { - DataType::Utf8 => Arc::new(StringArray::from(vals)), - DataType::Utf8View => Arc::new(StringViewArray::from(vals)), - DataType::LargeUtf8 => Arc::new(LargeStringArray::from(vals)), - _ => panic!("Unsupported data type"), + pub fn new(limit: usize, data_type: &DataType) -> Self { + let Ok(owned) = StringArrayType::try_from(data_type) else { + unreachable!("StringHashTable requires a string data type") }; - Self { owned, map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), - data_type, } } - - /// Extracts the string value at the given row index, handling nulls and different string types. - /// - /// Returns `None` if the value is null, otherwise `Some(value.to_string())`. - fn extract_string_value(&self, row_idx: usize) -> Option { - let is_null_and_value = match self.data_type { - DataType::Utf8 => { - let arr = self.owned.as_string::(); - (arr.is_null(row_idx), arr.value(row_idx)) - } - DataType::LargeUtf8 => { - let arr = self.owned.as_string::(); - (arr.is_null(row_idx), arr.value(row_idx)) - } - DataType::Utf8View => { - let arr = self.owned.as_string_view(); - (arr.is_null(row_idx), arr.value(row_idx)) - } - _ => panic!("Unsupported data type"), - }; - - let (is_null, value) = is_null_and_value; - if is_null { - None - } else { - Some(value.to_string()) - } - } - - /// Computes the id and its hash for the given row, for hash table lookups - fn id_and_hash(&self, row_idx: usize) -> (Option, u64) { - let id = self.extract_string_value(row_idx); - let hash = self.rnd.hash_one(id.as_deref()); - (id, hash) - } } impl ArrowHashTable for StringHashTable { fn set_batch(&mut self, ids: ArrayRef) { - self.owned = ids; + let Ok(owned) = StringArrayType::try_from(ids) else { + unreachable!("StringHashTable requires a string array") + }; + self.owned = owned; } fn len(&self) -> usize { @@ -211,11 +219,10 @@ impl ArrowHashTable for StringHashTable { fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - match self.data_type { - DataType::Utf8 => Arc::new(StringArray::from(ids)), - DataType::LargeUtf8 => Arc::new(LargeStringArray::from(ids)), - DataType::Utf8View => Arc::new(StringViewArray::from(ids)), - _ => unreachable!(), + match &self.owned { + StringArrayType::Utf8(_) => Arc::new(StringArray::from(ids)), + StringArrayType::Utf8View(_) => Arc::new(StringViewArray::from(ids)), + StringArrayType::LargeUtf8(_) => Arc::new(LargeStringArray::from(ids)), } } @@ -224,27 +231,26 @@ impl ArrowHashTable for StringHashTable { row_idx: usize, replace_idx: usize, ) -> (usize, InsertKind) { - let id = self.extract_string_value(row_idx); - - // Compute hash and create equality closure for hash table lookup. - let hash = self.rnd.hash_one(id.as_deref()); - let id_for_eq = id.clone(); - let eq = move |mi: &Option| id_for_eq.as_deref() == mi.as_deref(); + let id = self.owned.value(row_idx); + let hash = self.rnd.hash_one(id); + let eq = move |stored: &Option| id == stored.as_deref(); // Use entry API to avoid double lookup - self.map.find_or_insert(hash, id, replace_idx, eq) + self.map + .find_or_insert(hash, id.map(ToOwned::to_owned), replace_idx, eq) } fn insert_null(&mut self, row_idx: usize) -> bool { - let (id, hash) = self.id_and_hash(row_idx); - let id_for_eq = id.clone(); - let eq = move |mi: &Option| id_for_eq.as_deref() == mi.as_deref(); - self.map.insert_null(hash, id, eq) + let id = self.owned.value(row_idx); + let hash = self.rnd.hash_one(id); + let eq = move |stored: &Option| id == stored.as_deref(); + self.map.insert_null(hash, id.map(ToOwned::to_owned), eq) } fn remove_if_null(&mut self, row_idx: usize) -> bool { - let (id, hash) = self.id_and_hash(row_idx); - let eq = move |mi: &Option| id.as_deref() == mi.as_deref(); + let id = self.owned.value(row_idx); + let hash = self.rnd.hash_one(id); + let eq = move |stored: &Option| id == stored.as_deref(); self.map.remove_if_null(hash, eq) } @@ -255,30 +261,25 @@ impl ArrowHashTable for StringHashTable { impl PrimitiveHashTable where - Option<::Native>: Comparable, - Option<::Native>: HashValue, + Option<::Native>: Comparable + HashValue, { pub fn new(limit: usize, kt: DataType) -> Self { - let owned = Arc::new( - PrimitiveArray::::builder(0) - .with_data_type(kt.clone()) - .finish(), - ); + let owned = PrimitiveArray::::builder(0) + .with_data_type(kt) + .finish(); Self { owned, map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), - kt, } } /// Computes the id and its hash for the given row, for hash table lookups fn id_and_hash(&self, row_idx: usize) -> (Option, u64) { - let ids = self.owned.as_primitive::(); - let id: Option = if ids.is_null(row_idx) { + let id = if self.owned.is_null(row_idx) { None } else { - Some(ids.value(row_idx)) + Some(self.owned.value(row_idx)) }; let hash: u64 = id.hash(&self.rnd); (id, hash) @@ -287,11 +288,10 @@ where impl ArrowHashTable for PrimitiveHashTable where - Option<::Native>: Comparable, - Option<::Native>: HashValue, + Option<::Native>: Comparable + HashValue, { fn set_batch(&mut self, ids: ArrayRef) { - self.owned = ids; + self.owned = ids.as_primitive().clone(); } fn len(&self) -> usize { @@ -308,8 +308,8 @@ where fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - let mut builder: PrimitiveBuilder = - PrimitiveArray::builder(ids.len()).with_data_type(self.kt.clone()); + let mut builder: PrimitiveBuilder = PrimitiveArray::builder(ids.len()) + .with_data_type(self.owned.data_type().clone()); for id in ids.into_iter() { match id { None => builder.append_null(), @@ -325,11 +325,10 @@ where row_idx: usize, replace_idx: usize, ) -> (usize, InsertKind) { - let ids = self.owned.as_primitive::(); - let id: Option = if ids.is_null(row_idx) { + let id = if self.owned.is_null(row_idx) { None } else { - Some(ids.value(row_idx)) + Some(self.owned.value(row_idx)) }; // Compute hash and create equality closure for hash table lookup. let hash: u64 = id.hash(&self.rnd); @@ -597,9 +596,9 @@ pub fn new_hash_table( downcast_primitive! { kt => (downcast_helper, kt), - DataType::Utf8 => return Ok(Box::new(StringHashTable::new(limit, DataType::Utf8))), - DataType::LargeUtf8 => return Ok(Box::new(StringHashTable::new(limit, DataType::LargeUtf8))), - DataType::Utf8View => return Ok(Box::new(StringHashTable::new(limit, DataType::Utf8View))), + DataType::Utf8 => return Ok(Box::new(StringHashTable::new(limit, &DataType::Utf8))), + DataType::LargeUtf8 => return Ok(Box::new(StringHashTable::new(limit, &DataType::LargeUtf8))), + DataType::Utf8View => return Ok(Box::new(StringHashTable::new(limit, &DataType::Utf8View))), _ => {} } diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs index ca321cdf99784..3659c9757c7ae 100644 --- a/datafusion/physical-plan/src/aggregates/topk/heap.rs +++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs @@ -95,7 +95,7 @@ pub struct PrimitiveHeap where ::Native: Comparable, { - batch: ArrayRef, + batch: PrimitiveArray, heap: TopKHeap, desc: bool, data_type: DataType, @@ -106,9 +106,11 @@ where ::Native: Comparable, { pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { - let owned: ArrayRef = Arc::new(PrimitiveArray::::builder(0).finish()); + let batch = PrimitiveArray::::builder(0) + .with_data_type(data_type.clone()) + .finish(); Self { - batch: owned, + batch, heap: TopKHeap::new(limit, desc), desc, data_type, @@ -121,15 +123,14 @@ where ::Native: Comparable, { fn set_batch(&mut self, vals: ArrayRef) { - self.batch = vals; + self.batch = vals.as_primitive().clone(); } fn is_worse(&self, row_idx: usize) -> bool { if !self.heap.is_full() { return false; } - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); let worst_val = self.heap.worst_val().expect("Missing root"); (!self.desc && new_val > *worst_val) || (self.desc && new_val < *worst_val) } @@ -139,8 +140,7 @@ where } fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) { - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); self.heap.append_or_replace(new_val, map_idx, map); } @@ -150,8 +150,7 @@ where row_idx: usize, map: &mut Vec<(usize, usize)>, ) { - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); self.heap.replace_if_better(heap_idx, new_val, map); } @@ -443,14 +442,13 @@ impl TopKHeap { } fn swap(&mut self, a_idx: usize, b_idx: usize, mapper: &mut Vec<(usize, usize)>) { - let a_hi = self.heap[a_idx].take().expect("Missing heap entry"); - let b_hi = self.heap[b_idx].take().expect("Missing heap entry"); + self.heap.swap(a_idx, b_idx); - mapper.push((a_hi.map_idx, b_idx)); - mapper.push((b_hi.map_idx, a_idx)); + let b_hi = self.heap[b_idx].as_ref().expect("Missing heap entry"); + let a_hi = self.heap[a_idx].as_ref().expect("Missing heap entry"); - self.heap[a_idx] = Some(b_hi); - self.heap[b_idx] = Some(a_hi); + mapper.push((b_hi.map_idx, b_idx)); + mapper.push((a_hi.map_idx, a_idx)); } fn heapify_down(&mut self, node_idx: usize, mapper: &mut Vec<(usize, usize)>) { From aa9c213f7a96aa853a85d9b96fd795195749e1a5 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 24 Aug 2026 11:04:50 +0800 Subject: [PATCH 10/14] docs: require draining DF54 fork plans before upgrade Document the protobuf field-37 collision and missing dynamic-filter expr_id incompatibility, with an explicit stop, drain, coordinated-upgrade, and resume checklist. --- .../library-user-guide/upgrading/55.0.0.md | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index d64f287ea0b52..ac2dd219e3710 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -25,6 +25,30 @@ in this section pertains to features and changes that have already been merged to the main branch and are awaiting release in this version. +### Massive fork: drain serialized plans before upgrading + +This section applies only when upgrading from the Massive DataFusion 54 fork. +Treat serialized physical plans as incompatible across the upgrade boundary. + +The DataFusion 54 fork encoded `ParquetOptions.max_in_list_size` as protobuf +field 37. DataFusion 55 uses field 37 for `max_row_group_bytes` and field 38 for +`max_in_list_size`. Both field-37 values use the same protobuf wire type, so a +decoder cannot distinguish them. For example, the old default IN-list size of +20 would be decoded as a 20-byte row-group limit. + +Before starting any DataFusion 55 process: + +1. Stop all DataFusion 54 processes that can serialize physical plans. +2. Drain or cancel every queued, cached, or in-flight DataFusion 54 physical + plan. +3. Upgrade plan producers and consumers together. Do not send serialized plans + between DataFusion 54 and DataFusion 55 processes. +4. Resume plan production only after every consumer is running DataFusion 55. + +DataFusion 54 plans containing dynamic filters must also be drained. Plans +serialized without the deduplicating converter can omit the dynamic-filter +`expr_id`, which DataFusion 55 rejects during decoding. + ### `DataFrame::fill_null` now borrows its arguments `DataFrame::fill_null` previously took its arguments by value: From cbecdcde4e8b4820ae3195226aeb12af8d708f0d Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 24 Aug 2026 12:14:39 +0800 Subject: [PATCH 11/14] revert: keep fork rollout note out of upgrade guide --- .../library-user-guide/upgrading/55.0.0.md | 24 ------------------- 1 file changed, 24 deletions(-) diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index ac2dd219e3710..d64f287ea0b52 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -25,30 +25,6 @@ in this section pertains to features and changes that have already been merged to the main branch and are awaiting release in this version. -### Massive fork: drain serialized plans before upgrading - -This section applies only when upgrading from the Massive DataFusion 54 fork. -Treat serialized physical plans as incompatible across the upgrade boundary. - -The DataFusion 54 fork encoded `ParquetOptions.max_in_list_size` as protobuf -field 37. DataFusion 55 uses field 37 for `max_row_group_bytes` and field 38 for -`max_in_list_size`. Both field-37 values use the same protobuf wire type, so a -decoder cannot distinguish them. For example, the old default IN-list size of -20 would be decoded as a 20-byte row-group limit. - -Before starting any DataFusion 55 process: - -1. Stop all DataFusion 54 processes that can serialize physical plans. -2. Drain or cancel every queued, cached, or in-flight DataFusion 54 physical - plan. -3. Upgrade plan producers and consumers together. Do not send serialized plans - between DataFusion 54 and DataFusion 55 processes. -4. Resume plan production only after every consumer is running DataFusion 55. - -DataFusion 54 plans containing dynamic filters must also be drained. Plans -serialized without the deduplicating converter can omit the dynamic-filter -`expr_id`, which DataFusion 55 rejects during decoding. - ### `DataFrame::fill_null` now borrows its arguments `DataFrame::fill_null` previously took its arguments by value: From 7c67bad1ae7169620d2af8550402ad71e73925bd Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 24 Aug 2026 16:39:45 +0800 Subject: [PATCH 12/14] fix(proto): align max IN-list encoding with upstream Drop the fork-only presence oneof and restore DataFusion 55's scalar field 38 representation. This avoids carrying a divergent wire schema for ephemeral physical plans. --- .../datasource-parquet/src/file_format.rs | 6 +-- .../proto/datafusion_common.proto | 6 +-- datafusion/proto-common/src/from_proto/mod.rs | 37 +-------------- .../proto-common/src/generated/pbjson.rs | 46 +++++++++---------- .../proto-common/src/generated/prost.rs | 13 +----- datafusion/proto-common/src/to_proto/mod.rs | 2 +- datafusion/proto-models/src/from_proto.rs | 12 +---- .../src/generated/datafusion_proto_common.rs | 13 +----- .../proto/tests/cases/public_conversions.rs | 26 ----------- 9 files changed, 31 insertions(+), 130 deletions(-) diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 057f77815320c..6358201c06fa5 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -732,11 +732,7 @@ impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions { parquet_options::StatisticsEnabledOpt::StatisticsEnabled(enabled) }), max_row_group_size: global_options.global.max_row_group_size as u64, - max_in_list_size_opt: Some( - parquet_options::MaxInListSizeOpt::MaxInListSize( - global_options.global.max_in_list_size as u64, - ), - ), + max_in_list_size: global_options.global.max_in_list_size as u64, created_by: global_options.global.created_by.clone(), column_index_truncate_length_opt: global_options.global.column_index_truncate_length.map(|length| { parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length as u64) diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 71fca2d214100..748d77d63bf4c 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -617,11 +617,7 @@ message ParquetOptions { uint64 max_row_group_size = 15; - // Presence distinguishes an omitted value (default 20) from an explicit - // zero, which disables IN-list pruning. - oneof max_in_list_size_opt { - uint64 max_in_list_size = 38; - } + uint64 max_in_list_size = 38; string created_by = 16; diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 21a1614c969c3..169ff7f3d9ff2 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1081,12 +1081,7 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { }) .unwrap_or(None), max_row_group_size: value.max_row_group_size as usize, - // Plans written before this field was introduced do not carry a value. - // Preserve the documented default instead of treating absence as zero, - // which disables IN-list pruning. - max_in_list_size: value.max_in_list_size_opt.map(|opt| match opt { - protobuf::parquet_options::MaxInListSizeOpt::MaxInListSize(v) => v as usize, - }).unwrap_or(20), + max_in_list_size: value.max_in_list_size as usize, created_by: value.created_by.clone(), column_index_truncate_length: value .column_index_truncate_length_opt.as_ref() @@ -1398,36 +1393,6 @@ mod tests { ); } - #[test] - fn test_parquet_options_max_in_list_size_round_trip() { - let opts = ParquetOptions { - max_in_list_size: 64, - ..ParquetOptions::default() - }; - let recovered = parquet_options_proto_round_trip(opts); - assert_eq!(recovered.max_in_list_size, 64); - } - - #[test] - fn test_parquet_options_max_in_list_size_zero_round_trip() { - let opts = ParquetOptions { - max_in_list_size: 0, - ..ParquetOptions::default() - }; - let recovered = parquet_options_proto_round_trip(opts); - assert_eq!(recovered.max_in_list_size, 0); - } - - #[test] - fn test_parquet_options_max_in_list_size_absent_uses_default() { - let opts = ParquetOptions::default(); - let mut proto: crate::protobuf_common::ParquetOptions = - (&opts).try_into().expect("to_proto"); - proto.max_in_list_size_opt = None; - let recovered = ParquetOptions::try_from(&proto).expect("from_proto"); - assert_eq!(recovered.max_in_list_size, 20); - } - #[test] fn test_table_parquet_options_coerce_int96_tz_round_trip() { let mut opts = TableParquetOptions::default(); diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 6a29ad3400ff3..c222cd1cb8687 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -6409,6 +6409,9 @@ impl serde::Serialize for ParquetOptions { if self.max_row_group_size != 0 { len += 1; } + if self.max_in_list_size != 0 { + len += 1; + } if !self.created_by.is_empty() { len += 1; } @@ -6442,9 +6445,6 @@ impl serde::Serialize for ParquetOptions { if self.bloom_filter_ndv_opt.is_some() { len += 1; } - if self.max_in_list_size_opt.is_some() { - len += 1; - } if self.coerce_int96_opt.is_some() { len += 1; } @@ -6532,6 +6532,11 @@ impl serde::Serialize for ParquetOptions { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("maxRowGroupSize", ToString::to_string(&self.max_row_group_size).as_str())?; } + if self.max_in_list_size != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("maxInListSize", ToString::to_string(&self.max_in_list_size).as_str())?; + } if !self.created_by.is_empty() { struct_ser.serialize_field("createdBy", &self.created_by)?; } @@ -6609,15 +6614,6 @@ impl serde::Serialize for ParquetOptions { } } } - if let Some(v) = self.max_in_list_size_opt.as_ref() { - match v { - parquet_options::MaxInListSizeOpt::MaxInListSize(v) => { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("maxInListSize", ToString::to_string(&v).as_str())?; - } - } - } if let Some(v) = self.coerce_int96_opt.as_ref() { match v { parquet_options::CoerceInt96Opt::CoerceInt96(v) => { @@ -6699,6 +6695,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dataPageRowCountLimit", "max_row_group_size", "maxRowGroupSize", + "max_in_list_size", + "maxInListSize", "created_by", "createdBy", "content_defined_chunking", @@ -6719,8 +6717,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "bloomFilterFpp", "bloom_filter_ndv", "bloomFilterNdv", - "max_in_list_size", - "maxInListSize", "coerce_int96", "coerceInt96", "max_predicate_cache_size", @@ -6753,6 +6749,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { DictionaryPageSizeLimit, DataPageRowCountLimit, MaxRowGroupSize, + MaxInListSize, CreatedBy, ContentDefinedChunking, MetadataSizeHint, @@ -6764,7 +6761,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Encoding, BloomFilterFpp, BloomFilterNdv, - MaxInListSize, CoerceInt96, MaxPredicateCacheSize, MaxRowGroupBytes, @@ -6810,6 +6806,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dictionaryPageSizeLimit" | "dictionary_page_size_limit" => Ok(GeneratedField::DictionaryPageSizeLimit), "dataPageRowCountLimit" | "data_page_row_count_limit" => Ok(GeneratedField::DataPageRowCountLimit), "maxRowGroupSize" | "max_row_group_size" => Ok(GeneratedField::MaxRowGroupSize), + "maxInListSize" | "max_in_list_size" => Ok(GeneratedField::MaxInListSize), "createdBy" | "created_by" => Ok(GeneratedField::CreatedBy), "contentDefinedChunking" | "content_defined_chunking" => Ok(GeneratedField::ContentDefinedChunking), "metadataSizeHint" | "metadata_size_hint" => Ok(GeneratedField::MetadataSizeHint), @@ -6821,7 +6818,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "encoding" => Ok(GeneratedField::Encoding), "bloomFilterFpp" | "bloom_filter_fpp" => Ok(GeneratedField::BloomFilterFpp), "bloomFilterNdv" | "bloom_filter_ndv" => Ok(GeneratedField::BloomFilterNdv), - "maxInListSize" | "max_in_list_size" => Ok(GeneratedField::MaxInListSize), "coerceInt96" | "coerce_int96" => Ok(GeneratedField::CoerceInt96), "maxPredicateCacheSize" | "max_predicate_cache_size" => Ok(GeneratedField::MaxPredicateCacheSize), "maxRowGroupBytes" | "max_row_group_bytes" => Ok(GeneratedField::MaxRowGroupBytes), @@ -6865,6 +6861,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut dictionary_page_size_limit__ = None; let mut data_page_row_count_limit__ = None; let mut max_row_group_size__ = None; + let mut max_in_list_size__ = None; let mut created_by__ = None; let mut content_defined_chunking__ = None; let mut metadata_size_hint_opt__ = None; @@ -6876,7 +6873,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut encoding_opt__ = None; let mut bloom_filter_fpp_opt__ = None; let mut bloom_filter_ndv_opt__ = None; - let mut max_in_list_size_opt__ = None; let mut coerce_int96_opt__ = None; let mut max_predicate_cache_size_opt__ = None; let mut max_row_group_bytes_opt__ = None; @@ -7017,6 +7013,14 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::MaxInListSize => { + if max_in_list_size__.is_some() { + return Err(serde::de::Error::duplicate_field("maxInListSize")); + } + max_in_list_size__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } GeneratedField::CreatedBy => { if created_by__.is_some() { return Err(serde::de::Error::duplicate_field("createdBy")); @@ -7083,12 +7087,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { } bloom_filter_ndv_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::BloomFilterNdvOpt::BloomFilterNdv(x.0)); } - GeneratedField::MaxInListSize => { - if max_in_list_size_opt__.is_some() { - return Err(serde::de::Error::duplicate_field("maxInListSize")); - } - max_in_list_size_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::MaxInListSizeOpt::MaxInListSize(x.0)); - } GeneratedField::CoerceInt96 => { if coerce_int96_opt__.is_some() { return Err(serde::de::Error::duplicate_field("coerceInt96")); @@ -7136,6 +7134,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { dictionary_page_size_limit: dictionary_page_size_limit__.unwrap_or_default(), data_page_row_count_limit: data_page_row_count_limit__.unwrap_or_default(), max_row_group_size: max_row_group_size__.unwrap_or_default(), + max_in_list_size: max_in_list_size__.unwrap_or_default(), created_by: created_by__.unwrap_or_default(), content_defined_chunking: content_defined_chunking__, metadata_size_hint_opt: metadata_size_hint_opt__, @@ -7147,7 +7146,6 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { encoding_opt: encoding_opt__, bloom_filter_fpp_opt: bloom_filter_fpp_opt__, bloom_filter_ndv_opt: bloom_filter_ndv_opt__, - max_in_list_size_opt: max_in_list_size_opt__, coerce_int96_opt: coerce_int96_opt__, max_predicate_cache_size_opt: max_predicate_cache_size_opt__, max_row_group_bytes_opt: max_row_group_bytes_opt__, diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index 2c588aec829d6..bdbe38538e1d7 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -862,6 +862,8 @@ pub struct ParquetOptions { pub data_page_row_count_limit: u64, #[prost(uint64, tag = "15")] pub max_row_group_size: u64, + #[prost(uint64, tag = "38")] + pub max_in_list_size: u64, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] @@ -894,10 +896,6 @@ pub struct ParquetOptions { pub bloom_filter_fpp_opt: ::core::option::Option, #[prost(oneof = "parquet_options::BloomFilterNdvOpt", tags = "22")] pub bloom_filter_ndv_opt: ::core::option::Option, - /// Presence distinguishes an omitted value (default 20) from an explicit - /// zero, which disables IN-list pruning. - #[prost(oneof = "parquet_options::MaxInListSizeOpt", tags = "38")] - pub max_in_list_size_opt: ::core::option::Option, #[prost(oneof = "parquet_options::CoerceInt96Opt", tags = "32")] pub coerce_int96_opt: ::core::option::Option, #[prost(oneof = "parquet_options::MaxPredicateCacheSizeOpt", tags = "33")] @@ -962,13 +960,6 @@ pub mod parquet_options { #[prost(uint64, tag = "22")] BloomFilterNdv(u64), } - /// Presence distinguishes an omitted value (default 20) from an explicit - /// zero, which disables IN-list pruning. - #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] - pub enum MaxInListSizeOpt { - #[prost(uint64, tag = "38")] - MaxInListSize(u64), - } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] pub enum CoerceInt96Opt { #[prost(string, tag = "32")] diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 75a2ded4c9979..360981746585b 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -912,7 +912,7 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { dictionary_page_size_limit: value.dictionary_page_size_limit as u64, statistics_enabled_opt: value.statistics_enabled.clone().map(protobuf::parquet_options::StatisticsEnabledOpt::StatisticsEnabled), max_row_group_size: value.max_row_group_size as u64, - max_in_list_size_opt: Some(protobuf::parquet_options::MaxInListSizeOpt::MaxInListSize(value.max_in_list_size as u64)), + max_in_list_size: value.max_in_list_size as u64, created_by: value.created_by.clone(), column_index_truncate_length_opt: value.column_index_truncate_length.map(|v| protobuf::parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(v as u64)), statistics_truncate_length_opt: value.statistics_truncate_length.map(|v| protobuf::parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(v as u64)), diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index 2936143bb6dda..74ead8c52049b 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -374,17 +374,7 @@ impl TryFrom<&ParquetOptionsProto> for ParquetOptions { }, ), max_row_group_size: proto.max_row_group_size as usize, - // Preserve the documented default when an older plan omits the - // field; primitive zero disables IN-list pruning. - max_in_list_size: proto - .max_in_list_size_opt - .as_ref() - .map(|opt| match opt { - parquet_options::MaxInListSizeOpt::MaxInListSize(size) => { - *size as usize - } - }) - .unwrap_or(20), + max_in_list_size: proto.max_in_list_size as usize, created_by: proto.created_by.clone(), column_index_truncate_length: proto .column_index_truncate_length_opt diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index 2c588aec829d6..bdbe38538e1d7 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -862,6 +862,8 @@ pub struct ParquetOptions { pub data_page_row_count_limit: u64, #[prost(uint64, tag = "15")] pub max_row_group_size: u64, + #[prost(uint64, tag = "38")] + pub max_in_list_size: u64, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] @@ -894,10 +896,6 @@ pub struct ParquetOptions { pub bloom_filter_fpp_opt: ::core::option::Option, #[prost(oneof = "parquet_options::BloomFilterNdvOpt", tags = "22")] pub bloom_filter_ndv_opt: ::core::option::Option, - /// Presence distinguishes an omitted value (default 20) from an explicit - /// zero, which disables IN-list pruning. - #[prost(oneof = "parquet_options::MaxInListSizeOpt", tags = "38")] - pub max_in_list_size_opt: ::core::option::Option, #[prost(oneof = "parquet_options::CoerceInt96Opt", tags = "32")] pub coerce_int96_opt: ::core::option::Option, #[prost(oneof = "parquet_options::MaxPredicateCacheSizeOpt", tags = "33")] @@ -962,13 +960,6 @@ pub mod parquet_options { #[prost(uint64, tag = "22")] BloomFilterNdv(u64), } - /// Presence distinguishes an omitted value (default 20) from an explicit - /// zero, which disables IN-list pruning. - #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] - pub enum MaxInListSizeOpt { - #[prost(uint64, tag = "38")] - MaxInListSize(u64), - } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] pub enum CoerceInt96Opt { #[prost(string, tag = "32")] diff --git a/datafusion/proto/tests/cases/public_conversions.rs b/datafusion/proto/tests/cases/public_conversions.rs index f486ad06214ff..1cda8e01a0765 100644 --- a/datafusion/proto/tests/cases/public_conversions.rs +++ b/datafusion/proto/tests/cases/public_conversions.rs @@ -126,29 +126,3 @@ fn file_format_option_conversions_are_std_traits() { assert_from::<&JsonFormatFactory, protobuf::JsonOptions>(); assert_from::<&ParquetFormatFactory, protobuf::TableParquetOptions>(); } - -#[test] -fn parquet_max_in_list_size_preserves_zero_and_absent_default() { - let mut options = TableParquetOptions::default(); - options.global.max_in_list_size = 0; - let factory = ParquetFormatFactory::new_with_options(options); - let mut proto = protobuf::TableParquetOptions::from(&factory); - - let global = proto.global.as_ref().expect("global parquet options"); - assert!(matches!( - global.max_in_list_size_opt, - Some(protobuf::parquet_options::MaxInListSizeOpt::MaxInListSize( - 0 - )) - )); - let decoded = TableParquetOptions::try_from(&proto).expect("from_proto"); - assert_eq!(decoded.global.max_in_list_size, 0); - - proto - .global - .as_mut() - .expect("global parquet options") - .max_in_list_size_opt = None; - let decoded = TableParquetOptions::try_from(&proto).expect("from_proto"); - assert_eq!(decoded.global.max_in_list_size, 20); -} From 00951ea629485c39fccce0f1db7a96b68c7e2be2 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 24 Aug 2026 16:39:57 +0800 Subject: [PATCH 13/14] fix(topk): preserve declared hash key types Keep the constructor's Arrow type as the output source of truth and reject batches whose full type metadata differs. This prevents string variants, decimal metadata, or timestamp timezones from being silently adopted from the latest batch. --- .../src/aggregates/topk/hash_table.rs | 50 ++++++++++++++++--- 1 file changed, 43 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 771bf69182acf..5f4e6a70b37be 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -169,6 +169,7 @@ pub fn is_supported_hash_key_type(kt: &DataType) -> bool { // An implementation of ArrowHashTable for String keys pub struct StringHashTable { + data_type: DataType, owned: StringArrayType, map: TopKHashTable>, rnd: RandomState, @@ -179,6 +180,7 @@ struct PrimitiveHashTable where Option<::Native>: Comparable, { + data_type: DataType, owned: PrimitiveArray, map: TopKHashTable>, rnd: RandomState, @@ -190,6 +192,7 @@ impl StringHashTable { unreachable!("StringHashTable requires a string data type") }; Self { + data_type: data_type.clone(), owned, map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), @@ -199,6 +202,11 @@ impl StringHashTable { impl ArrowHashTable for StringHashTable { fn set_batch(&mut self, ids: ArrayRef) { + assert_eq!( + ids.data_type(), + &self.data_type, + "Top-K hash key batch type must match the declared type" + ); let Ok(owned) = StringArrayType::try_from(ids) else { unreachable!("StringHashTable requires a string array") }; @@ -219,10 +227,11 @@ impl ArrowHashTable for StringHashTable { fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - match &self.owned { - StringArrayType::Utf8(_) => Arc::new(StringArray::from(ids)), - StringArrayType::Utf8View(_) => Arc::new(StringViewArray::from(ids)), - StringArrayType::LargeUtf8(_) => Arc::new(LargeStringArray::from(ids)), + match &self.data_type { + DataType::Utf8 => Arc::new(StringArray::from(ids)), + DataType::Utf8View => Arc::new(StringViewArray::from(ids)), + DataType::LargeUtf8 => Arc::new(LargeStringArray::from(ids)), + _ => unreachable!("StringHashTable requires a string data type"), } } @@ -265,9 +274,10 @@ where { pub fn new(limit: usize, kt: DataType) -> Self { let owned = PrimitiveArray::::builder(0) - .with_data_type(kt) + .with_data_type(kt.clone()) .finish(); Self { + data_type: kt, owned, map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), @@ -291,6 +301,11 @@ where Option<::Native>: Comparable + HashValue, { fn set_batch(&mut self, ids: ArrayRef) { + assert_eq!( + ids.data_type(), + &self.data_type, + "Top-K hash key batch type must match the declared type" + ); self.owned = ids.as_primitive().clone(); } @@ -308,8 +323,8 @@ where fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - let mut builder: PrimitiveBuilder = PrimitiveArray::builder(ids.len()) - .with_data_type(self.owned.data_type().clone()); + let mut builder: PrimitiveBuilder = + PrimitiveArray::builder(ids.len()).with_data_type(self.data_type.clone()); for id in ids.into_iter() { match id { None => builder.append_null(), @@ -628,6 +643,27 @@ mod tests { Ok(()) } + #[test] + #[should_panic(expected = "Top-K hash key batch type must match the declared type")] + fn should_reject_primitive_key_metadata_mismatch() { + let ids = TimestampMillisecondArray::from(vec![1000]); + let data_type = DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())); + let mut hash_table = + new_hash_table(1, data_type).expect("create timestamp hash table"); + + hash_table.set_batch(Arc::new(ids)); + } + + #[test] + #[should_panic(expected = "Top-K hash key batch type must match the declared type")] + fn should_reject_string_key_type_mismatch() { + let ids = LargeStringArray::from(vec!["value"]); + let mut hash_table = + new_hash_table(1, DataType::Utf8).expect("create string hash table"); + + hash_table.set_batch(Arc::new(ids)); + } + #[test] fn should_resize_properly() -> Result<()> { let mut heap_to_map = BTreeMap::::new(); From 0437d4fc10ab72e60d5a041b77f29da08d6f825f Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Mon, 24 Aug 2026 16:40:04 +0800 Subject: [PATCH 14/14] test(expr): avoid shared extension rebuild state Use pointer identity alone to prove the expressionless extension was not rebuilt, removing the redundant process-global counter and its potential for cross-test interference. --- datafusion/expr/src/logical_plan/tree_node.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index 7e30568c450d2..fd3395cf5ba82 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -1038,7 +1038,6 @@ impl LogicalPlan { #[cfg(test)] mod tests { use std::fmt::Formatter; - use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::sync::{Arc, LazyLock}; use datafusion_common::{DFSchema, DFSchemaRef, Result}; @@ -1046,7 +1045,6 @@ mod tests { use super::*; use crate::UserDefinedLogicalNodeCore; - static REBUILD_COUNT: AtomicUsize = AtomicUsize::new(0); static EMPTY_SCHEMA: LazyLock = LazyLock::new(|| Arc::new(DFSchema::empty())); @@ -1081,14 +1079,12 @@ mod tests { ) -> Result { assert!(exprs.is_empty()); assert!(inputs.is_empty()); - REBUILD_COUNT.fetch_add(1, AtomicOrdering::Relaxed); Ok(Self) } } #[test] fn map_expressions_does_not_rebuild_expressionless_extension() -> Result<()> { - REBUILD_COUNT.store(0, AtomicOrdering::Relaxed); let original: Arc = Arc::new(EmptyExpressionNode); let plan = LogicalPlan::Extension(Extension { node: Arc::clone(&original), @@ -1100,7 +1096,6 @@ mod tests { }; assert!(Arc::ptr_eq(&node, &original)); - assert_eq!(REBUILD_COUNT.load(AtomicOrdering::Relaxed), 0); Ok(()) } }