From 075bcfdfc4cb93170a1c30597856f859fbee582c Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:33:28 +0800 Subject: [PATCH 01/10] feat(parquet): retain public reverse row-group controls Carry forward the downstream API from fork #73; row order within each row group is unchanged. --- 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 ac7e10f94e6f5..b698de590a1c2 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -511,13 +511,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 03f00512e5a21ce22faa9b5699c737e51818eeb5 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:34:42 +0800 Subject: [PATCH 02/10] perf(planner): retain left-input sort options for merge joins Port the fork #73 planner behavior and its four ordering regressions to DF56. --- 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 dd49e4ed6a673..602a8bdd9684e 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -49,7 +49,6 @@ use crate::physical_plan::{ use crate::schema_equivalence::schema_satisfied_by; use arrow::array::{ArrayRef, RecordBatch, UInt64Array}; -use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Schema}; use arrow_schema::Field; use datafusion_catalog::ScanArgs; @@ -1872,15 +1871,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 { @@ -3932,8 +3943,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; @@ -4464,6 +4475,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 4fd1bab8e8a923dcb128778e034dd576decc475c Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:36:49 +0800 Subject: [PATCH 03/10] fix(topk): retain declared primitive key type checks Keep DF56 typed storage and restore the fork review fix rejecting mismatched timezone or decimal metadata. Cover string variants using the upstream generic downcast guard. --- .../src/aggregates/topk/hash_table.rs | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 322ab7bb58dfc..06467e2224d0b 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -244,6 +244,11 @@ where Option: HashValue, { fn set_batch(&mut self, ids: ArrayRef) { + assert_eq!( + ids.data_type(), + &self.value_type, + "Top-K hash key batch type must match the declared type" + ); self.owned = ids.as_primitive().clone(); } @@ -608,6 +613,7 @@ pub fn new_hash_table( #[cfg(test)] mod tests { use super::*; + use arrow::array::Decimal128Array; use arrow::array::TimestampMillisecondArray; use arrow::datatypes::{Float16Type, Float32Type, Float64Type}; use arrow_schema::TimeUnit; @@ -696,6 +702,33 @@ mod tests { ) } + #[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).unwrap(); + 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_decimal_key_metadata_mismatch() { + let ids = Decimal128Array::from(vec![123]) + .with_precision_and_scale(10, 2) + .unwrap(); + let mut hash_table = new_hash_table(1, DataType::Decimal128(12, 3)).unwrap(); + hash_table.set_batch(Arc::new(ids)); + } + + #[test] + #[should_panic(expected = "Unsupported data type")] + fn should_reject_string_key_type_mismatch() { + let ids = LargeStringArray::from(vec!["value"]); + let mut hash_table = new_hash_table(1, DataType::Utf8).unwrap(); + hash_table.set_batch(Arc::new(ids)); + } + #[test] fn should_emit_correct_type() -> Result<()> { let ids = From 1808be52285ad74cb1fd31806d0462460a5c6fb5 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:36:49 +0800 Subject: [PATCH 04/10] fix(union): retain executable fallback for incompatible child rewrites Keep fork #73 behavior: replacing incompatible Interleave children returns a Union with every child partition. Update the upstream deferred-invariant regressions to require an executable fallback. --- datafusion/physical-plan/src/union.rs | 43 ++++++++------------------- 1 file changed, 13 insertions(+), 30 deletions(-) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 3b4cebc67ca86..55edc4716172f 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -783,9 +783,11 @@ impl ExecutionPlan for InterleaveExec { ..Self::clone(&*self) })), ChildrenPropertiesMode::Recompute => { - // The new children may no longer be interleavable; see - // `try_new_unchecked` for why this is not rejected here. - Ok(Arc::new(InterleaveExec::try_new_unchecked(children)?)) + if can_interleave(children.iter()) { + Ok(Arc::new(InterleaveExec::try_new(children)?)) + } else { + UnionExec::try_new(children) + } } } } @@ -1847,7 +1849,7 @@ mod tests { } #[test] - fn test_interleave_rebuild_defers_partitioning_invariant() -> Result<()> { + fn test_interleave_rebuild_falls_back_to_union() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("name", DataType::Int32, true)])); let hash = || make_hash_exec(&schema, vec!["name"], 3); @@ -1874,28 +1876,18 @@ mod tests { let children = rebuilt.children().into_iter().cloned().collect(); assert!(InterleaveExec::try_new(children).is_err()); - // The rebuilt node no longer claims a hash layout, is structurally - // sound, but is not executable until a distribution pass repairs it. + assert!(rebuilt.downcast_ref::().is_some()); assert!(matches!( rebuilt.output_partitioning(), - Partitioning::UnknownPartitioning(3) + Partitioning::UnknownPartitioning(6) )); rebuilt.check_invariants(InvariantLevel::Always)?; - let err = rebuilt - .check_invariants(InvariantLevel::Executable) - .unwrap_err() - .to_string(); - assert!( - err.contains( - "Not all InterleaveExec children have a consistent hash or range partitioning" - ), - "{err}" - ); + rebuilt.check_invariants(InvariantLevel::Executable)?; Ok(()) } #[test] - fn test_interleave_rebuild_reports_widest_partition_count() -> Result<()> { + fn test_interleave_rebuild_preserves_all_child_partitions() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("name", DataType::Int32, true)])); let interleave: Arc = @@ -1913,21 +1905,12 @@ mod tests { ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), )?; - // The widest child decides the reported count, so the partitions only - // it has are still visible to callers instead of being dropped. + assert!(rebuilt.downcast_ref::().is_some()); assert!(matches!( rebuilt.output_partitioning(), - Partitioning::UnknownPartitioning(5) + Partitioning::UnknownPartitioning(8) )); - // Executing one of them fails loudly rather than returning no rows. - let Err(err) = rebuilt.execute(4, Arc::new(TaskContext::default())) else { - panic!("executing a partition the narrow child lacks must fail"); - }; - let err = err.to_string(); - assert!( - err.contains("Partition 4 not found in InterleaveExec"), - "{err}" - ); + rebuilt.check_invariants(InvariantLevel::Executable)?; Ok(()) } From 7a9a47ff9a32af1afe0cf7d7b4c09551f1365cf5 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:36:49 +0800 Subject: [PATCH 05/10] test(parquet): retain fork schema adaptation regressions --- .../datasource-parquet/src/opener/mod.rs | 138 +++++++++++++++++- 1 file changed, 134 insertions(+), 4 deletions(-) diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index e188b4543f284..984a9ef375981 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -2152,8 +2152,11 @@ mod test { ParquetFileReaderFactory, ParquetFileSchemaProvider, ParquetRowSelection, RowGroupAccess, }; - use arrow::array::{Array, AsArray, RecordBatch, record_batch}; - use arrow::buffer::BooleanBuffer; + use arrow::array::{ + Array, ArrayRef, AsArray, Date32Array, Int32Array, ListArray, RecordBatch, + StringArray, record_batch, + }; + use arrow::buffer::{BooleanBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; use bytes::{BufMut, BytesMut}; use datafusion_common::{ @@ -2822,7 +2825,7 @@ mod test { let array = batch .column(0) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); for i in 0..array.len() { if !array.is_null(i) { @@ -2876,6 +2879,133 @@ mod test { (data_len, statistics) } + 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, @@ -5887,7 +6017,7 @@ mod test { let part = batch .column(1) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); assert!(part.iter().all(|v| v == Some(5))); From 6d8432b29e65bb5495a64420b0210e3143da1e1e Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:36:49 +0800 Subject: [PATCH 06/10] test(expr): retain expressionless extension Arc reuse regression --- datafusion/expr/src/logical_plan/tree_node.rs | 65 +++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index 7b9c4586bc4d1..7b524b08af84c 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -1114,3 +1114,68 @@ impl LogicalPlan { }) } } + +#[cfg(test)] +mod tests { + use std::fmt::Formatter; + use std::sync::{Arc, LazyLock}; + + use datafusion_common::{DFSchema, DFSchemaRef, Result}; + + use super::*; + use crate::UserDefinedLogicalNodeCore; + + 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()); + Ok(Self) + } + } + + #[test] + fn map_expressions_does_not_rebuild_expressionless_extension() -> Result<()> { + 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)); + Ok(()) + } +} From 1e4a9e76c02c6db4e0c74e61e3c34e1a85874f5d Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:40:08 +0800 Subject: [PATCH 07/10] feat(aggregate): retain preserving reads for state-heavy accumulators Port fork #77. Adapt Median through the shared DF56 percentile implementation, retaining decimal metadata and the per-group percentile state column. Preserve upstream ordered first/last fast paths and add regression coverage for the changed state layout. --- .../src/approx_distinct.rs | 115 ++++++- .../functions-aggregate/src/array_agg.rs | 167 +++++++++- .../functions-aggregate/src/first_last.rs | 255 ++++++++++++++- .../src/first_last/state.rs | 260 +++++++++++++++- .../src/percentile_cont.rs | 291 ++++++++++++++++-- 5 files changed, 1058 insertions(+), 30 deletions(-) diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index ded210502a501..1b51cfe251499 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -22,15 +22,15 @@ use arrow::datatypes::{ use datafusion_common::ScalarValue; use datafusion_common::hash_utils::create_hashes; use datafusion_common::{ - DataFusionError, Result, downcast_value, internal_datafusion_err, internal_err, - not_impl_err, + DataFusionError, Result, downcast_value, exec_datafusion_err, + internal_datafusion_err, internal_err, not_impl_err, }; use datafusion_expr::DistinctHandling; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, Signature, - Volatility, + Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupSelection, + GroupsAccumulator, Signature, Volatility, }; use datafusion_functions_aggregate_common::aggregate::count_distinct::{ Bitmap65536DistinctCountAccumulator, Bitmap65536DistinctCountAccumulatorI16, @@ -444,6 +444,16 @@ impl GroupHll { } } +fn checked_hll_state_data_length(current: usize, additional: usize) -> Result { + let total = current + .checked_add(additional) + .ok_or_else(|| exec_datafusion_err!("Approx distinct state length overflow"))?; + i32::try_from(total).map_err(|_| { + exec_datafusion_err!("Approx distinct state exceeds i32 offset capacity") + })?; + Ok(total) +} + /// A [`GroupsAccumulator`] for `approx_distinct` that keeps one adaptive /// (sparse → dense) HyperLogLog sketch per group. /// @@ -577,6 +587,18 @@ impl GroupsAccumulator for HllGroupsAccumulator { Ok(Arc::new(counts)) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.groups.len())?; + let counts = UInt64Array::from_iter_values( + selection.iter().map(|index| self.groups[index].count()), + ); + Ok(Arc::new(counts)) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: EmitTo) -> Result> { let mut groups = emit_to.take_needed(&mut self.groups); let mut builder = BinaryBuilder::new(); @@ -592,6 +614,26 @@ impl GroupsAccumulator for HllGroupsAccumulator { Ok(vec![Arc::new(builder.finish())]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + selection.validate_num_groups(self.groups.len())?; + let mut builder = BinaryBuilder::new(); + let mut scratch = Vec::new(); + let mut total_bytes = 0; + for index in selection.iter() { + self.groups[index].serialize(&mut scratch); + total_bytes = checked_hll_state_data_length(total_bytes, scratch.len())?; + builder.append_value(&scratch); + } + Ok(vec![Arc::new(builder.finish())]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn convert_to_state( &self, values: &[ArrayRef], @@ -1237,6 +1279,51 @@ mod tests { assert_eq!(counts.value(0), expected); } + #[test] + fn groups_preserving_reads_keep_hll_state() { + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 2, 3, 10, 10])); + let group_indices = vec![0usize, 0, 0, 1, 2, 2]; + let mut acc = HllGroupsAccumulator::new(); + acc.update_batch(&[values], &group_indices, None, 4) + .unwrap(); + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4).unwrap(); + let expected = UInt64Array::from(vec![1, 2, 1, 0]); + for _ in 0..2 { + let actual = acc.evaluate_preserving(selection).unwrap(); + assert_eq!(actual.as_primitive::(), &expected); + } + + let state = acc.state_preserving(selection).unwrap(); + let mut merged = HllGroupsAccumulator::new(); + merged.merge_batch(&state, &[0, 1, 2, 3], 4).unwrap(); + assert_eq!( + merged + .evaluate(EmitTo::All) + .unwrap() + .as_primitive::(), + &expected + ); + + let empty = GroupSelection::try_from_indices(&[], 4).unwrap(); + assert!(acc.evaluate_preserving(empty).unwrap().is_empty()); + let empty_state = acc.state_preserving(empty).unwrap(); + assert_eq!(empty_state.len(), 1); + assert_eq!(empty_state[0].data_type(), &DataType::Binary); + assert!(empty_state[0].is_empty()); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![4, 20])); + acc.update_batch(&[values], &[1, 3], None, 4).unwrap(); + assert_eq!( + acc.evaluate_preserving(GroupSelection::all(4)) + .unwrap() + .as_primitive::(), + &UInt64Array::from(vec![2, 2, 1, 1]) + ); + assert!(acc.supports_evaluate_preserving()); + assert!(acc.supports_state_preserving()); + } + #[test] fn groups_convert_to_state_roundtrips_through_merge() { let values: ArrayRef = Arc::new(Int64Array::from(vec![ @@ -1407,6 +1494,26 @@ mod tests { buf } + #[test] + fn hll_state_data_length_checks_offset_capacity() { + assert_eq!( + checked_hll_state_data_length(i32::MAX as usize, 0).unwrap(), + i32::MAX as usize + ); + assert!( + checked_hll_state_data_length(i32::MAX as usize, 1) + .unwrap_err() + .to_string() + .contains("i32 offset capacity") + ); + assert!( + checked_hll_state_data_length(usize::MAX, 1) + .unwrap_err() + .to_string() + .contains("length overflow") + ); + } + #[test] fn sparse_stays_sparse_for_small_groups() { let mut g = GroupHll::default(); diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 6663c63b518da..0b1effe5f16a7 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -27,14 +27,15 @@ use datafusion_common::utils::memory::estimate_memory_size; use datafusion_common::utils::proxy::HashTableAllocExt; use datafusion_common::utils::{SingleRowListArrayBuilder, offset_span}; use datafusion_common::{ - HashMap, HashSet, Result, ScalarValue, assert_eq_or_internal_err, exec_err, - internal_datafusion_err, internal_err, + HashMap, HashSet, Result, ScalarValue, assert_eq_or_internal_err, + exec_datafusion_err, exec_err, internal_datafusion_err, internal_err, }; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ Accumulator, AggregateMetric, AggregateMetricRecorder, AggregateMetrics, - AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, Signature, Volatility, + AggregateUDFImpl, Documentation, EmitTo, GroupSelection, GroupsAccumulator, + Signature, Volatility, }; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filter_to_nulls; use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity; @@ -541,6 +542,95 @@ impl ArrayAggGroupsAccumulator { } } + /// Materializes arbitrary groups without changing the stored batches or + /// group indices. + fn build_selected_groups(&self, selected_groups: &[usize]) -> Result { + let mut selected_positions = + HashMap::, RandomState>::with_capacity_and_hasher( + selected_groups.len(), + RandomState::default(), + ); + for (output_index, &group_index) in selected_groups.iter().enumerate() { + selected_positions + .entry(group_index) + .or_default() + .push(output_index); + } + + let mut counts = vec![0usize; selected_groups.len()]; + for entries in &self.batch_entries { + for &(group_index, _) in entries { + if let Some(output_indices) = + selected_positions.get(&(group_index as usize)) + { + for &output_index in output_indices { + counts[output_index] = + counts[output_index].checked_add(1).ok_or_else(|| { + exec_datafusion_err!( + "Array aggregation output length overflow" + ) + })?; + } + } + } + } + + let mut offsets = Vec::::with_capacity(selected_groups.len() + 1); + offsets.push(0); + let mut nulls_builder = NullBufferBuilder::new(selected_groups.len()); + let mut write_positions = Vec::with_capacity(selected_groups.len()); + let mut total_rows = 0usize; + for count in counts { + if count == 0 { + nulls_builder.append_null(); + } else { + nulls_builder.append_non_null(); + } + write_positions.push(total_rows); + total_rows = total_rows.checked_add(count).ok_or_else(|| { + exec_datafusion_err!("Array aggregation output length overflow") + })?; + offsets.push(i32::try_from(total_rows).map_err(|_| { + exec_datafusion_err!( + "Array aggregation output exceeds i32 offset capacity" + ) + })?); + } + + let flat_values = if total_rows == 0 { + new_empty_array(&self.datatype) + } else { + let mut interleave_indices = vec![(0usize, 0usize); total_rows]; + for (batch_index, entries) in self.batch_entries.iter().enumerate() { + for &(group_index, row_index) in entries { + if let Some(output_indices) = + selected_positions.get(&(group_index as usize)) + { + for &output_index in output_indices { + let write_position = write_positions[output_index]; + interleave_indices[write_position] = + (batch_index, row_index as usize); + write_positions[output_index] += 1; + } + } + } + } + + let sources: Vec<&dyn Array> = + self.batches.iter().map(|batch| batch.as_ref()).collect(); + arrow::compute::interleave(&sources, &interleave_indices)? + }; + + let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); + let field = Arc::new(Field::new_list_field(self.datatype.clone(), true)); + Ok(Arc::new(ListArray::new( + field, + offsets, + flat_values, + nulls_builder.finish(), + ))) + } + fn clear_state(&mut self) { // `size()` measures Vec capacity rather than len, so allocate new // buffers instead of using `clear()`. @@ -757,10 +847,31 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator { Ok(Arc::new(result)) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.num_groups)?; + let selected_groups = selection.iter().collect::>(); + self.build_selected_groups(&selected_groups) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: EmitTo) -> Result> { Ok(vec![self.evaluate(emit_to)?]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.evaluate_preserving(selection).map(|array| vec![array]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], @@ -4106,6 +4217,56 @@ mod tests { Ok(()) } + #[test] + fn groups_accumulator_preserving_reads() -> Result<()> { + let mut acc = ArrayAggGroupsAccumulator::new(DataType::Int32, false); + let values: ArrayRef = + Arc::new(Int32Array::from(vec![Some(10), None, Some(20), Some(30)])); + acc.update_batch(&[values], &[0, 1, 0, 2], None, 4)?; + let values: ArrayRef = Arc::new(Int32Array::from(vec![40, 50])); + acc.update_batch(&[values], &[2, 0], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3, 1], 4)?; + let expected = vec![ + Some(vec![Some(30), Some(40)]), + Some(vec![Some(10), Some(20), Some(50)]), + Some(vec![Some(30), Some(40)]), + None, + Some(vec![None]), + ]; + for _ in 0..2 { + let result = acc.evaluate_preserving(selection)?; + assert_eq!(list_array_to_i32_vecs(result.as_list::()), expected); + } + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 1); + assert_eq!(list_array_to_i32_vecs(state[0].as_list::()), expected); + + let empty = GroupSelection::try_from_indices(&[], 4)?; + let result = acc.evaluate_preserving(empty)?; + assert_eq!( + result.data_type(), + &DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))) + ); + assert!(result.is_empty()); + + let values: ArrayRef = Arc::new(Int32Array::from(vec![60])); + acc.update_batch(&[values], &[3], None, 4)?; + let all = acc.evaluate_preserving(GroupSelection::all(4))?; + assert_eq!( + list_array_to_i32_vecs(all.as_list::()), + vec![ + Some(vec![Some(10), Some(20), Some(50)]), + Some(vec![None]), + Some(vec![Some(30), Some(40)]), + Some(vec![Some(60)]), + ] + ); + assert!(acc.supports_evaluate_preserving()); + assert!(acc.supports_state_preserving()); + Ok(()) + } + #[test] fn groups_accumulator_multiple_batches() -> Result<()> { let mut acc = ArrayAggGroupsAccumulator::new(DataType::Int32, false); diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index e7f2d8b7eac47..fd2720b21988e 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -32,7 +32,7 @@ use datafusion_expr::utils::{ }; use datafusion_expr::{ Accumulator, AggregateUDFImpl, Documentation, EmitTo, Expr, ExprFunctionExt, - GroupsAccumulator, ReversedUDAF, Signature, SortExpr, Volatility, + GroupSelection, GroupsAccumulator, ReversedUDAF, Signature, SortExpr, Volatility, }; use datafusion_functions_aggregate_common::utils::get_sort_options; use datafusion_macros::user_doc; @@ -543,6 +543,12 @@ impl FirstLastGroupsAccumulator { Ok(()) } + fn validate_selection(&self, selection: GroupSelection<'_>) -> Result<()> { + selection.validate_num_groups(self.state.num_groups())?; + selection.validate_num_groups(self.orderings.len())?; + selection.validate_num_groups(self.is_sets.len()) + } + fn take_state( &mut self, emit_to: EmitTo, @@ -809,6 +815,15 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator) -> Result { + self.validate_selection(selection)?; + self.state.build_preserving(selection) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn state(&mut self, emit_to: EmitTo) -> Result> { let (val_arr, orderings, is_sets) = self.take_state(emit_to)?; let mut result = Vec::with_capacity(self.ordering_req.len() + 2); @@ -838,6 +853,41 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator, + ) -> Result> { + self.validate_selection(selection)?; + let mut result = Vec::with_capacity(self.ordering_req.len() + 2); + result.push(self.state.build_preserving(selection)?); + + for column_index in 0..self.ordering_req.len() { + let array = if selection.is_empty() { + self.default_orderings[column_index].to_array_of_size(0)? + } else { + ScalarValue::iter_to_array(selection.iter().map(|group_index| { + debug_assert_eq!( + self.orderings[group_index].len(), + self.ordering_req.len() + ); + self.orderings[group_index][column_index].clone() + }))? + }; + result.push(array); + } + + let mut is_sets = BooleanBufferBuilder::new(selection.len()); + for index in selection.iter() { + is_sets.append(self.is_sets.get_bit(index)); + } + result.push(Arc::new(BooleanArray::new(is_sets.finish(), None))); + Ok(result) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn merge_batch( &mut self, values: &[ArrayRef], @@ -1722,6 +1772,209 @@ mod tests { Ok(()) } + #[test] + fn first_groups_preserving_reads_keep_all_state() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "ordering", + DataType::Int64, + true, + )])); + let sort_keys = [PhysicalSortExpr { + expr: col("ordering", &schema).unwrap(), + options: SortOptions::default(), + }]; + let new_accumulator = || { + FirstLastGroupsAccumulator::try_new( + PrimitiveValueState::::new(DataType::Int64), + sort_keys.clone().into(), + false, + &[DataType::Int64], + true, + false, + ) + }; + let mut acc = new_accumulator()?; + + let values: ArrayRef = Arc::new(Int64Array::from(vec![10, 20, 30, 40])); + let ordering: ArrayRef = Arc::new(Int64Array::from(vec![4, 1, 3, 2])); + acc.update_batch(&[values, ordering], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4)?; + let expected = Int64Array::from(vec![Some(40), Some(20), Some(40), None]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + } + + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 3); + assert_eq!(state[0].as_primitive::(), &expected); + assert_eq!( + state[1].as_primitive::(), + &Int64Array::from(vec![Some(2), Some(1), Some(2), None]) + ); + assert_eq!( + state[2].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + + let mut merged = new_accumulator()?; + merged.merge_batch(&state, &[0, 1, 2, 3], 4)?; + assert_eq!( + merged + .evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &expected + ); + + let empty = GroupSelection::try_from_indices(&[], 4)?; + let result = acc.evaluate_preserving(empty)?; + assert_eq!(result.data_type(), &DataType::Int64); + assert!(result.is_empty()); + let state = acc.state_preserving(empty)?; + assert_eq!(state.len(), 3); + assert!(state.iter().all(|array| array.is_empty())); + assert_eq!(state[0].data_type(), &DataType::Int64); + assert_eq!(state[1].data_type(), &DataType::Int64); + assert_eq!(state[2].data_type(), &DataType::Boolean); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![15, 50])); + let ordering: ArrayRef = Arc::new(Int64Array::from(vec![0, 5])); + acc.update_batch(&[values, ordering], &[0, 3], None, 4)?; + assert_eq!( + acc.evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &Int64Array::from(vec![15, 30, 40, 50]) + ); + assert!(acc.supports_evaluate_preserving()); + assert!(acc.supports_state_preserving()); + Ok(()) + } + + #[test] + fn first_groups_preserving_reads_support_utf8_values() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "ordering", + DataType::Utf8, + true, + )])); + let sort_keys = [PhysicalSortExpr { + expr: col("ordering", &schema)?, + options: SortOptions::default(), + }]; + let mut acc = FirstLastGroupsAccumulator::try_new( + BytesValueState::try_new(DataType::Utf8)?, + sort_keys.into(), + false, + &[DataType::Utf8], + true, + false, + )?; + + let values: ArrayRef = Arc::new(StringArray::from(vec![ + "late", + "first", + "group-one", + "group-two", + ])); + let ordering: ArrayRef = Arc::new(StringArray::from(vec!["z", "a", "m", "b"])); + acc.update_batch(&[values, ordering], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4)?; + let expected_values = StringArray::from(vec![ + Some("group-two"), + Some("first"), + Some("group-two"), + None, + ]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)?.as_string::(), + &expected_values + ); + } + + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 3); + assert_eq!(state[0].as_string::(), &expected_values); + assert_eq!( + state[1].as_string::(), + &StringArray::from(vec![Some("b"), Some("a"), Some("b"), None]) + ); + assert_eq!( + state[2].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + Ok(()) + } + + #[test] + fn first_groups_preserving_reads_support_list_values() -> Result<()> { + let value_type = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let schema = Arc::new(Schema::new(vec![Field::new( + "ordering", + DataType::Int64, + true, + )])); + let sort_keys = [PhysicalSortExpr { + expr: col("ordering", &schema)?, + options: SortOptions::default(), + }]; + let mut acc = FirstLastGroupsAccumulator::try_new( + GenericValueState::new(value_type), + sort_keys.into(), + false, + &[DataType::Int64], + true, + false, + )?; + + let values = ListArray::from_iter_primitive::([ + Some(vec![Some(9)]), + Some(vec![Some(1), Some(2)]), + Some(vec![Some(3)]), + Some(vec![Some(4), Some(5)]), + ]); + let ordering = Int64Array::from(vec![9, 1, 3, 2]); + acc.update_batch( + &[Arc::new(values), Arc::new(ordering)], + &[0, 0, 1, 2], + None, + 4, + )?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4)?; + let expected_values = ListArray::from_iter_primitive::([ + Some(vec![Some(4), Some(5)]), + Some(vec![Some(1), Some(2)]), + Some(vec![Some(4), Some(5)]), + None, + ]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)?.as_list::(), + &expected_values + ); + } + + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 3); + assert_eq!(state[0].as_list::(), &expected_values); + assert_eq!( + state[1].as_primitive::(), + &Int64Array::from(vec![Some(2), Some(1), Some(2), None]) + ); + assert_eq!( + state[2].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + Ok(()) + } + #[test] fn test_first_group_acc() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/functions-aggregate/src/first_last/state.rs b/datafusion/functions-aggregate/src/first_last/state.rs index 52d752525e579..baf1b4ebf7c77 100644 --- a/datafusion/functions-aggregate/src/first_last/state.rs +++ b/datafusion/functions-aggregate/src/first_last/state.rs @@ -7,15 +7,33 @@ use std::sync::Arc; use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, AsArray, BinaryBuilder, BinaryViewBuilder, - BooleanBufferBuilder, LargeBinaryBuilder, LargeStringBuilder, PrimitiveArray, - StringBuilder, StringViewBuilder, + BooleanBufferBuilder, LargeBinaryBuilder, LargeStringBuilder, OffsetSizeTrait, + PrimitiveArray, StringBuilder, StringViewBuilder, }; use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::DataType; -use datafusion_common::{Result, ScalarValue, internal_err}; -use datafusion_expr::EmitTo; +use datafusion_common::{Result, ScalarValue, exec_datafusion_err, internal_err}; +use datafusion_expr::{EmitTo, GroupSelection}; + +fn checked_selected_value_length(current: usize, additional: usize) -> Result { + current + .checked_add(additional) + .ok_or_else(|| exec_datafusion_err!("First/last selected value length overflow")) +} + +fn validate_selected_value_offset(total_len: usize) -> Result<()> { + if total_len > O::MAX_OFFSET { + let offset_type = if O::IS_LARGE { "i64" } else { "i32" }; + return Err(exec_datafusion_err!( + "First/last selected value data exceeds {offset_type} offset capacity" + )); + } + Ok(()) +} pub(crate) trait ValueState: Send + Sync { + /// Returns the number of stored groups. + fn num_groups(&self) -> usize; /// Resizes the state to accommodate `new_size` groups. fn resize(&mut self, new_size: usize); /// Updates the state for the specified `group_idx` using the value at `idx` from the provided `array`. @@ -37,6 +55,8 @@ pub(crate) trait ValueState: Send + Sync { fn update(&mut self, group_idx: usize, array: &ArrayRef, idx: usize) -> Result<()>; /// Takes the accumulated state and returns it as an [`ArrayRef`], respecting the `emit_to` strategy. fn take(&mut self, emit_to: EmitTo) -> Result; + /// Materializes selected values without changing the stored state. + fn build_preserving(&self, selection: GroupSelection<'_>) -> Result; /// Returns the estimated memory size of the state in bytes. fn size(&self) -> usize; } @@ -59,6 +79,10 @@ impl PrimitiveValueState { } impl ValueState for PrimitiveValueState { + fn num_groups(&self) -> usize { + self.vals.len() + } + fn resize(&mut self, new_size: usize) { self.vals.resize(new_size, T::default_value()); self.nulls.resize(new_size); @@ -80,6 +104,25 @@ impl ValueState for PrimitiveValueState { Ok(Arc::new(array)) } + fn build_preserving(&self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.vals.len())?; + debug_assert_eq!(self.vals.len(), self.nulls.len()); + let values = selection + .iter() + .map(|index| self.vals[index]) + .collect::>(); + let mut valid = BooleanBufferBuilder::new(selection.len()); + for index in selection.iter() { + valid.append(self.nulls.get_bit(index)); + } + let array = PrimitiveArray::::new( + values.into(), + Some(NullBuffer::new(valid.finish())), + ) + .with_data_type(self.data_type.clone()); + Ok(Arc::new(array)) + } + fn size(&self) -> usize { self.vals.capacity() * size_of::() + self.nulls.capacity() / 8 } @@ -127,6 +170,10 @@ impl BytesValueState { } impl ValueState for BytesValueState { + fn num_groups(&self) -> usize { + self.vals.len() + } + fn resize(&mut self, new_size: usize) { if new_size < self.vals.len() { for v in self.vals[new_size..].iter().flatten() { @@ -264,6 +311,103 @@ impl ValueState for BytesValueState { } } + fn build_preserving(&self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.vals.len())?; + let total_len = selection.iter().try_fold(0usize, |total, index| { + let len = self.vals[index].as_ref().map_or(0, Vec::len); + checked_selected_value_length(total, len) + })?; + match self.data_type { + DataType::Utf8 | DataType::Binary => { + validate_selected_value_offset::(total_len)?; + } + DataType::LargeUtf8 | DataType::LargeBinary => { + validate_selected_value_offset::(total_len)?; + } + _ => {} + } + + match self.data_type { + DataType::Utf8 => { + let mut builder = + StringBuilder::with_capacity(selection.len(), total_len); + for index in selection.iter() { + match self.vals[index].as_deref() { + Some(value) => builder.append_value( + // SAFETY: values came from validated UTF-8 arrays. + unsafe { std::str::from_utf8_unchecked(value) }, + ), + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) + } + DataType::LargeUtf8 => { + let mut builder = + LargeStringBuilder::with_capacity(selection.len(), total_len); + for index in selection.iter() { + match self.vals[index].as_deref() { + Some(value) => builder.append_value( + // SAFETY: values came from validated UTF-8 arrays. + unsafe { std::str::from_utf8_unchecked(value) }, + ), + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) + } + DataType::Utf8View => { + let mut builder = StringViewBuilder::with_capacity(selection.len()); + for index in selection.iter() { + match self.vals[index].as_deref() { + Some(value) => builder.append_value( + // SAFETY: values came from validated UTF-8 arrays. + unsafe { std::str::from_utf8_unchecked(value) }, + ), + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) + } + DataType::Binary => { + let mut builder = + BinaryBuilder::with_capacity(selection.len(), total_len); + for index in selection.iter() { + match self.vals[index].as_deref() { + Some(value) => builder.append_value(value), + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) + } + DataType::LargeBinary => { + let mut builder = + LargeBinaryBuilder::with_capacity(selection.len(), total_len); + for index in selection.iter() { + match self.vals[index].as_deref() { + Some(value) => builder.append_value(value), + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) + } + DataType::BinaryView => { + let mut builder = BinaryViewBuilder::with_capacity(selection.len()); + for index in selection.iter() { + match self.vals[index].as_deref() { + Some(value) => builder.append_value(value), + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) + } + _ => internal_err!( + "Unsupported data type for BytesValueState: {}", + self.data_type + ), + } + } + fn size(&self) -> usize { self.vals.capacity() * size_of::>>() + self.total_capacity } @@ -304,6 +448,10 @@ impl GenericValueState { } impl ValueState for GenericValueState { + fn num_groups(&self) -> usize { + self.vals.len() + } + fn resize(&mut self, new_size: usize) { if new_size < self.vals.len() { for v in self.vals[new_size..].iter().flatten() { @@ -344,6 +492,20 @@ impl ValueState for GenericValueState { ScalarValue::iter_to_array(scalars) } + fn build_preserving(&self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.vals.len())?; + if selection.is_empty() { + return Ok(arrow::array::new_empty_array(&self.data_type)); + } + + let default = ScalarValue::try_from(&self.data_type)?; + let scalars = selection + .iter() + .map(|index| self.vals[index].clone().unwrap_or_else(|| default.clone())) + .collect::>(); + ScalarValue::iter_to_array(scalars) + } + fn size(&self) -> usize { self.vals.capacity() * size_of::>() + self.total_size } @@ -378,6 +540,37 @@ mod tests { use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{DataType, Field, Fields}; + #[test] + fn selected_value_length_checks_offset_capacity() { + assert_eq!( + checked_selected_value_length(usize::MAX, 0).unwrap(), + usize::MAX + ); + assert!( + checked_selected_value_length(usize::MAX, 1) + .unwrap_err() + .to_string() + .contains("length overflow") + ); + + assert!(validate_selected_value_offset::(i32::MAX as usize).is_ok()); + assert!( + validate_selected_value_offset::(i32::MAX as usize + 1) + .unwrap_err() + .to_string() + .contains("i32 offset capacity") + ); + assert!(validate_selected_value_offset::(i64::MAX as usize).is_ok()); + if let Some(too_large) = (i64::MAX as usize).checked_add(1) { + assert!( + validate_selected_value_offset::(too_large) + .unwrap_err() + .to_string() + .contains("i64 offset capacity") + ); + } + } + #[test] fn test_bytes_value_state_utf8() -> Result<()> { let mut state = BytesValueState::try_new(DataType::Utf8)?; @@ -398,6 +591,21 @@ mod tests { state.update(0, &array, 2)?; assert_eq!(state.total_capacity, state.total_capacity_calculated()); + let selected = state.build_preserving(GroupSelection::try_from_indices( + &[1, 0, 1], + state.num_groups(), + )?)?; + assert_eq!( + selected.as_string::(), + &StringArray::from(vec!["world", "longer_string_than_hello", "world"]) + ); + let empty = state.build_preserving(GroupSelection::try_from_indices( + &[], + state.num_groups(), + )?)?; + assert_eq!(empty.data_type(), &DataType::Utf8); + assert!(empty.is_empty()); + let result = state.take(EmitTo::All)?; let result = result.as_string::(); assert_eq!(result.len(), 2); @@ -417,6 +625,8 @@ mod tests { state.resize(1); let array: ArrayRef = Arc::new(LargeStringArray::from(vec!["large_utf8"])); state.update(0, &array, 0)?; + let selected = state.build_preserving(GroupSelection::all(1))?; + assert_eq!(selected.as_string::().value(0), "large_utf8"); let result = state.take(EmitTo::All)?; assert_eq!(result.as_string::().value(0), "large_utf8"); Ok(()) @@ -428,6 +638,8 @@ mod tests { state.resize(1); let array: ArrayRef = Arc::new(StringViewArray::from(vec!["Utf8View"])); state.update(0, &array, 0)?; + let selected = state.build_preserving(GroupSelection::all(1))?; + assert_eq!(selected.as_string_view().value(0), "Utf8View"); let result = state.take(EmitTo::All)?; assert_eq!(result.as_string_view().value(0), "Utf8View"); Ok(()) @@ -439,6 +651,8 @@ mod tests { state.resize(1); let array: ArrayRef = Arc::new(BinaryArray::from(vec![b"binary" as &[u8]])); state.update(0, &array, 0)?; + let selected = state.build_preserving(GroupSelection::all(1))?; + assert_eq!(selected.as_binary::().value(0), b"binary"); let result = state.take(EmitTo::All)?; assert_eq!(result.as_binary::().value(0), b"binary"); Ok(()) @@ -451,6 +665,8 @@ mod tests { let array: ArrayRef = Arc::new(LargeBinaryArray::from(vec![b"large_binary" as &[u8]])); state.update(0, &array, 0)?; + let selected = state.build_preserving(GroupSelection::all(1))?; + assert_eq!(selected.as_binary::().value(0), b"large_binary"); let result = state.take(EmitTo::All)?; assert_eq!(result.as_binary::().value(0), b"large_binary"); Ok(()) @@ -466,6 +682,11 @@ mod tests { state.update(0, &array, 0)?; + let selected = state.build_preserving(GroupSelection::all(1))?; + assert_eq!( + selected.as_binary_view().value(0), + b"long_binary_value_to_test_view" + ); let result = state.take(EmitTo::All)?; let result = result.as_binary_view(); assert_eq!(result.value(0), b"long_binary_value_to_test_view"); @@ -583,6 +804,37 @@ mod tests { "size accounting drifted after overwrite" ); + let selected = state.build_preserving(GroupSelection::try_from_indices( + &[1, 0, 1], + state.num_groups(), + )?)?; + let selected = selected.as_list::(); + assert_eq!(selected.len(), 3); + assert_eq!( + selected + .value(0) + .as_any() + .downcast_ref::() + .unwrap(), + &StringArray::from(vec!["b", "c"]) + ); + assert_eq!( + selected + .value(1) + .as_any() + .downcast_ref::() + .unwrap(), + &StringArray::from(vec!["d", "e", "f"]) + ); + assert_eq!( + selected + .value(2) + .as_any() + .downcast_ref::() + .unwrap(), + &StringArray::from(vec!["b", "c"]) + ); + let result = state.take(EmitTo::All)?; let result = result.as_list::(); assert_eq!(result.len(), 2); diff --git a/datafusion/functions-aggregate/src/percentile_cont.rs b/datafusion/functions-aggregate/src/percentile_cont.rs index d719a082af6cc..4ebc28c6f5c65 100644 --- a/datafusion/functions-aggregate/src/percentile_cont.rs +++ b/datafusion/functions-aggregate/src/percentile_cont.rs @@ -9,8 +9,8 @@ use std::mem::{size_of, size_of_val}; use std::sync::Arc; use arrow::array::{ - ArrowNumericType, ArrowPrimitiveType, BooleanArray, ListArray, PrimitiveArray, - PrimitiveBuilder, + ArrowNumericType, ArrowPrimitiveType, BooleanArray, Float64Array, ListArray, + PrimitiveArray, PrimitiveBuilder, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::{ @@ -42,7 +42,7 @@ use datafusion_expr::{ Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, Signature, TypeSignature, TypeSignatureClass, Volatility, }; -use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; use datafusion_expr::{ expr::{AggregateFunction, Sort}, function::{AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs}, @@ -645,6 +645,37 @@ impl> } Ok(()) } + + fn build_state_preserving(&self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.group_values.len())?; + let mut offsets = Vec::with_capacity(selection.len() + 1); + offsets.push(0); + let mut total_values = 0usize; + for index in selection.iter() { + total_values = total_values + .checked_add(self.group_values[index].len()) + .ok_or_else(|| { + exec_datafusion_err!("Percentile state length overflow") + })?; + offsets.push(i32::try_from(total_values).map_err(|_| { + exec_datafusion_err!("Percentile state exceeds i32 offset capacity") + })?); + } + + let flattened = selection + .iter() + .flat_map(|index| self.group_values[index].iter().copied()) + .collect::>(); + let values = PrimitiveArray::::new(ScalarBuffer::from(flattened), None) + .with_data_type(self.data_type.clone()); + let array = ListArray::new( + Arc::new(Field::new_list_field(self.data_type.clone(), true)), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + Arc::new(values), + None, + ); + Ok(Arc::new(array)) + } } impl GroupsAccumulator for PercentileContGroupsAccumulator @@ -750,13 +781,11 @@ where ); let percentile_array: ArrayRef = match self.group_percentiles.as_mut() { - Some(group_percentiles) => Arc::new(arrow::array::Float64Array::from( - emit_to.take_needed(group_percentiles), - )), - None => Arc::new(arrow::array::Float64Array::from(vec![ - self.percentile - .get() - .ok(); + Some(group_percentiles) => { + Arc::new(Float64Array::from(emit_to.take_needed(group_percentiles))) + } + None => Arc::new(Float64Array::from(vec![ + self.percentile.get().ok(); num_groups ])), }; @@ -764,6 +793,26 @@ where Ok(vec![Arc::new(result_list_array), percentile_array]) } + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + let values = self.build_state_preserving(selection)?; + let percentiles = match &self.group_percentiles { + Some(percentiles) => Float64Array::from_iter( + selection + .iter() + .map(|index| percentiles.get(index).copied().flatten()), + ), + None => Float64Array::from(vec![self.percentile.get().ok(); selection.len()]), + }; + Ok(vec![values, Arc::new(percentiles)]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + fn evaluate(&mut self, emit_to: EmitTo) -> Result { // Emit values let mut emit_group_values = emit_to.take_needed(&mut self.group_values); @@ -804,6 +853,38 @@ where Ok(Arc::new(evaluate_result_builder.finish())) } + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + selection.validate_num_groups(self.group_values.len())?; + let mut builder = PrimitiveBuilder::::with_capacity(selection.len()) + .with_data_type(self.data_type.clone()); + let constant = self + .group_percentiles + .is_none() + .then(|| self.percentile.effective_percentile()) + .transpose()?; + for index in selection.iter() { + let mut values = self.group_values[index].clone(); + if values.is_empty() { + builder.append_null(); + continue; + } + let percentile = match constant { + Some(percentile) => percentile, + None => self.percentile.apply_order(self.percentile.get_group_value( + self.group_percentiles.as_ref().and_then(|percentiles| { + percentiles.get(index).copied().flatten() + }), + )?), + }; + builder.append_option(calculate_percentile::(&mut values, percentile)?); + } + Ok(Arc::new(builder.finish())) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + fn convert_to_state( &self, values: &[ArrayRef], @@ -858,16 +939,10 @@ where let percentile_values = self.percentile.to_float64_array(percentile_values)?; let nulls = filtered_null_mask(opt_filter, &percentile_values); - Arc::new(arrow::array::Float64Array::new( - percentile_values.values().clone(), - nulls, - )) + Arc::new(Float64Array::new(percentile_values.values().clone(), nulls)) } else { let percentile = self.percentile.get().ok(); - Arc::new(arrow::array::Float64Array::from(vec![ - percentile; - input_array.len() - ])) + Arc::new(Float64Array::from(vec![percentile; input_array.len()])) }; Ok(vec![Arc::new(converted_list_array), percentile_array]) @@ -1244,6 +1319,186 @@ mod tests { ) } + #[test] + fn preserving_reads_keep_per_group_percentiles() -> Result<()> { + let new_accumulator = || { + PercentileContGroupsAccumulator::::new( + PercentileParam { + aggregate_fn_name: "PERCENTILE_CONT".to_string(), + state: PercentileParamState::Pending, + is_desc: false, + }, + DataType::Float64, + ) + }; + let mut acc = new_accumulator(); + acc.update_batch( + &[ + Arc::new(Float64Array::from(vec![0.0, 4.0, 10.0, 20.0])), + Arc::new(Float64Array::from(vec![0.25, 0.25, 0.75, 0.75])), + ], + &[0, 0, 1, 1], + None, + 3, + )?; + let selection = GroupSelection::try_from_indices(&[1, 0, 1, 2], 3)?; + let expected = Float64Array::from(vec![Some(17.5), Some(1.0), Some(17.5), None]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + } + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 2); + assert_eq!( + state[1].as_primitive::(), + &Float64Array::from(vec![Some(0.75), Some(0.25), Some(0.75), None]) + ); + let mut merged = new_accumulator(); + merged.merge_batch(&state, &[0, 1, 2, 3], 4)?; + assert_eq!( + merged.evaluate(EmitTo::All)?.as_primitive::(), + &expected + ); + acc.update_batch( + &[ + Arc::new(Float64Array::from(vec![6.0])), + Arc::new(Float64Array::from(vec![0.5])), + ], + &[2], + None, + 3, + )?; + assert_eq!( + acc.evaluate(EmitTo::All)?.as_primitive::(), + &Float64Array::from(vec![1.0, 17.5, 6.0]) + ); + Ok(()) + } + + #[test] + fn preserving_reads_keep_median_decimal_metadata() -> Result<()> { + let data_type = DataType::Decimal128(12, 2); + let mut acc = create_percentile_groups_accumulator( + "median", + PercentileParam { + aggregate_fn_name: "MEDIAN".to_string(), + state: PercentileParamState::Resolved(0.5), + is_desc: false, + }, + &data_type, + )?; + let values = PrimitiveArray::::from(vec![100, 300, 900]) + .with_data_type(data_type.clone()); + acc.update_batch(&[Arc::new(values)], &[0, 0, 1], None, 3)?; + let selection = GroupSelection::try_from_indices(&[1, 0, 1, 2], 3)?; + let expected = PrimitiveArray::::from(vec![ + Some(900), + Some(200), + Some(900), + None, + ]) + .with_data_type(data_type.clone()); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + } + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 2); + assert_eq!(state[0].as_list::().values().data_type(), &data_type); + assert_eq!( + state[1].as_primitive::(), + &Float64Array::from(vec![0.5; 4]) + ); + let all = acc.evaluate(EmitTo::All)?; + assert_eq!(all.data_type(), &data_type); + assert_eq!( + all.as_primitive::() + .iter() + .collect::>(), + vec![Some(200), Some(900), None] + ); + Ok(()) + } + + #[test] + fn groups_preserving_reads_keep_percentile_state() -> Result<()> { + let mut acc = + PercentileContGroupsAccumulator::::new( + PercentileParam { + aggregate_fn_name: "PERCENTILE_CONT".to_string(), + state: PercentileParamState::Resolved(0.25), + is_desc: false, + }, + DataType::Float64, + ); + let values: ArrayRef = + Arc::new(Float64Array::from(vec![4.0, 8.0, 0.0, 20.0, 0.0, 10.0])); + acc.update_batch(&[values], &[0, 1, 0, 3, 3, 3], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[3, 0, 3, 2], 4)?; + let expected = Float64Array::from(vec![Some(5.0), Some(1.0), Some(5.0), None]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)? + .as_primitive::(), + &expected + ); + } + + let state = acc.state_preserving(selection)?; + let state = state[0].as_list::(); + let state_values = state + .iter() + .map(|values| { + values + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!( + state_values, + vec![ + vec![20.0, 0.0, 10.0], + vec![4.0, 0.0], + vec![20.0, 0.0, 10.0], + vec![], + ] + ); + + let empty = GroupSelection::try_from_indices(&[], 4)?; + let result = acc.evaluate_preserving(empty)?; + assert_eq!(result.data_type(), &DataType::Float64); + assert!(result.is_empty()); + let state = acc.state_preserving(empty)?; + assert_eq!(state.len(), 2); + assert_eq!(state[1].data_type(), &DataType::Float64); + assert!(state[1].is_empty()); + assert_eq!( + state[0].data_type(), + &DataType::List(Arc::new(Field::new_list_field(DataType::Float64, true))) + ); + assert!(state[0].is_empty()); + + let values: ArrayRef = Arc::new(Float64Array::from(vec![4.0, 12.0])); + acc.update_batch(&[values], &[2, 2], None, 4)?; + assert_eq!( + acc.evaluate_preserving(GroupSelection::all(4))? + .as_primitive::(), + &Float64Array::from(vec![1.0, 8.0, 6.0, 5.0]) + ); + assert!(acc.supports_evaluate_preserving()); + assert!(acc.supports_state_preserving()); + Ok(()) + } + #[test] fn retract_batch_errors_on_untracked_value() { let mut acc = create_accumulator(); From 2f6fe0f5fe1a4fae8b4e32b152cc6ae6da0c9b40 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:41:31 +0800 Subject: [PATCH 08/10] feat(aggregate): retain native unordered first_value groups Port fork #91, including preserving reads and filter/null/merge regressions, while retaining DF56 ordered-aggregate fast paths. --- .../functions-aggregate/src/first_last.rs | 549 ++++++++++++++++-- 1 file changed, 512 insertions(+), 37 deletions(-) diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index fd2720b21988e..ad8a5de9f3e4e 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -72,6 +72,12 @@ fn create_groups_accumulator_helper( state: S, ) -> Result> { let Some(ordering) = LexOrdering::new(args.order_bys.to_vec()) else { + if is_first { + return Ok(Box::new(UnorderedFirstValueGroupsAccumulator::new( + state, + args.ignore_nulls, + ))); + } return internal_err!("Groups accumulator must have an ordering."); }; @@ -191,44 +197,47 @@ fn create_groups_accumulator( } } +/// Returns true when a value state exists for the aggregate's return type. +/// +/// `LAST_VALUE` also needs an ordering; `FIRST_VALUE` has a grouped +/// accumulator with or without one. fn groups_accumulator_supported(args: &AccumulatorArgs) -> bool { use DataType::*; - !args.order_bys.is_empty() - && matches!( - args.return_field.data_type(), - Int8 | Int16 - | Int32 - | Int64 - | UInt8 - | UInt16 - | UInt32 - | UInt64 - | Float16 - | Float32 - | Float64 - | Decimal32(_, _) - | Decimal64(_, _) - | Decimal128(_, _) - | Decimal256(_, _) - | Date32 - | Date64 - | Time32(_) - | Time64(_) - | Timestamp(_, _) - | Utf8 - | LargeUtf8 - | Utf8View - | Binary - | LargeBinary - | BinaryView - | List(_) - | LargeList(_) - | ListView(_) - | LargeListView(_) - | FixedSizeList(_, _) - | Struct(_) - | Map(_, _) - ) + matches!( + args.return_field.data_type(), + Int8 | Int16 + | Int32 + | Int64 + | UInt8 + | UInt16 + | UInt32 + | UInt64 + | Float16 + | Float32 + | Float64 + | Decimal32(_, _) + | Decimal64(_, _) + | Decimal128(_, _) + | Decimal256(_, _) + | Date32 + | Date64 + | Time32(_) + | Time64(_) + | Timestamp(_, _) + | Utf8 + | LargeUtf8 + | Utf8View + | Binary + | LargeBinary + | BinaryView + | List(_) + | LargeList(_) + | ListView(_) + | LargeListView(_) + | FixedSizeList(_, _) + | Struct(_) + | Map(_, _) + ) } #[user_doc( @@ -958,6 +967,169 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator { + state: S, + /// `is_sets[group_idx]` is true once the group holds its first value. + is_sets: BooleanBufferBuilder, + ignore_nulls: bool, +} + +impl UnorderedFirstValueGroupsAccumulator { + fn new(state: S, ignore_nulls: bool) -> Self { + Self { + state, + is_sets: BooleanBufferBuilder::new(0), + ignore_nulls, + } + } + + fn resize(&mut self, total_num_groups: usize) { + if self.is_sets.len() < total_num_groups { + self.state.resize(total_num_groups); + self.is_sets.resize(total_num_groups); + } + } + + /// Stores row `idx` of `values` as the group's value unless the group + /// already has one. + fn set_if_unset( + &mut self, + group_idx: usize, + values: &ArrayRef, + idx: usize, + ) -> Result<()> { + if !self.is_sets.get_bit(group_idx) { + self.state.update(group_idx, values, idx)?; + self.is_sets.set_bit(group_idx, true); + } + Ok(()) + } + + fn validate_selection(&self, selection: GroupSelection<'_>) -> Result<()> { + selection.validate_num_groups(self.state.num_groups())?; + selection.validate_num_groups(self.is_sets.len()) + } +} + +impl GroupsAccumulator + for UnorderedFirstValueGroupsAccumulator +{ + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + self.resize(total_num_groups); + let values = &values[0]; + for (idx, &group_idx) in group_indices.iter().enumerate() { + // A row passes the FILTER clause only when the predicate is + // `true`; rows whose predicate evaluates to `null` are excluded. + if let Some(filter) = opt_filter + && !(filter.is_valid(idx) && filter.value(idx)) + { + continue; + } + if self.ignore_nulls && values.is_null(idx) { + continue; + } + self.set_if_unset(group_idx, values, idx)?; + } + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result { + let values = self.state.take(emit_to)?; + state::take_need(&mut self.is_sets, emit_to); + Ok(values) + } + + fn evaluate_preserving(&mut self, selection: GroupSelection<'_>) -> Result { + self.validate_selection(selection)?; + self.state.build_preserving(selection) + } + + fn supports_evaluate_preserving(&self) -> bool { + true + } + + fn state(&mut self, emit_to: EmitTo) -> Result> { + let values = self.state.take(emit_to)?; + let is_sets = state::take_need(&mut self.is_sets, emit_to); + Ok(vec![values, Arc::new(BooleanArray::new(is_sets, None))]) + } + + fn state_preserving( + &mut self, + selection: GroupSelection<'_>, + ) -> Result> { + self.validate_selection(selection)?; + let values = self.state.build_preserving(selection)?; + let mut is_sets = BooleanBufferBuilder::new(selection.len()); + for group_idx in selection.iter() { + is_sets.append(self.is_sets.get_bit(group_idx)); + } + Ok(vec![ + values, + Arc::new(BooleanArray::new(is_sets.finish(), None)), + ]) + } + + fn supports_state_preserving(&self) -> bool { + true + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + let [values, is_set] = values else { + return internal_err!( + "first_value: expected 2 state columns, got {}", + values.len() + ); + }; + let is_set = as_boolean_array(is_set)?; + validate_is_set_flags(is_set, "first_value")?; + self.resize(total_num_groups); + for (idx, &group_idx) in group_indices.iter().enumerate() { + if is_set.value(idx) { + self.set_if_unset(group_idx, values, idx)?; + } + } + Ok(()) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + let values = &values[0]; + let is_set = BooleanArray::from_iter((0..values.len()).map(|idx| { + let passed_filter = + opt_filter.is_none_or(|f| f.is_valid(idx) && f.value(idx)); + Some(passed_filter && !(self.ignore_nulls && values.is_null(idx))) + })); + Ok(vec![Arc::clone(values), Arc::new(is_set)]) + } + + fn size(&self) -> usize { + self.state.size() + self.is_sets.capacity() / 8 + } +} + /// This accumulator is used when there is no ordering specified for the /// `FIRST_VALUE` aggregation. It simply returns the first value it sees /// according to the pre-existing ordering of the input data, and provides @@ -1339,7 +1511,7 @@ impl AggregateUDFImpl for LastValue { } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - groups_accumulator_supported(&args) + !args.order_bys.is_empty() && groups_accumulator_supported(&args) } fn create_groups_accumulator( @@ -1642,6 +1814,7 @@ fn convert_to_sort_cols(arrs: &[ArrayRef], sort_exprs: &LexOrdering) -> Vec Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(1), Some(2)])); + group_acc.update_batch(&[values], &[0, 0, 1], None, 3)?; + + // RESPECT NULLS keeps the leading null as group 0's first value. + let state = group_acc.state_preserving(GroupSelection::all(3))?; + let expected_state: Vec = vec![ + Arc::new(Int64Array::from(vec![None, Some(2), None])), + Arc::new(BooleanArray::from(vec![true, true, false])), + ]; + assert_eq!(state, expected_state); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![5, 6, 7])); + group_acc.update_batch(&[values], &[0, 1, 2], None, 4)?; + + // Groups 0 and 1 already have a first value, so only group 2 changes. + let binding = group_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + let expect = Int64Array::from(vec![None, Some(2), Some(7), None]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_ignore_nulls() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + true, + ); + + let values: ArrayRef = + Arc::new(Int64Array::from(vec![None, Some(1), None, Some(3)])); + group_acc.update_batch(&[values], &[0, 0, 1, 0], None, 2)?; + + // Group 1 only saw a null, so it stays unset. + let state = group_acc.state(EmitTo::All)?; + let expected_state: Vec = vec![ + Arc::new(Int64Array::from(vec![Some(1), None])), + Arc::new(BooleanArray::from(vec![true, false])), + ]; + assert_eq!(state, expected_state); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_filter_null_predicate() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![10, 20, 30])); + + // Row 0: predicate is null with its value bit set + // Row 1: predicate is false + // Row 2: predicate is true + let filter = BooleanArray::new( + BooleanBuffer::from(vec![false, true, false, true]), + Some(NullBuffer::from(BooleanBuffer::from(vec![ + true, false, true, true, + ]))), + ) + .slice(1, 3); + assert_eq!(filter.offset(), 1); + + group_acc.update_batch(&[values], &[0, 0, 1], Some(&filter), 2)?; + + let binding = group_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + + // Group 0 has no row with a `true` predicate, so it must stay unset. + let expect = Int64Array::from(vec![None, Some(30)]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_merge() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![5])); + group_acc.update_batch(&[values], &[0], None, 1)?; + + let state: Vec = vec![ + Arc::new(Int64Array::from(vec![100, 200, 300, 400])), + Arc::new(BooleanArray::from(vec![true, true, false, true])), + ]; + group_acc.merge_batch(&state, &[2, 0, 1, 1], 3)?; + + let binding = group_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + + // group 0 keeps its value from update_batch. + // group 1 skips 300 (is_set=false) and takes 400. + // group 2 takes 100. + let expect = Int64Array::from(vec![5, 400, 100]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + /// `convert_to_state` folds the FILTER clause and IGNORE NULLS into + /// `is_set`, so `merge_batch` keeps only rows `update_batch` would keep. + #[test] + fn test_unordered_first_group_acc_convert_to_state() -> Result<()> { + let group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + true, + ); + + let values: ArrayRef = + Arc::new(Int64Array::from(vec![None, Some(10), Some(20), Some(30)])); + + // Row 0: null value, skipped by IGNORE NULLS + // Row 1: predicate is null with its value bit set + // Row 2: predicate is false + // Row 3: predicate is true + let filter = BooleanArray::new( + BooleanBuffer::from(vec![true, true, false, true]), + Some(NullBuffer::from(BooleanBuffer::from(vec![ + true, false, true, true, + ]))), + ); + + let state = group_acc.convert_to_state(&[values], Some(&filter))?; + assert_eq!(state.len(), 2); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![false, false, false, true]) + ); + + let mut merging_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + true, + ); + merging_acc.merge_batch(&state, &[0, 0, 0, 0], 1)?; + + let binding = merging_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_primitive::(); + let expect = Int64Array::from(vec![Some(30)]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_emit_first() -> Result<()> { + let mut group_acc = UnorderedFirstValueGroupsAccumulator::new( + PrimitiveValueState::::new(DataType::Int64), + false, + ); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + group_acc.update_batch(&[values], &[0, 1, 2], None, 3)?; + + let binding = group_acc.evaluate(EmitTo::First(2))?; + assert_eq!( + binding.as_primitive::(), + &Int64Array::from(vec![1, 2]) + ); + + // Old group 2 is now group 0 and keeps its value; group 1 is new. + let values: ArrayRef = Arc::new(Int64Array::from(vec![9, 8])); + group_acc.update_batch(&[values], &[0, 1], None, 2)?; + + let binding = group_acc.evaluate(EmitTo::All)?; + assert_eq!( + binding.as_primitive::(), + &Int64Array::from(vec![3, 8]) + ); + + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_preserving_reads_utf8() -> Result<()> { + let mut acc = UnorderedFirstValueGroupsAccumulator::new( + BytesValueState::try_new(DataType::Utf8)?, + false, + ); + + let values: ArrayRef = Arc::new(StringArray::from(vec![ + "first", + "late", + "group-one", + "group-two", + ])); + acc.update_batch(&[values], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4)?; + let expected_values = StringArray::from(vec![ + Some("group-two"), + Some("first"), + Some("group-two"), + None, + ]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)?.as_string::(), + &expected_values + ); + } + + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 2); + assert_eq!(state[0].as_string::(), &expected_values); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + assert!(acc.supports_evaluate_preserving()); + assert!(acc.supports_state_preserving()); + Ok(()) + } + + #[test] + fn test_unordered_first_group_acc_preserving_reads_list() -> Result<()> { + let value_type = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let mut acc = UnorderedFirstValueGroupsAccumulator::new( + GenericValueState::new(value_type), + false, + ); + + let values = ListArray::from_iter_primitive::([ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(9)]), + Some(vec![Some(3)]), + Some(vec![Some(4), Some(5)]), + ]); + acc.update_batch(&[Arc::new(values)], &[0, 0, 1, 2], None, 4)?; + + let selection = GroupSelection::try_from_indices(&[2, 0, 2, 3], 4)?; + let expected_values = ListArray::from_iter_primitive::([ + Some(vec![Some(4), Some(5)]), + Some(vec![Some(1), Some(2)]), + Some(vec![Some(4), Some(5)]), + None, + ]); + for _ in 0..2 { + assert_eq!( + acc.evaluate_preserving(selection)?.as_list::(), + &expected_values + ); + } + + let state = acc.state_preserving(selection)?; + assert_eq!(state.len(), 2); + assert_eq!(state[0].as_list::(), &expected_values); + assert_eq!( + state[1].as_boolean(), + &BooleanArray::from(vec![true, true, true, false]) + ); + Ok(()) + } + #[test] fn test_first_value_merge_compacts_list_state() -> Result<()> { let values: ArrayRef = @@ -3247,4 +3688,38 @@ mod tests { assert_eq!(merged.size(), updated.size()); Ok(()) } + + #[test] + fn test_first_value_unordered_groups_accumulator() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]); + let exprs = [col("a", &schema)?]; + let expr_fields = [Arc::new(Field::new("a", DataType::Int64, true))]; + let args = || AccumulatorArgs { + return_field: Arc::new(Field::new("f", DataType::Int64, true)), + schema: &schema, + expr_fields: &expr_fields, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "first_value(a)", + is_distinct: false, + exprs: &exprs, + }; + + // Without ORDER BY, FIRST_VALUE uses its own grouped accumulator + // rather than falling back to `GroupsAccumulatorAdapter`. + let first_value = FirstValue::new(); + assert!(first_value.groups_accumulator_supported(args())); + let group_acc = first_value.create_groups_accumulator(args())?; + let group_acc: &dyn Any = group_acc.as_ref(); + assert!( + group_acc + .is::>>() + ); + + // LAST_VALUE still requires an ordering. + assert!(!LastValue::new().groups_accumulator_supported(args())); + + Ok(()) + } } From 9a7be3e6d2f0d42a9a9ad732f66fe9019e8d398f Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:46:41 +0800 Subject: [PATCH 09/10] feat(csv): retain matching file columns by header Port fork #84 while preserving DF56 null-regex and custom terminator support. CSV scan field 9 now belongs to upstream terminator; use field 10 for header matching and regenerate protobuf bindings. Serialized CSV plans must not be exchanged between DF55-fork and DF56-fork processes. --- Cargo.lock | 2 + datafusion/common/src/config.rs | 17 + .../core/src/datasource/physical_plan/csv.rs | 280 +++++++++ datafusion/datasource-csv/Cargo.toml | 2 + datafusion/datasource-csv/src/file_format.rs | 11 + datafusion/datasource-csv/src/source.rs | 557 ++++++++++++++++-- .../proto/datafusion_common.proto | 2 + datafusion/proto-common/src/from_proto/mod.rs | 4 + .../proto-common/src/generated/pbjson.rs | 22 + .../proto-common/src/generated/prost.rs | 3 + datafusion/proto-common/src/to_proto/mod.rs | 3 + .../proto-models/proto/datafusion.proto | 1 + datafusion/proto-models/src/from_proto.rs | 5 + .../src/generated/datafusion_proto_common.rs | 3 + .../proto-models/src/generated/pbjson.rs | 18 + .../proto-models/src/generated/prost.rs | 2 + datafusion/proto/tests/cases/plans/sources.rs | 41 +- .../sqllogictest/test_files/csv_files.slt | 40 ++ docs/source/user-guide/sql/format_options.md | 39 +- 19 files changed, 997 insertions(+), 55 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e7fd5e05d03b3..ce0efd62a6d5c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2042,11 +2042,13 @@ dependencies = [ "datafusion-datasource", "datafusion-execution", "datafusion-expr", + "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", "futures", + "log", "object_store", "regex", "tokio", diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 825536e165ddd..651e681cbe97b 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -4102,6 +4102,16 @@ config_namespace! { /// of columns, it creates a union schema containing all columns found across the files, and will /// pad any files missing columns with null values for their rows. pub truncated_rows: Option, default = None + /// Whether to map file columns to the table schema by header name instead of by position. + /// + /// Only applies when `has_header` is true. When enabled, each file's header row is read + /// and columns are matched by name: table columns missing from a file are filled with + /// nulls, and file columns not declared in the table are ignored. Names are matched exactly + /// (case-sensitive) and duplicate header names resolve to the first occurrence. Header + /// names containing quoted newlines are not supported. + /// + /// Defaults to `false` when `None`. + pub match_columns_by_header: Option, default = None } } @@ -4227,6 +4237,13 @@ impl CsvOptions { self } + /// Whether to map file columns to the table schema by header name instead of by position. + /// Only applies when the file has a header. + pub fn with_match_columns_by_header(mut self, match_by_header: bool) -> Self { + self.match_columns_by_header = Some(match_by_header); + self + } + /// Set the compression level for the output file. /// The valid range depends on the compression algorithm. /// If not specified, the default level for the algorithm is used. diff --git a/datafusion/core/src/datasource/physical_plan/csv.rs b/datafusion/core/src/datasource/physical_plan/csv.rs index 004e1d60cb22d..a2367b4ba1856 100644 --- a/datafusion/core/src/datasource/physical_plan/csv.rs +++ b/datafusion/core/src/datasource/physical_plan/csv.rs @@ -802,4 +802,284 @@ mod tests { Arc::new(schema) } + + /// Files sharing the table `(a INT, b TEXT, c INT, d TEXT)` but written + /// with different header layouts + const HEADER_LAYOUT_FILES: [(&str, &str); 3] = [ + ("full.csv", "a,b,c,d\n1,x,10,p\n4,y,40,s\n"), + ("missing_middle.csv", "a,d\n2,q\n5,t\n"), + ("reordered.csv", "d,extra,c,a\nr,ignored,30,3\n"), + ]; + + #[derive(Debug, Clone, Copy)] + enum HeaderLayoutRead { + Local, + /// Tiny object store chunks so the header spans several chunks + Chunked, + /// Byte-range partitions that start mid-file + Repartitioned, + } + + fn write_header_layout_files(dir: &std::path::Path) -> Result<()> { + for (name, contents) in HEADER_LAYOUT_FILES { + fs::write(dir.join(name), contents)?; + } + Ok(()) + } + + async fn header_layout_table( + ctx: &SessionContext, + location: &str, + options: &str, + ) -> Result<()> { + ctx.sql(&format!( + "CREATE EXTERNAL TABLE t (a INT, b TEXT, c INT, d TEXT) \ + STORED AS CSV LOCATION '{location}' \ + OPTIONS ('format.has_header' 'true'{options})" + )) + .await? + .collect() + .await?; + Ok(()) + } + + #[rstest] + #[case::local(HeaderLayoutRead::Local)] + #[case::chunked(HeaderLayoutRead::Chunked)] + #[case::repartitioned(HeaderLayoutRead::Repartitioned)] + #[tokio::test] + async fn csv_match_columns_by_header(#[case] read: HeaderLayoutRead) -> Result<()> { + let tmp_dir = TempDir::new()?; + write_header_layout_files(tmp_dir.path())?; + + let ctx = match read { + HeaderLayoutRead::Local | HeaderLayoutRead::Chunked => SessionContext::new(), + HeaderLayoutRead::Repartitioned => SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_repartition_file_min_size(1), + ), + }; + if matches!(read, HeaderLayoutRead::Chunked) { + ctx.register_object_store( + &Url::parse("file://").unwrap(), + Arc::new(ChunkedStore::new(Arc::new(LocalFileSystem::new()), 3)), + ); + } + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true'", + ) + .await?; + + let plan = ctx + .sql("SELECT * FROM t") + .await? + .create_physical_plan() + .await?; + let plan = datafusion_physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!(plan.contains("match_columns_by_header=true"), "{plan}"); + if matches!(read, HeaderLayoutRead::Repartitioned) { + assert!(plan.contains("missing_middle.csv:"), "{plan}"); + } + + let all = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + let projected = ctx + .sql("SELECT d, c FROM t ORDER BY d") + .await? + .collect() + .await?; + let missing_only = ctx + .sql("SELECT count(b), count(*) FROM t") + .await? + .collect() + .await?; + + insta::allow_duplicates! { + assert_snapshot!(batches_to_string(&all), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 1 | x | 10 | p | + | 2 | | | q | + | 3 | | 30 | r | + | 4 | y | 40 | s | + | 5 | | | t | + +---+---+----+---+ + "); + assert_snapshot!(batches_to_string(&projected), @r" + +---+----+ + | d | c | + +---+----+ + | p | 10 | + | q | | + | r | 30 | + | s | 40 | + | t | | + +---+----+ + "); + assert_snapshot!(batches_to_string(&missing_only), @r" + +------------+----------+ + | count(t.b) | count(*) | + +------------+----------+ + | 2 | 5 | + +------------+----------+ + "); + } + Ok(()) + } + + #[rstest] + #[case::local(HeaderLayoutRead::Local)] + #[case::chunked(HeaderLayoutRead::Chunked)] + #[case::repartitioned(HeaderLayoutRead::Repartitioned)] + #[tokio::test] + async fn csv_match_columns_by_header_with_comments( + #[case] read: HeaderLayoutRead, + ) -> Result<()> { + let tmp_dir = TempDir::new()?; + fs::write( + tmp_dir.path().join("a.csv"), + "# generated\n# by test\nd,c,a\nr,30,3\n# trailing\ns,40,4\n", + )?; + + let ctx = match read { + HeaderLayoutRead::Local | HeaderLayoutRead::Chunked => SessionContext::new(), + HeaderLayoutRead::Repartitioned => SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_repartition_file_min_size(1), + ), + }; + if matches!(read, HeaderLayoutRead::Chunked) { + ctx.register_object_store( + &Url::parse("file://").unwrap(), + Arc::new(ChunkedStore::new(Arc::new(LocalFileSystem::new()), 3)), + ); + } + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true', 'format.comment' '#'", + ) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + insta::allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 3 | | 30 | r | + | 4 | | 40 | s | + +---+---+----+---+ + "); + } + Ok(()) + } + + #[tokio::test] + async fn csv_match_columns_by_header_with_truncated_rows() -> Result<()> { + let tmp_dir = TempDir::new()?; + fs::write(tmp_dir.path().join("a.csv"), "d,a,c\nr,1,10\ns,2\n")?; + + let ctx = SessionContext::new(); + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true', 'format.truncated_rows' 'true'", + ) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 1 | | 10 | r | + | 2 | | | s | + +---+---+----+---+ + "); + Ok(()) + } + + #[tokio::test] + async fn csv_columns_mapped_by_position_by_default() -> Result<()> { + let tmp_dir = TempDir::new()?; + fs::write(tmp_dir.path().join("a.csv"), "a,d\n2,q\n")?; + + let ctx = SessionContext::new(); + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.truncated_rows' 'true'", + ) + .await?; + + let batches = ctx.sql("SELECT * FROM t").await?.collect().await?; + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+---+---+ + | a | b | c | d | + +---+---+---+---+ + | 2 | q | | | + +---+---+---+---+ + "); + Ok(()) + } + + #[cfg(feature = "compression")] + #[tokio::test] + async fn csv_match_columns_by_header_gzip() -> Result<()> { + use flate2::Compression; + use flate2::write::GzEncoder; + + let tmp_dir = TempDir::new()?; + for (name, contents) in HEADER_LAYOUT_FILES { + let file = File::create(tmp_dir.path().join(format!("{name}.gz")))?; + let mut encoder = GzEncoder::new(file, Compression::default()); + encoder.write_all(contents.as_bytes())?; + encoder.finish()?; + } + + let ctx = SessionContext::new(); + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true', 'format.compression' 'gzip'", + ) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 1 | x | 10 | p | + | 2 | | | q | + | 3 | | 30 | r | + | 4 | y | 40 | s | + | 5 | | | t | + +---+---+----+---+ + "); + Ok(()) + } } diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index 7e7195dfda9d5..5ed78ea04a8c0 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -47,11 +47,13 @@ datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } +datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } +log = { workspace = true } object_store = { workspace = true } regex = { workspace = true } tokio = { workspace = true } diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 578694c18a283..07a75591a3a16 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -293,6 +293,14 @@ impl CsvFormat { self } + /// Set whether file columns are matched to the table schema by header name + /// instead of by position + /// - defaults to false + pub fn with_match_columns_by_header(mut self, match_by_header: bool) -> Self { + self.options.match_columns_by_header = Some(match_by_header); + self + } + /// The delimiter character. pub fn delimiter(&self) -> u8 { self.options.delimiter @@ -998,6 +1006,9 @@ impl From<&CsvFormatFactory> for datafusion_proto_models::protobuf::CsvOptions { ignore_trailing_whitespace: options .ignore_trailing_whitespace .map_or(vec![], |v| vec![v as u8]), + match_columns_by_header: options + .match_columns_by_header + .map_or(vec![], |v| vec![v as u8]), } } else { datafusion_proto_models::protobuf::CsvOptions::default() diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index cbc8df820f222..4fb659d115faa 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -8,7 +8,8 @@ use datafusion_datasource::boundary_stream::AlignedBoundaryStream; use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use datafusion_physical_plan::projection::ProjectionExprs; use std::fmt; -use std::io::Read; +use std::io::{BufRead, BufReader, Cursor, Read}; +use std::ops::Range; use std::sync::Arc; use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream}; @@ -18,7 +19,10 @@ use datafusion_datasource::{ FileRange, ListingTableUrl, PartitionedFile, TableSchema, as_file_source, }; +use arrow::array::RecordBatch; use arrow::csv; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; use datafusion_common::config::CsvOptions; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{DataFusionError, Result, exec_datafusion_err}; @@ -26,15 +30,21 @@ use datafusion_common_runtime::JoinSet; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_execution::TaskContext; +use datafusion_physical_expr_adapter::{ + BatchAdapter, BatchAdapterFactory, DefaultPhysicalExprAdapterFactory, + PhysicalExprAdapterFactory, +}; use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; use datafusion_physical_plan::{ DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, }; use crate::file_format::CsvDecoder; +use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use object_store::buffered::BufWriter; -use object_store::{GetOptions, GetResultPayload, ObjectStore}; +use object_store::path::Path; +use object_store::{GetOptions, GetResultPayload, ObjectStore, ObjectStoreExt}; use regex::Regex; use tokio::io::AsyncWriteExt; @@ -174,28 +184,53 @@ impl CsvSource { pub fn newlines_in_values(&self) -> bool { self.options.newlines_in_values.unwrap_or(false) } + + /// true if file columns are matched to the table schema by header name + pub fn match_columns_by_header(&self) -> bool { + self.options.match_columns_by_header.unwrap_or(false) && self.has_header() + } + + /// Whether to match file columns to the table schema by header name + pub fn with_match_columns_by_header(&self, match_by_header: bool) -> Self { + let mut conf = self.clone(); + conf.options.match_columns_by_header = Some(match_by_header); + conf + } } impl CsvSource { - fn open(&self, reader: R) -> Result> { - Ok(self.builder()?.build(reader)?) + fn open( + &self, + reader: R, + mapping: Option<&HeaderMapping>, + ) -> Result> { + Ok(self.builder(mapping)?.build(reader)?) } - fn builder(&self) -> Result { - let mut builder = - csv::ReaderBuilder::new(Arc::clone(self.table_schema.file_schema())) - .with_delimiter(self.delimiter()) - .with_batch_size( - self.batch_size - .expect("Batch size must be set before initializing builder"), - ) - .with_header(self.has_header()) - .with_quote(self.quote()) - .with_truncated_rows(self.truncate_rows()); + fn builder(&self, mapping: Option<&HeaderMapping>) -> Result { + let (schema, projection) = match mapping { + Some(mapping) => ( + Arc::clone(&mapping.physical_schema), + mapping.physical_projection.clone(), + ), + None => ( + Arc::clone(self.table_schema.file_schema()), + self.projection.file_indices.clone(), + ), + }; + let mut builder = csv::ReaderBuilder::new(schema) + .with_delimiter(self.delimiter()) + .with_batch_size( + self.batch_size + .expect("Batch size must be set before initializing builder"), + ) + .with_header(self.has_header()) + .with_quote(self.quote()) + .with_truncated_rows(self.truncate_rows()); if let Some(terminator) = self.terminator() { builder = builder.with_terminator(terminator); } - builder = builder.with_projection(self.projection.file_indices.clone()); + builder = builder.with_projection(projection); if let Some(escape) = self.escape() { builder = builder.with_escape(escape) } @@ -211,6 +246,194 @@ impl CsvSource { Ok(builder) } + + fn parse_header(&self, line: &[u8]) -> Result> { + let mut format = csv::reader::Format::default() + .with_header(true) + .with_delimiter(self.delimiter()) + .with_quote(self.quote()); + if let Some(terminator) = self.terminator() { + format = format.with_terminator(terminator); + } + if let Some(escape) = self.escape() { + format = format.with_escape(escape); + } + if let Some(comment) = self.comment() { + format = format.with_comment(comment); + } + let (schema, _) = format.infer_schema(line, Some(0))?; + Ok(schema.fields().iter().map(|f| f.name().clone()).collect()) + } + + /// Builds the mapping for a file whose header row is `line`, or `None` if + /// the file can be read positionally. + fn header_mapping( + &self, + line: &[u8], + expr_adapter_factory: Arc, + ) -> Result>> { + if line.is_empty() { + return Ok(None); + } + let header = self.parse_header(line)?; + HeaderMapping::try_new( + &header, + self.table_schema.file_schema(), + &self.projection.file_indices, + expr_adapter_factory, + ) + .map(|mapping| mapping.map(Arc::new)) + } +} + +/// Maps the columns of a single CSV file onto the table's file schema by +/// header name. +/// +/// Only the header columns that are declared in the table are decoded; the +/// decoded batch is then adapted to the projected file schema with the scan's +/// [`PhysicalExprAdapterFactory`], which is the same mechanism Parquet uses +/// to fill missing columns and reorder by name. +#[derive(Debug)] +struct HeaderMapping { + /// One field per header column, typed from the table schema where the + /// name is declared there + physical_schema: SchemaRef, + /// Header positions to decode, in the order of the table projection + physical_projection: Vec, + adapter: BatchAdapter, +} + +impl HeaderMapping { + fn try_new( + header: &[String], + file_schema: &Schema, + file_indices: &[usize], + expr_adapter_factory: Arc, + ) -> Result> { + let matches_positionally = header.len() == file_schema.fields().len() + && header + .iter() + .zip(file_schema.fields()) + .all(|(name, field)| name == field.name()); + if matches_positionally { + return Ok(None); + } + + let physical_fields: Vec = header + .iter() + .map(|name| match file_schema.field_with_name(name) { + Ok(field) => field.clone(), + Err(_) => Field::new(name, DataType::Utf8, true), + }) + .collect(); + + let physical_projection: Vec = file_indices + .iter() + .filter_map(|&index| { + let name = file_schema.field(index).name(); + header.iter().position(|column| column == name) + }) + .collect(); + + let physical_schema = Arc::new(Schema::new(physical_fields)); + let decoded_schema = Arc::new(physical_schema.project(&physical_projection)?); + let target_schema = Arc::new(file_schema.project(file_indices)?); + let adapter = BatchAdapterFactory::new(target_schema) + .with_adapter_factory(expr_adapter_factory) + .make_adapter(&decoded_schema)?; + + Ok(Some(Self { + physical_schema, + physical_projection, + adapter, + })) + } + + fn remap(&self, batch: &RecordBatch) -> Result { + self.adapter.adapt_batch(batch) + } +} + +/// Returns the byte range of the header row in `buf`: the first line that is +/// not a comment line, matching how the CSV decoder skips comments. Returns +/// `None` when that line is not yet complete and more input may follow. +fn find_header_line( + buf: &[u8], + terminator: u8, + comment: Option, + eof: bool, +) -> Option> { + let is_comment = |line: &[u8]| comment.is_some_and(|c| line.first() == Some(&c)); + let mut start = 0; + loop { + let rest = &buf[start..]; + match rest.iter().position(|b| *b == terminator) { + Some(end) if is_comment(&rest[..end]) => start += end + 1, + Some(end) => return Some(start..start + end), + None if !eof => return None, + None if is_comment(rest) => return Some(buf.len()..buf.len()), + None => return Some(start..buf.len()), + } + } +} + +/// Reads the header row of `input`, returning it together with a stream that +/// still yields every byte of `input`. +async fn peek_header_line( + mut input: BoxStream<'static, Result>, + terminator: u8, + comment: Option, +) -> Result<(Vec, BoxStream<'static, Result>)> { + let mut prefix = Vec::new(); + let mut buffered = Vec::new(); + loop { + let chunk = input.next().await.transpose()?; + let eof = chunk.is_none(); + if let Some(chunk) = chunk { + prefix.extend_from_slice(&chunk); + buffered.push(Ok(chunk)); + } + if let Some(header) = find_header_line(&prefix, terminator, comment, eof) { + let header = prefix[header].to_vec(); + let buffered = futures::stream::iter(buffered); + let stream = if eof { + buffered.boxed() + } else { + buffered.chain(input).boxed() + }; + return Ok((header, stream)); + } + } +} + +/// Reads the header row of an uncompressed file, growing the requested range +/// until it is complete. Also returns the byte offset where the header starts. +async fn fetch_header_line( + store: &dyn ObjectStore, + location: &Path, + file_size: u64, + terminator: u8, + comment: Option, +) -> Result<(Vec, u64)> { + let mut len = file_size.min(64 * 1024); + loop { + let bytes = store.get_range(location, 0..len).await?; + let eof = len >= file_size; + if let Some(header) = find_header_line(&bytes, terminator, comment, eof) { + return Ok((bytes[header.clone()].to_vec(), header.start as u64)); + } + len = (len * 2).min(file_size); + } +} + +fn remap_stream( + stream: BoxStream<'static, Result>, + mapping: Option>, +) -> BoxStream<'static, Result> { + match mapping { + Some(mapping) => stream.map(move |batch| mapping.remap(&batch?)).boxed(), + None => stream, + } } /// A [`FileOpener`] that opens a CSV file and yields a [`FileOpenFuture`] @@ -219,6 +442,7 @@ pub struct CsvOpener { file_compression_type: FileCompressionType, object_store: Arc, partition_index: usize, + expr_adapter_factory: Arc, } impl CsvOpener { @@ -233,8 +457,19 @@ impl CsvOpener { file_compression_type, object_store, partition_index: 0, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), } } + + /// Set the [`PhysicalExprAdapterFactory`] used to adapt files whose + /// columns are matched to the table schema by header name + pub fn with_expr_adapter_factory( + mut self, + expr_adapter_factory: Arc, + ) -> Self { + self.expr_adapter_factory = expr_adapter_factory; + self + } } impl From for Arc { @@ -250,11 +485,21 @@ impl FileSource for CsvSource { base_config: &FileScanConfig, partition_index: usize, ) -> Result> { + if self.options.match_columns_by_header == Some(true) && !self.has_header() { + log::warn!( + "CSV option match_columns_by_header is ignored because has_header is false; columns are matched by position" + ); + } + let expr_adapter_factory = base_config + .expr_adapter_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); let mut opener = Arc::new(CsvOpener { config: Arc::new(self.clone()), file_compression_type: base_config.file_compression_type, object_store, partition_index, + expr_adapter_factory, }) as Arc; opener = ProjectionOpener::try_new( self.projection.clone(), @@ -307,7 +552,11 @@ impl FileSource for CsvSource { fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, ", has_header={}", self.has_header()) + write!(f, ", has_header={}", self.has_header())?; + if self.match_columns_by_header() { + write!(f, ", match_columns_by_header=true")?; + } + Ok(()) } DisplayFormatType::TreeRender => Ok(()), } @@ -360,6 +609,7 @@ impl FileSource for CsvSource { newlines_in_values: self.options.newlines_in_values.unwrap_or(false), truncate_rows: self.truncate_rows(), terminator: self.terminator().map(|terminator| vec![terminator]), + match_columns_by_header: self.match_columns_by_header(), }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::CsvScan(node)), @@ -395,6 +645,8 @@ impl FileOpener for CsvOpener { // `self.config.has_header` controls whether to skip reading the 1st line header // If the .csv file is read in parallel and this `CsvOpener` is only reading some middle // partition, then don't skip first line + let match_by_header = self.config.match_columns_by_header(); + let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); let mut csv_has_header = self.config.has_header(); if let Some(FileRange { start, .. }) = partitioned_file.range && start != 0 @@ -417,6 +669,8 @@ impl FileOpener for CsvOpener { let store = Arc::clone(&self.object_store); let terminator = self.config.terminator(); + let line_end = terminator.unwrap_or(b'\n'); + let comment = self.config.comment(); let baseline_metrics = BaselineMetrics::new(&self.config.metrics, self.partition_index); @@ -440,18 +694,40 @@ impl FileOpener for CsvOpener { ) })?; + // A partition may start mid-file, so the header is fetched + // separately. Every byte-range partition of a file repeats this + // small request; cache per file if that becomes costly. + let mapping = if match_by_header { + let (header, header_start) = fetch_header_line( + store.as_ref(), + &location, + file_size, + line_end, + comment, + ) + .await?; + // A partition starting within leading comment lines owns + // the header row, so it has to skip it. + if raw_start <= header_start { + config.options.has_header = Some(true); + } + config.header_mapping(&header, Arc::clone(&expr_adapter_factory))? + } else { + None + }; + let aligned_stream = AlignedBoundaryStream::new( Arc::clone(&store), location.clone(), raw_start, raw_end, file_size, - terminator.unwrap_or(b'\n'), + line_end, ) .await? .map_err(DataFusionError::from); - let decoder = config.builder()?.build_decoder(); + let decoder = config.builder(mapping.as_deref())?.build_decoder(); let input = file_compression_type .convert_stream(aligned_stream.boxed())? .fuse(); @@ -459,7 +735,7 @@ impl FileOpener for CsvOpener { input, DecoderDeserializer::new(CsvDecoder::new(decoder)), ); - return Ok(stream.map_err(Into::into).boxed()); + return Ok(remap_stream(stream.map_err(Into::into).boxed(), mapping)); } // No range specified — read the entire file @@ -470,7 +746,25 @@ impl FileOpener for CsvOpener { #[cfg(not(target_arch = "wasm32"))] GetResultPayload::File(file, _) => { let decoder = file_compression_type.convert_read(file)?; - let mut reader = config.open(decoder)?; + let (decoder, mapping): (Box, _) = if match_by_header + { + let mut decoder = BufReader::new(decoder); + let mut prefix = Vec::new(); + let header = loop { + let eof = decoder.read_until(line_end, &mut prefix)? == 0; + if let Some(header) = + find_header_line(&prefix, line_end, comment, eof) + { + break prefix[header].to_vec(); + } + }; + let mapping = config + .header_mapping(&header, Arc::clone(&expr_adapter_factory))?; + (Box::new(Cursor::new(prefix).chain(decoder)), mapping) + } else { + (decoder, None) + }; + let mut reader = config.open(decoder, mapping.as_deref())?; // Use std::iter::from_fn to wrap execution of iterator's next() method. let iterator = std::iter::from_fn(move || { @@ -480,20 +774,29 @@ impl FileOpener for CsvOpener { result }); - Ok(futures::stream::iter(iterator) + let stream = futures::stream::iter(iterator) .map(|r| r.map_err(Into::into)) - .boxed()) + .boxed(); + Ok(remap_stream(stream, mapping)) } GetResultPayload::Stream(s) => { - let decoder = config.builder()?.build_decoder(); let s = s.map_err(DataFusionError::from); - let input = file_compression_type.convert_stream(s.boxed())?.fuse(); - + let mut input = file_compression_type.convert_stream(s.boxed())?; + let mut mapping = None; + if match_by_header { + let (header, rest) = + peek_header_line(input, line_end, comment).await?; + mapping = config + .header_mapping(&header, Arc::clone(&expr_adapter_factory))?; + input = rest; + } + + let decoder = config.builder(mapping.as_deref())?.build_decoder(); let stream = deserialize_stream( - input, + input.fuse(), DecoderDeserializer::new(CsvDecoder::new(decoder)), ); - Ok(stream.map_err(Into::into).boxed()) + Ok(remap_stream(stream.map_err(Into::into).boxed(), mapping)) } } })) @@ -519,7 +822,7 @@ pub async fn plan_to_csv( let storeref = Arc::clone(&store); let plan: Arc = Arc::clone(&plan); let filename = format!("{}/part-{i}.csv", parsed.prefix()); - let file = object_store::path::Path::parse(filename)?; + let file = Path::parse(filename)?; let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; join_set.spawn(async move { @@ -643,6 +946,7 @@ impl CsvSource { quote: proto_str_to_byte(&scan.quote, "quote")?, newlines_in_values: Some(scan.newlines_in_values), truncated_rows: Some(scan.truncate_rows), + match_columns_by_header: Some(scan.match_columns_by_header), ..Default::default() }; let source = Arc::new( @@ -661,7 +965,7 @@ impl CsvSource { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Array, Int64Array, StringArray}; + use arrow::array::{Array, Int32Array, Int64Array, RecordBatchOptions, StringArray}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; fn csv_source(schema: SchemaRef, null_regex: Option<&str>) -> CsvSource { @@ -685,7 +989,7 @@ mod tests { let data = "id,name\n1,alice\n2,N/A\n3,carol\n"; let batch = source - .open(data.as_bytes()) + .open(data.as_bytes(), None) .unwrap() .next() .unwrap() @@ -713,7 +1017,7 @@ mod tests { let data = "id,value\n1,10\n2,N/A\n3,30\n"; let batch = source - .open(data.as_bytes()) + .open(data.as_bytes(), None) .unwrap() .next() .unwrap() @@ -739,7 +1043,7 @@ mod tests { let data = "id,name\n1,alice\n2,N/A\n3,carol\n"; let batch = source - .open(data.as_bytes()) + .open(data.as_bytes(), None) .unwrap() .next() .unwrap() @@ -758,10 +1062,195 @@ mod tests { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, true)])); let source = csv_source(schema, Some("(")); - let err = source.open(b"id\n1\n".as_slice()).unwrap_err(); + let err = source.open(b"id\n1\n".as_slice(), None).unwrap_err(); assert!( err.to_string().contains("Unable to parse CSV null regex"), "unexpected error: {err}" ); } + + fn table_schema() -> Schema { + Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Int32, true), + ]) + } + + fn header(names: &[&str]) -> Vec { + names.iter().map(|name| name.to_string()).collect() + } + + fn adapter_factory() -> Arc { + Arc::new(DefaultPhysicalExprAdapterFactory) + } + + #[test] + fn header_mapping_preserves_null_regex_and_terminator() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("name", DataType::Utf8, true), + ])); + let source = csv_source(schema, Some("^N/A$")) + .with_match_columns_by_header(true) + .with_terminator(Some(b';')); + let mapping = source + .header_mapping(b"name,id;", adapter_factory())? + .unwrap(); + let mut reader = + source.open(b"name,id;N/A,1;alice,2;".as_slice(), Some(&mapping))?; + let batch = mapping.remap(&reader.next().unwrap()?)?; + assert_eq!(batch.num_rows(), 2); + assert_eq!( + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(), + &Int64Array::from(vec![1, 2]) + ); + let names = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(names.is_null(0)); + assert_eq!(names.value(1), "alice"); + Ok(()) + } + + #[test] + fn header_mapping_skipped_when_header_matches() -> Result<()> { + let mapping = HeaderMapping::try_new( + &header(&["a", "b", "c"]), + &table_schema(), + &[0, 2], + adapter_factory(), + )?; + assert!(mapping.is_none()); + Ok(()) + } + + #[test] + fn header_mapping_by_name() -> Result<()> { + let mapping = HeaderMapping::try_new( + &header(&["c", "extra", "a"]), + &table_schema(), + &[0, 1, 2], + adapter_factory(), + )? + .unwrap(); + + let physical_types: Vec<_> = mapping + .physical_schema + .fields() + .iter() + .map(|f| f.data_type().clone()) + .collect(); + assert_eq!( + physical_types, + vec![DataType::Int32, DataType::Utf8, DataType::Int32] + ); + assert_eq!(mapping.physical_projection, vec![2, 0]); + + let decoded_schema = mapping.physical_schema.project(&[2, 0])?; + let decoded = RecordBatch::try_new( + Arc::new(decoded_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![10, 20])), + ], + )?; + let remapped = mapping.remap(&decoded)?; + assert_eq!(remapped.schema().as_ref(), &table_schema()); + assert_eq!(remapped.column(1).null_count(), 2); + assert_eq!( + remapped.column(0).as_ref(), + &Int32Array::from(vec![1, 2]) as &dyn Array + ); + assert_eq!( + remapped.column(2).as_ref(), + &Int32Array::from(vec![10, 20]) as &dyn Array + ); + Ok(()) + } + + #[test] + fn header_mapping_only_missing_columns() -> Result<()> { + let mapping = HeaderMapping::try_new( + &header(&["a", "c"]), + &table_schema(), + &[1], + adapter_factory(), + )? + .unwrap(); + assert!(mapping.physical_projection.is_empty()); + + let decoded = RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(3)), + )?; + let remapped = mapping.remap(&decoded)?; + assert_eq!(remapped.num_rows(), 3); + assert_eq!(remapped.column(0).null_count(), 3); + Ok(()) + } + + #[test] + fn header_mapping_missing_non_nullable_column() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, false), + ]); + let err = + HeaderMapping::try_new(&header(&["a"]), &schema, &[0, 1], adapter_factory()) + .unwrap_err(); + assert!( + err.to_string() + .contains("Non-nullable column 'b' is missing from the physical schema"), + "{err}" + ); + } + + #[test] + fn parse_header_handles_quotes() -> Result<()> { + let source = CsvSource::new(Arc::new(table_schema())); + assert_eq!( + source.parse_header(b"\"a,1\",b\r")?, + vec!["a,1".to_string(), "b".to_string()] + ); + Ok(()) + } + + #[test] + fn find_header_line_skips_comments() { + let comment = Some(b'#'); + assert_eq!( + find_header_line(b"# c1\n#c2\na,b\n1,2\n", b'\n', comment, false), + Some(9..12) + ); + assert_eq!( + find_header_line(b"#c\na,b\n", b'\n', None, false), + Some(0..2) + ); + assert_eq!(find_header_line(b"# c1\na,", b'\n', comment, false), None); + assert_eq!( + find_header_line(b"# c1\na,b", b'\n', comment, true), + Some(5..8) + ); + assert_eq!( + find_header_line(b"# c1\n# c2", b'\n', comment, true), + Some(9..9) + ); + assert_eq!(find_header_line(b"", b'\n', comment, true), Some(0..0)); + } + + #[test] + fn match_columns_by_header_requires_header() { + let source = CsvSource::new(Arc::new(table_schema())) + .with_csv_options(CsvOptions::default().with_has_header(false)) + .with_match_columns_by_header(true); + assert!(!source.match_columns_by_header()); + } } diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index ba6e6937daf38..4dedf36b03432 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -507,6 +507,8 @@ message CsvOptions { bytes ignore_leading_whitespace = 21; // Whether to ignore trailing whitespace in string values bytes ignore_trailing_whitespace = 22; + // Whether to map file columns to the table schema by header name + bytes match_columns_by_header = 23; } // Options controlling CSV format diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index fd94fc38d5901..2c7598273535f 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1163,6 +1163,10 @@ impl TryFrom<&protobuf::CsvOptions> for CsvOptions { .ignore_trailing_whitespace .first() .map(|h| *h != 0), + match_columns_by_header: proto_opts + .match_columns_by_header + .first() + .map(|h| *h != 0), }) } } diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 09386a01eb556..0b06993145d0e 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -1755,6 +1755,9 @@ impl serde::Serialize for CsvOptions { if !self.ignore_trailing_whitespace.is_empty() { len += 1; } + if !self.match_columns_by_header.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion_common.CsvOptions", len)?; if !self.has_header.is_empty() { #[allow(clippy::needless_borrow)] @@ -1850,6 +1853,11 @@ impl serde::Serialize for CsvOptions { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("ignoreTrailingWhitespace", pbjson::private::base64::encode(&self.ignore_trailing_whitespace).as_str())?; } + if !self.match_columns_by_header.is_empty() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("matchColumnsByHeader", pbjson::private::base64::encode(&self.match_columns_by_header).as_str())?; + } struct_ser.end() } } @@ -1898,6 +1906,8 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { "ignoreLeadingWhitespace", "ignore_trailing_whitespace", "ignoreTrailingWhitespace", + "match_columns_by_header", + "matchColumnsByHeader", ]; #[allow(clippy::enum_variant_names)] @@ -1924,6 +1934,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { QuoteStyle, IgnoreLeadingWhitespace, IgnoreTrailingWhitespace, + MatchColumnsByHeader, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -1967,6 +1978,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { "quoteStyle" | "quote_style" => Ok(GeneratedField::QuoteStyle), "ignoreLeadingWhitespace" | "ignore_leading_whitespace" => Ok(GeneratedField::IgnoreLeadingWhitespace), "ignoreTrailingWhitespace" | "ignore_trailing_whitespace" => Ok(GeneratedField::IgnoreTrailingWhitespace), + "matchColumnsByHeader" | "match_columns_by_header" => Ok(GeneratedField::MatchColumnsByHeader), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -2008,6 +2020,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { let mut quote_style__ = None; let mut ignore_leading_whitespace__ = None; let mut ignore_trailing_whitespace__ = None; + let mut match_columns_by_header__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::HasHeader => { @@ -2168,6 +2181,14 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { Some(map_.next_value::<::pbjson::private::BytesDeserialize<_>>()?.0) ; } + GeneratedField::MatchColumnsByHeader => { + if match_columns_by_header__.is_some() { + return Err(serde::de::Error::duplicate_field("matchColumnsByHeader")); + } + match_columns_by_header__ = + Some(map_.next_value::<::pbjson::private::BytesDeserialize<_>>()?.0) + ; + } } } Ok(CsvOptions { @@ -2193,6 +2214,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { quote_style: quote_style__.unwrap_or_default(), ignore_leading_whitespace: ignore_leading_whitespace__.unwrap_or_default(), ignore_trailing_whitespace: ignore_trailing_whitespace__.unwrap_or_default(), + match_columns_by_header: match_columns_by_header__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index ee586fa75ee37..8cfb07c02afbb 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -711,6 +711,9 @@ pub struct CsvOptions { /// Whether to ignore trailing whitespace in string values #[prost(bytes = "vec", tag = "22")] pub ignore_trailing_whitespace: ::prost::alloc::vec::Vec, + /// Whether to map file columns to the table schema by header name + #[prost(bytes = "vec", tag = "23")] + pub match_columns_by_header: ::prost::alloc::vec::Vec, } /// Options controlling CSV format #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 80a9dc662db5e..2404bba41fff9 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -1076,6 +1076,9 @@ impl TryFrom<&CsvOptions> for protobuf::CsvOptions { ignore_trailing_whitespace: opts .ignore_trailing_whitespace .map_or_else(Vec::new, |h| vec![h as u8]), + match_columns_by_header: opts + .match_columns_by_header + .map_or_else(Vec::new, |h| vec![h as u8]), }) } } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index df39cccca2371..01749542e864d 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1414,6 +1414,7 @@ message CsvScanExecNode { bool truncate_rows = 8; // Custom one-byte line terminator. Absent means the default newline terminator. optional bytes terminator = 9; + bool match_columns_by_header = 10; } message JsonScanExecNode { diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index 5ec8939f8b530..4f34657984b7c 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -283,6 +283,11 @@ impl TryFrom<&CsvOptionsProto> for CsvOptions { } else { Some(proto.ignore_trailing_whitespace[0] != 0) }, + match_columns_by_header: if proto.match_columns_by_header.is_empty() { + None + } else { + Some(proto.match_columns_by_header[0] != 0) + }, }) } } diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index ee586fa75ee37..8cfb07c02afbb 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -711,6 +711,9 @@ pub struct CsvOptions { /// Whether to ignore trailing whitespace in string values #[prost(bytes = "vec", tag = "22")] pub ignore_trailing_whitespace: ::prost::alloc::vec::Vec, + /// Whether to map file columns to the table schema by header name + #[prost(bytes = "vec", tag = "23")] + pub match_columns_by_header: ::prost::alloc::vec::Vec, } /// Options controlling CSV format #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 725cceeb75970..fb93e524f1025 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -5015,6 +5015,9 @@ impl serde::Serialize for CsvScanExecNode { if self.terminator.is_some() { len += 1; } + if self.match_columns_by_header { + len += 1; + } if self.optional_escape.is_some() { len += 1; } @@ -5045,6 +5048,9 @@ impl serde::Serialize for CsvScanExecNode { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("terminator", pbjson::private::base64::encode(&v).as_str())?; } + if self.match_columns_by_header { + struct_ser.serialize_field("matchColumnsByHeader", &self.match_columns_by_header)?; + } if let Some(v) = self.optional_escape.as_ref() { match v { csv_scan_exec_node::OptionalEscape::Escape(v) => { @@ -5080,6 +5086,8 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { "truncate_rows", "truncateRows", "terminator", + "match_columns_by_header", + "matchColumnsByHeader", "escape", "comment", ]; @@ -5093,6 +5101,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { NewlinesInValues, TruncateRows, Terminator, + MatchColumnsByHeader, Escape, Comment, } @@ -5123,6 +5132,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { "newlinesInValues" | "newlines_in_values" => Ok(GeneratedField::NewlinesInValues), "truncateRows" | "truncate_rows" => Ok(GeneratedField::TruncateRows), "terminator" => Ok(GeneratedField::Terminator), + "matchColumnsByHeader" | "match_columns_by_header" => Ok(GeneratedField::MatchColumnsByHeader), "escape" => Ok(GeneratedField::Escape), "comment" => Ok(GeneratedField::Comment), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), @@ -5151,6 +5161,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { let mut newlines_in_values__ = None; let mut truncate_rows__ = None; let mut terminator__ = None; + let mut match_columns_by_header__ = None; let mut optional_escape__ = None; let mut optional_comment__ = None; while let Some(k) = map_.next_key()? { @@ -5199,6 +5210,12 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| x.0) ; } + GeneratedField::MatchColumnsByHeader => { + if match_columns_by_header__.is_some() { + return Err(serde::de::Error::duplicate_field("matchColumnsByHeader")); + } + match_columns_by_header__ = Some(map_.next_value()?); + } GeneratedField::Escape => { if optional_escape__.is_some() { return Err(serde::de::Error::duplicate_field("escape")); @@ -5221,6 +5238,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { newlines_in_values: newlines_in_values__.unwrap_or_default(), truncate_rows: truncate_rows__.unwrap_or_default(), terminator: terminator__, + match_columns_by_header: match_columns_by_header__.unwrap_or_default(), optional_escape: optional_escape__, optional_comment: optional_comment__, }) diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index bb0da62d84c71..abb256d41caa1 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2145,6 +2145,8 @@ pub struct CsvScanExecNode { /// Custom one-byte line terminator. Absent means the default newline terminator. #[prost(bytes = "vec", optional, tag = "9")] pub terminator: ::core::option::Option<::prost::alloc::vec::Vec>, + #[prost(bool, tag = "10")] + pub match_columns_by_header: bool, #[prost(oneof = "csv_scan_exec_node::OptionalEscape", tags = "5")] pub optional_escape: ::core::option::Option, #[prost(oneof = "csv_scan_exec_node::OptionalComment", tags = "6")] diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index 32909afddf83b..dd68ee78b686a 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -13,6 +13,7 @@ use arrow::array::RecordBatch; use arrow::datatypes::Fields; use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::config::CsvOptions; use datafusion::datasource::empty::EmptyTable; use datafusion::datasource::file_format::json::JsonFormat; use datafusion::datasource::listing::{ @@ -518,8 +519,6 @@ fn roundtrip_avro_scan() -> Result<()> { #[test] fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { - use datafusion::common::config::CsvOptions; - let file_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); for newlines_in_values in [Some(true), Some(false), None] { @@ -609,6 +608,44 @@ fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { Ok(()) } +#[test] +fn roundtrip_csv_scan_preserves_match_columns_by_header() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, true)])); + let file_source = Arc::new( + CsvSource::new(TableSchema::from(&file_schema)).with_csv_options( + CsvOptions::default() + .with_has_header(true) + .with_match_columns_by_header(true) + .with_terminator(Some(b';')), + ), + ); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.csv".to_string(), + 1024, + )])]) + .build(); + + let ctx = SessionContext::new(); + let roundtripped = roundtrip_test_and_return( + DataSourceExec::from_data_source(scan_config), + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let csv_source = roundtripped + .downcast_ref::() + .and_then(|exec| exec.data_source().downcast_ref::()) + .and_then(|scan| scan.file_source().downcast_ref::()) + .ok_or_else(|| internal_datafusion_err!("Expected CsvSource"))?; + + assert!(csv_source.match_columns_by_header()); + assert_eq!(csv_source.terminator(), Some(b';')); + Ok(()) +} + #[tokio::test] async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { let mut file_group = diff --git a/datafusion/sqllogictest/test_files/csv_files.slt b/datafusion/sqllogictest/test_files/csv_files.slt index 549a74c4eb54a..d8ce7e91567dd 100644 --- a/datafusion/sqllogictest/test_files/csv_files.slt +++ b/datafusion/sqllogictest/test_files/csv_files.slt @@ -597,3 +597,43 @@ DROP TABLE stored_trim_both; statement ok DROP TABLE whitespace_source; + +# Match file columns to the table schema by header name +query I +COPY (SELECT 1 AS a, 'x' AS b, 10 AS c) +TO 'test_files/scratch/csv_files/header_layouts/full.csv' +STORED AS csv +OPTIONS ('format.has_header' 'true'); +---- +1 + +query I +COPY (SELECT 20 AS c, 'ignored' AS undeclared, 'y' AS b) +TO 'test_files/scratch/csv_files/header_layouts/reordered.csv' +STORED AS csv +OPTIONS ('format.has_header' 'true'); +---- +1 + +statement ok +CREATE EXTERNAL TABLE header_layouts ( + a INT, + b TEXT, + c INT +) STORED AS CSV +LOCATION 'test_files/scratch/csv_files/header_layouts/' +OPTIONS ('format.has_header' 'true', 'format.match_columns_by_header' 'true'); + +query ITI +SELECT a, b, c FROM header_layouts ORDER BY c; +---- +1 x 10 +NULL y 20 + +query I +SELECT c FROM header_layouts WHERE a IS NULL; +---- +20 + +statement ok +DROP TABLE header_layouts; diff --git a/docs/source/user-guide/sql/format_options.md b/docs/source/user-guide/sql/format_options.md index 3d188c771d312..8deda045ee4f8 100644 --- a/docs/source/user-guide/sql/format_options.md +++ b/docs/source/user-guide/sql/format_options.md @@ -118,25 +118,26 @@ OPTIONS('COMPRESSION' 'gzip'); The following options are available when reading or writing CSV files. Note: If any unsupported option is specified, an error will be raised and the query will fail. -| Option | Description | Default Value | -| -------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------ | -| COMPRESSION | Sets the compression that should be applied to the entire CSV file. Supported values are GZIP, BZIP2, XZ, ZSTD, and UNCOMPRESSED. | UNCOMPRESSED | -| HAS_HEADER | Sets if the CSV file should include column headers. If not set, uses session or system default. | None | -| DELIMITER | Sets the character which should be used as the column delimiter within the CSV file. | `,` (comma) | -| QUOTE | Sets the character which should be used for quoting values within the CSV file. | `"` (double quote) | -| TERMINATOR | Sets the character which should be used as the line terminator within the CSV file. | None | -| ESCAPE | Sets the character which should be used for escaping special characters within the CSV file. | None | -| DOUBLE_QUOTE | Sets if quotes within quoted fields should be escaped by doubling them (e.g., `"aaa""bbb"`). | None | -| NEWLINES_IN_VALUES | Sets if newlines in quoted values are supported. If not set, uses session or system default. | None | -| DATE_FORMAT | Sets the format that dates should be encoded in within the CSV file. | None | -| DATETIME_FORMAT | Sets the format that datetimes should be encoded in within the CSV file. | None | -| TIMESTAMP_FORMAT | Sets the format that timestamps should be encoded in within the CSV file. | None | -| TIMESTAMP_TZ_FORMAT | Sets the format that timestamps with timezone should be encoded in within the CSV file. | None | -| TIME_FORMAT | Sets the format that times should be encoded in within the CSV file. | None | -| NULL_VALUE | Sets the string which should be used to indicate null values within the CSV file. | None | -| NULL_REGEX | Sets the regex pattern to match null values when loading CSVs. | None | -| SCHEMA_INFER_MAX_REC | Sets the maximum number of records to scan to infer the schema. If set to 0, schema inference is disabled and all fields will be inferred as Utf8 (string) type. | None | -| COMMENT | Sets the character which should be used to indicate comment lines in the CSV file. | None | +| Option | Description | Default Value | +| ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------ | +| COMPRESSION | Sets the compression that should be applied to the entire CSV file. Supported values are GZIP, BZIP2, XZ, ZSTD, and UNCOMPRESSED. | UNCOMPRESSED | +| HAS_HEADER | Sets if the CSV file should include column headers. If not set, uses session or system default. | None | +| DELIMITER | Sets the character which should be used as the column delimiter within the CSV file. | `,` (comma) | +| QUOTE | Sets the character which should be used for quoting values within the CSV file. | `"` (double quote) | +| TERMINATOR | Sets the character which should be used as the line terminator within the CSV file. | None | +| ESCAPE | Sets the character which should be used for escaping special characters within the CSV file. | None | +| DOUBLE_QUOTE | Sets if quotes within quoted fields should be escaped by doubling them (e.g., `"aaa""bbb"`). | None | +| NEWLINES_IN_VALUES | Sets if newlines in quoted values are supported. If not set, uses session or system default. | None | +| DATE_FORMAT | Sets the format that dates should be encoded in within the CSV file. | None | +| DATETIME_FORMAT | Sets the format that datetimes should be encoded in within the CSV file. | None | +| TIMESTAMP_FORMAT | Sets the format that timestamps should be encoded in within the CSV file. | None | +| TIMESTAMP_TZ_FORMAT | Sets the format that timestamps with timezone should be encoded in within the CSV file. | None | +| TIME_FORMAT | Sets the format that times should be encoded in within the CSV file. | None | +| NULL_VALUE | Sets the string which should be used to indicate null values within the CSV file. | None | +| NULL_REGEX | Sets the regex pattern to match null values when loading CSVs. | None | +| SCHEMA_INFER_MAX_REC | Sets the maximum number of records to scan to infer the schema. If set to 0, schema inference is disabled and all fields will be inferred as Utf8 (string) type. | None | +| COMMENT | Sets the character which should be used to indicate comment lines in the CSV file. | None | +| MATCH_COLUMNS_BY_HEADER | Only applies when `HAS_HEADER` is true. Matches columns to the table schema by header name instead of by position. Missing columns are read as null and undeclared columns are ignored. Names are matched exactly (case-sensitive); duplicate header names resolve to the first occurrence. | None | **Example:** From db5f3c2b23d17923d2dd1f48a8b23a760f6282a1 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 10 Nov 2026 09:54:32 +0800 Subject: [PATCH 10/10] test(optimizer): verify immediate Interleave fallback across passes --- .../enforce_distribution.rs | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index e728a5512e330..a7653c079a742 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -4227,7 +4227,7 @@ fn existing_interleave_is_kept_when_children_stay_interleavable() -> Result<()> } #[test] -fn interleave_broken_by_later_rewrite_is_repaired_on_next_pass() -> Result<()> { +fn interleave_fallback_remains_valid_across_distribution_passes() -> Result<()> { // The chain reported in https://github.com/apache/datafusion/issues/21826: // a distribution pass builds an interleave, a later rule changes one // child's partitioning while rebuilding the tree, and another distribution @@ -4250,14 +4250,14 @@ fn interleave_broken_by_later_rewrite_is_repaired_on_next_pass() -> Result<()> { vec![Arc::clone(children[0]), coalesced], ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), )?; + assert!(rewritten.is::()); assert!(matches!( rewritten.output_partitioning(), - Partitioning::UnknownPartitioning(10) + Partitioning::UnknownPartitioning(11) )); - // Without a repair, the plan is rejected rather than executed. - assert!(check_plan_tree_sanity(&rewritten, &config.optimizer).is_err()); + check_plan_tree_sanity(&rewritten, &config.optimizer)?; - // The next distribution pass restores a valid interleave. + // The next distribution pass can recover the original interleave. let pass2 = EnsureRequirements::new().optimize(rewritten, &config)?; check_plan_tree_sanity(&pass2, &config.optimizer)?; assert_plan!(pass2, @@ -4324,15 +4324,15 @@ fn issue_21826_join_selection_after_distribution_pass() -> Result<()> { // Rebuilding the interleave over the swapped join used to fail here with // "Can not create InterleaveExec: new children can not be interleaved". let reordered = JoinSelection::new().optimize(pass1, &config)?; + assert!(reordered.is::()); assert!(matches!( reordered.output_partitioning(), - Partitioning::UnknownPartitioning(10) + Partitioning::UnknownPartitioning(20) )); - // Without a repair, the plan is rejected rather than executed. - assert!(check_plan_tree_sanity(&reordered, &config.optimizer).is_err()); + check_plan_tree_sanity(&reordered, &config.optimizer)?; - // The second distribution pass repairs it. The two joins no longer share - // a partitioning, so the result is a plain union. + // The second distribution pass keeps the union because the two joins no + // longer share a partitioning. let pass2 = EnsureRequirements::new().optimize(reordered, &config)?; check_plan_tree_sanity(&pass2, &config.optimizer)?; assert_plan!(pass2,