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/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() 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, 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/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))); 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 } } 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(()) + } +} 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..ad8a5de9f3e4e 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; @@ -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( @@ -543,6 +552,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 +824,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 +862,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], @@ -908,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 @@ -1289,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( @@ -1592,6 +1814,7 @@ fn convert_to_sort_cols(arrs: &[ArrayRef], sort_exprs: &LexOrdering) -> Vec 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![ @@ -2973,6 +3399,274 @@ mod tests { Ok(()) } + #[test] + fn test_unordered_first_group_acc() -> 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 = @@ -2994,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(()) + } } 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(); 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 = 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(()) } 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:**