From a323585e8411250c1c05afc7c81be2b6b592e0df Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 20 Aug 2026 15:53:11 -0400 Subject: [PATCH 01/16] [branch-55] fix: make UnnestExec respect datafusion.execution.batch_size (#24384) (#24529) ## Which issue does this PR close? - Part of https://github.com/apache/datafusion/issues/24462 - Backport of #24384 to `branch-55` (for 55.1.0, tracked in #24462). - Fixes #24383 ## Rationale for this change `UnnestExec` emitted exactly one output batch per input batch, however many rows the unnesting produced, never consulting `datafusion.execution.batch_size`. This let downstream operators receive arbitrarily large batches and made peak memory scale with input batch size times list length instead of `batch_size`. No API change, so it fits the backport criteria. ## What changes are included in this PR? Clean cherry-pick of #24384 (commit 6187d47ec5658190e729d4d3cad2cc10805aecdb). No adaptation required. ## Are these changes tested? Yes. Carries the original regression coverage, all tests pass. ## Are there any user-facing changes? `UnnestExec` output batches now respect `datafusion.execution.batch_size` as an upper bound. No API changes. Co-authored-by: Andy Grove --- datafusion/physical-plan/src/unnest.rs | 472 ++++++++++++++++-- datafusion/sqllogictest/test_files/unnest.slt | 40 ++ 2 files changed, 476 insertions(+), 36 deletions(-) diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 3fa274b27a7bd..6dfc2a0e537e7 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -23,10 +23,10 @@ use std::task::{Poll, ready}; use super::metrics::{ self, BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, - MetricsSet, RecordOutput, + MetricsSet, RecordOutput, SplitMetrics, }; use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{BatchSplitStream, EmptyRecordBatchStream}; use crate::{ ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, @@ -293,17 +293,29 @@ impl ExecutionPlan for UnnestExec { partition: usize, context: Arc, ) -> Result { + let batch_size = context.session_config().batch_size(); let input = self.input.execute(partition, context)?; let metrics = UnnestMetrics::new(partition, &self.metrics); - Ok(Box::pin(UnnestStream { + let stream = Box::pin(UnnestStream { input, schema: Arc::clone(&self.schema), list_type_columns: self.list_column_indices.clone(), struct_column_indices: self.struct_column_indices.iter().copied().collect(), options: self.options.clone(), metrics, - })) + batch_size, + pending_input: None, + }); + + // Chunking the input bounds each build to roughly `batch_size` rows, but two cases + // can still produce an oversized batch (see `predict_output_lens`), so the output + // goes through the shared splitter to make the bound unconditional. + Ok(Box::pin(BatchSplitStream::new( + stream, + batch_size, + SplitMetrics::new(&self.metrics, partition), + ))) } fn metrics(&self) -> Option { @@ -518,6 +530,67 @@ struct UnnestStream { options: UnnestOptions, /// Metrics metrics: UnnestMetrics, + /// Target number of rows per output batch, from `datafusion.execution.batch_size`. + batch_size: usize, + /// Rows of the current input batch that have not been unnested yet. Unnesting one + /// input batch can produce arbitrarily many output rows, so the input is consumed in + /// chunks small enough that each chunk's output stays near `batch_size`. + /// + /// Note the scope of the memory bound this buys: chunking removes the input batch size + /// from the peak, but not the length of an individual list. A single row whose list is + /// longer than `batch_size`, and recursive unnesting (where the expansion cannot be + /// predicted up front), both still materialize their full expansion in one build. + pending_input: Option, +} + +/// An input batch being unnested incrementally, a chunk of rows at a time. +struct PendingInput { + /// The full input batch. Rows before `row_offset` have already been unnested. + batch: RecordBatch, + /// Index of the next input row to unnest. + row_offset: usize, + /// How many output rows each input row expands into, indexed by input row. + /// + /// `None` when the expansion cannot be predicted from the input alone, in which case + /// the whole remaining input is unnested in one call and only the output is split. + /// See [`UnnestStream::predict_output_lens`]. + output_lens: Option>, +} + +impl PendingInput { + fn remaining_rows(&self) -> usize { + self.batch.num_rows() - self.row_offset + } + + /// How many input rows to unnest next so the resulting batch holds at most + /// `batch_size` rows. + fn next_chunk_rows(&self, batch_size: usize) -> usize { + let Some(output_lens) = &self.output_lens else { + return self.remaining_rows(); + }; + + let lens = &output_lens.values()[self.row_offset..]; + let batch_size = batch_size as i64; + let mut output_rows = 0i64; + for (rows, len) in lens.iter().enumerate() { + // The first row is always taken, even if it alone overshoots `batch_size`: an + // input row is never split across builds, so this is what guarantees progress. + // An oversized build is sliced down by `BatchSplitStream` on the way out. + if rows > 0 && output_rows + len > batch_size { + return rows; + } + output_rows += len; + } + lens.len() + } + + /// The per-row output lengths covering the next `rows` input rows, so the unnesting + /// does not have to recompute what `predict_output_lens` already derived. + fn chunk_lengths(&self, rows: usize) -> Option> { + self.output_lens + .as_ref() + .map(|lens| lens.slice(self.row_offset, rows)) + } } impl RecordBatchStream for UnnestStream { @@ -546,30 +619,60 @@ impl UnnestStream { cx: &mut std::task::Context<'_>, ) -> Poll>> { loop { - return Poll::Ready(match ready!(self.input.poll_next_unpin(cx)) { + // Unnest the next chunk of the input batch already in hand. + if let Some(pending) = self.pending_input.as_mut() { + // `PendingInput` is only built from a non-empty batch and `next_chunk_rows` + // always consumes at least one row, so it is dropped the moment it drains. + debug_assert!(pending.remaining_rows() > 0); + + let rows = pending.next_chunk_rows(self.batch_size); + let chunk = pending.batch.slice(pending.row_offset, rows); + let chunk_lengths = pending.chunk_lengths(rows); + pending.row_offset += rows; + let drained = pending.remaining_rows() == 0; + + let timer = self.metrics.baseline_metrics.elapsed_compute().timer(); + let result = build_batch( + &chunk, + &self.schema, + &self.list_type_columns, + &self.struct_column_indices, + &self.options, + chunk_lengths.as_ref(), + ); + timer.done(); + + if drained { + self.pending_input = None; + } + + // A chunk can legitimately produce no rows at all, for example when every + // list in it is empty under `NullHandling::Drop`; `build_batch` signals + // that with `None` rather than an empty batch. + if let Some(batch) = result? { + debug_assert!(batch.num_rows() > 0); + (&batch).record_output(&self.metrics.baseline_metrics); + return Poll::Ready(Some(Ok(batch))); + } + continue; + } + + // Otherwise pull the next input batch. + match ready!(self.input.poll_next_unpin(cx)) { Some(Ok(batch)) => { - let elapsed_compute = - self.metrics.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); self.metrics.input_batches.add(1); self.metrics.input_rows.add(batch.num_rows()); - let result = build_batch( - &batch, - &self.schema, - &self.list_type_columns, - &self.struct_column_indices, - &self.options, - )?; - timer.done(); - let Some(result_batch) = result else { - continue; - }; - (&result_batch).record_output(&self.metrics.baseline_metrics); - - // Empty record batches should not be emitted. - // They need to be treated as [`Option`]es and handled separately - debug_assert!(result_batch.num_rows() > 0); - Some(Ok(result_batch)) + if batch.num_rows() > 0 { + let timer = + self.metrics.baseline_metrics.elapsed_compute().timer(); + let output_lens = self.predict_output_lens(&batch); + timer.done(); + self.pending_input = Some(PendingInput { + batch, + row_offset: 0, + output_lens: output_lens?, + }); + } } // If the stream is depleted or returned an error, log the finish message: other => { @@ -590,11 +693,48 @@ impl UnnestStream { self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); } - other + return Poll::Ready(other); } - }); + } } } + + /// Compute how many output rows each input row of `batch` will expand into, so the + /// input can be chunked to keep each build bounded. + /// + /// Returns `None` when the count cannot be derived from the input alone, which is the + /// signal to unnest the whole batch in one call: + /// + /// * With no list columns, unnesting only widens structs and leaves the row count + /// alone, so the output is already bounded by the input batch size. + /// * With recursion (`depth > 1`), a row's expansion depends on the lengths of inner + /// lists that only exist after the outer levels have been unnested, so it cannot be + /// predicted up front. + fn predict_output_lens( + &self, + batch: &RecordBatch, + ) -> Result>> { + if self.list_type_columns.is_empty() + || self + .list_type_columns + .iter() + .any(|unnest| unnest.depth != 1) + { + return Ok(None); + } + + let list_arrays: Vec = self + .list_type_columns + .iter() + .map(|unnest| Arc::clone(batch.column(unnest.index_in_input_schema))) + .collect(); + + // This is exactly the per-row length that `list_unnest_at_level` derives when it + // actually unnests, so the chunk boundaries are exact rather than estimated, and + // each chunk's slice of it is handed back to `build_batch` instead of recomputed. + let longest_length = find_longest_length(&list_arrays, &self.options)?; + Ok(Some(longest_length.as_primitive::().clone())) + } } /// Given a set of struct column indices to flatten @@ -686,6 +826,7 @@ fn list_unnest_at_level( temp_unnested_arrs: &mut HashMap, level_to_unnest: usize, options: &UnnestOptions, + precomputed_lengths: Option<&PrimitiveArray>, ) -> Result>> { // Extract unnestable columns at this level let (arrs_to_unnest, list_unnest_specs): (Vec>, Vec<_>) = @@ -712,8 +853,17 @@ fn list_unnest_at_level( // Filter out so that list_arrays only contain column with the highest depth // at the same time, during iteration remove this depth so next time we don't have to unnest them again - let longest_length = find_longest_length(&arrs_to_unnest, options)?; - let unnested_length = longest_length.as_primitive::(); + // + // The caller may already have computed these lengths to decide how many input rows to + // feed us; reusing them avoids running the kernel chain twice over the same rows. + // Cloning is an `Arc` bump on the underlying buffer, not a copy. + let longest_length = match precomputed_lengths { + Some(lengths) => lengths.clone(), + None => find_longest_length(&arrs_to_unnest, options)? + .as_primitive::() + .clone(), + }; + let unnested_length = &longest_length; let total_length = if unnested_length.is_empty() { 0 } else { @@ -830,6 +980,7 @@ fn build_batch( list_type_columns: &[ListUnnest], struct_column_indices: &HashSet, options: &UnnestOptions, + precomputed_lengths: Option<&PrimitiveArray>, ) -> Result> { let transformed = match list_type_columns.len() { 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), @@ -851,12 +1002,21 @@ fn build_batch( true => batch.columns(), false => &flatten_arrs, }; + // Only sound for a single non-recursive level: with recursion the deeper + // levels' lengths depend on arrays that do not exist yet, which is also why + // the caller does not predict lengths in that case. + let level_lengths = if max_recursion == 1 { + precomputed_lengths + } else { + None + }; let Some(temp_result) = list_unnest_at_level( input, list_type_columns, &mut temp_unnested_result, depth, options, + level_lengths, )? else { return Ok(None); @@ -1269,7 +1429,7 @@ fn repeat_arrs_from_indices( mod tests { use super::*; use arrow::array::{ - GenericListArray, NullBufferBuilder, OffsetSizeTrait, StringArray, + GenericListArray, Int32Array, NullBufferBuilder, OffsetSizeTrait, StringArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{Field, Int32Type}; @@ -1459,6 +1619,7 @@ mod tests { null_handling: NullHandling::Preserve, recursions: vec![], }, + None, )? .unwrap(); @@ -1503,8 +1664,7 @@ mod tests { // NULL -> one row with c2 = 5 and unnested value NULL // [NULL, F] -> two rows with c2 = 6, 6 let list_array = Arc::new(make_generic_array::()) as ArrayRef; - let other = - Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let other = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; let in_schema = Arc::new(Schema::new(vec![ Field::new( "c1", @@ -1535,6 +1695,7 @@ mod tests { null_handling: NullHandling::PreserveAndExpandEmpty, recursions: vec![], }, + None, )? .unwrap(); @@ -1562,8 +1723,7 @@ mod tests { #[test] fn test_build_batch_preserve_and_expand_empty_largelist() -> Result<()> { let list_array = Arc::new(make_generic_array::()) as ArrayRef; - let other = - Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let other = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; let in_schema = Arc::new(Schema::new(vec![ Field::new( "c1", @@ -1597,6 +1757,7 @@ mod tests { null_handling: NullHandling::PreserveAndExpandEmpty, recursions: vec![], }, + None, )? .unwrap(); @@ -1647,8 +1808,7 @@ mod tests { b.append(false); b.finish() }; - let id = - Arc::new(arrow::array::Int32Array::from(vec![10, 20, 30, 40])) as ArrayRef; + let id = Arc::new(Int32Array::from(vec![10, 20, 30, 40])) as ArrayRef; let in_schema = Arc::new(Schema::new(vec![ Field::new( @@ -1696,6 +1856,7 @@ mod tests { null_handling: NullHandling::PreserveAndExpandEmpty, recursions: vec![], }, + None, )? .unwrap(); @@ -1817,6 +1978,7 @@ mod tests { null_handling: NullHandling::PreserveAndExpandEmpty, recursions: vec![], }, + None, )? .unwrap(); @@ -2005,4 +2167,242 @@ mod tests { assert_eq!(take_indices, expected); Ok(()) } + + /// Build a single-column `List` batch where row `i` holds `lens[i]` elements, + /// numbered consecutively from 0 across the whole batch. A `None` length is a NULL + /// list. + fn list_batch(lens: &[Option]) -> RecordBatch { + let mut next = 0i32; + let rows: Vec>>> = lens + .iter() + .map(|len| { + len.map(|len| { + (0..len) + .map(|_| { + next += 1; + Some(next - 1) + }) + .collect() + }) + }) + .collect(); + let list = ListArray::from_iter_primitive::(rows); + let schema = Arc::new(Schema::new(vec![Field::new( + "l", + list.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(list)]).unwrap() + } + + /// Run a depth-1 unnest of column "l" over `input`, with the given + /// `datafusion.execution.batch_size`, and return the output batches. + async fn unnest_with_batch_size( + input: Vec, + batch_size: usize, + options: UnnestOptions, + ) -> Result> { + unnest_at_depth(input, batch_size, options, 1).await + } + + /// Unnest column "l" of `input` to `depth`, with the given + /// `datafusion.execution.batch_size`, and return the output batches. + async fn unnest_at_depth( + input: Vec, + batch_size: usize, + options: UnnestOptions, + depth: usize, + ) -> Result> { + let input_schema = input[0].schema(); + let output_schema = + Arc::new(Schema::new(vec![Field::new("l", DataType::Int32, true)])); + let source = + crate::test::TestMemoryExec::try_new_exec(&[input], input_schema, None)?; + let unnest = UnnestExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth, + }], + vec![], + output_schema, + options, + )?; + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + datafusion_execution::config::SessionConfig::new() + .with_batch_size(batch_size), + ), + ); + crate::common::collect(unnest.execute(0, task_ctx)?).await + } + + /// The values an unnest produces, flattened across all output batches. + fn output_values(batches: &[RecordBatch]) -> Vec> { + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .iter() + .collect::>() + }) + .collect() + } + + /// Output batch sizes are fully determined by the input lengths and `batch_size`, so + /// assert the exact shapes rather than just the `<= batch_size` bound. Each case pins a + /// distinct path through `next_chunk_rows`. + #[tokio::test] + async fn test_unnest_stream_output_batch_shapes() -> Result<()> { + struct Case { + /// One inner slice per input batch, each of that batch's per-row list lengths. + lens_per_batch: &'static [&'static [Option]], + batch_size: usize, + expected_sizes: &'static [usize], + } + let cases: &[Case] = &[ + // Chunks pack several input rows. This is the case that distinguishes chunking + // the input from building everything and slicing: slicing a single 30-row build + // would give [8, 8, 8, 6]. + Case { + lens_per_batch: &[&[Some(3); 10]], + batch_size: 8, + expected_sizes: &[6, 6, 6, 6, 6], + }, + // Output smaller than batch_size comes back as one batch. + Case { + lens_per_batch: &[&[Some(3), Some(2)]], + batch_size: 1024, + expected_sizes: &[5], + }, + // One row expanding past batch_size cannot be chunked on the input side, so the + // oversized build is sliced on the way out instead. + Case { + lens_per_batch: &[&[Some(25)]], + batch_size: 10, + expected_sizes: &[10, 10, 5], + }, + // Chunk boundaries are per input batch, so each batch contributes a short tail. + Case { + lens_per_batch: &[&[Some(5), Some(5)], &[Some(1)], &[Some(7), Some(2)]], + batch_size: 4, + expected_sizes: &[4, 1, 4, 1, 1, 4, 3, 2], + }, + ]; + + for case in cases { + let input: Vec = case + .lens_per_batch + .iter() + .map(|lens| list_batch(lens)) + .collect(); + let batches = + unnest_with_batch_size(input, case.batch_size, UnnestOptions::default()) + .await?; + + let sizes: Vec = batches.iter().map(|b| b.num_rows()).collect(); + assert_eq!( + sizes, case.expected_sizes, + "lens={:?} batch_size={}", + case.lens_per_batch, case.batch_size + ); + + // `list_batch` numbers each batch's elements from 0, so the expected values are + // one run per input batch. Splitting must not perturb values or their order. + let expected_values: Vec> = case + .lens_per_batch + .iter() + .flat_map(|lens| { + (0..lens.iter().flatten().sum::() as i32).map(Some) + }) + .collect(); + assert_eq!( + output_values(&batches), + expected_values, + "lens={:?} batch_size={}", + case.lens_per_batch, + case.batch_size + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_unnest_stream_chunking_preserves_null_handling() -> Result<()> { + // NULL and empty lists each contribute one NULL output row under + // PreserveAndExpandEmpty, and the per-row output counts that drive chunking must + // agree with that or chunk boundaries would drift out of step with the unnesting. + let lens = &[Some(3), Some(0), None, Some(2), None, Some(0)]; + let options = + UnnestOptions::new().with_null_handling(NullHandling::PreserveAndExpandEmpty); + + let chunked = + unnest_with_batch_size(vec![list_batch(lens)], 2, options.clone()).await?; + let whole = unnest_with_batch_size(vec![list_batch(lens)], 1024, options).await?; + + assert!(chunked.iter().all(|b| b.num_rows() <= 2)); + // 3 + 1 + 1 + 2 + 1 + 1 + assert_eq!(chunked.iter().map(|b| b.num_rows()).sum::(), 9); + assert_eq!(output_values(&chunked), output_values(&whole)); + Ok(()) + } + + #[tokio::test] + async fn test_unnest_stream_drop_null_handling() -> Result<()> { + // Under Drop, NULL and empty lists produce nothing. Chunks made up entirely of + // such rows yield no batch at all, and must not stall the stream or leak an + // empty batch into the output. + let lens = &[None, Some(0), None, Some(4), Some(0), None]; + let options = UnnestOptions::new().with_null_handling(NullHandling::Drop); + + let batches = unnest_with_batch_size(vec![list_batch(lens)], 2, options).await?; + + assert!(batches.iter().all(|b| b.num_rows() > 0)); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 4); + Ok(()) + } + + #[tokio::test] + async fn test_unnest_stream_recursive_respects_batch_size() -> Result<()> { + // Recursive unnest cannot have its expansion predicted from the input, so it falls + // back to unnesting a whole input batch and slicing the output. The batch_size + // guarantee has to hold on that path too. + let inner = Field::new_list_field(DataType::Int32, true); + let outer = + Field::new_list_field(DataType::new_list(DataType::Int32, true), true); + let values = Int32Array::from((0..24).collect::>()); + // 12 inner lists of 2 elements each... + let inner_list = ListArray::new( + Arc::new(inner), + OffsetBuffer::new((0..=12).map(|i| i * 2).collect::>().into()), + Arc::new(values), + None, + ); + // ...grouped 3 to a row, so 4 input rows expand to 24 output rows at depth 2. + let outer_list = ListArray::new( + Arc::new(outer), + OffsetBuffer::new((0..=4).map(|i| i * 3).collect::>().into()), + Arc::new(inner_list), + None, + ); + let input_schema = Arc::new(Schema::new(vec![Field::new( + "l", + outer_list.data_type().clone(), + true, + )])); + let input = RecordBatch::try_new(input_schema, vec![Arc::new(outer_list)])?; + + let batches = + unnest_at_depth(vec![input], 7, UnnestOptions::default(), 2).await?; + + let sizes: Vec = batches.iter().map(|b| b.num_rows()).collect(); + assert_eq!(sizes, vec![7, 7, 7, 3]); + assert_eq!( + output_values(&batches), + (0..24).map(Some).collect::>() + ); + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index a3385b81d70d1..8c6cc3fa4922f 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -1646,3 +1646,43 @@ DROP TABLE varied_lists; statement ok DROP TABLE dense_lists; + +## Unnest respects `datafusion.execution.batch_size` (apache/datafusion#24383). Chunking +## the input must not perturb row order or drop rows; batch shapes are asserted in the +## Rust unit tests, which is where they are actually observable. + +statement ok +CREATE TABLE fanout AS + SELECT i AS id, range(0, i % 7) AS xs FROM (SELECT unnest(range(1, 40)) AS i); + +statement ok +set datafusion.execution.batch_size = 3; + +query II +SELECT id, unnest(xs) AS x FROM fanout WHERE id IN (5, 6, 7, 8) ORDER BY id, x; +---- +5 0 +5 1 +5 2 +5 3 +5 4 +6 0 +6 1 +6 2 +6 3 +6 4 +6 5 +8 0 + +# Recursive unnest cannot have its expansion predicted from the input, so it takes the +# output-slicing path instead of input chunking. +query I +SELECT count(*) FROM (SELECT unnest(unnest(make_array(range(0, 20), range(0, 20)))) AS x); +---- +40 + +statement ok +set datafusion.execution.batch_size = 8192; + +statement ok +DROP TABLE fanout; From 045b590c4822fdbd2ce47eac2f7a5a55b28cde45 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Fri, 21 Aug 2026 14:43:25 -0400 Subject: [PATCH 02/16] [branch-55] fix: apply struct field filters when the file schema needs adaptation (#24125) (#24530) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Part of https://github.com/apache/datafusion/issues/24462 - Backport of #24125 to `branch-55` (for 55.1.0, tracked in #24462). - Fixes #24109 ## Rationale for this change With `datafusion.execution.parquet.pushdown_filters = true`, a predicate on a struct field was reported as fully handled by the scan whenever the file needed schema adaptation, so `FilterExec` was removed from the plan and the predicate was silently dropped — returning every row instead of the filtered set. This is a correctness bug (wrong results), not specific to 55.0.0, so it fits the backport criteria. ## What changes are included in this PR? Cherry-pick of #24125 (commit 40c208e5d65b931ef3cc3e121e85e228045a1a17). Git's recursive merge auto-resolved surrounding context differences in `datafusion/physical-expr-adapter/src/schema_rewriter.rs` and `datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt`; no manual conflict resolution or adaptation of the fix itself was required. One follow-up commit adapts a test expectation: `branch-55` doesn't have #24130/#24315, which taught nested schema pruning to union the leaves needed by mixed whole-column + field-access reads (e.g. `select s, s['y'] from narrow`). Without that optimization, the mixed-access case falls back to reading every physical leaf, so `bytes_scanned` is `219` here instead of the `146` the original PR's test expects on `main`. This is a pre-existing difference in pruning capability, not a correctness regression from this fix. ## Are these changes tested? Yes. Carries the original regression coverage, all tests pass. ## Are there any user-facing changes? `WHERE s['field'] = ...` predicates on struct columns now filter correctly when Parquet filter pushdown requires schema adaptation. No API changes. --------- Co-authored-by: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Co-authored-by: Claude --- .../src/schema_rewriter.rs | 533 +++++++++++++++++- .../test_files/parquet_filter_pushdown.slt | 113 ++++ .../parquet_nested_schema_pruning.slt | 20 +- 3 files changed, 654 insertions(+), 12 deletions(-) diff --git a/datafusion/physical-expr-adapter/src/schema_rewriter.rs b/datafusion/physical-expr-adapter/src/schema_rewriter.rs index ef25af7d920fb..a5d1c494c8ae4 100644 --- a/datafusion/physical-expr-adapter/src/schema_rewriter.rs +++ b/datafusion/physical-expr-adapter/src/schema_rewriter.rs @@ -25,7 +25,7 @@ use std::hash::Hash; use std::sync::Arc; use arrow::array::RecordBatch; -use arrow::datatypes::{DataType, FieldRef, SchemaRef}; +use arrow::datatypes::{DataType, FieldRef, Fields, SchemaRef}; use datafusion_common::{ DataFusionError, Result, ScalarValue, exec_err, metadata::FieldMetadata, @@ -273,6 +273,42 @@ struct DefaultPhysicalExprAdapterRewriter { physical_file_schema: SchemaRef, } +/// Outcome of walking a `get_field` key path through nested struct fields. +enum FieldPathResolution<'a> { + /// The leaf field the path points at. + Found(&'a FieldRef), + /// Some key along the path does not exist, so the access reads as null. + Missing, + /// An intermediate field is not a struct, so the path cannot be resolved + /// statically. + NotAStruct, +} + +/// Follow a `get_field` key path (`['a', 'b']` for `s['a']['b']`) through +/// nested struct fields. +/// +/// The first key is taken separately from the rest so that the "at least one +/// key" invariant is carried by the signature: there is no empty path to +/// resolve. +fn resolve_field_path<'a>( + fields: &'a Fields, + field_name: &str, + rest: &[&str], +) -> FieldPathResolution<'a> { + let Some(field) = fields.iter().find(|f| f.name() == field_name) else { + return FieldPathResolution::Missing; + }; + let Some((next_field_name, rest)) = rest.split_first() else { + return FieldPathResolution::Found(field); + }; + match field.data_type() { + DataType::Struct(nested_fields) => { + resolve_field_path(nested_fields, next_field_name, rest) + } + _ => FieldPathResolution::NotAStruct, + } +} + impl DefaultPhysicalExprAdapterRewriter { fn rewrite_expr( &self, @@ -282,6 +318,10 @@ impl DefaultPhysicalExprAdapterRewriter { return Ok(Transformed::yes(transformed)); } + if let Some(transformed) = self.try_narrow_struct_cast(&expr)? { + return Ok(Transformed::yes(transformed)); + } + if let Some(column) = expr.downcast_ref::() { return self.rewrite_column(Arc::clone(&expr), column); } @@ -289,6 +329,135 @@ impl DefaultPhysicalExprAdapterRewriter { Ok(Transformed::no(expr)) } + /// Rewrite `get_field(cast(s AS Struct<..>), 'f')` into + /// `cast(get_field(s, 'f') AS )`. + /// + /// Expressions are rewritten bottom-up, so by the time we reach a + /// `get_field` node its struct argument has already been wrapped in a cast + /// by [`Self::rewrite_column`] whenever the logical and physical struct + /// types differ. + /// + /// Narrowing that cast is worthwhile for two reasons: + /// + /// 1. Reading one field should not cost a whole struct. The wide form + /// casts every field of the column — including ones the query never + /// reads — to produce a value that is immediately discarded except for + /// one field. + /// 2. It keeps the column visible. Consumers throughout the codebase + /// pattern match on `get_field(column, 'f')` to recognise a struct + /// field access; a cast between the `get_field` and its column defeats + /// that match, and each such consumer then falls back to whatever it + /// does for an unrecognised expression. + /// + /// The Parquet scan is one such consumer, and the reason this is a + /// correctness fix rather than only an optimisation: it decides at + /// planning time, against the table schema, that a struct-field predicate + /// can be evaluated as a row filter, and reports the predicate as fully + /// handled. See . + /// + /// Fixing it here rather than teaching that one consumer to see through + /// casts is deliberate: the adapter is where the obscuring cast is + /// introduced, so every consumer benefits, and no consumer has to loosen + /// its pattern to accept arbitrary casts between a `get_field` and its + /// column. + /// + /// `get_field` also has a flattened multi-key form: `s['a']['b']` is + /// simplified to `get_field(s, 'a', 'b')`, so the whole field path is + /// resolved here rather than just the first key. + /// + /// Only struct casts are narrowed. `get_field` on a Map column performs a + /// runtime key lookup rather than a schema-level field access, so the map + /// value must keep its cast. + fn try_narrow_struct_cast( + &self, + expr: &Arc, + ) -> Result>> { + let Some(get_field_expr) = + ScalarFunctionExpr::try_downcast_func::(expr.as_ref()) + else { + return Ok(None); + }; + let Some((source_expr, field_name_exprs)) = get_field_expr.args().split_first() + else { + return Ok(None); + }; + let Some(cast) = source_expr.downcast_ref::() else { + return Ok(None); + }; + + // Every key has to be a string literal, otherwise the leaf field + // cannot be resolved statically. + let mut field_path = Vec::with_capacity(field_name_exprs.len()); + for field_name_expr in field_name_exprs { + let Some(field_name) = field_name_expr + .downcast_ref::() + .and_then(|lit| lit.value().try_as_str().flatten()) + else { + return Ok(None); + }; + field_path.push(field_name); + } + // A `get_field` with no keys is not a field access we can narrow. + let Some((first_key, rest_keys)) = field_path.split_first() else { + return Ok(None); + }; + + let DataType::Struct(logical_struct_fields) = cast.target_field().data_type() + else { + return Ok(None); + }; + let FieldPathResolution::Found(logical_struct_field) = + resolve_field_path(logical_struct_fields, first_key, rest_keys) + else { + return Ok(None); + }; + + let inner = cast.expr(); + let DataType::Struct(physical_struct_fields) = + inner.data_type(&self.physical_file_schema)? + else { + return Ok(None); + }; + let physical_struct_field = + match resolve_field_path(&physical_struct_fields, first_key, rest_keys) { + FieldPathResolution::Found(field) => field, + FieldPathResolution::Missing => { + // The file does not have this field at all, so reading it + // yields null. Note that the cast would have produced the + // same value: struct casts fill missing target fields with + // nulls. + let null_value = + ScalarValue::Null.cast_to(logical_struct_field.data_type())?; + return Ok(Some(Arc::new(Literal::new_with_metadata( + null_value, + Some(FieldMetadata::from(logical_struct_field.as_ref())), + )))); + } + FieldPathResolution::NotAStruct => return Ok(None), + }; + + // Rebuild `get_field` over the uncast struct so its return field is + // recomputed from the physical field type. + let mut args = Vec::with_capacity(get_field_expr.args().len()); + args.push(Arc::clone(inner)); + args.extend(field_name_exprs.iter().map(Arc::clone)); + let extracted = Arc::new(ScalarFunctionExpr::try_new( + Arc::new(get_field_expr.fun().clone()), + args, + &self.physical_file_schema, + Arc::new(get_field_expr.config_options().clone()), + )?) as Arc; + + if physical_struct_field == logical_struct_field { + return Ok(Some(extracted)); + } + Ok(Some(Arc::new(CastExpr::new_with_target_field( + extracted, + Arc::clone(logical_struct_field), + Some(cast.cast_options().clone()), + )))) + } + /// Attempt to rewrite struct field access expressions to return null if the field does not exist in the physical schema. /// Note that this does *not* handle nested struct fields, only top-level struct field access. /// See for more details. @@ -1426,6 +1595,368 @@ mod tests { // datafusion/core/tests/parquet/schema_adapter.rs provide better coverage for this functionality. } + /// Build `get_field(column, 'field')` against `schema`. + fn get_field_expr( + schema: &Schema, + column: &str, + field: &str, + ) -> Arc { + let index = schema.index_of(column).unwrap(); + Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())), + vec![ + Arc::new(Column::new(column, index)), + Arc::new(Literal::new(ScalarValue::from(field))), + ], + schema, + Arc::new(datafusion_common::config::ConfigOptions::default()), + ) + .unwrap(), + ) + } + + fn struct_schemas( + physical_fields: Vec, + logical_fields: Vec, + ) -> (SchemaRef, SchemaRef) { + let physical = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(physical_fields.into()), + true, + )])); + let logical = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(logical_fields.into()), + true, + )])); + (logical, physical) + } + + /// `s['x']` where the file stores `x` as `Int32` and the table declares + /// `Int64` must cast the extracted field, not the whole struct, so that + /// the column stays visible under the `get_field`. + /// + /// See . + #[test] + fn test_narrow_struct_cast_to_field_access() { + let (logical_schema, physical_schema) = struct_schemas( + vec![Field::new("x", DataType::Int32, true)], + vec![Field::new("x", DataType::Int64, true)], + ); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let rewritten = adapter + .rewrite(get_field_expr(&logical_schema, "s", "x")) + .unwrap(); + + let cast = assert_cast_expr(&rewritten); + assert_eq!(cast.cast_type(), &DataType::Int64); + let get_field = cast + .expr() + .downcast_ref::() + .expect("Expected get_field under the cast"); + assert_eq!(get_field.return_type(), &DataType::Int32); + assert!( + get_field.args()[0].downcast_ref::().is_some(), + "the struct column must not be hidden behind a cast, got: {rewritten}" + ); + } + + /// A struct field that only differs in a nested leaf type still ends up + /// with a single cast on the extracted field. + #[test] + fn test_narrow_struct_cast_nested_field_access() { + let (logical_schema, physical_schema) = struct_schemas( + vec![Field::new( + "inner", + DataType::Struct(vec![Field::new("x", DataType::Utf8, true)].into()), + true, + )], + vec![Field::new( + "inner", + DataType::Struct(vec![Field::new("x", DataType::Utf8View, true)].into()), + true, + )], + ); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let outer = get_field_expr(&logical_schema, "s", "inner"); + let expr = Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())), + vec![outer, Arc::new(Literal::new(ScalarValue::from("x")))], + &logical_schema, + Arc::new(datafusion_common::config::ConfigOptions::default()), + ) + .unwrap(), + ) as Arc; + + let rewritten = adapter.rewrite(expr).unwrap(); + + let cast = assert_cast_expr(&rewritten); + assert_eq!(cast.cast_type(), &DataType::Utf8View); + let outer_get_field = cast + .expr() + .downcast_ref::() + .expect("Expected get_field under the cast"); + let inner_get_field = outer_get_field.args()[0] + .downcast_ref::() + .expect("Expected a nested get_field"); + assert!( + inner_get_field.args()[0].downcast_ref::().is_some(), + "the struct column must not be hidden behind a cast, got: {rewritten}" + ); + } + + /// A struct column that needs no adaptation at all is left completely + /// alone — the narrowing must not disturb the common case. + #[test] + fn test_narrow_struct_cast_leaves_matching_schema_alone() { + let (logical_schema, physical_schema) = struct_schemas( + vec![Field::new("x", DataType::Int32, true)], + vec![Field::new("x", DataType::Int32, true)], + ); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let expr = get_field_expr(&logical_schema, "s", "x"); + + let rewritten = adapter.rewrite(Arc::clone(&expr)).unwrap(); + + assert_eq!( + rewritten.to_string(), + expr.to_string(), + "an unadapted struct column must pass through untouched" + ); + } + + /// When the accessed field has the same type in both schemas, the struct + /// cast disappears entirely rather than being replaced by a field cast: + /// only a sibling field forced the column-level cast in the first place. + #[test] + fn test_narrow_struct_cast_drops_cast_when_field_types_match() { + let (logical_schema, physical_schema) = struct_schemas( + vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ], + vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int64, true), + ], + ); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let expr = Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())), + vec![ + Arc::new(Column::new("s", 0)), + Arc::new(Literal::new(ScalarValue::from("x"))), + ], + &logical_schema, + Arc::new(datafusion_common::config::ConfigOptions::default()), + ) + .unwrap(), + ) as Arc; + + let rewritten = adapter.rewrite(expr).unwrap(); + + assert!( + rewritten.downcast_ref::().is_none(), + "`x` has the same type in both schemas, so no cast is needed, got: {rewritten}" + ); + let get_field = rewritten + .downcast_ref::() + .expect("Expected a bare get_field"); + assert_eq!(get_field.return_type(), &DataType::Int32); + assert!( + get_field.args()[0].downcast_ref::().is_some(), + "the struct column must not be hidden behind a cast, got: {rewritten}" + ); + } + + /// `s['inner']['x']` is simplified to the flattened `get_field(s, 'inner', + /// 'x')`, so the whole key path has to be resolved. + #[test] + fn test_narrow_struct_cast_flattened_field_path() { + let (logical_schema, physical_schema) = struct_schemas( + vec![Field::new( + "inner", + DataType::Struct(vec![Field::new("x", DataType::Utf8, true)].into()), + true, + )], + vec![Field::new( + "inner", + DataType::Struct(vec![Field::new("x", DataType::Utf8View, true)].into()), + true, + )], + ); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let expr = Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())), + vec![ + Arc::new(Column::new("s", 0)), + Arc::new(Literal::new(ScalarValue::from("inner"))), + Arc::new(Literal::new(ScalarValue::from("x"))), + ], + &logical_schema, + Arc::new(datafusion_common::config::ConfigOptions::default()), + ) + .unwrap(), + ) as Arc; + + let rewritten = adapter.rewrite(expr).unwrap(); + + let cast = assert_cast_expr(&rewritten); + assert_eq!(cast.cast_type(), &DataType::Utf8View); + let get_field = cast + .expr() + .downcast_ref::() + .expect("Expected get_field under the cast"); + assert_eq!(get_field.return_type(), &DataType::Utf8); + assert_eq!( + get_field.args().len(), + 3, + "the full key path must be preserved, got: {rewritten}" + ); + assert!( + get_field.args()[0].downcast_ref::().is_some(), + "the struct column must not be hidden behind a cast, got: {rewritten}" + ); + } + + /// A key path whose leaf is missing from the file still resolves to a + /// typed null literal. + #[test] + fn test_narrow_struct_cast_flattened_field_path_missing_leaf() { + let (logical_schema, physical_schema) = struct_schemas( + vec![Field::new( + "inner", + DataType::Struct(vec![Field::new("x", DataType::Int32, true)].into()), + true, + )], + vec![Field::new( + "inner", + DataType::Struct( + vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Utf8, true), + ] + .into(), + ), + true, + )], + ); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let expr = Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(datafusion_expr::ScalarUDF::from(GetFieldFunc::new())), + vec![ + Arc::new(Column::new("s", 0)), + Arc::new(Literal::new(ScalarValue::from("inner"))), + Arc::new(Literal::new(ScalarValue::from("y"))), + ], + &logical_schema, + Arc::new(datafusion_common::config::ConfigOptions::default()), + ) + .unwrap(), + ) as Arc; + + let rewritten = adapter.rewrite(expr).unwrap(); + + let literal = rewritten + .downcast_ref::() + .expect("Expected a null literal"); + assert_eq!(*literal.value(), ScalarValue::Utf8(None)); + } + + /// Accessing a field the file does not have yields a typed null literal. + #[test] + fn test_narrow_struct_cast_missing_field() { + let (logical_schema, physical_schema) = struct_schemas( + vec![Field::new("x", DataType::Int32, true)], + vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Utf8, true), + ], + ); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let rewritten = adapter + .rewrite(get_field_expr(&logical_schema, "s", "y")) + .unwrap(); + + let literal = rewritten + .downcast_ref::() + .expect("Expected a null literal"); + assert_eq!(*literal.value(), ScalarValue::Utf8(None)); + } + + /// `get_field` on a Map column is a runtime key lookup, not a schema-level + /// field access, so the map value must keep its cast. + #[test] + fn test_map_field_access_keeps_cast() { + let map_type = |value_type: DataType| { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", value_type, true), + ] + .into(), + ), + false, + )), + false, + ) + }; + let physical_schema = Arc::new(Schema::new(vec![Field::new( + "s", + map_type(DataType::Int32), + true, + )])); + let logical_schema = Arc::new(Schema::new(vec![Field::new( + "s", + map_type(DataType::Int64), + true, + )])); + + let adapter = DefaultPhysicalExprAdapterFactory + .create(Arc::clone(&logical_schema), physical_schema) + .unwrap(); + let rewritten = adapter + .rewrite(get_field_expr(&logical_schema, "s", "k")) + .unwrap(); + + let get_field = rewritten + .downcast_ref::() + .expect("Expected the get_field to be preserved"); + assert!( + get_field.args()[0].downcast_ref::().is_some(), + "map columns must keep the whole-column cast, got: {rewritten}" + ); + } + // ============================================================================ // BatchAdapterFactory and BatchAdapter tests // ============================================================================ diff --git a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt index cb3be93191fb2..f7d23001fcf50 100644 --- a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt +++ b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt @@ -890,6 +890,119 @@ set datafusion.execution.parquet.pushdown_filters = false; statement ok DROP TABLE t_struct_filter; +########## +# Filters on struct fields, with pushdown enabled, where the declared table +# schema differs from the physical file schema and the scan therefore has to +# adapt the struct column. +# +# See https://github.com/apache/datafusion/issues/24109. +########## + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +COPY ( + SELECT + column1 as id, + named_struct('x', arrow_cast(column2, 'Int32')) as s + FROM VALUES (1, 100), (2, 200), (3, 300) +) TO 'test_files/scratch/parquet_filter_pushdown/struct_schema_cast.parquet' +STORED AS PARQUET; + +# `x` is stored as Int32 in the file but declared as BIGINT here, which forces +# the scan to adapt the struct column. +statement ok +CREATE EXTERNAL TABLE t_struct_schema_cast (id BIGINT, s STRUCT) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/struct_schema_cast.parquet'; + +query II +SELECT id, s['x'] FROM t_struct_schema_cast WHERE s['x'] = 200; +---- +2 200 + +# Conjunction of a struct-field filter and a primitive filter. +query II +SELECT id, s['x'] FROM t_struct_schema_cast WHERE s['x'] > 100 AND id > 2; +---- +3 300 + +# Control: the same file read through a schema that matches it exactly, so no +# cast is inserted. This path has always worked. +statement ok +CREATE EXTERNAL TABLE t_struct_no_schema_cast (id BIGINT, s STRUCT) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/struct_schema_cast.parquet'; + +query II +SELECT id, s['x'] FROM t_struct_no_schema_cast WHERE s['x'] = 200; +---- +2 200 + +query II +SELECT id, s['x'] FROM t_struct_no_schema_cast WHERE s['x'] > 100 AND id > 2; +---- +3 300 + +# Declaring a field the file does not have must not disturb the fields it does +# have. +statement ok +CREATE EXTERNAL TABLE t_struct_missing_field (id BIGINT, s STRUCT) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/struct_schema_cast.parquet'; + +query II +SELECT id, s['x'] FROM t_struct_missing_field WHERE s['x'] = 200; +---- +2 200 + +query II +SELECT id, s['x'] FROM t_struct_missing_field WHERE s['x'] > 100 AND id > 2; +---- +3 300 + +# The absent field itself reads as null, so it matches nothing. +query II +SELECT id, s['x'] FROM t_struct_missing_field WHERE s['missing'] = 200; +---- + +# Same, for a nested field path. +statement ok +COPY ( + SELECT + column1 as id, + named_struct('inner', named_struct('x', arrow_cast(column2, 'Int32'))) as s + FROM VALUES (1, 100), (2, 200), (3, 300) +) TO 'test_files/scratch/parquet_filter_pushdown/struct_nested_schema_cast.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE t_struct_nested_schema_cast (id BIGINT, s STRUCT>) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/struct_nested_schema_cast.parquet'; + +query II +SELECT id, s['inner']['x'] FROM t_struct_nested_schema_cast WHERE s['inner']['x'] = 200; +---- +2 200 + +# Clean up +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +statement ok +DROP TABLE t_struct_schema_cast; + +statement ok +DROP TABLE t_struct_no_schema_cast; + +statement ok +DROP TABLE t_struct_missing_field; + +statement ok +DROP TABLE t_struct_nested_schema_cast; + ########## # Regression test for https://github.com/apache/datafusion/issues/20937 # diff --git a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt index d936a89beb9f7..c5af117403c13 100644 --- a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt +++ b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt @@ -125,24 +125,22 @@ explain analyze select s from full_schema; ---- Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=219] -# `get_field` on a schema-narrowed struct becomes `get_field(CAST(s), 'x')`; -# the read clips to the cast target (every field the *narrow* schema -# declares), not further down to just `x`. The fair "nothing was clipped" -# baseline is therefore reading every physical leaf of `s` -# (`select s from full_schema` above), not the same `get_field` query against -# `full_schema` -- that one needs no cast at all and takes `get_field`'s own, -# more precise, single-leaf pushdown path. +# Selecting a single field of `s` reads fewer bytes than selecting `s` itself +# (above): the read clips down to `x` rather than to every field the narrow +# schema declares. query TT explain analyze select s['x'] from narrow; ---- -Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=146] +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=75] -# Mixed access -- the whole (narrowed) column and a subfield of it -- still -# reads only the narrow schema's leaves. +# Mixed access -- the whole (narrowed) column and a subfield of it. A root +# carrying both access kinds falls back to reading every physical leaf, so +# this costs the same as the unclipped full_schema read below rather than +# clipping to the narrow schema's leaves. query TT explain analyze select s, s['y'] from narrow; ---- -Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=146] +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=219] query TT explain analyze select s, s['y'] from full_schema; From 20268a7a5212e3364609137d957b3625849f3cea Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 26 Aug 2026 15:47:24 -0400 Subject: [PATCH 03/16] [branch-55] fix(ffi): use the session's LogicalExtensionCodec in create_physical_plan (#24690) (#24694) This PR is a cherry-pick of #24690 onto `branch-55`. See the original PR for a description of the issue. Co-authored-by: Claude Opus 5 (1M context) --- datafusion/ffi/src/session/mod.rs | 57 ++++++++++++++++++++++++++----- 1 file changed, 49 insertions(+), 8 deletions(-) diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index 83f842508ab2c..f215a6ba5a568 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -53,8 +53,8 @@ use datafusion_expr::{ use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; use datafusion_proto::bytes::{ - logical_plan_from_bytes, logical_plan_from_bytes_with_extension_codec, - logical_plan_to_bytes, logical_plan_to_bytes_with_extension_codec, + logical_plan_from_bytes_with_extension_codec, + logical_plan_to_bytes_with_extension_codec, }; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::logical_plan::from_proto::parse_expr; @@ -246,13 +246,17 @@ unsafe extern "C" fn create_physical_plan_fn_wrapper( let runtime = session.runtime().clone(); let session = session.clone(); async move { + let logical_codec: Arc = + (&session.logical_codec).into(); let session = session.inner(); let task_ctx = session.task_ctx(); - let logical_plan = sresult_return!(logical_plan_from_bytes( - logical_plan_serialized.as_slice(), - task_ctx.as_ref(), - )); + let logical_plan = + sresult_return!(logical_plan_from_bytes_with_extension_codec( + logical_plan_serialized.as_slice(), + task_ctx.as_ref(), + logical_codec.as_ref(), + )); let physical_plan = session.create_physical_plan(&logical_plan).await; @@ -746,7 +750,10 @@ impl Session for ForeignSession { logical_plan: &LogicalPlan, ) -> datafusion_common::Result> { unsafe { - let logical_plan = logical_plan_to_bytes(logical_plan)?; + let codec: Arc = + (&self.session.logical_codec).into(); + let logical_plan = + logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?; let physical_plan = df_result!( (self.session.create_physical_plan)( &self.session, @@ -851,8 +858,9 @@ mod tests { use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; + use arrow::array::record_batch; use arrow_schema::{DataType, Field, Schema}; - use datafusion::catalog::MemoryCatalogProvider; + use datafusion::catalog::{MemTable, MemoryCatalogProvider}; use datafusion::execution::SessionStateBuilder; use datafusion_common::DataFusionError; use datafusion_expr::col; @@ -860,6 +868,7 @@ mod tests { use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; use super::*; + use crate::proto::physical_extension_codec::tests::TestExtensionCodec; static QUERY_PLANNER_CALLS: AtomicUsize = AtomicUsize::new(0); static PHYSICAL_OPTIMIZER_CALLS: AtomicUsize = AtomicUsize::new(0); @@ -1004,4 +1013,36 @@ mod tests { Ok(()) } + + /// `create_physical_plan` must serialize with the session's logical codec on + /// both sides of the boundary. A plan that scans a custom table provider is + /// unserializable without it. + #[tokio::test] + async fn test_create_physical_plan_uses_logical_codec() -> Result<(), DataFusionError> + { + let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch = record_batch!(("a", Int32, [1, 2, 3]))?; + let table = MemTable::try_new(schema, vec![vec![batch]])?; + ctx.register_table("test_table", Arc::new(table))?; + + let logical_codec = FFI_LogicalExtensionCodec::new( + Arc::new(TestExtensionCodec), + None, + task_ctx_provider, + ); + + let state = ctx.state(); + let local_session = FFI_SessionRef::new(&state, None, logical_codec); + let foreign_session = ForeignSession::try_from(&local_session)?; + + let logical_plan = ctx.table("test_table").await?.into_optimized_plan()?; + let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?; + + assert_eq!(physical_plan.name(), "DataSourceExec"); + assert_eq!(physical_plan.schema().field(0).name(), "a"); + + Ok(()) + } } From 1a944f43983188712aa1934d034f3087d82172e1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 28 Aug 2026 10:03:31 -0400 Subject: [PATCH 04/16] [branch-55] fix: preserve provided arguments during FFI object construction (#24723) (#24752) This is a back port of #24723 onto `branch-55` to support `datafusion-python` upgrade to 55.1.0. The details can be found in the linked PR. This is needed for https://github.com/apache/datafusion-python/pull/1677 Co-authored-by: Claude Opus 5 (1M context) --- .../ffi/src/proto/logical_extension_codec.rs | 74 ++++++++++++++++++- .../ffi/src/proto/physical_extension_codec.rs | 47 +++++++++++- datafusion/ffi/src/query_planner.rs | 36 +++++++++ datafusion/ffi/src/table_provider.rs | 50 ++++++++++++- 4 files changed, 202 insertions(+), 5 deletions(-) diff --git a/datafusion/ffi/src/proto/logical_extension_codec.rs b/datafusion/ffi/src/proto/logical_extension_codec.rs index ed2c594f1bc02..2668421a93254 100644 --- a/datafusion/ffi/src/proto/logical_extension_codec.rs +++ b/datafusion/ffi/src/proto/logical_extension_codec.rs @@ -294,6 +294,15 @@ impl Drop for FFI_LogicalExtensionCodec { impl FFI_LogicalExtensionCodec { /// Creates a new [`FFI_LogicalExtensionCodec`]. + /// + /// If `codec` is already foreign, this re-exports its original FFI handle + /// rather than adding another wrapper layer. The handle still adopts the + /// `task_ctx_provider` supplied here, so it is never silently discarded and + /// an imported codec can be rebound to a different session. + /// + /// `runtime` is only honored when a new wrapper is created. An + /// already-foreign handle keeps the runtime of the library that owns it, + /// because that value lives in private data this side cannot reach. pub fn new( codec: Arc, runtime: Option, @@ -302,7 +311,9 @@ impl FFI_LogicalExtensionCodec { if let Some(codec) = (Arc::clone(&codec) as Arc) .downcast_ref::() { - return codec.0.clone(); + let mut codec = codec.0.clone(); + codec.task_ctx_provider = task_ctx_provider.into(); + return codec; } let task_ctx_provider = task_ctx_provider.into(); @@ -507,7 +518,9 @@ mod tests { use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; - use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; + use crate::proto::logical_extension_codec::{ + FFI_LogicalExtensionCodec, ForeignLogicalExtensionCodec, + }; use crate::proto::physical_extension_codec::tests::TestExtensionCodec; fn create_test_table() -> MemTable { @@ -727,4 +740,61 @@ mod tests { let foreign_codec: Arc = (&ffi_codec).into(); assert!(!arc_ptr_eq(&foreign_codec, &codec)); } + + /// Importing a codec and re-wrapping it with a different task context + /// provider must rebind the handle. See + /// . + #[test] + fn ffi_logical_extension_codec_rebind_adopts_task_ctx_provider() { + let (_ctx_a, provider_a) = crate::util::tests::test_session_and_ctx(); + let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx(); + + let mut ffi_codec = FFI_LogicalExtensionCodec::new( + Arc::new(TestExtensionCodec {}) as Arc, + None, + provider_a, + ); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + + let imported: Arc = (&ffi_codec).into(); + assert!( + (Arc::clone(&imported) as Arc) + .downcast_ref::() + .is_some() + ); + + let rebound = FFI_LogicalExtensionCodec::new(imported, None, provider_b); + + let task_ctx: Arc = (&rebound.task_ctx_provider) + .try_into() + .expect("rebound codec resolves"); + assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); + } + + /// Because the provider is held as a `Weak`, a codec that cannot be rebound + /// forces callers to keep the original session alive. Once rebinding works, + /// dropping it must not invalidate the handle. + #[test] + fn ffi_logical_extension_codec_rebind_releases_original_session() { + let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx(); + + let rebound = { + let (ctx_a, provider_a) = crate::util::tests::test_session_and_ctx(); + let mut ffi_codec = FFI_LogicalExtensionCodec::new( + Arc::new(TestExtensionCodec {}) as Arc, + None, + provider_a, + ); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + let imported: Arc = (&ffi_codec).into(); + drop(ctx_a); + + FFI_LogicalExtensionCodec::new(imported, None, provider_b) + }; + + let task_ctx: Arc = (&rebound.task_ctx_provider) + .try_into() + .expect("rebound codec must not depend on the original session"); + assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); + } } diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index 95d2ed68a6ea3..ad8a208967c36 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -280,6 +280,15 @@ impl Drop for FFI_PhysicalExtensionCodec { impl FFI_PhysicalExtensionCodec { /// Creates a new [`FFI_PhysicalExtensionCodec`]. + /// + /// If `codec` is already foreign, this re-exports its original FFI handle + /// rather than adding another wrapper layer. The handle still adopts the + /// `task_ctx_provider` supplied here, so it is never silently discarded and + /// an imported codec can be rebound to a different session. + /// + /// `runtime` is only honored when a new wrapper is created. An + /// already-foreign handle keeps the runtime of the library that owns it, + /// because that value lives in private data this side cannot reach. pub fn new( codec: Arc, runtime: Option, @@ -288,7 +297,9 @@ impl FFI_PhysicalExtensionCodec { if let Some(codec) = (Arc::clone(&codec) as Arc) .downcast_ref::() { - return codec.0.clone(); + let mut codec = codec.0.clone(); + codec.task_ctx_provider = task_ctx_provider.into(); + return codec; } let task_ctx_provider = task_ctx_provider.into(); @@ -449,7 +460,9 @@ pub(crate) mod tests { }; use crate::execution_plan::tests::EmptyExec; - use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; + use crate::proto::physical_extension_codec::{ + FFI_PhysicalExtensionCodec, ForeignPhysicalExtensionCodec, + }; #[derive(Debug)] pub(crate) struct TestExtensionCodec; @@ -710,4 +723,34 @@ pub(crate) mod tests { let foreign_codec: Arc = (&ffi_codec).into(); assert!(!arc_ptr_eq(&foreign_codec, &codec)); } + + /// Importing a codec and re-wrapping it with a different task context + /// provider must rebind the handle. See + /// . + #[test] + fn ffi_physical_extension_codec_rebind_adopts_task_ctx_provider() { + let (_ctx_a, provider_a) = crate::util::tests::test_session_and_ctx(); + let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx(); + + let mut ffi_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(TestExtensionCodec {}) as Arc, + None, + provider_a, + ); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + + let imported: Arc = (&ffi_codec).into(); + assert!( + (Arc::clone(&imported) as Arc) + .downcast_ref::() + .is_some() + ); + + let rebound = FFI_PhysicalExtensionCodec::new(imported, None, provider_b); + + let task_ctx: Arc = (&rebound.task_ctx_provider) + .try_into() + .expect("rebound codec resolves"); + assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); + } } diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 6d895d65c5fc1..727e1e57819a7 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -442,4 +442,40 @@ mod tests { Ok(()) } + + // Control for https://github.com/apache/datafusion/issues/24722: this + // constructor adopts the supplied codecs on the already-foreign path. + #[test] + fn test_rebind_foreign_query_planner_adopts_codecs() { + use datafusion_execution::TaskContext; + + let ctx_a = Arc::new(SessionContext::new()); + let ctx_b = Arc::new(SessionContext::new()); + let provider_b = Arc::clone(&ctx_b) as Arc; + + let mut ffi_a = create_ffi_query_planner(Arc::clone(&ctx_a)); + ffi_a.library_marker_id = crate::mock_foreign_marker_id; + let imported: Arc = (&ffi_a).into(); + let any_ref: &dyn std::any::Any = imported.as_ref(); + assert!(any_ref.downcast_ref::().is_some()); + + let rebound = FFI_QueryPlanner::new_with_ffi_codecs( + imported, + FFI_LogicalExtensionCodec::new( + Arc::new(DefaultLogicalExtensionCodec {}), + None, + &provider_b, + ), + FFI_PhysicalExtensionCodec::new( + Arc::new(DefaultPhysicalExtensionCodec {}), + None, + &provider_b, + ), + ); + + let bound_to: Arc = (&rebound.logical_codec.task_ctx_provider) + .try_into() + .unwrap(); + assert_eq!(bound_to.session_id(), ctx_b.task_ctx().session_id()); + } } diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index ee9377bff064e..f1eb6076b6783 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -402,6 +402,16 @@ impl FFI_TableProvider { ) } + /// Creates an [`FFI_TableProvider`] using a prebuilt FFI logical codec. + /// + /// If `provider` is already foreign, this re-exports its original FFI + /// handle rather than adding another wrapper layer. The handle still adopts + /// the `logical_codec` supplied here, so it is never silently discarded and + /// an imported provider can be rebound to a different session. + /// + /// `runtime` is only honored when a new wrapper is created. An + /// already-foreign handle keeps the runtime of the library that owns it, + /// because that value lives in private data this side cannot reach. pub fn new_with_ffi_codec( provider: Arc, can_support_pushdown_filters: bool, @@ -409,7 +419,9 @@ impl FFI_TableProvider { logical_codec: FFI_LogicalExtensionCodec, ) -> Self { if let Some(provider) = provider.downcast_ref::() { - return provider.0.clone(); + let mut provider = provider.0.clone(); + provider.logical_codec = logical_codec; + return provider; } let private_data = Box::new(ProviderPrivateData { provider, runtime }); @@ -903,4 +915,40 @@ mod tests { Ok(()) } + + /// Re-wrapping an imported provider with a rebuilt logical codec must adopt + /// that codec. See . + #[test] + fn test_rebind_foreign_table_provider_adopts_logical_codec() -> Result<()> { + let (_ctx_a, provider_a) = crate::util::tests::test_session_and_ctx(); + let (ctx_b, provider_b) = crate::util::tests::test_session_and_ctx(); + + let mut ffi_provider = FFI_TableProvider::new( + create_test_table_provider()?, + true, + None, + provider_a, + None, + ); + ffi_provider.library_marker_id = crate::mock_foreign_marker_id; + + let imported: Arc = (&ffi_provider).into(); + assert!(imported.downcast_ref::().is_some()); + + // Rebuild the codec against session B and re-wrap. + let codec_b = FFI_LogicalExtensionCodec::new( + Arc::new(DefaultLogicalExtensionCodec {}), + None, + provider_b, + ); + let rebound = + FFI_TableProvider::new_with_ffi_codec(imported, true, None, codec_b); + + let task_ctx: Arc = (&rebound.logical_codec.task_ctx_provider) + .try_into() + .expect("rebound provider's codec resolves"); + assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); + + Ok(()) + } } From 795c861a5b5f38387022a3276cd9bc335ef50594 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 28 Aug 2026 11:32:53 -0400 Subject: [PATCH 05/16] [branch-55] fix: adapt input batches with stricter nested nullability to planned schema in aggregation (#24394) (#24699) This is a back port of #24394 for `branch-55`. The original issue is https://github.com/apache/datafusion/issues/24069 Note that the original PR was marked as an auto-detected api change, but I believe that is a false positive. --------- Co-authored-by: Patrick Ribbsaeter Co-authored-by: Claude Opus 5 (1M context) --- datafusion/common/src/nested_struct.rs | 713 +++++++++++++++++- datafusion/core/tests/sql/aggregates/mod.rs | 1 + .../sql/aggregates/nested_nullability.rs | 246 ++++++ datafusion/physical-plan/src/memory.rs | 317 ++++++++ 4 files changed, 1275 insertions(+), 2 deletions(-) create mode 100644 datafusion/core/tests/sql/aggregates/nested_nullability.rs diff --git a/datafusion/common/src/nested_struct.rs b/datafusion/common/src/nested_struct.rs index e915b91b911cc..a6eb7df22a2cf 100644 --- a/datafusion/common/src/nested_struct.rs +++ b/datafusion/common/src/nested_struct.rs @@ -19,11 +19,14 @@ use crate::error::{_plan_err, Result}; use arrow::{ array::{ Array, ArrayRef, AsArray, DictionaryArray, FixedSizeListArray, GenericListArray, - GenericListViewArray, StructArray, downcast_integer, make_array, new_null_array, + GenericListViewArray, RecordBatch, StructArray, UnionArray, downcast_integer, + make_array, new_null_array, }, buffer::NullBuffer, compute::{CastOptions, can_cast_types, cast_with_options}, - datatypes::{DataType, DataType::Struct, Field, FieldRef}, + datatypes::{ + DataType, DataType::Struct, Field, FieldRef, SchemaRef, UnionFields, UnionMode, + }, }; use std::{collections::HashSet, sync::Arc}; @@ -121,6 +124,63 @@ fn cast_struct_column( } } +/// Cast a union column to match target union fields, handling child fields recursively. +/// +/// ## Casting Behavior +/// - Preserves union mode (sparse or dense). Incompatible modes are rejected. +/// - Requires exact matching union type ID sets (order may differ). +/// - Recursively adapts each matching child array using `cast_column`. +/// - Preserves row-level `type_ids` and dense `offsets` buffers without copying primitive data. +fn cast_union_column( + source_col: &ArrayRef, + source_fields: &UnionFields, + source_mode: &UnionMode, + target_fields: &UnionFields, + target_mode: &UnionMode, + cast_options: &CastOptions, +) -> Result { + validate_union_schema_compatibility( + source_fields, + source_mode, + target_fields, + target_mode, + )?; + + let source_union = source_col + .as_any() + .downcast_ref::() + .ok_or_else(|| { + crate::error::DataFusionError::Plan(format!( + "Expected UnionArray for Union data type, got {}", + source_col.data_type() + )) + })?; + + let mut children = Vec::with_capacity(target_fields.len()); + + for (target_type_id, target_field) in target_fields.iter() { + let source_child = source_union.child(target_type_id); + + children.push( + cast_column(source_child, target_field.data_type(), cast_options).map_err( + |e| { + e.context(format!( + "While adapting Union child type ID {target_type_id} ('{}')", + target_field.name() + )) + }, + )?, + ); + } + + Ok(Arc::new(UnionArray::try_new( + target_fields.clone(), + source_union.type_ids().clone(), + source_union.offsets().cloned(), + children, + )?)) +} + /// Cast a column to match the target field type, with special handling for nested structs. /// /// This function serves as the main entry point for column casting operations. For struct @@ -215,6 +275,17 @@ pub fn cast_column( target_value_type, cast_options, ), + ( + DataType::Union(source_fields, source_mode), + DataType::Union(target_fields, target_mode), + ) => cast_union_column( + source_col, + source_fields, + source_mode, + target_fields, + target_mode, + cast_options, + ), _ => Ok(cast_with_options(source_col, target_type, cast_options)?), } } @@ -490,6 +561,48 @@ fn validate_field_compatibility( ) } +fn validate_union_schema_compatibility( + source_fields: &UnionFields, + source_mode: &UnionMode, + target_fields: &UnionFields, + target_mode: &UnionMode, +) -> Result<()> { + if source_mode != target_mode { + return _plan_err!( + "Cannot adapt Union from mode {source_mode:?} to {target_mode:?}" + ); + } + + // This adapter is for schema conformance, not general Union variant-set evolution. + if source_fields.len() != target_fields.len() { + return _plan_err!( + "Cannot adapt Union schema with different field sets: source has {} fields, target has {}", + source_fields.len(), + target_fields.len() + ); + } + + for (target_type_id, target_field) in target_fields.iter() { + let Some((_, source_field)) = source_fields + .iter() + .find(|(source_type_id, _)| *source_type_id == target_type_id) + else { + return _plan_err!( + "Cannot adapt Union schema: target type ID {target_type_id} ('{}') is missing from source", + target_field.name() + ); + }; + + if !target_field.contains(source_field) { + return _plan_err!( + "Cannot adapt Union child with type ID {target_type_id}: source field {source_field} is not contained by target field {target_field}" + ); + } + } + + Ok(()) +} + /// Validates that `source_type` can be cast to `target_type`, recursively /// handling container types that wrap structs. pub fn validate_data_type_compatibility( @@ -524,6 +637,17 @@ pub fn validate_data_type_compatibility( } validate_data_type_compatibility(field_name, s_val, t_val)?; } + ( + DataType::Union(source_fields, source_mode), + DataType::Union(target_fields, target_mode), + ) => { + validate_union_schema_compatibility( + source_fields, + source_mode, + target_fields, + target_mode, + )?; + } _ => { if !can_cast_types(source_type, target_type) { return _plan_err!( @@ -1703,3 +1827,588 @@ mod tests { )); } } + +/// Adapts a `RecordBatch` to conform to `target_schema`, verifying that each target field +/// type contains the incoming column data type (as verified by [`arrow::datatypes::DataType::contains`]) +/// and transforms the metadata/types of differing columns to match `target_schema` +/// without copying primitive buffer data. +/// +/// If `batch` has an incompatible column count or incompatible column data types, +/// an error is returned. +pub fn adapt_batch_to_schema( + batch: RecordBatch, + target_schema: &SchemaRef, +) -> Result { + if Arc::ptr_eq(batch.schema_ref(), target_schema) + || batch.schema().as_ref() == target_schema.as_ref() + { + return Ok(batch); + } + + if batch.num_columns() != target_schema.fields().len() { + return _plan_err!( + "Batch schema does not conform to expected schema (column count mismatch). Expected: {target_schema}, got: {}", + batch.schema() + ); + } + + let mut columns = Vec::with_capacity(batch.num_columns()); + let mut needs_column_adaptation = false; + let cast_options = CastOptions::default(); + + for (target_field, col) in target_schema.fields().iter().zip(batch.columns()) { + if target_field.data_type() != col.data_type() { + // If data types differ, verify that target_field's data type contains + // the column's data type (e.g. stricter nested struct / list field nullability). + if !target_field.data_type().contains(col.data_type()) { + return _plan_err!( + "Batch column '{}' with type {} cannot be adapted to expected type {}", + target_field.name(), + col.data_type(), + target_field.data_type() + ); + } + needs_column_adaptation = true; + let adapted_col = cast_column(col, target_field.data_type(), &cast_options)?; + columns.push(adapted_col); + } else { + columns.push(Arc::clone(col)); + } + } + + if needs_column_adaptation { + Ok(RecordBatch::try_new(Arc::clone(target_schema), columns)?) + } else { + // Schema differs only in top-level metadata or field nullability, while + // column data types match exactly. Replace the schema on the batch. + Ok(RecordBatch::try_new( + Arc::clone(target_schema), + batch.columns().to_vec(), + )?) + } +} + +#[cfg(test)] +mod adapt_schema_tests { + use super::*; + use arrow::array::{Int32Array, StringArray}; + use arrow::datatypes::{Field, Fields, Schema}; + + #[test] + fn test_adapt_batch_to_schema_identical() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ])); + + let a = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let b = Arc::new(StringArray::from(vec![Some("x"), None, Some("z")])) as ArrayRef; + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![a, b])?; + + let adapted = adapt_batch_to_schema(batch.clone(), &schema)?; + assert_eq!(adapted, batch); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_stricter_nested_struct() -> Result<()> { + // Declared table schema: {a: Struct({x: Int32 (nullable), y: Utf8 (nullable)})} + let declared_inner_fields = Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Utf8, true), + ]); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "a", + Struct(declared_inner_fields), + false, + )])); + + // Runtime batch schema: {a: Struct({x: Int32 (NON-nullable), y: Utf8 (NON-nullable)})} + let runtime_inner_fields = Fields::from(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Utf8, false), + ]); + let runtime_schema = Arc::new(Schema::new(vec![Field::new( + "a", + Struct(runtime_inner_fields.clone()), + false, + )])); + + let x = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let y = Arc::new(StringArray::from(vec!["x", "y", "z"])) as ArrayRef; + let struct_array = + Arc::new(StructArray::new(runtime_inner_fields, vec![x, y], None)) + as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![struct_array])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + assert_eq!(adapted.num_rows(), 3); + + // Verify nested fields now have the declared nullability + let Struct(fields) = adapted.column(0).data_type() else { + panic!("expected struct"); + }; + assert!(fields[0].is_nullable()); + assert!(fields[1].is_nullable()); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_top_level_nullability_only() -> Result<()> { + // Declared schema has nullable column 'a', runtime batch has non-nullable 'a' + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let runtime_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + let a = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![a])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + assert!(adapted.schema().field(0).is_nullable()); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_null_into_non_nullable_rejected() { + // Declared schema is non-nullable, but runtime batch is nullable + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let runtime_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let a = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])) as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![a]).unwrap(); + + // Must reject because nullable is not contained by non-nullable + let result = adapt_batch_to_schema(batch, &declared_schema); + assert!(result.is_err()); + } + + #[test] + fn test_adapt_batch_to_schema_incompatible_type_rejected() { + let declared_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let runtime_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + + let a = Arc::new(StringArray::from(vec!["1", "2"])) as ArrayRef; + let batch = RecordBatch::try_new(runtime_schema, vec![a]).unwrap(); + + let result = adapt_batch_to_schema(batch, &declared_schema); + assert!(result.is_err()); + } + + fn test_two_field_union(nullable: bool) -> UnionFields { + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, nullable), + Field::new("str", DataType::Utf8, nullable), + ], + ) + .unwrap() + } + + #[test] + fn test_adapt_batch_to_schema_stricter_sparse_union() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let target_union_fields = test_two_field_union(true); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Sparse), + false, + )])); + + let source_union_fields = test_two_field_union(false); + let runtime_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Sparse), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); + let str_array: ArrayRef = + Arc::new(StringArray::from(vec!["hello", "world", "!"])); + let type_ids = [0, 0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + None, + vec![int_array, str_array], + )?; + let batch = RecordBatch::try_new(runtime_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let DataType::Union(fields, mode) = adapted_union.data_type() else { + panic!("expected union"); + }; + assert_eq!(*mode, UnionMode::Sparse); + assert!(fields.iter().all(|(_, f)| f.is_nullable())); + + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_stricter_dense_union() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let target_union_fields = test_two_field_union(true); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + )])); + + let source_union_fields = test_two_field_union(false); + let runtime_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 30])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["hello"])); + let type_ids = [0, 1, 0].into_iter().collect::>(); + let offsets = [0, 0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + Some(offsets), + vec![int_array, str_array], + )?; + let batch = RecordBatch::try_new(runtime_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let DataType::Union(fields, mode) = adapted_union.data_type() else { + panic!("expected union"); + }; + assert_eq!(*mode, UnionMode::Dense); + assert!(fields.iter().all(|(_, f)| f.is_nullable())); + + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_union_reordered_and_non_contiguous_type_ids() + -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let target_union_fields = UnionFields::try_new( + vec![3, 1], + vec![ + Field::new("str", DataType::Utf8, true), + Field::new("int", DataType::Int32, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + )])); + + let source_union_fields = UnionFields::try_new( + vec![1, 3], + vec![ + Field::new("int", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + ], + )?; + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 30])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["b"])); + let type_ids = [1, 3, 1].into_iter().collect::>(); + let offsets = [0, 0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids.clone(), + Some(offsets.clone()), + vec![int_array, str_array], + )?; + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(source_batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + adapted_union.data_type(), + declared_schema.field(0).data_type() + ); + assert_eq!(adapted_union.type_ids(), &type_ids); + assert_eq!(adapted_union.offsets(), Some(&offsets)); + + // Child 1 is int, Child 3 is str (accessed by type ID) + let int_child = adapted_union + .child(1) + .as_any() + .downcast_ref::() + .unwrap(); + let str_child = adapted_union + .child(3) + .as_any() + .downcast_ref::() + .unwrap(); + + // Row 0: type_id 1 -> int value 10 + assert_eq!(adapted_union.type_id(0), 1); + assert_eq!(int_child.value(adapted_union.value_offset(0)), 10); + + // Row 1: type_id 3 -> str value "b" + assert_eq!(adapted_union.type_id(1), 3); + assert_eq!(str_child.value(adapted_union.value_offset(1)), "b"); + + // Row 2: type_id 1 -> int value 30 + assert_eq!(adapted_union.type_id(2), 1); + assert_eq!(int_child.value(adapted_union.value_offset(2)), 30); + + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_union_nested_struct() -> Result<()> { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_struct_fields = vec![Field::new("x", DataType::Int32, true)]; + let target_union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("s", Struct(target_struct_fields.into()), true)], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Dense), + false, + )])); + + let source_struct_fields = vec![Field::new("x", DataType::Int32, false)]; + let source_union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("s", Struct(source_struct_fields.into()), false)], + )?; + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let struct_child: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("x", DataType::Int32, false)].into(), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + None, + )); + let type_ids = [0, 0].into_iter().collect::>(); + let offsets = [0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids.clone(), + Some(offsets.clone()), + vec![struct_child], + )?; + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)])?; + + let adapted = adapt_batch_to_schema(source_batch, &declared_schema)?; + assert_eq!(adapted.schema().as_ref(), declared_schema.as_ref()); + let adapted_union = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let adapted_child = adapted_union.child(0); + let struct_arr = adapted_child + .as_any() + .downcast_ref::() + .unwrap(); + assert!(struct_arr.fields()[0].is_nullable()); + Ok(()) + } + + #[test] + fn test_adapt_batch_to_schema_union_incompatible_mode_rejected() { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::UnionMode; + + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(test_two_field_union(true), UnionMode::Dense), + false, + )])); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(test_two_field_union(false), UnionMode::Sparse), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let type_ids = [0, 0].into_iter().collect::>(); + let source_union = UnionArray::try_new( + test_two_field_union(false), + type_ids, + None, + vec![int_array, str_array], + ) + .unwrap(); + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)]).unwrap(); + + let res = adapt_batch_to_schema(source_batch, &declared_schema); + assert!(res.is_err()); + } + + #[test] + fn test_adapt_batch_to_schema_union_field_set_mismatch_rejected() { + use arrow::array::UnionArray; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{UnionFields, UnionMode}; + + // Target has type ID [0] + let target_union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("value", DataType::Int32, true)], + ) + .unwrap(); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(target_union_fields, UnionMode::Sparse), + false, + )])); + + // Source has type IDs [0, 1] (where ID 0 is compatible) + let source_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("value", DataType::Int32, false), + Field::new("extra", DataType::Utf8, false), + ], + ) + .unwrap(); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Sparse), + false, + )])); + + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b"])); + let type_ids = [0, 1].into_iter().collect::>(); + let source_union = UnionArray::try_new( + source_union_fields, + type_ids, + None, + vec![int_array, str_array], + ) + .unwrap(); + + let source_batch = + RecordBatch::try_new(source_schema, vec![Arc::new(source_union)]).unwrap(); + + let res = adapt_batch_to_schema(source_batch, &declared_schema); + assert!(res.is_err()); + let err = res.unwrap_err().to_string(); + assert!( + err.contains("different field sets") + || err.contains("cannot be adapted to expected type"), + "unexpected error message: {err}" + ); + } + + #[test] + fn test_validate_data_type_compatibility_union() { + use arrow::datatypes::{UnionFields, UnionMode}; + + let target_type = DataType::Union(test_two_field_union(true), UnionMode::Dense); + + // Compatible: exact same type IDs in different order with stricter nullability + let reordered_source_fields = UnionFields::try_new( + vec![1, 0], + vec![ + Field::new("str", DataType::Utf8, false), + Field::new("value", DataType::Int32, false), + ], + ) + .unwrap(); + let source_type = DataType::Union(reordered_source_fields, UnionMode::Dense); + assert!( + validate_data_type_compatibility("u", &source_type, &target_type).is_ok() + ); + + // Incompatible: mismatched mode + let sparse_source_type = + DataType::Union(test_two_field_union(false), UnionMode::Sparse); + assert!( + validate_data_type_compatibility("u", &sparse_source_type, &target_type) + .is_err() + ); + + // Incompatible: field-set mismatch (extra source ID 2) + let extra_id_source = DataType::Union( + UnionFields::try_new( + vec![0, 1, 2], + vec![ + Field::new("value", DataType::Int32, false), + Field::new("str", DataType::Utf8, false), + Field::new("extra", DataType::Int32, false), + ], + ) + .unwrap(), + UnionMode::Dense, + ); + assert!( + validate_data_type_compatibility("u", &extra_id_source, &target_type) + .is_err() + ); + + // Incompatible: field-set mismatch (missing source ID 1) + let missing_id_source = DataType::Union( + UnionFields::try_new( + vec![0], + vec![Field::new("value", DataType::Int32, false)], + ) + .unwrap(), + UnionMode::Dense, + ); + assert!( + validate_data_type_compatibility("u", &missing_id_source, &target_type) + .is_err() + ); + } +} diff --git a/datafusion/core/tests/sql/aggregates/mod.rs b/datafusion/core/tests/sql/aggregates/mod.rs index b209e91cc81e7..186297b639cbd 100644 --- a/datafusion/core/tests/sql/aggregates/mod.rs +++ b/datafusion/core/tests/sql/aggregates/mod.rs @@ -1021,3 +1021,4 @@ pub fn split_fuzz_timestamp_data_into_batches( pub mod basic; pub mod dict_nulls; +mod nested_nullability; diff --git a/datafusion/core/tests/sql/aggregates/nested_nullability.rs b/datafusion/core/tests/sql/aggregates/nested_nullability.rs new file mode 100644 index 0000000000000..448759ad74c54 --- /dev/null +++ b/datafusion/core/tests/sql/aggregates/nested_nullability.rs @@ -0,0 +1,246 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Regression tests for aggregating batches whose data types are *stricter* +//! than the table's declared schema. See +//! . +//! +//! Builds on the end-to-end reproducer from #24278 by @alamb. +//! +//! A `RecordBatch` is a valid instance of a schema that is a superset of its +//! own (see [`Schema::contains`] / `Field::contains`): most commonly the +//! schema declares a (possibly nested) field as nullable while the batch's +//! arrays mark it non-nullable. `MemTable::try_new` accepts such batches via +//! exactly that check, and engines embedding DataFusion (e.g. Comet) feed +//! such batches over FFI. Aggregations must therefore not fail when the +//! runtime arrays are stricter than the planned schema. +//! +//! [`Schema::contains`]: arrow::datatypes::Schema::contains + +use std::sync::Arc; + +use arrow::array::{BooleanArray, RecordBatch, StructArray, UInt32Array}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion::physical_plan::collect; +use datafusion::physical_plan::expressions::col; +use datafusion::prelude::*; +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::FairSpillPool; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_functions_aggregate::array_agg::array_agg_udaf; + +/// Returns the fields of the struct column `b`: a single `colA Boolean`. +/// +/// `col_a_nullable` controls whether `colA` is declared nullable — the only +/// difference between the table's declared schema (`true`) and the actual +/// batches (`false`). +fn make_struct_fields(col_a_nullable: bool) -> Fields { + Fields::from(vec![Field::new("colA", DataType::Boolean, col_a_nullable)]) +} + +/// Returns the schema `(a UInt32 NOT NULL, b Struct("colA" Boolean) NOT NULL)` +/// with the nested field `b.colA` nullable per `col_a_nullable`. +/// +/// See [`make_struct_fields`]. +fn make_schema(col_a_nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new( + "b", + DataType::Struct(make_struct_fields(col_a_nullable)), + false, + ), + ])) +} + +/// Runs a SQL aggregation over a table whose batches are stricter than its +/// declared schema. +/// +/// [`Self::run`] registers table `t(a UInt32, b Struct("colA" Boolean))` +/// where the declared schema marks the nested field `colA` as nullable, but +/// the batches carry a stricter, non-nullable `colA`, then runs the query +/// and returns the collected result. +struct AggregateBatchesTest { + /// Number of rows in the table. `a` is `0..num_rows` (so also the number + /// of groups for `GROUP BY a`) and `b.colA` alternates `true` / `false`. + num_rows: u32, + /// If set, the context uses a [`FairSpillPool`] of this size (and a small + /// batch size) so the aggregation is forced to spill. + memory_limit: Option, +} + +impl AggregateBatchesTest { + fn new() -> Self { + Self { + num_rows: 100, + memory_limit: None, + } + } + + fn with_num_rows(mut self, num_rows: u32) -> Self { + self.num_rows = num_rows; + self + } + + fn with_memory_limit(mut self, memory_limit: usize) -> Self { + self.memory_limit = Some(memory_limit); + self + } + + /// Runs `sql` against the table described above and asserts the result + /// has one output row per group (i.e. [`Self::num_rows`] rows in total). + async fn run(self, sql: &str) -> Result<()> { + // The table's declared schema: the nested field `b.colA` is + // nullable ... + let declared_schema = make_schema(true); + + // ... while the batches are stricter: `b.colA` is non-nullable. + // `MemTable::try_new` accepts this combination via + // `Schema::contains`. + let batch_struct_fields = make_struct_fields(false); + let batch = RecordBatch::try_new( + make_schema(false), + vec![ + Arc::new(UInt32Array::from_iter_values(0..self.num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..self.num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let table = MemTable::try_new(declared_schema, vec![vec![batch]])?; + + let ctx = match self.memory_limit { + Some(limit) => { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(limit))) + .build_arc()?; + SessionContext::new_with_config_rt( + SessionConfig::new().with_batch_size(100), + runtime, + ) + } + None => SessionContext::new(), + }; + ctx.register_table("t", Arc::new(table))?; + + let result = ctx.sql(sql).await?.collect().await?; + + let total_rows: usize = result.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, self.num_rows as usize); + Ok(()) + } +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches() -> Result<()> { + AggregateBatchesTest::new() + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(b) FROM t GROUP BY a") + .await +} + +#[tokio::test] +async fn array_agg_distinct_struct_from_stricter_batches_with_spilling() -> Result<()> { + AggregateBatchesTest::new() + .with_num_rows(10_000) + .with_memory_limit(4_000_000) + .run("SELECT a, array_agg(DISTINCT b) FROM t GROUP BY a") + .await +} + +/// Direct unit test for `AggregateExec` boundary adaptation: +/// Feeds `AggregateExec` directly from a `MemorySourceConfig` whose batches carry +/// a stricter nested struct nullability than the plan schema without going +/// through `MemTable`. +#[tokio::test] +async fn test_aggregate_exec_direct_input_adaptation() -> Result<()> { + let declared_schema = make_schema(true); + let batch_struct_fields = make_struct_fields(false); + let num_rows = 100_u32; + let stricter_batch = RecordBatch::try_new( + make_schema(false), + vec![ + Arc::new(UInt32Array::from_iter_values(0..num_rows)), + Arc::new(StructArray::new( + batch_struct_fields, + vec![Arc::new(BooleanArray::from_iter( + (0..num_rows).map(|i| Some(i % 2 == 0)), + ))], + None, + )), + ], + )?; + + let input_plan: Arc = MemorySourceConfig::try_new_exec( + &[vec![stricter_batch]], + Arc::clone(&declared_schema), + None, + )?; + + let grouping_set = + PhysicalGroupBy::new_single(vec![(col("a", &declared_schema)?, "a".to_string())]); + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("b", &declared_schema)?]) + .schema(Arc::clone(&declared_schema)) + .alias("array_agg(b)") + .build()?, + )]; + + let agg_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + grouping_set, + aggregates, + vec![None], + input_plan, + Arc::clone(&declared_schema), + )?); + + let task_ctx = Arc::new(TaskContext::default()); + let results = collect(agg_exec, task_ctx).await?; + + let total_rows: usize = results.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, num_rows as usize); + Ok(()) +} diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index efe42c7ebc5f0..0c77d7e773265 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -106,6 +106,17 @@ impl Stream for MemoryStream { None => batch.clone(), }; + // MemoryStream advertises `self.schema`, therefore emitted RecordBatches + // must conform to it when batches were provided with stricter nested types + // (e.g. MemTable accepts stricter batches via Schema::contains). + let batch = if batch.schema().as_ref() != self.schema.as_ref() + && self.schema.contains(batch.schema().as_ref()) + { + datafusion_common::nested_struct::adapt_batch_to_schema(batch, &self.schema)? + } else { + batch + }; + let Some(&fetch) = self.fetch.as_ref() else { return Poll::Ready(Some(Ok(batch))); }; @@ -673,4 +684,310 @@ mod lazy_memory_tests { Ok(()) } + + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema() -> Result<()> { + use arrow::array::{ArrayRef, BooleanArray, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + // Declared schema expects nullable struct field colA + let declared_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, true)]); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "b", + DataType::Struct(declared_fields), + false, + )])); + + // Runtime batch has stricter non-nullable struct field colA + let source_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, false)]); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "b", + DataType::Struct(source_fields.clone()), + false, + )])); + + let struct_array: ArrayRef = Arc::new(StructArray::new( + source_fields, + vec![Arc::new(BooleanArray::from(vec![true, false]))], + None, + )); + let stricter_batch = RecordBatch::try_new(source_schema, vec![struct_array])?; + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), declared_schema); + + let struct_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(struct_col.fields()[0].is_nullable()); + let bool_child = struct_col + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(bool_child.value(0)); + assert!(!bool_child.value(1)); + + Ok(()) + } + + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_with_projection() + -> Result<()> { + use arrow::array::{ArrayRef, BooleanArray, Int32Array, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + // Declared full schema: col a (Int32), col b (Struct) + let declared_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, true)]); + let full_declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Struct(declared_fields), false), + ])); + + // Projected schema for column "b" (projection = [1]) + let projected_schema = Arc::new(full_declared_schema.project(&[1])?); + + // Runtime batch has stricter struct + let source_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, false)]); + let source_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Struct(source_fields.clone()), false), + ])); + + let struct_array: ArrayRef = Arc::new(StructArray::new( + source_fields, + vec![Arc::new(BooleanArray::from(vec![true, false]))], + None, + )); + let stricter_batch = RecordBatch::try_new( + source_schema, + vec![Arc::new(Int32Array::from(vec![10, 20])), struct_array], + )?; + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&projected_schema), + Some(vec![1]), + )?; + + assert_eq!(stream.schema(), projected_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), projected_schema); + assert_eq!(emitted_batch.num_columns(), 1); + + let struct_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(struct_col.fields()[0].is_nullable()); + let bool_child = struct_col + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(bool_child.value(0)); + assert!(!bool_child.value(1)); + + Ok(()) + } + + /// Regression for the Union reconstruction path at the `MemoryStream` + /// producer boundary: a declared nullable Union child vs a stricter + /// non-nullable runtime child. + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_union() -> Result<()> + { + use arrow::array::{Array, ArrayRef, Float64Array, Int32Array, UnionArray}; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{DataType, Field, Schema, UnionFields, UnionMode}; + use futures::StreamExt; + + let declared_union_fields = UnionFields::try_new( + vec![0_i8, 1], + vec![ + Field::new("i", DataType::Int32, true), + Field::new("f", DataType::Float64, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(declared_union_fields, UnionMode::Dense), + false, + )])); + + let source_union_fields = UnionFields::try_new( + vec![0_i8, 1], + vec![ + Field::new("i", DataType::Int32, false), + Field::new("f", DataType::Float64, false), + ], + )?; + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let type_ids = ScalarBuffer::from(vec![0_i8, 1, 0]); + let offsets = ScalarBuffer::from(vec![0_i32, 0, 1]); + let union_array: ArrayRef = Arc::new(UnionArray::try_new( + source_union_fields, + type_ids, + Some(offsets), + vec![ + Arc::new(Int32Array::from(vec![10, 20])), + Arc::new(Float64Array::from(vec![1.5])), + ], + )?); + let stricter_batch = RecordBatch::try_new(source_schema, vec![union_array])?; + + assert!(declared_schema.contains(stricter_batch.schema().as_ref())); + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), stream.schema()); + assert_eq!(emitted_batch.schema(), declared_schema); + + let union_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(union_col.len(), 3); + assert_eq!(union_col.type_id(0), 0); + assert_eq!(union_col.type_id(1), 1); + assert_eq!(union_col.type_id(2), 0); + let i_child = union_col + .child(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(i_child.values(), &[10, 20]); + + Ok(()) + } + + /// Regression for a contained `Map<.., Struct>` whose runtime nested field + /// is non-nullable while the declared nested field is nullable. + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_map_of_struct() + -> Result<()> { + use arrow::array::{ + Array, ArrayRef, Int32Array, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + fn map_field(value_child_nullable: bool) -> Field { + let value_struct = DataType::Struct(Fields::from(vec![Field::new( + "v", + DataType::Int32, + value_child_nullable, + )])); + let entries = Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", value_struct, true), + ])), + false, + ); + Field::new("m", DataType::Map(Arc::new(entries), false), true) + } + + let declared_schema = Arc::new(Schema::new(vec![map_field(true)])); + let source_schema = Arc::new(Schema::new(vec![map_field(false)])); + + let value_fields = Fields::from(vec![Field::new("v", DataType::Int32, false)]); + let values_struct = StructArray::new( + value_fields, + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + None, + ); + let entries = StructArray::new( + Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", values_struct.data_type().clone(), true), + ]), + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef, + Arc::new(values_struct) as ArrayRef, + ], + None, + ); + let DataType::Map(source_entries_field, _) = source_schema.field(0).data_type() + else { + unreachable!("map field") + }; + let map_array: ArrayRef = Arc::new(MapArray::try_new( + Arc::clone(source_entries_field), + OffsetBuffer::new(vec![0, 2, 3].into()), + entries, + None, + false, + )?); + let stricter_batch = RecordBatch::try_new(source_schema, vec![map_array])?; + + // The stricter batch is accepted by `MemTable::try_new`-style checks. + assert!(declared_schema.contains(stricter_batch.schema().as_ref())); + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), stream.schema()); + assert_eq!(emitted_batch.schema(), declared_schema); + + let map_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(map_col.len(), 2); + let values = map_col + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert!(values.fields()[0].is_nullable()); + let ints = values + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ints.values(), &[1, 2, 3]); + + Ok(()) + } } From 15408533b208c196e0534ad7a85615738fbda0c7 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Sat, 29 Aug 2026 12:00:56 -0400 Subject: [PATCH 06/16] [branch-55] Fix panic on RightMark hash joins when propagating ordering (#24718) (#24759) --- .../physical_optimizer/enforce_sorting.rs | 34 +++++++++++++++++++ .../enforce_sorting/sort_pushdown.rs | 4 +-- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index a8162f137ed0a..76af9b0c29218 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -243,6 +243,40 @@ async fn test_remove_unnecessary_sort5() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_propagate_in_right_mark_join() -> Result<()> { + let left_schema = create_test_schema2()?; + let right_schema = create_test_schema3()?; + let left_input = memory_exec(&left_schema); + let parquet_ordering = [sort_expr("a", &right_schema)].into(); + let right_input = + parquet_exec_with_sort(right_schema.clone(), vec![parquet_ordering]); + let on = vec![( + Arc::new(Column::new_with_schema("col_a", &left_schema)?) as _, + Arc::new(Column::new_with_schema("c", &right_schema)?) as _, + )]; + let join = hash_join_exec(left_input, right_input, on, None, &JoinType::RightMark)?; + let physical_plan = sort_exec([sort_expr("a", &join.schema())].into(), join); + + let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); + assert_snapshot!(test.run(), @r" + Input Plan: + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(col_a@0, c@2)] + DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [a@0 ASC] + HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(col_a@0, c@2)] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + "); + Ok(()) +} + #[tokio::test] async fn test_hash_join_interleaved_projection_preserves_parent_sort() -> Result<()> { let left_schema = create_test_schema()?; diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 5c17ffbd1e7db..03a28e5f647fc 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -954,7 +954,7 @@ fn handle_custom_pushdown( } // For hash join we only maintain the input order for the right child -// for join type: Inner, Right, RightSemi, RightAnti +// for join type: Inner, Right, RightSemi, RightAnti, RightMark fn handle_hash_join( plan: &HashJoinExec, parent_required: OrderingRequirements, @@ -1045,7 +1045,7 @@ fn build_join_column_index(plan: &HashJoinExec) -> Vec { .chain(map_fields(plan.right().schema(), JoinSide::Right)) .collect::>() } - JoinType::RightSemi | JoinType::RightAnti => { + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { map_fields(plan.right().schema(), JoinSide::Right) } _ => unreachable!("unexpected join type: {}", plan.join_type()), From 573f81614541d828ad3c117b8d393c6d11d4a480 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 2 Sep 2026 09:40:43 -0400 Subject: [PATCH 07/16] [branch-55] fix(common): support empty struct in ScalarValue::compact and new_default (#24582) (#24876) This is a back port of #24582 into `branch-55` for inclusion in 55.1.0. Please see the original PR for details. Co-authored-by: Unik Dahal <61407386+unikdahal@users.noreply.github.com> --- datafusion/common/src/scalar/mod.rs | 137 +++++++++++++++++++++++----- 1 file changed, 115 insertions(+), 22 deletions(-) diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index cb0442392ad21..90982ea58ab13 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -1743,18 +1743,24 @@ impl ScalarValue { // Struct types DataType::Struct(fields) => { - let values = fields - .iter() - .map(|f| ScalarValue::new_default(f.data_type())) - .collect::>>()?; - Ok(ScalarValue::Struct(Arc::new(StructArray::new( - fields.clone(), - values - .into_iter() - .map(|v| v.to_array()) - .collect::>()?, - None, - )))) + if fields.is_empty() { + Ok(ScalarValue::Struct(Arc::new( + StructArray::new_empty_fields(1, None), + ))) + } else { + let values = fields + .iter() + .map(|f| ScalarValue::new_default(f.data_type())) + .collect::>>()?; + Ok(ScalarValue::Struct(Arc::new(StructArray::new( + fields.clone(), + values + .into_iter() + .map(|v| v.to_array()) + .collect::>()?, + None, + )))) + } } // Dictionary types @@ -5035,16 +5041,22 @@ impl ScalarValue { DataType::BinaryView => Arc::new(array.as_binary_view().gc()), DataType::Struct(_) => { let s = array.as_struct(); - let columns = s - .columns() - .iter() - .map(|c| ScalarValue::compact_view_buffers(Arc::clone(c))) - .collect(); - Arc::new(StructArray::new( - s.fields().clone(), - columns, - s.nulls().cloned(), - )) + if s.fields().is_empty() { + // Zero-field structs carry no child buffers to compact, so + // return the input array unchanged. + array + } else { + let columns = s + .columns() + .iter() + .map(|c| ScalarValue::compact_view_buffers(Arc::clone(c))) + .collect(); + Arc::new(StructArray::new( + s.fields().clone(), + columns, + s.nulls().cloned(), + )) + } } DataType::List(field) => gc_list!(field, i32, ListArray), DataType::LargeList(field) => gc_list!(field, i64, LargeListArray), @@ -11610,4 +11622,85 @@ mod tests { ); assert_eq!(&large_list.value(0), &expected_array); } + + #[test] + fn test_compact_empty_struct() { + // A struct scalar wraps a single-row StructArray; use a null row to also + // exercise null-buffer preservation. + let nulls = NullBuffer::from(vec![false]); + let empty_struct = Arc::new(StructArray::new_empty_fields(1, Some(nulls))); + let mut scalar = ScalarValue::Struct(empty_struct); + + // Before fix: panics inside compact_view_buffers calling StructArray::new on 0 fields + scalar.compact(); + + let ScalarValue::Struct(arr) = &scalar else { + panic!("expected Struct") + }; + assert_eq!(arr.len(), 1); + assert_eq!(arr.num_columns(), 0); + assert_eq!(arr.null_count(), 1); + assert!(arr.is_null(0)); + } + + #[test] + fn test_compact_nested_empty_struct() { + // 1. List of empty structs + let inner_struct_field = + Arc::new(Field::new("item", DataType::Struct(Fields::empty()), true)); + let inner_struct_arr = + Arc::new(StructArray::new_empty_fields(2, None)) as ArrayRef; + let list_arr = ListArray::new( + inner_struct_field, + OffsetBuffer::new(vec![0i32, 2].into()), + inner_struct_arr, + None, + ); + let mut list_scalar = ScalarValue::List(Arc::new(list_arr)); + list_scalar.compact(); + + let ScalarValue::List(res_list) = &list_scalar else { + panic!("expected List") + }; + assert_eq!(res_list.len(), 1); + assert_eq!(res_list.values().len(), 2); + + // 2. Struct containing an empty struct field + let empty_field = Arc::new(Field::new( + "empty_child", + DataType::Struct(Fields::empty()), + true, + )); + let int_field = Arc::new(Field::new("int_child", DataType::Int32, true)); + let outer_struct = StructArray::new( + Fields::from(vec![Arc::clone(&empty_field), Arc::clone(&int_field)]), + vec![ + Arc::new(StructArray::new_empty_fields(2, None)) as ArrayRef, + Arc::new(Int32Array::from(vec![10, 20])) as ArrayRef, + ], + None, + ); + let mut outer_scalar = ScalarValue::Struct(Arc::new(outer_struct)); + outer_scalar.compact(); + + let ScalarValue::Struct(res_outer) = &outer_scalar else { + panic!("expected Struct") + }; + assert_eq!(res_outer.len(), 2); + let child_empty = res_outer.column(0).as_struct(); + assert_eq!(child_empty.len(), 2); + } + + #[test] + fn test_new_default_empty_struct() { + let empty_struct_type = DataType::Struct(Fields::empty()); + let scalar = ScalarValue::new_default(&empty_struct_type).unwrap(); + + let ScalarValue::Struct(arr) = &scalar else { + panic!("expected Struct") + }; + assert_eq!(arr.len(), 1); + assert_eq!(arr.null_count(), 0); + assert_eq!(arr.num_columns(), 0); + } } From 4d524a140efc3432c88e3e391549766ca062aa87 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 2 Sep 2026 11:32:48 -0400 Subject: [PATCH 08/16] [branch-55] chore: Minor dependency updates (#24879) This is a small update to the Cargo.lock file to update - chacha20 v0.10.0 -> v0.10.2 - h2 v0.4.13 -> v0.4.19 The first was due to a yanked crate. The second was due to a security vulnerability, RUSTSEC-2026-0258. The yanked crate is a warning in cargo audit and will auto-resolve on building to the newer 0.10.2, and the security vulnerability is already updated in `main`. --- Cargo.lock | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 62c04d332b98e..c9f916386df71 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -112,7 +112,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -123,7 +123,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -1251,9 +1251,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chacha20" -version = "0.10.0" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", "cpufeatures 0.3.0", @@ -2801,7 +2801,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -2940,7 +2940,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3284,9 +3284,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.13" +version = "0.4.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "ef8e5e5a340588f4452631496976cf8636d4a7ecf600239fdc27615d2530bc16" dependencies = [ "atomic-waker", "bytes", @@ -4212,7 +4212,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -5373,7 +5373,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5845,7 +5845,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -5945,7 +5945,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -6125,7 +6125,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -7066,7 +7066,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] From cf1ccaef18b3aefbf2bd4351711e2bbdb78e7590 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 2 Sep 2026 14:41:53 -0400 Subject: [PATCH 09/16] [branch-55] Align metadata propagation through Physical and Logical casts (#23169) (#24875) This PR is a backport of #23169 onto `branch-55`. I made one update the `use` statement in `planner.rs`. Co-authored-by: Dewey Dunnington Co-authored-by: Andrew Lamb --- Cargo.lock | 1 + .../custom_data_source/custom_file_casts.rs | 5 +- datafusion/expr/src/expr_schema.rs | 167 +++++++- datafusion/functions/src/core/arrow_cast.rs | 26 +- .../functions/src/core/arrow_try_cast.rs | 24 +- .../physical-expr-adapter/src/rewrite.rs | 7 +- datafusion/physical-expr/Cargo.toml | 1 + .../physical-expr/src/expressions/cast.rs | 391 ++++++++++++++---- .../physical-expr/src/expressions/mod.rs | 2 +- .../physical-expr/src/expressions/try_cast.rs | 354 +++++++++++++++- datafusion/physical-expr/src/planner.rs | 161 ++++++-- .../proto/src/logical_plan/from_proto.rs | 6 +- .../tests/cases/roundtrip_logical_plan.rs | 26 ++ .../cast_extension_type_metadata.slt | 54 ++- 14 files changed, 1070 insertions(+), 155 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c9f916386df71..2483c2f325843 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2406,6 +2406,7 @@ name = "datafusion-physical-expr" version = "55.0.0" dependencies = [ "arrow", + "arrow-schema", "criterion", "datafusion-common", "datafusion-expr", diff --git a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs index 71addc6d1bcb0..202c0a71257e9 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs @@ -188,10 +188,11 @@ impl PhysicalExprAdapter for CustomCastsPhysicalExprAdapter { if let Some(cast) = expr.downcast_ref::() { let input_data_type = cast.expr().data_type(&self.physical_file_schema)?; - let output_data_type = cast.target_field().data_type(); + let output_field = cast.target_field(); if !cast.is_bigger_cast(&input_data_type) { return not_impl_err!( - "Unsupported CAST from {input_data_type} to {output_data_type}" + "Unsupported CAST from {input_data_type} to {}", + output_field.data_type() ); } } diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 8927fcf4d0bbe..cdd66f99ad3b7 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -31,6 +31,7 @@ use crate::{LogicalPlan, Projection, Subquery, WindowFunctionDefinition, utils}; use arrow::compute::can_cast_types; use arrow::datatypes::FieldRef; use arrow::datatypes::{DataType, Field}; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::FieldExt; use datafusion_common::{ Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference, @@ -69,18 +70,46 @@ pub trait ExprSchemable { -> Result<(DataType, bool)>; } -/// Derives the output field for a cast expression from the source field. +/// Derives the output field for a cast expression from the source and target fields. +/// +/// Metadata handling: +/// - Type-only casts (i.e., target_field == DataType::SomeDataType.into_nullable_field()) +/// propagate non extension-type metadata from the source. This is for backward compatibility +/// (casts have propagated source metadata for many if not all previous versions), recognizing +/// that the return type of `::` should have the +/// return type of `` (e.g., casting arrow.json to utf8). +/// - All other casts preserve target metadata exactly. This ensures in particular that output +/// metadata when casting to an extension type contains the extension information in the +/// output field. Callers that wish to have some mix of source and target metadata can use +/// Alias or construct an output field themselves (whose metadata will be used directly). +/// /// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL. fn cast_output_field( source_field: &FieldRef, - target_type: &DataType, + target_field: &FieldRef, force_nullable: bool, ) -> Arc { + // Check if this is a "type-only" cast (target_field == DataType::X.into_nullable_field()) + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + let metadata = if is_type_only { + // Type-only cast: propagate source metadata, stripping extension type keys + let mut meta = source_field.metadata().clone(); + meta.remove(EXTENSION_TYPE_NAME_KEY); + meta.remove(EXTENSION_TYPE_METADATA_KEY); + meta + } else { + // Explicit target field: use target metadata exactly + target_field.metadata().clone() + }; + let mut f = source_field .as_ref() .clone() - .with_data_type(target_type.clone()) - .with_metadata(source_field.metadata().clone()); + .with_data_type(target_field.data_type().clone()) + .with_metadata(metadata); if force_nullable { f = f.with_nullable(true); } @@ -462,7 +491,8 @@ impl ExprSchemable for Expr { /// - **Aliases**: Merge underlying expr metadata with alias-specific metadata, preferring the alias metadata /// - **Binary expressions**: field metadata is empty /// - **Boolean expressions**: field metadata is empty - /// - **Cast expressions**: determined by the input expression's field metadata handling + /// - **Cast expressions**: Type-only casts pass through source metadata (stripping extension + /// type keys); casts with explicit target fields use target metadata exactly /// - **Scalar functions**: Generate metadata via function's [`return_field_from_args`] method, /// with the default implementation returning empty field metadata /// - **Aggregate functions**: Generate metadata via function's [`return_field`] method, @@ -602,20 +632,16 @@ impl ExprSchemable for Expr { func.return_field_from_args(args) } // _ => Ok((self.get_type(schema)?, self.nullable(schema)?)), - Expr::Cast(Cast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), false) - }) - } + Expr::Cast(Cast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, false)), Expr::Placeholder(Placeholder { id: _, field: Some(field), }) => Ok(Arc::clone(field).renamed(&schema_name)), - Expr::TryCast(TryCast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), true) - }) - } + Expr::TryCast(TryCast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, true)), Expr::LambdaVariable(LambdaVariable { field: Some(field), .. }) => Ok(Arc::clone(field).renamed(&schema_name)), @@ -1290,4 +1316,115 @@ mod tests { assert_eq!(meta, expr.metadata(&schema).unwrap()); } + + #[test] + fn test_cast_and_try_cast_extension_type_metadata() { + use crate::expr::{Cast, TryCast}; + use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, + }; + + // Helper to build either Cast or TryCast expression + fn make_cast_expr( + expr: Expr, + target_field: FieldRef, + use_try_cast: bool, + ) -> Expr { + if use_try_cast { + Expr::TryCast(TryCast { + expr: Box::new(expr), + field: target_field, + }) + } else { + Expr::Cast(Cast { + expr: Box::new(expr), + field: target_field, + }) + } + } + + // Run the same test logic for both Cast and TryCast + for use_try_cast in [false, true] { + let cast_name = if use_try_cast { "TryCast" } else { "Cast" }; + + // Create a schema with a field that has extension type metadata + let mut source_meta = HashMap::new(); + source_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ); + source_meta.insert("custom_key".to_string(), "custom_value".to_string()); + + let source_field = Field::new("foo", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta); + + let schema = MockExprSchema::new() + .with_data_type(DataType::FixedSizeBinary(16)) + .with_metadata(FieldMetadata::from(source_field.metadata().clone())); + + // Test 1: Cast to a type without extension metadata strips extension metadata + // but preserves non-extension metadata + let cast_expr = make_cast_expr( + col("foo"), + Arc::new(Field::new("", DataType::Utf8, true)), + use_try_cast, + ); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert!( + result_field + .metadata() + .get(EXTENSION_TYPE_NAME_KEY) + .is_none(), + "{cast_name}: Extension type name should be stripped when target has no extension metadata" + ); + assert_eq!( + result_field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "{cast_name}: Non-extension metadata should be preserved" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + + // Test 2: Cast to a field with explicit metadata uses target metadata exactly + let mut target_meta = HashMap::new(); + target_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ); + target_meta.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()); + + let target_field = + Field::new("", DataType::Utf8, true).with_metadata(target_meta); + + let cast_expr = + make_cast_expr(col("foo"), Arc::new(target_field), use_try_cast); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()), + "{cast_name}: Extension type name should come from target field" + ); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()), + "{cast_name}: Extension type metadata should come from target field" + ); + assert!( + result_field.metadata().get("custom_key").is_none(), + "{cast_name}: Source metadata should NOT propagate when target has explicit metadata" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + } + } } diff --git a/datafusion/functions/src/core/arrow_cast.rs b/datafusion/functions/src/core/arrow_cast.rs index 0b67883c17c87..929a92abcca23 100644 --- a/datafusion/functions/src/core/arrow_cast.rs +++ b/datafusion/functions/src/core/arrow_cast.rs @@ -27,8 +27,8 @@ use datafusion_common::{ use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, Expr, ExprSchemable, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -161,9 +161,27 @@ impl ScalarUDFImpl for ArrowCastFunc { let [source_arg, type_arg] = take_function_args(self.name(), args)?; let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; + + // We can skip the cast only if: + // 1. The source and target types are the same + // 2. The source has no extension metadata that needs to be stripped let new_expr = if source_type == target_type { - // the argument's data type is already the correct type - source_arg + // Check if source has extension metadata + let source_field = source_arg.to_field(info.schema())?; + let has_extension_metadata = source_field + .1 + .metadata() + .contains_key("ARROW:extension:name"); + if has_extension_metadata { + // Need to create a cast to strip extension metadata + Expr::Cast(datafusion_expr::Cast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + // the argument's data type is already the correct type + source_arg + } } else { // Use an actual cast to get the correct type Expr::Cast(datafusion_expr::Cast { diff --git a/datafusion/functions/src/core/arrow_try_cast.rs b/datafusion/functions/src/core/arrow_try_cast.rs index d27b29ba5736d..0914695b60370 100644 --- a/datafusion/functions/src/core/arrow_try_cast.rs +++ b/datafusion/functions/src/core/arrow_try_cast.rs @@ -26,8 +26,8 @@ use datafusion_common::{ use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, Expr, ExprSchemable, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -134,8 +134,26 @@ impl ScalarUDFImpl for ArrowTryCastFunc { let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; + + // We can skip the cast only if: + // 1. The source and target types are the same + // 2. The source has no extension metadata that needs to be stripped let new_expr = if source_type == target_type { - source_arg + // Check if source has extension metadata + let source_field = source_arg.to_field(info.schema())?; + let has_extension_metadata = source_field + .1 + .metadata() + .contains_key("ARROW:extension:name"); + if has_extension_metadata { + // Need to create a try_cast to strip extension metadata + Expr::TryCast(datafusion_expr::TryCast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + source_arg + } } else { Expr::TryCast(datafusion_expr::TryCast { expr: Box::new(source_arg), diff --git a/datafusion/physical-expr-adapter/src/rewrite.rs b/datafusion/physical-expr-adapter/src/rewrite.rs index 7345a587ee6a4..3d31feca96031 100644 --- a/datafusion/physical-expr-adapter/src/rewrite.rs +++ b/datafusion/physical-expr-adapter/src/rewrite.rs @@ -319,6 +319,9 @@ mod tests { assert_eq!(source.name(), "__datafusion_file_row_index"); assert_eq!(source.index(), 2); + // The row index column is at index 2, beyond the user-visible schema. + // When the source column lookup fails, the field name is empty. The + // correct field name would be provided by a parent projection/alias. let input_schema = Schema::new(vec![ Field::new("value", DataType::Int64, true), Field::new("__datafusion_file_row_index", DataType::Int64, false) @@ -328,9 +331,11 @@ mod tests { )])), ]); let return_field = expr.return_field(&input_schema)?; - assert_eq!(return_field.name(), "file_row_index"); + // Field name is empty because column index 2 is beyond the schema + assert_eq!(return_field.name(), ""); assert_eq!(return_field.data_type(), &DataType::Int64); assert!(return_field.is_nullable()); + // Exact target field does not preserve source metadata assert!(return_field.metadata().is_empty()); Ok(()) } diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 65ef2a3ceb216..0588a777230fb 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -51,6 +51,7 @@ proto = [ [dependencies] arrow = { workspace = true } +arrow-schema = { workspace = true } datafusion-common = { workspace = true } datafusion-expr = { workspace = true } datafusion-expr-common = { workspace = true } diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index dbb91e365af90..d9ee0798433c1 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,8 +23,9 @@ use std::sync::Arc; use crate::physical_expr::PhysicalExpr; use arrow::compute::{CastOptions, can_cast_types}; -use arrow::datatypes::{DataType, DataType::*, FieldRef, Schema}; +use arrow::datatypes::{DataType, DataType::*, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::nested_struct::{ @@ -59,8 +61,18 @@ fn can_cast_named_struct_types(source: &DataType, target: &DataType) -> bool { pub struct CastExpr { /// The expression to cast pub expr: Arc, - /// Field metadata describing the desired output after casting + /// The target field. + /// + /// For a type-only cast (see [`CastExpr::new`]) this is a field synthesized + /// from the target data type alone and only its data type is meaningful. + /// For a cast built from an explicit field (see + /// [`CastExpr::new_with_target_field`]) its metadata and nullability are + /// applied to the output field as-is. target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata and + /// nullability describe the output field exactly. + explicit_target: bool, /// Cast options cast_options: CastOptions<'static>, } @@ -68,8 +80,12 @@ pub struct CastExpr { // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for CastExpr { fn eq(&self, other: &Self) -> bool { + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. self.expr.eq(&other.expr) - && self.target_field.eq(&other.target_field) + && self.cast_type().eq(other.cast_type()) + && self.target_metadata().eq(&other.target_metadata()) + && self.target_nullable().eq(&other.target_nullable()) && self.cast_options.eq(&other.cast_options) } } @@ -77,7 +93,17 @@ impl PartialEq for CastExpr { impl Hash for CastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.target_field.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } + self.target_nullable().hash(state); self.cast_options.hash(state); } } @@ -85,39 +111,39 @@ impl Hash for CastExpr { impl CastExpr { /// Create a new `CastExpr` using only a `DataType`. /// - /// This constructor is provided for compatibility with existing call sites - /// that only know the target type. It synthesizes a ``Field`` with the - /// given type (**nullable by default**) and no name metadata. Callers that - /// already have a `FieldRef` (for example, coming from schema inference or a - /// resolved column) should prefer [`CastExpr::new_with_target_field`], which - /// preserves the field's name, nullability, and other metadata. In other - /// words: + /// This constructor creates a type-only cast where metadata and nullability + /// are passed through from the source expression (with extension type keys + /// stripped from metadata). This is the most common use case when you only + /// need to change the data type. /// - /// * use `new()` when only a `DataType` is available and you want the legacy - /// semantics of a type-only cast - /// * use `new_with_target_field()` when you need explicit field - /// metadata/name/nullability preserved + /// For explicit control over the output field's metadata and nullability, + /// use [`CastExpr::new_with_target_field`] or the individual builder methods. pub fn new( expr: Arc, cast_type: DataType, cast_options: Option>, ) -> Self { - Self::new_with_target_field( + Self { expr, - cast_type.into_nullable_field_ref(), - cast_options, - ) + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), + } } /// Create a new `CastExpr` with an explicit target `FieldRef`. /// - /// The provided `target_field` is used verbatim for the expression's - /// return schema, so the field's name, nullability, and other metadata are - /// preserved. This is the preferred constructor when the caller already - /// has field information (for example, during logical-to-physical planning). + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// - The field's nullability is preserved + /// + /// This is the preferred constructor when the caller has explicit field + /// information that should be used exactly (for example, during schema + /// enforcement or adapter layers). /// - /// See [`CastExpr::new`] for the compatibility constructor that only accepts - /// a `DataType`. + /// See [`CastExpr::new`] for type-only casts where source metadata should + /// pass through. pub fn new_with_target_field( expr: Arc, target_field: FieldRef, @@ -126,6 +152,7 @@ impl CastExpr { Self { expr, target_field, + explicit_target: true, cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), } } @@ -140,7 +167,30 @@ impl CastExpr { self.target_field.data_type() } - /// Field metadata describing the output column after casting. + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// Explicit nullability for the output field, or `None` to pass through source nullability. + pub fn target_nullable(&self) -> Option { + self.explicit_target + .then(|| self.target_field.is_nullable()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. Note that the returned + /// field may not match what `return_field()` returns when evaluated against + /// a schema, since `return_field()` may incorporate source field information. + /// + /// Prefer [`cast_type()`], [`target_metadata()`], and [`target_nullable()`] + /// for direct access to the individual components. + /// + /// [`cast_type()`]: CastExpr::cast_type + /// [`target_metadata()`]: CastExpr::target_metadata + /// [`target_nullable()`]: CastExpr::target_nullable pub fn target_field(&self) -> &FieldRef { &self.target_field } @@ -150,19 +200,53 @@ impl CastExpr { &self.cast_options } + /// Whether this cast has explicit metadata (vs pass-through from source). + pub fn has_explicit_metadata(&self) -> bool { + self.explicit_target + } + + /// Whether this cast has explicit nullability (vs pass-through from source). + pub fn has_explicit_nullability(&self) -> bool { + self.explicit_target + } + fn resolved_target_field(&self, input_schema: &Schema) -> Result { - if is_default_target_field(&self.target_field) { - self.expr.return_field(input_schema).map(|field| { - Arc::new( - field - .as_ref() - .clone() - .with_data_type(self.cast_type().clone()), + // Try to get the source field for the name. If the target field is + // explicit, we can fall back to an empty name if the source lookup fails + // (e.g., for virtual row-index columns appended at scan time). + let source_result = self.expr.return_field(input_schema); + + if self.explicit_target { + // Metadata and nullability come from the target field verbatim + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new( + name, + self.cast_type().clone(), + self.target_field.is_nullable(), ) - }) - } else { - Ok(Arc::clone(&self.target_field)) + .with_metadata(self.target_field.metadata().clone()), + )); } + + // Type-only cast: pass through the source metadata and nullability, + // stripping extension type keys (the cast is to a plain storage type). + source_result.map(|source_field| { + let mut metadata = source_field.metadata().clone(); + metadata.remove(EXTENSION_TYPE_NAME_KEY); + metadata.remove(EXTENSION_TYPE_METADATA_KEY); + + Arc::new( + source_field + .as_ref() + .clone() + .with_data_type(self.cast_type().clone()) + .with_metadata(metadata), + ) + }) } /// Check if casting from the specified source type to the target type is a @@ -191,12 +275,6 @@ impl CastExpr { } } -fn is_default_target_field(target_field: &FieldRef) -> bool { - target_field.name().is_empty() - && target_field.is_nullable() - && target_field.metadata().is_empty() -} - pub(crate) fn is_order_preserving_cast_family( source_type: &DataType, target_type: &DataType, @@ -267,11 +345,12 @@ impl PhysicalExpr for CastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(CastExpr::new_with_target_field( - Arc::clone(&children[0]), - Arc::clone(&self.target_field), - Some(self.cast_options.clone()), - ))) + Ok(Arc::new(CastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + cast_options: self.cast_options.clone(), + })) } fn evaluate_bounds(&self, children: &[&Interval]) -> Result { @@ -372,32 +451,61 @@ pub fn cast_with_options( cast_type: DataType, cast_options: Option>, ) -> Result> { - cast_with_target_field( - expr, - input_schema, - cast_type.into_nullable_field_ref(), - cast_options, - ) + let expr_type = expr.data_type(input_schema)?; + + // If the types match, no cast is needed for a type-only cast + if expr_type == cast_type { + return Ok(Arc::clone(&expr)); + } + + let can_build_cast = if requires_nested_struct_cast(&expr_type, &cast_type) { + can_cast_named_struct_types(&expr_type, &cast_type) + } else { + can_cast_types(&expr_type, &cast_type) + }; + + if !can_build_cast { + return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); + } + + Ok(Arc::new(CastExpr::new(expr, cast_type, cast_options))) } /// Return a PhysicalExpression representing `expr` casted to `target_field`, /// preserving any explicit field semantics such as name, nullability, and /// metadata. /// -/// If the input expression already has the same data type, this helper still -/// preserves an explicit `target_field` by constructing a field-aware -/// [`CastExpr`]. Only the default synthesized field created by the legacy -/// type-only API is elided back to the original child expression. +/// If the input expression already has the same data type and the target field +/// has no explicit metadata or nullability constraints, the original expression +/// is returned unchanged. pub fn cast_with_target_field( expr: Arc, input_schema: &Schema, - target_field: FieldRef, + target_field: &FieldRef, cast_options: Option>, ) -> Result> { let expr_type = expr.data_type(input_schema)?; let cast_type = target_field.data_type(); - if expr_type == *cast_type && is_default_target_field(&target_field) { - return Ok(Arc::clone(&expr)); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata or nullability constraints). This is the field created by + // `into_nullable_field_ref()` when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + // For same-type casts, we can skip creating a CastExpr only if: + // 1. The target is type-only (no explicit metadata) + // 2. The source has no extension metadata that needs to be stripped + // Otherwise we need the CastExpr to strip extension metadata from the source. + if expr_type == *cast_type && is_type_only { + let source_field = expr.return_field(input_schema)?; + let has_extension_metadata = source_field + .metadata() + .contains_key(EXTENSION_TYPE_NAME_KEY); + if !has_extension_metadata { + return Ok(Arc::clone(&expr)); + } } let can_build_cast = if requires_nested_struct_cast(&expr_type, cast_type) { @@ -415,11 +523,22 @@ pub fn cast_with_target_field( return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); } - Ok(Arc::new(CastExpr::new_with_target_field( - expr, - target_field, - cast_options, - ))) + // For type-only casts, use CastExpr::new which preserves source metadata/nullability. + // For explicit target fields, use new_with_target_field which applies the target's + // extension metadata and nullability. + if is_type_only { + Ok(Arc::new(CastExpr::new( + expr, + cast_type.clone(), + cast_options, + ))) + } else { + Ok(Arc::new(CastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + cast_options, + ))) + } } /// Return a PhysicalExpression representing `expr` casted to @@ -990,26 +1109,32 @@ mod tests { #[test] fn field_aware_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); for (child_nullable, target_nullable) in [(true, false), (false, true)] { let schema = Schema::new(vec![Field::new("a", Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", Int64, target_nullable) + .with_metadata(metadata.clone()), + ); let expr = CastExpr::new_with_target_field( col("a", &schema)?, - Arc::new( - Field::new("cast_target", Int64, target_nullable) - .with_metadata(metadata.clone()), - ), + Arc::clone(&target_field), None, ); let field = expr.return_field(&schema)?; - assert_eq!(field.name(), "cast_target"); + // Field name comes from source + assert_eq!(field.name(), "a"); assert_eq!(field.data_type(), &Int64); + // Nullability comes from target assert_eq!(field.is_nullable(), target_nullable); + // Target metadata should be preserved exactly assert_eq!( - field.metadata().get("target_meta").map(String::as_str), - Some("1") + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" ); assert_eq!(expr.nullable(&schema)?, child_nullable || target_nullable); } @@ -1017,6 +1142,38 @@ mod tests { Ok(()) } + #[test] + fn target_field_accessor_returns_the_constructed_field() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", Int32, true)]); + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + let target_field = + Arc::new(Field::new("cast_target", Int64, false).with_metadata(metadata)); + + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + // The field is returned verbatim, including its name. + assert_eq!(expr.target_field(), &target_field); + assert_eq!(expr.cast_type(), &Int64); + assert_eq!(expr.target_metadata(), Some(target_field.metadata())); + assert_eq!(expr.target_nullable(), Some(false)); + assert!(expr.has_explicit_metadata()); + assert!(expr.has_explicit_nullability()); + + // A type-only cast reports no explicit target. + let type_only = CastExpr::new(col("a", &schema)?, Int64, None); + assert_eq!(type_only.cast_type(), &Int64); + assert_eq!(type_only.target_metadata(), None); + assert_eq!(type_only.target_nullable(), None); + assert!(!type_only.has_explicit_metadata()); + assert!(!type_only.has_explicit_nullability()); + + Ok(()) + } + #[test] fn type_only_cast_preserves_legacy_field_name_and_nullability() -> Result<()> { let schema = Schema::new(vec![Field::new("a", Int32, false)]); @@ -1170,7 +1327,8 @@ mod tests { let literal = Arc::new(crate::expressions::Literal::new(ScalarValue::Struct( Arc::new(scalar_struct), ))); - let expr = CastExpr::new_with_target_field(literal, Arc::new(target_field), None); + let target_field = Arc::new(target_field); + let expr = CastExpr::new_with_target_field(literal, target_field, None); let batch = RecordBatch::new_empty(schema); let result = expr.evaluate(&batch)?; @@ -1216,6 +1374,93 @@ mod tests { Ok(()) } + #[test] + fn type_only_cast_strips_extension_metadata() -> Result<()> { + // When using type-only cast (new()), extension metadata from source should NOT propagate + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let expr = CastExpr::new(col("a", &schema)?, Utf8, None); + + let field = expr.return_field(&schema)?; + assert!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "Type-only cast should strip extension type name from source" + ); + assert_eq!( + field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "Type-only cast should preserve non-extension metadata" + ); + + Ok(()) + } + + #[test] + fn field_aware_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", Utf8, true).with_metadata(target_meta)); + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware cast should preserve target's non-extension metadata" + ); + + Ok(()) + } + #[test] fn test_check_bigger_cast_precision_loss() { use DataType::*; diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 035dd5d5072b0..f5784a10a38e8 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -58,7 +58,7 @@ pub use literal::{Literal, lit}; pub use negative::{NegativeExpr, negative}; pub use no_op::NoOp; pub use not::{NotExpr, not}; -pub use try_cast::{TryCastExpr, try_cast}; +pub use try_cast::{TryCastExpr, try_cast, try_cast_with_target_field}; pub use unknown_column::UnKnownColumn; pub(crate) use cast::cast_with_target_field; diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index 65b953fd181b7..c054026724fb1 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,40 +23,94 @@ use std::sync::Arc; use crate::PhysicalExpr; use arrow::compute; use arrow::compute::CastOptions; -use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use compute::can_cast_types; +use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::{Result, not_impl_err}; use datafusion_expr::ColumnarValue; /// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast -#[derive(Debug, Eq)] +#[derive(Debug, Clone, Eq)] pub struct TryCastExpr { /// The expression to cast expr: Arc, - /// The data type to cast to - cast_type: DataType, + /// The target field. + /// + /// For a type-only cast (see [`TryCastExpr::new`]) this is a field + /// synthesized from the target data type alone and only its data type is + /// meaningful. For a cast built from an explicit field (see + /// [`TryCastExpr::new_with_target_field`]) its metadata is applied to the + /// output field as-is. + target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata + /// describes the output field exactly. + explicit_target: bool, } // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for TryCastExpr { fn eq(&self, other: &Self) -> bool { - self.expr.eq(&other.expr) && self.cast_type == other.cast_type + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. + self.expr.eq(&other.expr) + && self.cast_type() == other.cast_type() + && self.target_metadata() == other.target_metadata() } } impl Hash for TryCastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.cast_type.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } } } impl TryCastExpr { - /// Create a new CastExpr + /// Create a new `TryCastExpr` using only a `DataType`. + /// + /// This constructor creates a type-only cast where metadata is passed through + /// from the source expression (with extension type keys stripped). + /// TRY_CAST results are always nullable since failed casts return NULL. pub fn new(expr: Arc, cast_type: DataType) -> Self { - Self { expr, cast_type } + Self { + expr, + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + } + } + + /// Create a new `TryCastExpr` with an explicit target `FieldRef`. + /// + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// + /// TRY_CAST results are always nullable since failed casts return NULL. + /// + /// See [`TryCastExpr::new`] for type-only casts where source metadata should + /// pass through. + pub fn new_with_target_field( + expr: Arc, + target_field: FieldRef, + ) -> Self { + Self { + expr, + target_field, + explicit_target: true, + } } /// The expression to cast @@ -65,19 +120,33 @@ impl TryCastExpr { /// The data type to cast to pub fn cast_type(&self) -> &DataType { - &self.cast_type + self.target_field.data_type() + } + + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. TRY_CAST results are + /// always nullable regardless of the target field's nullability. + pub fn target_field(&self) -> &FieldRef { + &self.target_field } } impl fmt::Display for TryCastExpr { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type) + write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type()) } } impl PhysicalExpr for TryCastExpr { fn data_type(&self, _input_schema: &Schema) -> Result { - Ok(self.cast_type.clone()) + Ok(self.cast_type().clone()) } fn nullable(&self, _input_schema: &Schema) -> Result { @@ -90,14 +159,41 @@ impl PhysicalExpr for TryCastExpr { safe: true, format_options: DEFAULT_FORMAT_OPTIONS, }; - value.cast_to(&self.cast_type, Some(&options)) + value.cast_to(self.cast_type(), Some(&options)) } fn return_field(&self, input_schema: &Schema) -> Result { - self.expr - .return_field(input_schema) - .map(|f| f.as_ref().clone().with_data_type(self.cast_type.clone())) - .map(Arc::new) + // If metadata is explicit, we can build the field without source + // (though we still try to get source for the name) + let source_result = self.expr.return_field(input_schema); + + if let Some(metadata) = self.target_metadata() { + // Explicit metadata: use it exactly, TRY_CAST is always nullable + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new(name, self.cast_type().clone(), true) + .with_metadata(metadata.clone()), + )); + } + + // Pass-through metadata from source (stripping extension keys) + source_result.map(|source_field| { + let mut metadata = source_field.metadata().clone(); + metadata.remove(EXTENSION_TYPE_NAME_KEY); + metadata.remove(EXTENSION_TYPE_METADATA_KEY); + + Arc::new( + source_field + .as_ref() + .clone() + .with_data_type(self.cast_type().clone()) + .with_nullable(true) // TRY_CAST is always nullable + .with_metadata(metadata), + ) + }) } fn children(&self) -> Vec<&Arc> { @@ -108,16 +204,17 @@ impl PhysicalExpr for TryCastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(TryCastExpr::new( - Arc::clone(&children[0]), - self.cast_type.clone(), - ))) + Ok(Arc::new(TryCastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + })) } fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "TRY_CAST(")?; self.expr.fmt_sql(f)?; - write!(f, " AS {:?})", self.cast_type) + write!(f, " AS {:?})", self.cast_type()) } #[cfg(feature = "proto")] @@ -190,6 +287,60 @@ pub fn try_cast( } } +/// Return a PhysicalExpression representing `expr` casted to `target_field`, +/// preserving any explicit field semantics such as metadata. +/// +/// TRY_CAST results are always nullable since failed casts return NULL. +/// +/// If the input expression already has the same data type, the target field +/// has no explicit metadata constraints, and the source has no extension +/// metadata to strip, the original expression is returned unchanged. +pub fn try_cast_with_target_field( + expr: Arc, + input_schema: &Schema, + target_field: &FieldRef, +) -> Result> { + let expr_type = expr.data_type(input_schema)?; + let cast_type = target_field.data_type(); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata constraints). This is the field created by `into_nullable_field_ref()` + // when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + // For same-type casts, we can skip creating a TryCastExpr only if: + // 1. The target is type-only (no explicit metadata) + // 2. The source has no extension metadata that needs to be stripped + // Otherwise we need the TryCastExpr to strip extension metadata from the source. + if expr_type == *cast_type && is_type_only { + let source_field = expr.return_field(input_schema)?; + let has_extension_metadata = source_field + .metadata() + .contains_key(EXTENSION_TYPE_NAME_KEY); + if !has_extension_metadata { + return Ok(Arc::clone(&expr)); + } + } + + if !can_cast_types(&expr_type, cast_type) { + return not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}"); + } + + // For type-only casts, use TryCastExpr::new which preserves source metadata. + // For explicit target fields, use new_with_target_field which applies the target's + // metadata exactly. + if is_type_only { + Ok(Arc::new(TryCastExpr::new(expr, cast_type.clone()))) + } else { + Ok(Arc::new(TryCastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + ))) + } +} + #[cfg(test)] mod tests { use super::*; @@ -642,6 +793,167 @@ mod tests { Ok(()) } + + #[test] + fn field_aware_try_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", DataType::Utf8, true).with_metadata(target_meta)); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware try_cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware try_cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware try_cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware try_cast should preserve target's non-extension metadata" + ); + // TRY_CAST is always nullable + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn field_aware_try_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). + // TRY_CAST is always nullable regardless of target field's nullability. + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + + for child_nullable in [true, false] { + let schema = + Schema::new(vec![Field::new("a", DataType::Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", DataType::Int64, false) // target says non-nullable + .with_metadata(metadata.clone()), + ); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + // Field name comes from source + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + // TRY_CAST is ALWAYS nullable (ignores target field's nullability) + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + // Target metadata should be preserved exactly + assert_eq!( + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" + ); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + } + + Ok(()) + } + + #[test] + fn type_only_try_cast_strips_extension_keys() -> Result<()> { + // Type-only cast should strip extension keys but preserve other source metadata + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.extension".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "ext_meta".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false).with_metadata(source_meta), + ]); + + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + let field = expr.return_field(&schema)?; + + // Extension keys should be stripped + assert!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "Type-only try_cast should strip extension type name" + ); + assert!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY).is_none(), + "Type-only try_cast should strip extension type metadata" + ); + // Non-extension metadata should pass through + assert_eq!( + field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "Type-only try_cast should preserve non-extension metadata" + ); + // Field name preserved, type changed, always nullable + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn type_only_try_cast_is_always_nullable() -> Result<()> { + // TRY_CAST is always nullable even when source is non-nullable + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + + let field = expr.return_field(&schema)?; + + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + + Ok(()) + } } #[cfg(all(test, feature = "proto"))] diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index f80d1b15bdc59..91c36cb09cd35 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -27,7 +27,7 @@ use crate::{ use arrow::datatypes::Schema; use datafusion_common::config::ConfigOptions; use datafusion_common::datatype::FieldExt; -use datafusion_common::metadata::{FieldMetadata, format_type_and_metadata}; +use datafusion_common::metadata::FieldMetadata; use datafusion_common::{ DFSchema, Result, ScalarValue, TableReference, ToDFSchema, exec_err, internal_datafusion_err, not_impl_err, plan_datafusion_err, plan_err, @@ -41,7 +41,7 @@ use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::var_provider::VarType; use datafusion_expr::var_provider::is_system_variables; use datafusion_expr::{ - Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, TryCast, binary_expr, lit, + Between, BinaryExpr, Expr, Like, Operator, TryCast, binary_expr, lit, }; /// [PhysicalExpr] evaluate DataFusion expressions such as `A + 1`, or `CAST(c1 @@ -387,23 +387,11 @@ pub fn create_physical_expr( Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, - Arc::clone(field), + field, None, ), Expr::TryCast(TryCast { expr, field }) => { - if !field.metadata().is_empty() { - let (_, src_field) = expr.to_field(input_dfschema)?; - return plan_err!( - "TryCast from {} to {} is not supported", - format_type_and_metadata( - src_field.data_type(), - Some(src_field.metadata()), - ), - format_type_and_metadata(field.data_type(), Some(field.metadata())) - ); - } - - expressions::try_cast( + expressions::try_cast_with_target_field( create_physical_expr( expr, input_dfschema, @@ -411,7 +399,7 @@ pub fn create_physical_expr( planning_ctx, )?, input_schema, - field.data_type().clone(), + field, ) } Expr::Not(expr) => expressions::not(create_physical_expr( @@ -744,6 +732,7 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc { mod tests { use arrow::array::{ArrayRef, BooleanArray, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field}; + use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_expr::col; use super::*; @@ -768,6 +757,14 @@ mod tests { .expect("planner should lower logical CAST to CastExpr") } + fn as_planner_try_cast( + physical: &Arc, + ) -> &expressions::TryCastExpr { + physical + .downcast_ref::() + .expect("planner should lower logical TRY_CAST to TryCastExpr") + } + #[test] fn test_create_physical_expr_scalar_input_output() -> Result<()> { let expr = col("letter").eq(lit("A")); @@ -801,9 +798,21 @@ mod tests { #[test] fn test_cast_lowering_preserves_target_field_metadata() -> Result<()> { let schema = test_cast_schema(); + + // Target field with both extension metadata and custom metadata. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( - Field::new("cast_target", DataType::Int64, true) - .with_metadata([("target_meta".to_string(), "1".to_string())].into()), + Field::new("cast_target", DataType::Int64, true).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ), + (EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()), + ("custom_target_meta".to_string(), "custom_value".to_string()), + ] + .into(), + ), ); let cast_expr = Expr::Cast(Cast::new_from_field( Box::new(col("a")), @@ -813,8 +822,35 @@ mod tests { let physical = lower_cast_expr(&cast_expr, &schema)?; let cast = as_planner_cast(&physical); - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); + // The CastExpr stores the target type and all target metadata + assert_eq!(cast.cast_type(), &DataType::Int64); + let target_metadata = cast.target_metadata().expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + assert_eq!(cast.target_nullable(), Some(true)); + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_target_meta"), + Some(&"custom_value".to_string()), + "All target metadata should propagate with exact semantics" + ); assert!(physical.nullable(&schema)?); Ok(()) @@ -840,22 +876,85 @@ mod tests { #[test] fn test_cast_lowering_preserves_same_type_field_semantics() -> Result<()> { let schema = test_cast_schema(); + + // Same-type cast with extension metadata on target. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( Field::new("same_type_cast", DataType::Int32, true).with_metadata( - [("target_meta".to_string(), "same-type".to_string())].into(), + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.opaque".to_string(), + ), + ("custom_meta".to_string(), "custom_value".to_string()), + ] + .into(), ), ); - let cast_expr = Expr::Cast(Cast::new_from_field( - Box::new(col("a")), - Arc::clone(&target_field), - )); - let physical = lower_cast_expr(&cast_expr, &schema)?; - let cast = as_planner_cast(&physical); + for use_try_cast in [false, true] { + // For error labelling + let cast_name = if use_try_cast { "TRY_CAST" } else { "CAST" }; - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); - assert!(physical.nullable(&schema)?); + let cast_expr = if use_try_cast { + Expr::TryCast(TryCast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + } else { + Expr::Cast(Cast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + }; + + let physical = lower_cast_expr(&cast_expr, &schema)?; + + // Extract common fields - both CastExpr and TryCastExpr have these + let (cast_type, target_metadata, target_nullable) = if use_try_cast { + let cast = as_planner_try_cast(&physical); + (cast.cast_type(), cast.target_metadata(), None) + } else { + let cast = as_planner_cast(&physical); + ( + cast.cast_type(), + cast.target_metadata(), + cast.target_nullable(), + ) + }; + + // Verify the physical expression stores correct metadata (same for both) + assert_eq!(cast_type, &DataType::Int32, "{cast_name}: cast_type"); + let target_metadata = target_metadata.expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: extension type name" + ); + + // Only CastExpr tracks target_nullable (TryCast is always nullable) + if !use_try_cast { + assert_eq!(target_nullable, Some(true), "{cast_name}: target_nullable"); + } + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: return_field extension type name" + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_meta"), + Some(&"custom_value".to_string()), + "{cast_name}: All target metadata should propagate with exact semantics" + ); + assert!( + physical.nullable(&schema)?, + "{cast_name}: should be nullable" + ); + } Ok(()) } diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index d4d0ea7292ffe..2830edd35760f 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -429,7 +429,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::Cast(Cast::new_from_field(expr, Arc::new(field)))) } ExprType::TryCast(cast) => { @@ -442,7 +443,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::TryCast(TryCast::new_from_field( expr, Arc::new(field), diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index a450f7a7e888f..4dc3291f31d57 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -2693,6 +2693,19 @@ fn roundtrip_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::Cast(Cast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); } #[test] @@ -2703,6 +2716,19 @@ fn roundtrip_try_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::TryCast(TryCast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); + let test_expr = Expr::TryCast(TryCast::new(Box::new(lit("not a bool")), DataType::Boolean)); diff --git a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt index 425d8ac16eaee..01a19454e9a80 100644 --- a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt +++ b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt @@ -45,5 +45,55 @@ FROM ( ---- 00010203040506070809000102030506 arrow.uuid -statement error DataFusion error: Optimizer rule 'simplify_expressions' failed[\s\S]*TryCast from FixedSizeBinary\(16\) to FixedSizeBinary\(16\)<\{"ARROW:extension:name": "arrow\.uuid"\}> is not supported -SELECT TRY_CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID); +# TRY_CAST to extension type should also preserve extension metadata +query ?T +SELECT + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + arrow_metadata( + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + 'ARROW:extension:name' + ); +---- +00010203040506070809000102030506 arrow.uuid + +# TRY_CAST to UUID from a subquery +query ?T +SELECT + TRY_CAST(raw AS UUID), + arrow_metadata(TRY_CAST(raw AS UUID), 'ARROW:extension:name') +FROM ( + VALUES ( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + ) +) AS uuids(raw); +---- +00010203040506070809000102030506 arrow.uuid + +# arrow_cast from UUID to same underlying type (FixedSizeBinary(16)) strips +# extension metadata (type-only cast semantics) +query ?T +SELECT + arrow_cast(uuid_val, 'FixedSizeBinary(16)'), + arrow_metadata(arrow_cast(uuid_val, 'FixedSizeBinary(16)'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL + +# arrow_cast to a different type strips extension metadata (type-only cast semantics) +query ?T +SELECT + arrow_cast(uuid_val, 'Binary'), + arrow_metadata(arrow_cast(uuid_val, 'Binary'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL From 66ed9b55a3b4e1f873ff64a6d59074068ee1fa9e Mon Sep 17 00:00:00 2001 From: Gene Bordegaray Date: Tue, 8 Sep 2026 06:46:46 -0500 Subject: [PATCH 10/16] [branch-55] fix: preserve projection metadata during optimization (#24670) (#24992) Backports #24670 to branch-55. Co-authored-by: Tim Saucer Co-authored-by: Claude Opus 5 (1M context) --- datafusion/physical-plan/src/projection.rs | 333 ++++++++++++++++++++- 1 file changed, 319 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index cf362cdee55d3..1e672bb8b98e7 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -82,6 +82,9 @@ pub struct ProjectionExec { metrics: ExecutionPlanMetricsSet, /// Cache holding plan properties like equivalences, output partitioning etc. cache: Arc, + /// Whether the output metadata differs from the metadata derived from the + /// projection expressions and input schema. + overrides_metadata: bool, } impl ProjectionExec { @@ -143,7 +146,7 @@ impl ProjectionExec { let expr_arc = expr.into_iter().map(Into::into).collect::>(); let projection = ProjectionExprs::from_expressions(expr_arc); let projector = projection.make_projector(&input_schema)?; - Self::try_from_projector(projector, input) + Self::try_from_projector(projector, input, false) } /// Create a projection using field and schema metadata from @@ -171,12 +174,15 @@ impl ProjectionExec { let projection = ProjectionExprs::from_expressions(expr_arc); let projector = projection .make_projector_with_schema_metadata(&input_schema, projected_schema)?; - Self::try_from_projector(projector, input) + let overrides_metadata = + Self::compute_overrides_metadata(&projector, &input_schema)?; + Self::try_from_projector(projector, input, overrides_metadata) } fn try_from_projector( projector: Projector, input: Arc, + overrides_metadata: bool, ) -> Result { // Construct a map from the input expressions to the output expression of the Projection let projection_mapping = @@ -191,6 +197,7 @@ impl ProjectionExec { input, metrics: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), + overrides_metadata, }) } @@ -231,6 +238,33 @@ impl ProjectionExec { )) } + /// Returns whether `projector`'s output metadata differs from the metadata + /// derived from its expressions and `input_schema`. + fn compute_overrides_metadata( + projector: &Projector, + input_schema: &Schema, + ) -> Result { + let output_schema = projector.output_schema(); + if input_schema.metadata() != output_schema.metadata() { + return Ok(true); + } + for (projection, output_field) in + projector.projection().iter().zip(output_schema.fields()) + { + let derived_field = projection.expr.return_field(input_schema)?; + if derived_field.metadata() != output_field.metadata() { + return Ok(true); + } + } + Ok(false) + } + + /// Returns whether this projection's output metadata differs from the + /// metadata derived when the projection was constructed. + fn overrides_metadata(&self) -> bool { + self.overrides_metadata + } + /// Collect reverse alias mapping from projection expressions. /// The result hash map is a map from aliased Column in parent to original expr. fn collect_reverse_alias( @@ -352,11 +386,16 @@ impl ExecutionPlan for ProjectionExec { metrics: ExecutionPlanMetricsSet::new(), ..Self::clone(&*self) })), - ChildrenPropertiesMode::Recompute => ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), - ) - .map(|p| Arc::new(p) as _), + ChildrenPropertiesMode::Recompute => { + let input = children.swap_remove(0); + let projector = self.projector.clone(); + let overrides_metadata = ProjectionExec::compute_overrides_metadata( + &projector, + input.schema().as_ref(), + )?; + ProjectionExec::try_from_projector(projector, input, overrides_metadata) + .map(|p| Arc::new(p) as _) + } } } @@ -946,7 +985,12 @@ pub fn remove_unnecessary_projections( if is_projection_removable(projection) { return Ok(Transformed::yes(Arc::clone(projection.input()))); } - // If it does, check if we can push it under its child(ren): + // Swapping a projection with observable metadata can change query results + // by changing the metadata visible to its child expressions. + if projection.overrides_metadata() { + return Ok(Transformed::no(plan)); + } + // Otherwise, check if we can push it under its child(ren): projection .input() .try_swapping_with_projection(projection)? @@ -958,6 +1002,7 @@ pub fn remove_unnecessary_projections( /// Compare the inputs and outputs of the projection. All expressions must be /// columns without alias, and projection does not change the order of fields. +/// The input and output schemas must also match exactly to preserve metadata. /// For example, if the input schema is `a, b`, `SELECT a, b` is removable, /// but `SELECT b, a` and `SELECT a+1, b` and `SELECT a AS c, b` are not. fn is_projection_removable(projection: &ProjectionExec) -> bool { @@ -968,6 +1013,7 @@ fn is_projection_removable(projection: &ProjectionExec) -> bool { }; col.name() == proj_expr.alias && col.index() == idx }) && exprs.len() == projection.input().schema().fields().len() + && projection.schema() == projection.input().schema() } /// Given the expression set of a projection, checks if the projection causes @@ -1001,13 +1047,17 @@ pub fn new_projections_for_columns( } /// Creates a new [`ProjectionExec`] instance with the given child plan and -/// projected expressions. +/// projected expressions, preserving the original output metadata. pub fn make_with_child( projection: &ProjectionExec, child: &Arc, ) -> Result> { - ProjectionExec::try_new(projection.expr().to_vec(), Arc::clone(child)) - .map(|e| Arc::new(e) as _) + ProjectionExec::try_new_with_schema_metadata( + projection.expr().to_vec(), + Arc::clone(child), + projection.schema().as_ref(), + ) + .map(|e| Arc::new(e) as _) } /// Returns `true` if all the expressions in the argument are `Column`s. @@ -1255,6 +1305,11 @@ pub fn update_join_filter( /// Collapse a chain of consecutive [`ProjectionExec`]s into one. Returns /// `None` if nothing could be merged. +/// +/// The projection-removal optimizer checks `outer.overrides_metadata()` before +/// reaching this helper. The unified projection also keeps `outer`'s schema, so +/// collapsing cannot lose its output metadata. Inner projections still need the +/// check below because outer expressions may observe their metadata. fn try_collapse_projection_chain( outer: &ProjectionExec, ) -> Result>> { @@ -1264,6 +1319,10 @@ fn try_collapse_projection_chain( let mut collapsed_any = false; 'outer: while let Some(inner_proj) = current_input.downcast_ref::() { + if inner_proj.overrides_metadata() { + break; + } + // Collect the column references usage in the outer projection. column_ref_map.clear(); for proj_expr in ¤t_exprs { @@ -1313,8 +1372,13 @@ fn try_collapse_projection_chain( } // To unify 3 or more sequential projections: + // Preserve the outer projection's output metadata. let unified: Arc = - Arc::new(ProjectionExec::try_new(current_exprs, current_input)?); + Arc::new(ProjectionExec::try_new_with_schema_metadata( + current_exprs, + current_input, + outer.schema().as_ref(), + )?); remove_unnecessary_projections(unified).data().map(Some) } @@ -1436,19 +1500,23 @@ mod tests { use crate::common::collect; use crate::empty::EmptyExec; + use crate::filter::FilterExec; use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; use crate::test::exec::StatisticsExec; + use arrow::array::StringArray; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::ScalarValue; use datafusion_common::stats::{ColumnStatistics, Precision, Statistics}; - use datafusion_expr::Operator; + use datafusion_expr::{Operator, ScalarUDF}; + use datafusion_functions::core::arrow_metadata::ArrowMetadataFunc; + use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::expressions::{ - BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit, + BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, is_null, lit, }; #[test] @@ -1491,6 +1559,243 @@ mod tests { Ok(()) } + fn identity_projection_with_metadata( + input: Arc, + field_metadata: HashMap, + schema_metadata: HashMap, + ) -> Result> { + let metadata_schema = Schema::new_with_metadata( + vec![Field::new("i", DataType::Int32, true).with_metadata(field_metadata)], + schema_metadata, + ); + Ok(Arc::new(ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }], + input, + &metadata_schema, + )?)) + } + + #[test] + fn test_field_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert!(optimized.downcast_ref::().is_some()); + assert_eq!(optimized.schema(), expected_schema); + Ok(()) + } + + #[test] + fn test_schema_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::new(), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert!(optimized.downcast_ref::().is_some()); + assert_eq!(optimized.schema(), expected_schema); + Ok(()) + } + + #[test] + fn test_replace_children_recomputes_metadata_override() -> Result<()> { + let field_metadata = + HashMap::from([("event_field".to_string(), "true".to_string())]); + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + field_metadata.clone(), + HashMap::new(), + )?; + assert!( + projection + .downcast_ref::() + .expect("test plan should be a ProjectionExec") + .overrides_metadata() + ); + + let replacement_schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true).with_metadata(field_metadata), + ])); + let replacement: Arc = + Arc::new(EmptyExec::new(replacement_schema)); + let replaced = projection.replace_children( + vec![replacement], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + assert!( + !replaced + .downcast_ref::() + .expect("replaced plan should be a ProjectionExec") + .overrides_metadata() + ); + Ok(()) + } + + #[test] + fn test_make_with_child_preserves_output_metadata() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + let projection = projection + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + + let rebuilt = make_with_child(projection, &test::scan_partitioned(1))?; + + assert_eq!(rebuilt.schema(), projection.schema()); + Ok(()) + } + + #[tokio::test] + async fn test_metadata_observing_parent_blocks_projection_collapse() -> Result<()> { + let inner = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + let arrow_metadata = ScalarFunctionExpr::new( + "arrow_metadata", + Arc::new(ScalarUDF::new_from_impl(ArrowMetadataFunc::new())), + vec![ + Arc::new(Column::new("i", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "event_field".to_string(), + )))), + ], + Arc::new(Field::new("arrow_metadata", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + ); + let outer: Arc = Arc::new(ProjectionExec::try_new( + [ProjectionExpr { + expr: Arc::new(arrow_metadata), + alias: "metadata".to_string(), + }], + inner, + )?); + + let outer_projection = outer + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + assert!(try_collapse_projection_chain(outer_projection)?.is_none()); + + let optimized = remove_unnecessary_projections(outer)?.data; + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("metadata expression should return Utf8"); + assert_eq!(values.value(0), "true"); + Ok(()) + } + + #[tokio::test] + async fn test_metadata_observing_filter_blocks_projection_pushdown() -> Result<()> { + let widened: Arc = Arc::new(ProjectionExec::try_new( + [ + ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "j".to_string(), + }, + ], + test::scan_partitioned(1), + )?); + let arrow_metadata = Arc::new(ScalarFunctionExpr::new( + "arrow_metadata", + Arc::new(ScalarUDF::new_from_impl(ArrowMetadataFunc::new())), + vec![ + Arc::new(Column::new("i", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "event_field".to_string(), + )))), + ], + Arc::new(Field::new("arrow_metadata", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )); + let filter: Arc = + Arc::new(FilterExec::try_new(is_null(arrow_metadata)?, widened)?); + let projection = identity_projection_with_metadata( + filter, + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + assert_eq!(optimized.schema(), expected_schema); + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 100 + ); + Ok(()) + } + + // A schema-only metadata override must block projection embedding. The filter + // rebuilds the schema from expressions and would otherwise drop this metadata. + #[tokio::test] + async fn test_schema_level_metadata_blocks_projection_embedding() -> Result<()> { + let scan = test::scan_partitioned(1); + let predicate = binary( + col("i", &scan.schema())?, + Operator::Gt, + lit(ScalarValue::Int32(Some(-1))), + &scan.schema(), + )?; + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, scan)?); + let projection = identity_projection_with_metadata( + filter, + HashMap::new(), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + // Field metadata matches, so this checks the schema-level comparison. + let projection_exec = projection + .downcast_ref::() + .expect("test plan should be a ProjectionExec"); + assert!(projection_exec.overrides_metadata()); + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert_eq!(optimized.schema(), expected_schema); + assert_eq!( + optimized.schema().metadata(), + &HashMap::from([("schema-key".to_string(), "schema-value".to_string())]) + ); + + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 100 + ); + Ok(()) + } + #[test] fn test_collect_column_indices() -> Result<()> { let expr = Arc::new(BinaryExpr::new( From 7d3835c71f30cbd3c3ae4041732267f1f453097a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 8 Sep 2026 08:29:54 -0400 Subject: [PATCH 11/16] chore: update version 55.1.0 (#25070) This PR updates the version number and changelog for the 55.1.0 release. --- Cargo.lock | 86 +++++++++---------- Cargo.toml | 78 ++++++++--------- dev/changelog/55.1.0.md | 52 +++++++++++ docs/source/download.md | 2 +- docs/source/user-guide/configs.md | 2 +- docs/source/user-guide/crate-configuration.md | 2 +- docs/source/user-guide/example-usage.md | 2 +- 7 files changed, 138 insertions(+), 86 deletions(-) create mode 100644 dev/changelog/55.1.0.md diff --git a/Cargo.lock b/Cargo.lock index 2483c2f325843..141d5b6fdddb6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1707,7 +1707,7 @@ dependencies = [ [[package]] name = "datafusion" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-schema", @@ -1781,7 +1781,7 @@ dependencies = [ [[package]] name = "datafusion-benchmarks" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -1812,7 +1812,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -1835,7 +1835,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -1859,7 +1859,7 @@ dependencies = [ [[package]] name = "datafusion-cli" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -1891,7 +1891,7 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-ipc", @@ -1920,7 +1920,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "55.0.0" +version = "55.1.0" dependencies = [ "futures", "log", @@ -1929,7 +1929,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-compression", @@ -1968,7 +1968,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-ipc", @@ -1992,7 +1992,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-avro", @@ -2010,7 +2010,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -2032,7 +2032,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -2055,7 +2055,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-schema", @@ -2090,11 +2090,11 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "55.0.0" +version = "55.1.0" [[package]] name = "datafusion-examples" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-flight", @@ -2135,7 +2135,7 @@ dependencies = [ [[package]] name = "datafusion-execution" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-buffer", @@ -2162,7 +2162,7 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-schema", @@ -2188,7 +2188,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "datafusion-common", @@ -2199,7 +2199,7 @@ dependencies = [ [[package]] name = "datafusion-ffi" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-schema", @@ -2236,7 +2236,7 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-buffer", @@ -2270,7 +2270,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "criterion", @@ -2291,7 +2291,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "criterion", @@ -2303,7 +2303,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-ord", @@ -2329,7 +2329,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -2343,7 +2343,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "criterion", @@ -2359,7 +2359,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "55.0.0" +version = "55.1.0" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2367,7 +2367,7 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "55.0.0" +version = "55.1.0" dependencies = [ "datafusion-doc", "quote", @@ -2376,7 +2376,7 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -2403,7 +2403,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-schema", @@ -2430,7 +2430,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "datafusion-common", @@ -2443,7 +2443,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "chrono", @@ -2461,7 +2461,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "datafusion-common", @@ -2483,7 +2483,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "arrow-data", @@ -2525,7 +2525,7 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -2561,7 +2561,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "datafusion-common", @@ -2573,7 +2573,7 @@ dependencies = [ [[package]] name = "datafusion-proto-models" -version = "55.0.0" +version = "55.1.0" dependencies = [ "datafusion-common", "datafusion-proto-common", @@ -2584,7 +2584,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "datafusion-common", @@ -2602,7 +2602,7 @@ dependencies = [ [[package]] name = "datafusion-session" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow-schema", "async-trait", @@ -2615,7 +2615,7 @@ dependencies = [ [[package]] name = "datafusion-spark" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "bigdecimal", @@ -2645,7 +2645,7 @@ dependencies = [ [[package]] name = "datafusion-sql" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "bigdecimal", @@ -2671,7 +2671,7 @@ dependencies = [ [[package]] name = "datafusion-sqllogictest" -version = "55.0.0" +version = "55.1.0" dependencies = [ "arrow", "async-trait", @@ -2703,7 +2703,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "55.0.0" +version = "55.1.0" dependencies = [ "async-recursion", "async-trait", @@ -2724,7 +2724,7 @@ dependencies = [ [[package]] name = "datafusion-wasmtest" -version = "55.0.0" +version = "55.1.0" dependencies = [ "bytes", "chrono", diff --git a/Cargo.toml b/Cargo.toml index 4526fa0c58934..7060b1df93f26 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -80,7 +80,7 @@ repository = "https://github.com/apache/datafusion" # Define Minimum Supported Rust Version (MSRV) rust-version = "1.94.0" # Define DataFusion version -version = "55.0.0" +version = "55.1.0" [workspace.dependencies] # We turn off default-features for some dependencies here so the workspaces which inherit them can @@ -121,44 +121,44 @@ chrono = { version = "0.4.45", default-features = false } criterion = "0.8" ctor = "1.0.7" dashmap = "6.2.1" -datafusion = { path = "datafusion/core", version = "55.0.0", default-features = false } -datafusion-catalog = { path = "datafusion/catalog", version = "55.0.0" } -datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "55.0.0" } -datafusion-common = { path = "datafusion/common", version = "55.0.0", default-features = false } -datafusion-common-runtime = { path = "datafusion/common-runtime", version = "55.0.0" } -datafusion-datasource = { path = "datafusion/datasource", version = "55.0.0", default-features = false } -datafusion-datasource-arrow = { path = "datafusion/datasource-arrow", version = "55.0.0", default-features = false } -datafusion-datasource-avro = { path = "datafusion/datasource-avro", version = "55.0.0", default-features = false } -datafusion-datasource-csv = { path = "datafusion/datasource-csv", version = "55.0.0", default-features = false } -datafusion-datasource-json = { path = "datafusion/datasource-json", version = "55.0.0", default-features = false } -datafusion-datasource-parquet = { path = "datafusion/datasource-parquet", version = "55.0.0", default-features = false } -datafusion-doc = { path = "datafusion/doc", version = "55.0.0" } -datafusion-execution = { path = "datafusion/execution", version = "55.0.0", default-features = false } -datafusion-expr = { path = "datafusion/expr", version = "55.0.0", default-features = false } -datafusion-expr-common = { path = "datafusion/expr-common", version = "55.0.0" } -datafusion-ffi = { path = "datafusion/ffi", version = "55.0.0" } -datafusion-functions = { path = "datafusion/functions", version = "55.0.0" } -datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "55.0.0" } -datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "55.0.0" } -datafusion-functions-nested = { path = "datafusion/functions-nested", version = "55.0.0", default-features = false } -datafusion-functions-table = { path = "datafusion/functions-table", version = "55.0.0" } -datafusion-functions-window = { path = "datafusion/functions-window", version = "55.0.0" } -datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "55.0.0" } -datafusion-macros = { path = "datafusion/macros", version = "55.0.0" } -datafusion-optimizer = { path = "datafusion/optimizer", version = "55.0.0", default-features = false } -datafusion-physical-expr = { path = "datafusion/physical-expr", version = "55.0.0", default-features = false } -datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter", version = "55.0.0", default-features = false } -datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "55.0.0", default-features = false } -datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "55.0.0" } -datafusion-physical-plan = { path = "datafusion/physical-plan", version = "55.0.0" } -datafusion-proto = { path = "datafusion/proto", version = "55.0.0", default-features = false } -datafusion-proto-common = { path = "datafusion/proto-common", version = "55.0.0" } -datafusion-proto-models = { path = "datafusion/proto-models", version = "55.0.0" } -datafusion-pruning = { path = "datafusion/pruning", version = "55.0.0" } -datafusion-session = { path = "datafusion/session", version = "55.0.0" } -datafusion-spark = { path = "datafusion/spark", version = "55.0.0" } -datafusion-sql = { path = "datafusion/sql", version = "55.0.0" } -datafusion-substrait = { path = "datafusion/substrait", version = "55.0.0" } +datafusion = { path = "datafusion/core", version = "55.1.0", default-features = false } +datafusion-catalog = { path = "datafusion/catalog", version = "55.1.0" } +datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "55.1.0" } +datafusion-common = { path = "datafusion/common", version = "55.1.0", default-features = false } +datafusion-common-runtime = { path = "datafusion/common-runtime", version = "55.1.0" } +datafusion-datasource = { path = "datafusion/datasource", version = "55.1.0", default-features = false } +datafusion-datasource-arrow = { path = "datafusion/datasource-arrow", version = "55.1.0", default-features = false } +datafusion-datasource-avro = { path = "datafusion/datasource-avro", version = "55.1.0", default-features = false } +datafusion-datasource-csv = { path = "datafusion/datasource-csv", version = "55.1.0", default-features = false } +datafusion-datasource-json = { path = "datafusion/datasource-json", version = "55.1.0", default-features = false } +datafusion-datasource-parquet = { path = "datafusion/datasource-parquet", version = "55.1.0", default-features = false } +datafusion-doc = { path = "datafusion/doc", version = "55.1.0" } +datafusion-execution = { path = "datafusion/execution", version = "55.1.0", default-features = false } +datafusion-expr = { path = "datafusion/expr", version = "55.1.0", default-features = false } +datafusion-expr-common = { path = "datafusion/expr-common", version = "55.1.0" } +datafusion-ffi = { path = "datafusion/ffi", version = "55.1.0" } +datafusion-functions = { path = "datafusion/functions", version = "55.1.0" } +datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "55.1.0" } +datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "55.1.0" } +datafusion-functions-nested = { path = "datafusion/functions-nested", version = "55.1.0", default-features = false } +datafusion-functions-table = { path = "datafusion/functions-table", version = "55.1.0" } +datafusion-functions-window = { path = "datafusion/functions-window", version = "55.1.0" } +datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "55.1.0" } +datafusion-macros = { path = "datafusion/macros", version = "55.1.0" } +datafusion-optimizer = { path = "datafusion/optimizer", version = "55.1.0", default-features = false } +datafusion-physical-expr = { path = "datafusion/physical-expr", version = "55.1.0", default-features = false } +datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter", version = "55.1.0", default-features = false } +datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "55.1.0", default-features = false } +datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "55.1.0" } +datafusion-physical-plan = { path = "datafusion/physical-plan", version = "55.1.0" } +datafusion-proto = { path = "datafusion/proto", version = "55.1.0", default-features = false } +datafusion-proto-common = { path = "datafusion/proto-common", version = "55.1.0" } +datafusion-proto-models = { path = "datafusion/proto-models", version = "55.1.0" } +datafusion-pruning = { path = "datafusion/pruning", version = "55.1.0" } +datafusion-session = { path = "datafusion/session", version = "55.1.0" } +datafusion-spark = { path = "datafusion/spark", version = "55.1.0" } +datafusion-sql = { path = "datafusion/sql", version = "55.1.0" } +datafusion-substrait = { path = "datafusion/substrait", version = "55.1.0" } doc-comment = "0.3" env_logger = "0.11" diff --git a/dev/changelog/55.1.0.md b/dev/changelog/55.1.0.md new file mode 100644 index 0000000000000..93a9f5b9a8ebf --- /dev/null +++ b/dev/changelog/55.1.0.md @@ -0,0 +1,52 @@ + + +# Apache DataFusion 55.1.0 Changelog + +This release consists of 10 commits from 3 contributors. See credits at the end of this changelog for more information. + +See the [upgrade guide](https://datafusion.apache.org/library-user-guide/upgrading.html) for information on how to upgrade from previous versions. + +**Fixed bugs:** + +- [branch-55] fix(ffi): use the session's LogicalExtensionCodec in create_physical_plan (#24690) [#24694](https://github.com/apache/datafusion/pull/24694) (timsaucer) + +**Other:** + +- [branch-55] fix: make UnnestExec respect datafusion.execution.batch_size (#24384) [#24529](https://github.com/apache/datafusion/pull/24529) (alamb) +- [branch-55] fix: apply struct field filters when the file schema needs adaptation (#24125) [#24530](https://github.com/apache/datafusion/pull/24530) (alamb) +- [branch-55] fix: preserve provided arguments during FFI object construction (#24723) [#24752](https://github.com/apache/datafusion/pull/24752) (timsaucer) +- [branch-55] fix: adapt input batches with stricter nested nullability to planned schema in aggregation (#24394) [#24699](https://github.com/apache/datafusion/pull/24699) (timsaucer) +- [branch-55] Fix panic on RightMark hash joins when propagating ordering (#24718) [#24759](https://github.com/apache/datafusion/pull/24759) (timsaucer) +- [branch-55] fix(common): support empty struct in ScalarValue::compact and new_default (#24582) [#24876](https://github.com/apache/datafusion/pull/24876) (timsaucer) +- [branch-55] chore: Minor dependency updates [#24879](https://github.com/apache/datafusion/pull/24879) (timsaucer) +- [branch-55] Align metadata propagation through Physical and Logical casts (#23169) [#24875](https://github.com/apache/datafusion/pull/24875) (timsaucer) +- [branch-55] fix: preserve projection metadata during optimization (#24670) [#24992](https://github.com/apache/datafusion/pull/24992) (gene-bordegaray) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 7 Tim Saucer + 2 Andrew Lamb + 1 Gene Bordegaray +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/docs/source/download.md b/docs/source/download.md index 34296262071c8..2a0cc8e6e3f6e 100644 --- a/docs/source/download.md +++ b/docs/source/download.md @@ -26,7 +26,7 @@ For example: ```toml [dependencies] -datafusion = "55.0.0" +datafusion = "55.1.0" ``` While DataFusion is distributed via [crates.io] as a convenience, the diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index e02ada03fc413..c750a3e689951 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -104,7 +104,7 @@ The following configuration settings are available: | datafusion.execution.parquet.statistics_enabled | page | (writing) Sets if statistics are enabled for any column Valid values are: "none", "chunk", and "page" These values are not case sensitive. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.max_row_group_size | 1048576 | (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. | | datafusion.execution.parquet.max_row_group_bytes | NULL | (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. | -| datafusion.execution.parquet.created_by | datafusion version 55.0.0 | (writing) Sets "created by" property | +| datafusion.execution.parquet.created_by | datafusion version 55.1.0 | (writing) Sets "created by" property | | datafusion.execution.parquet.column_index_truncate_length | 64 | (writing) Sets column index truncate length | | datafusion.execution.parquet.statistics_truncate_length | 64 | (writing) Sets statistics truncate length. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.data_page_row_count_limit | 20000 | (writing) Sets best effort maximum number of rows in data page | diff --git a/docs/source/user-guide/crate-configuration.md b/docs/source/user-guide/crate-configuration.md index 09c65107e58c8..75ac2fbdde939 100644 --- a/docs/source/user-guide/crate-configuration.md +++ b/docs/source/user-guide/crate-configuration.md @@ -156,7 +156,7 @@ By default, Datafusion returns errors as a plain text message. You can enable mo such as backtraces by enabling the `backtrace` feature to your `Cargo.toml` file like this: ```toml -datafusion = { version = "55.0.0", features = ["backtrace"]} +datafusion = { version = "55.1.0", features = ["backtrace"]} ``` Set environment [variables](https://doc.rust-lang.org/std/backtrace/index.html#environment-variables) diff --git a/docs/source/user-guide/example-usage.md b/docs/source/user-guide/example-usage.md index 6f2419b9cd182..59495434a150e 100644 --- a/docs/source/user-guide/example-usage.md +++ b/docs/source/user-guide/example-usage.md @@ -29,7 +29,7 @@ Find latest available Datafusion version on [DataFusion's crates.io] page. Add the dependency to your `Cargo.toml` file: ```toml -datafusion = "55.0.0" +datafusion = "55.1.0" tokio = { version = "1.0", features = ["rt-multi-thread"] } ``` From 180db4c1dde0d891b8d5a3ac3a836d4e7ba5992d Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 19 Sep 2026 05:41:25 -0400 Subject: [PATCH 12/16] [branch-55] Pull the MinIO test image from quay.io (#25092, #25216) (#25485) - Fixes https://github.com/apache/datafusion/issues/25215 on branch-55 This PR: - Backports https://github.com/apache/datafusion/pull/25092 and https://github.com/apache/datafusion/pull/25216 from @comphead to the branch-55 line ## Rationale MinIO withdrew `minio/minio` from Docker Hub on 2026-09-11, so every `datafusion-cli` storage integration test on `branch-55` now panics while pulling the image (for example the CI failure on https://github.com/apache/datafusion/pull/25277). --------- Co-authored-by: Oleks V --- .github/workflows/rust.yml | 16 ++ datafusion-cli/tests/cli_integration.rs | 314 ++++++++++++++++-------- 2 files changed, 225 insertions(+), 105 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index eaa9b21a1b343..03e5fbe45b595 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -370,6 +370,22 @@ jobs: with: save-if: false # set in linux-test shared-key: "amd-ci" + # The storage integration tests start MinIO containers. Pulling the image + # once up front keeps the pull off the critical path of the tests, which + # would otherwise pull it several times concurrently and occasionally fail + # with transient Docker transport errors. The tests retry the pull + # themselves, so a failure here is only a warning. + # + # MINIO_IMAGE must match what the tests start: the tag comes from the + # `minio` module of `testcontainers-modules`, the registry from + # MINIO_IMAGE_NAME in `datafusion-cli/tests/cli_integration.rs`. The + # `minio_image_matches_ci_prepull` test fails if either drifts. + - name: Pre-pull MinIO image + env: + MINIO_IMAGE: quay.io/minio/minio:RELEASE.2025-02-28T09-55-16Z + run: | + ci/scripts/retry timeout 120 docker pull "$MINIO_IMAGE" \ + || echo "::warning::Could not pre-pull $MINIO_IMAGE, the tests will pull it themselves" - name: Run tests (excluding doctests) env: RUST_BACKTRACE: 1 diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index 4dc244445a2eb..53f2ab89442ff 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/datafusion-cli/tests/cli_integration.rs @@ -24,12 +24,13 @@ use insta::internals::SettingsBindDropGuard; use insta::{Settings, glob}; use insta_cmd::{assert_cmd_snapshot, get_cargo_bin}; use std::path::PathBuf; +use std::time::Duration; use std::{env, fs}; use testcontainers_modules::minio; use testcontainers_modules::testcontainers::core::{CmdWaitFor, ExecCommand, Mount}; use testcontainers_modules::testcontainers::runners::AsyncRunner; use testcontainers_modules::testcontainers::{ - ContainerAsync, ImageExt, TestcontainersError, + ContainerAsync, Image, ImageExt, TestcontainersError, }; fn cli() -> Command { @@ -45,10 +46,117 @@ fn make_settings() -> Settings { settings } +const MINIO_ROOT_USER: &str = "TEST-DataFusionLogin"; +const MINIO_ROOT_PASSWORD: &str = "TEST-DataFusionPassword"; + +/// Registry override for the image pinned by `testcontainers-modules`. +/// +/// MinIO withdrew `minio/minio` from Docker Hub on 2026-09-11. quay.io still +/// serves the same tag, so only the registry changes here. An unblock, not a +/// fix: see . +const MINIO_IMAGE_NAME: &str = "quay.io/minio/minio"; + +/// How many times to try bringing up the MinIO container before failing. +/// +/// Both the image pull and the `mc` calls that provision the bucket fail +/// intermittently on CI with transient errors such as +/// `bytes remaining on stream`. Retrying is much cheaper than a flaky run. +const MINIO_SETUP_ATTEMPTS: u32 = 3; + +/// Delay before the first retry of the MinIO setup, doubled on each attempt. +const MINIO_SETUP_RETRY_DELAY: Duration = Duration::from_secs(5); + +/// Time budget for a single MinIO setup attempt. A stalled image pull or `mc` +/// invocation is retried instead of hanging the whole test run. +const MINIO_SETUP_TIMEOUT: Duration = Duration::from_mins(3); + +/// Starts a MinIO container preloaded with the test data, retrying transient +/// Docker failures. +/// +/// Returns `None` when the test should be skipped, that is when +/// `TEST_STORAGE_INTEGRATION` is unset or the registry is rate limiting the +/// image pull. Panics if the container cannot be started for any other reason. +async fn start_minio_or_skip() -> Option> { + if env::var("TEST_STORAGE_INTEGRATION").is_err() { + eprintln!("Skipping external storages integration tests"); + return None; + } + + match setup_minio_container().await { + Ok(container) => Some(container), + Err(e) if is_docker_pull_rate_limit(&e) => { + eprintln!("Skipping test: Docker pull rate limit reached: {e}"); + None + } + Err(e) => panic!("{e}"), + } +} + +/// A registry pull rate limit does not clear up within a test run, so the +/// affected tests are skipped rather than retried. +fn is_docker_pull_rate_limit(error: &str) -> bool { + error.contains("toomanyrequests") +} + +/// Retrying only pays off for transient failures. An exhausted pull quota or a +/// Docker daemon that cannot be reached at all stays broken for the whole run. +fn is_retryable(error: &str) -> bool { + !is_docker_pull_rate_limit(error) + && !error.contains("failed to initialize a docker client") +} + async fn setup_minio_container() -> Result, String> { - const MINIO_ROOT_USER: &str = "TEST-DataFusionLogin"; - const MINIO_ROOT_PASSWORD: &str = "TEST-DataFusionPassword"; + let mut delay = MINIO_SETUP_RETRY_DELAY; + let mut last_error = String::from("MinIO container setup was not attempted at all"); + + for attempt in 1..=MINIO_SETUP_ATTEMPTS { + last_error = match tokio::time::timeout( + MINIO_SETUP_TIMEOUT, + try_setup_minio_container(), + ) + .await + { + Ok(Ok(container)) => return Ok(container), + Ok(Err(e)) => e, + Err(_) => format!( + "Timed out after {MINIO_SETUP_TIMEOUT:?} while starting the MinIO container" + ), + }; + + if attempt == MINIO_SETUP_ATTEMPTS || !is_retryable(&last_error) { + break; + } + + eprintln!( + "MinIO container setup failed (attempt {attempt}/{MINIO_SETUP_ATTEMPTS}), \ + retrying in {delay:?}: {last_error}" + ); + tokio::time::sleep(delay).await; + delay *= 2; + } + + Err(last_error) +} + +/// A single attempt at starting and provisioning a MinIO container. +/// +/// The container is removed again if provisioning fails, so that the next +/// attempt starts from a clean state. +async fn try_setup_minio_container() -> Result, String> { + let container = start_minio_container().await?; + + match provision_minio_container(&container).await { + Ok(()) => Ok(container), + Err(e) => { + if let Err(rm_error) = container.rm().await { + eprintln!("Failed to remove the MinIO container: {rm_error}"); + } + Err(e) + } + } +} +async fn start_minio_container() -> Result, String> { let data_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../datafusion/core/tests/data"); @@ -56,7 +164,8 @@ async fn setup_minio_container() -> Result, String> .canonicalize() .expect("Failed to get absolute path for test data"); - let container = minio::MinIO::default() + minio::MinIO::default() + .with_name(MINIO_IMAGE_NAME) .with_env_var("MINIO_ROOT_USER", MINIO_ROOT_USER) .with_env_var("MINIO_ROOT_PASSWORD", MINIO_ROOT_PASSWORD) .with_mount(Mount::bind_mount( @@ -64,60 +173,93 @@ async fn setup_minio_container() -> Result, String> "/source", )) .start() - .await; - - match container { - Ok(container) => { - // We wait for MinIO to be healthy and prepare test files. We do it via CLI to avoid s3 dependency - let commands = [ - ExecCommand::new(["/usr/bin/mc", "ready", "local"]), - ExecCommand::new([ - "/usr/bin/mc", - "alias", - "set", - "localminio", - "http://localhost:9000", - MINIO_ROOT_USER, - MINIO_ROOT_PASSWORD, - ]), - ExecCommand::new(["/usr/bin/mc", "mb", "localminio/data"]), - ExecCommand::new([ - "/usr/bin/mc", - "cp", - "-r", - "/source/", - "localminio/data/", - ]), - ]; - - for command in commands { - let command = - command.with_cmd_ready_condition(CmdWaitFor::Exit { code: Some(0) }); - - let cmd_ref = format!("{command:?}"); - - if let Err(e) = container.exec(command).await { - let stdout = container.stdout_to_vec().await.unwrap_or_default(); - let stderr = container.stderr_to_vec().await.unwrap_or_default(); - - return Err(format!( - "Failed to execute command: {}\nError: {}\nStdout: {:?}\nStderr: {:?}", - cmd_ref, - e, - String::from_utf8_lossy(&stdout), - String::from_utf8_lossy(&stderr) - )); - } - } + .await + .map_err(|e| match e { + TestcontainersError::Client(e) => format!( + "Failed to start MinIO container. Ensure Docker is running and accessible: {e}" + ), + e => format!("Failed to start MinIO container: {e}"), + }) +} - Ok(container) +/// Waits for MinIO to be healthy and uploads the test files. +/// +/// This is done via the `mc` CLI shipped in the image to avoid an s3 dependency. +async fn provision_minio_container( + container: &ContainerAsync, +) -> Result<(), String> { + let commands = [ + ExecCommand::new(["/usr/bin/mc", "ready", "local"]), + ExecCommand::new([ + "/usr/bin/mc", + "alias", + "set", + "localminio", + "http://localhost:9000", + MINIO_ROOT_USER, + MINIO_ROOT_PASSWORD, + ]), + ExecCommand::new(["/usr/bin/mc", "mb", "localminio/data"]), + ExecCommand::new(["/usr/bin/mc", "cp", "-r", "/source/", "localminio/data/"]), + ]; + + for command in commands { + let command = + command.with_cmd_ready_condition(CmdWaitFor::Exit { code: Some(0) }); + + let cmd_ref = format!("{command:?}"); + + if let Err(e) = container.exec(command).await { + let stdout = container.stdout_to_vec().await.unwrap_or_default(); + let stderr = container.stderr_to_vec().await.unwrap_or_default(); + + return Err(format!( + "Failed to execute command: {}\nError: {}\nStdout: {:?}\nStderr: {:?}", + cmd_ref, + e, + String::from_utf8_lossy(&stdout), + String::from_utf8_lossy(&stderr) + )); } - - Err(TestcontainersError::Client(e)) => Err(format!( - "Failed to start MinIO container. Ensure Docker is running and accessible: {e}" - )), - Err(e) => Err(format!("Failed to start MinIO container: {e}")), } + + Ok(()) +} + +/// CI pre-pulls the MinIO image so that the storage integration tests do not +/// have to pull it themselves. Guard against that pre-pull going stale when +/// `testcontainers-modules` bumps the image it uses. +#[test] +fn minio_image_matches_ci_prepull() { + let workflow = + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../.github/workflows/rust.yml"); + + // The workflow is not shipped with the published crate. + let Ok(contents) = fs::read_to_string(&workflow) else { + return; + }; + + let image = minio::MinIO::default(); + + // The override only redirects the registry, so it would silently stop + // tracking upstream if the crate ever pinned a different image. + assert!( + MINIO_IMAGE_NAME.ends_with(&format!("/{}", image.name())), + "`testcontainers-modules` now uses `{}`, which MINIO_IMAGE_NAME \ + (`{MINIO_IMAGE_NAME}`) no longer mirrors.", + image.name() + ); + + // Match the assignment: `minio/minio:` is a substring of the quay + // reference and would pass either way. + let image_ref = format!("{MINIO_IMAGE_NAME}:{}", image.tag()); + + assert!( + contents.contains(&format!("MINIO_IMAGE: {image_ref}")), + "{} does not pre-pull `{image_ref}`. Update MINIO_IMAGE in the \ + `Pre-pull MinIO image` step to match the image used by the tests.", + workflow.display() + ); } #[cfg(test)] @@ -522,18 +664,8 @@ fn test_cli_wide_result_set_no_crash() { #[tokio::test] async fn test_cli() { - if env::var("TEST_STORAGE_INTEGRATION").is_err() { - eprintln!("Skipping external storages integration tests"); + let Some(container) = start_minio_or_skip().await else { return; - } - - let container = match setup_minio_container().await { - Ok(c) => c, - Err(e) if e.contains("toomanyrequests") => { - eprintln!("Skipping test: Docker pull rate limit reached: {e}"); - return; - } - e @ Err(_) => e.unwrap(), }; let settings = make_settings(); @@ -546,8 +678,8 @@ async fn test_cli() { assert_cmd_snapshot!( cli() .env_clear() - .env("AWS_ACCESS_KEY_ID", "TEST-DataFusionLogin") - .env("AWS_SECRET_ACCESS_KEY", "TEST-DataFusionPassword") + .env("AWS_ACCESS_KEY_ID", MINIO_ROOT_USER) + .env("AWS_SECRET_ACCESS_KEY", MINIO_ROOT_PASSWORD) .env("AWS_ENDPOINT", format!("http://localhost:{port}")) .env("AWS_ALLOW_HTTP", "true") .pass_stdin(input) @@ -559,22 +691,13 @@ async fn test_cli() { async fn test_aws_options() { // Separate test is needed to pass aws as options in sql and not via env - if env::var("TEST_STORAGE_INTEGRATION").is_err() { - eprintln!("Skipping external storages integration tests"); + let Some(container) = start_minio_or_skip().await else { return; - } + }; let settings = make_settings(); let _bound = settings.bind_to_scope(); - let container = match setup_minio_container().await { - Ok(c) => c, - Err(e) if e.contains("toomanyrequests") => { - eprintln!("Skipping test: Docker pull rate limit reached: {e}"); - return; - } - e @ Err(_) => e.unwrap(), - }; let port = container.get_host_port_ipv4(9000).await.unwrap(); let input = format!( @@ -582,8 +705,8 @@ async fn test_aws_options() { STORED AS CSV LOCATION 's3://data/cars.csv' OPTIONS( - 'aws.access_key_id' 'TEST-DataFusionLogin', - 'aws.secret_access_key' 'TEST-DataFusionPassword', + 'aws.access_key_id' '{MINIO_ROOT_USER}', + 'aws.secret_access_key' '{MINIO_ROOT_PASSWORD}', 'aws.endpoint' 'http://localhost:{port}', 'aws.allow_http' 'true' ); @@ -658,18 +781,8 @@ fn test_backtrace_output(#[case] query: &str) { #[tokio::test] async fn test_s3_url_fallback() { - if env::var("TEST_STORAGE_INTEGRATION").is_err() { - eprintln!("Skipping external storages integration tests"); + let Some(container) = start_minio_or_skip().await else { return; - } - - let container = match setup_minio_container().await { - Ok(c) => c, - Err(e) if e.contains("toomanyrequests") => { - eprintln!("Skipping test: Docker pull rate limit reached: {e}"); - return; - } - e @ Err(_) => e.unwrap(), }; let mut settings = make_settings(); @@ -695,19 +808,10 @@ SELECT * FROM partitioned_data ORDER BY column_1, column_2 LIMIT 5; /// Validate object store profiling output #[tokio::test] async fn test_object_store_profiling() { - if env::var("TEST_STORAGE_INTEGRATION").is_err() { - eprintln!("Skipping external storages integration tests"); + let Some(container) = start_minio_or_skip().await else { return; - } - - let container = match setup_minio_container().await { - Ok(c) => c, - Err(e) if e.contains("toomanyrequests") => { - eprintln!("Skipping test: Docker pull rate limit reached: {e}"); - return; - } - e @ Err(_) => e.unwrap(), }; + let mut settings = make_settings(); // as the object store profiling contains timestamps and durations, we must @@ -769,8 +873,8 @@ impl MinioCommandExt for Command { let port = container.get_host_port_ipv4(9000).await.unwrap(); self.env_clear() - .env("AWS_ACCESS_KEY_ID", "TEST-DataFusionLogin") - .env("AWS_SECRET_ACCESS_KEY", "TEST-DataFusionPassword") + .env("AWS_ACCESS_KEY_ID", MINIO_ROOT_USER) + .env("AWS_SECRET_ACCESS_KEY", MINIO_ROOT_PASSWORD) .env("AWS_ENDPOINT", format!("http://localhost:{port}")) .env("AWS_ALLOW_HTTP", "true") } From cf7cfd188755e7f8c5a8fcf8ff2af3c41fb23f3b Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 19 Sep 2026 17:10:23 -0400 Subject: [PATCH 13/16] [branch-55] fix: update rustls to address RUSTSEC-2026-0285 (#25309) (#25514) - Fixes the `security_audit` failure on `branch-55` (for example https://github.com/apache/datafusion/actions/runs/35435408307/job/105877099456 on https://github.com/apache/datafusion/pull/25277) This PR: - Backports https://github.com/apache/datafusion/pull/25309 from @Xuanwo to the branch-55 line ## Rationale `branch-55` locks `rustls` 0.23.39, which is affected by [RUSTSEC-2026-0285](https://rustsec.org/advisories/RUSTSEC-2026-0285.html) ("TLS 1.3 handshake messages incorrectly accepted across encryption level boundaries", published 2026-09-14). Every PR targeting `branch-55` now fails the `security_audit` job until the lockfile is updated. `rustls` 0.23.45 fixes the vulnerability. The cherry-pick applied cleanly: `branch-55` and `main` had identical `rustls`, `aws-lc-rs`, `aws-lc-sys`, and `rustls-webpki` entries before the fix. ## Testing `cargo audit` with the same ignores as the CI job passes on this branch (it fails on `branch-55` without this change): ```sh cargo audit --ignore RUSTSEC-2026-0194 --ignore RUSTSEC-2026-0195 ``` `cargo fetch --locked` and `cargo check --locked -p datafusion-cli` also succeed, confirming the updated lockfile resolves and builds. Co-authored-by: Xuanwo --- Cargo.lock | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 141d5b6fdddb6..5ea0f9673cb50 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -586,9 +586,9 @@ dependencies = [ [[package]] name = "aws-lc-rs" -version = "1.16.3" +version = "1.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ec6fb3fe69024a75fa7e1bfb48aa6cf59706a101658ea01bfd33b2b248a038f" +checksum = "b281d307588d634de920874890732659e2e7672f72b5e10e81badc1a8a83621e" dependencies = [ "aws-lc-sys", "zeroize", @@ -596,14 +596,15 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.40.0" +version = "0.45.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f50037ee5e1e41e7b8f9d161680a725bd1626cb6f8c7e901f91f942850852fe7" +checksum = "9bff6c3b54fad79a2e60b8102caf565819711497c1f5f092f49508e2f5c31b27" dependencies = [ "cc", "cmake", "dunce", "fs_extra", + "pkg-config", ] [[package]] @@ -5379,9 +5380,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.39" +version = "0.23.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c2c118cb077cca2822033836dfb1b975355dfb784b5e8da48f7b6c5db74e60e" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" dependencies = [ "aws-lc-rs", "log", @@ -5417,9 +5418,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.103.13" +version = "0.103.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" dependencies = [ "aws-lc-rs", "ring", From 5c4db5a904c34a805bdc8b83b4cd9f102c9243a5 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 19 Sep 2026 19:22:47 -0400 Subject: [PATCH 14/16] [branch-55] Reduce datafusion-spark dependency footprint (#24351) (#25277) - Part of https://github.com/apache/datafusion/issues/24462 - Fixes https://github.com/apache/datafusion/issues/24474 on branch-55 This PR: - Backports https://github.com/apache/datafusion/pull/24351 from @dentiny to the branch-55 line ## Rationale As noted in https://github.com/apache/datafusion/issues/24462#issuecomment-5637098663, publishing `datafusion-spark` for the 55.1.0 release required a manual local change because the crate does not compile without its optional `core` feature (see https://github.com/apache/datafusion/issues/24474). Backporting this fix means a future 55.2.0 release can be published without manual intervention. Verified locally that `cargo check -p datafusion-spark --no-default-features` now succeeds on this branch (it fails on `branch-55` without this change). Co-authored-by: dentiny --- Cargo.lock | 3 +-- datafusion/spark/Cargo.toml | 5 ++--- datafusion/spark/src/function/string/quote.rs | 6 ++++-- datafusion/spark/src/function/table/mod.rs | 2 +- datafusion/spark/src/lib.rs | 6 +++--- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5ea0f9673cb50..f2705f6d0ed6e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2624,14 +2624,13 @@ dependencies = [ "crc32fast", "criterion", "datafusion", - "datafusion-catalog", "datafusion-common", - "datafusion-execution", "datafusion-expr", "datafusion-functions", "datafusion-functions-aggregate", "datafusion-functions-aggregate-common", "datafusion-functions-nested", + "datafusion-session", "log", "num-traits", "percent-encoding", diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index 93987b553f2f5..01fd9ec9b3f5b 100644 --- a/datafusion/spark/Cargo.toml +++ b/datafusion/spark/Cargo.toml @@ -49,14 +49,13 @@ chrono = { workspace = true } crc32fast = "1.4" # Optional dependency for SessionStateBuilderSpark extension trait datafusion = { workspace = true, optional = true, default-features = false } -datafusion-catalog = { workspace = true } datafusion-common = { workspace = true } -datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } -datafusion-functions = { workspace = true, features = ["crypto_expressions"] } +datafusion-functions = { workspace = true } datafusion-functions-aggregate = { workspace = true } datafusion-functions-aggregate-common = { workspace = true } datafusion-functions-nested = { workspace = true } +datafusion-session = { workspace = true } log = { workspace = true } num-traits = { workspace = true } percent-encoding = "2.3.2" diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs index 39ad8bf841764..55d827fb7a280 100644 --- a/datafusion/spark/src/function/string/quote.rs +++ b/datafusion/spark/src/function/string/quote.rs @@ -17,12 +17,14 @@ use arrow::array::{ArrayRef, OffsetSizeTrait, StringArray}; use arrow::datatypes::DataType; -use datafusion::logical_expr::{Coercion, ColumnarValue, Signature, TypeSignatureClass}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; use datafusion_common::types::{NativeType, logical_string}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, exec_err}; -use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Volatility}; +use datafusion_expr::{ + Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignatureClass, Volatility, +}; use datafusion_functions::utils::make_scalar_function; use std::sync::Arc; diff --git a/datafusion/spark/src/function/table/mod.rs b/datafusion/spark/src/function/table/mod.rs index aba7b7ceb78ea..5ccf8d91df045 100644 --- a/datafusion/spark/src/function/table/mod.rs +++ b/datafusion/spark/src/function/table/mod.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use datafusion_catalog::TableFunction; +use datafusion_session::TableFunction; use std::sync::Arc; pub mod expr_fn {} diff --git a/datafusion/spark/src/lib.rs b/datafusion/spark/src/lib.rs index 6cd4678da7560..c75f94142ff8a 100644 --- a/datafusion/spark/src/lib.rs +++ b/datafusion/spark/src/lib.rs @@ -42,7 +42,7 @@ //! Spark functions taking priority. //! //! ``` -//! # use datafusion_execution::FunctionRegistry; +//! # use datafusion_expr::registry::FunctionRegistry; //! # use datafusion_expr::{ScalarUDF, AggregateUDF, WindowUDF, HigherOrderUDF}; //! # use datafusion_expr::planner::ExprPlanner; //! # use datafusion_common::Result; @@ -137,10 +137,10 @@ mod session_state; #[cfg(feature = "core")] pub use session_state::SessionStateBuilderSpark; -use datafusion_catalog::TableFunction; use datafusion_common::Result; -use datafusion_execution::FunctionRegistry; +use datafusion_expr::registry::FunctionRegistry; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_session::TableFunction; use log::debug; use std::sync::Arc; From 9f5714b0ce725b6af2d27cef4a1dc2aaa4e76cc1 Mon Sep 17 00:00:00 2001 From: Oleks V Date: Fri, 25 Sep 2026 12:20:42 -0700 Subject: [PATCH 15/16] [branch-55] fix datafusion-cli CI: switch from minio to rustfs (#25706) (#25759) ## Which issue does this PR close? - Fixes #25705 on `branch-55`. ## Rationale for this change This PR backports #25706 from @a10y to the `branch-55` line. `quay.io/minio/minio` no longer allows anonymous pulls. `branch-55` pulls that image for the `datafusion-cli` storage integration tests (it moved there from Docker Hub in #25485), so the `cargo test datafusion-cli (amd64)` job is expected to hit the same failure that #25705 reported on `main`. The last `branch-55` CI run was on 2026-09-19, before the image became unavailable, so no run on this branch has failed yet. ## What changes are included in this PR? A clean `git cherry-pick -x` of e2ca7f38051744b2010cf09b80db8e9dfa4b5d38 with no conflicts. The patch is identical to the original (same `git patch-id`). - `datafusion-cli/tests/cli_integration.rs`: run the S3 tests against a RustFS container (`docker.io/rustfs/rustfs:1.0.0`) instead of MinIO. - `.github/workflows/rust.yml`: pre-pull the RustFS image instead of the MinIO one. - `datafusion-cli/Cargo.toml`: drop the `minio` feature of `testcontainers-modules`. - `datafusion-cli/CONTRIBUTING.md`: update the storage test docs. ## What is the testing strategy for this PR? The existing `datafusion-cli` storage integration tests cover this change, and CI runs them in the `cargo test datafusion-cli (amd64)` job. On this branch, `datafusion-cli/tests/` and its snapshots match `main` apart from one unrelated Spark test. `testcontainers` and `testcontainers-modules` also resolve to the same versions as on `main`. Locally, `cargo fmt --all -- --check`, `cargo clippy -p datafusion-cli --all-targets --all-features -- -D warnings`, `taplo format --check` and the prettier docs check pass. The Docker-based storage tests were not run locally. CI runs them. ## Are there any user-facing changes? No. The change only affects tests, CI and contributor docs. Signed-off-by: Andrew Duffy Co-authored-by: Andrew Duffy --- .github/workflows/rust.yml | 17 +- datafusion-cli/CONTRIBUTING.md | 6 +- datafusion-cli/Cargo.toml | 2 +- datafusion-cli/tests/cli_integration.rs | 246 +++++++++++++----------- 4 files changed, 149 insertions(+), 122 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 03e5fbe45b595..a2916779fca87 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -370,22 +370,21 @@ jobs: with: save-if: false # set in linux-test shared-key: "amd-ci" - # The storage integration tests start MinIO containers. Pulling the image + # The storage integration tests start RustFS containers. Pulling the image # once up front keeps the pull off the critical path of the tests, which # would otherwise pull it several times concurrently and occasionally fail # with transient Docker transport errors. The tests retry the pull # themselves, so a failure here is only a warning. # - # MINIO_IMAGE must match what the tests start: the tag comes from the - # `minio` module of `testcontainers-modules`, the registry from - # MINIO_IMAGE_NAME in `datafusion-cli/tests/cli_integration.rs`. The - # `minio_image_matches_ci_prepull` test fails if either drifts. - - name: Pre-pull MinIO image + # RUSTFS_IMAGE must match RUSTFS_IMAGE_NAME and RUSTFS_IMAGE_TAG in + # `datafusion-cli/tests/cli_integration.rs`. The + # `rustfs_image_matches_ci_prepull` test fails if they differ. + - name: Pre-pull RustFS image env: - MINIO_IMAGE: quay.io/minio/minio:RELEASE.2025-02-28T09-55-16Z + RUSTFS_IMAGE: docker.io/rustfs/rustfs:1.0.0 run: | - ci/scripts/retry timeout 120 docker pull "$MINIO_IMAGE" \ - || echo "::warning::Could not pre-pull $MINIO_IMAGE, the tests will pull it themselves" + ci/scripts/retry timeout 120 docker pull "$RUSTFS_IMAGE" \ + || echo "::warning::Could not pre-pull $RUSTFS_IMAGE, the tests will pull it themselves" - name: Run tests (excluding doctests) env: RUST_BACKTRACE: 1 diff --git a/datafusion-cli/CONTRIBUTING.md b/datafusion-cli/CONTRIBUTING.md index 8be656ec4ee34..fa615028a4ddd 100644 --- a/datafusion-cli/CONTRIBUTING.md +++ b/datafusion-cli/CONTRIBUTING.md @@ -35,7 +35,7 @@ cargo test --all-targets ## Running Storage Integration Tests -By default, storage integration tests are not run. These tests use the `testcontainers` crate to start up a local MinIO server using Docker on port 9000. +By default, storage integration tests are not run. These tests use the `testcontainers` crate to start up a local RustFS server using Docker with a dynamically assigned host port. To run them you will need to set `TEST_STORAGE_INTEGRATION`: @@ -47,8 +47,8 @@ For some of the tests, [snapshots](https://datafusion.apache.org/contributor-gui ### AWS -S3 integration is tested against [Minio](https://github.com/minio/minio) with [TestContainers](https://github.com/testcontainers/testcontainers-rs) -This requires Docker to be running on your machine and port 9000 to be free. +S3 integration is tested against [RustFS](https://github.com/rustfs/rustfs) with [TestContainers](https://github.com/testcontainers/testcontainers-rs) +This requires Docker to be running on your machine. If you see an error mentioning "failed to load IMDS session token" such as diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index 62eedafe798d4..c21bc972161d6 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -77,7 +77,7 @@ ctor = { workspace = true } insta = { workspace = true } insta-cmd = "0.7.0" rstest = { workspace = true } -testcontainers-modules = { workspace = true, features = ["minio"] } +testcontainers-modules = { workspace = true } # Makes sure `test_display_pg_json` behaves in a consistent way regardless of # feature unification with dependencies serde_json = { workspace = true, features = ["preserve_order"] } diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index 53f2ab89442ff..92de3ad68e92f 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/datafusion-cli/tests/cli_integration.rs @@ -20,17 +20,22 @@ use std::process::Command; use rstest::rstest; use async_trait::async_trait; +use futures::TryStreamExt; use insta::internals::SettingsBindDropGuard; use insta::{Settings, glob}; use insta_cmd::{assert_cmd_snapshot, get_cargo_bin}; +use object_store::{ + ObjectStore, ObjectStoreExt, aws::AmazonS3Builder, local::LocalFileSystem, +}; use std::path::PathBuf; use std::time::Duration; use std::{env, fs}; -use testcontainers_modules::minio; -use testcontainers_modules::testcontainers::core::{CmdWaitFor, ExecCommand, Mount}; +use testcontainers_modules::testcontainers::core::{ + CmdWaitFor, ExecCommand, IntoContainerPort, +}; use testcontainers_modules::testcontainers::runners::AsyncRunner; use testcontainers_modules::testcontainers::{ - ContainerAsync, Image, ImageExt, TestcontainersError, + ContainerAsync, GenericImage, ImageExt, TestcontainersError, }; fn cli() -> Command { @@ -46,43 +51,39 @@ fn make_settings() -> Settings { settings } -const MINIO_ROOT_USER: &str = "TEST-DataFusionLogin"; -const MINIO_ROOT_PASSWORD: &str = "TEST-DataFusionPassword"; +const RUSTFS_ACCESS_KEY: &str = "TEST-DataFusionLogin"; +const RUSTFS_SECRET_KEY: &str = "TEST-DataFusionPassword"; -/// Registry override for the image pinned by `testcontainers-modules`. -/// -/// MinIO withdrew `minio/minio` from Docker Hub on 2026-09-11. quay.io still -/// serves the same tag, so only the registry changes here. An unblock, not a -/// fix: see . -const MINIO_IMAGE_NAME: &str = "quay.io/minio/minio"; +/// Pinned RustFS image, also pulled by the CLI CI job. +const RUSTFS_IMAGE_NAME: &str = "docker.io/rustfs/rustfs"; +const RUSTFS_IMAGE_TAG: &str = "1.0.0"; -/// How many times to try bringing up the MinIO container before failing. +/// How many times to try bringing up the RustFS container before failing. /// -/// Both the image pull and the `mc` calls that provision the bucket fail -/// intermittently on CI with transient errors such as -/// `bytes remaining on stream`. Retrying is much cheaper than a flaky run. -const MINIO_SETUP_ATTEMPTS: u32 = 3; +/// Image pulls and fixture uploads can fail with transient errors such as +/// `bytes remaining on stream`. Retry these failures before failing the test. +const RUSTFS_SETUP_ATTEMPTS: u32 = 3; -/// Delay before the first retry of the MinIO setup, doubled on each attempt. -const MINIO_SETUP_RETRY_DELAY: Duration = Duration::from_secs(5); +/// Delay before the first retry of the RustFS setup, doubled on each attempt. +const RUSTFS_SETUP_RETRY_DELAY: Duration = Duration::from_secs(5); -/// Time budget for a single MinIO setup attempt. A stalled image pull or `mc` +/// Time budget for a single RustFS setup attempt. A stalled image pull or `curl` /// invocation is retried instead of hanging the whole test run. -const MINIO_SETUP_TIMEOUT: Duration = Duration::from_mins(3); +const RUSTFS_SETUP_TIMEOUT: Duration = Duration::from_mins(3); -/// Starts a MinIO container preloaded with the test data, retrying transient +/// Starts a RustFS container preloaded with the test data, retrying transient /// Docker failures. /// /// Returns `None` when the test should be skipped, that is when /// `TEST_STORAGE_INTEGRATION` is unset or the registry is rate limiting the /// image pull. Panics if the container cannot be started for any other reason. -async fn start_minio_or_skip() -> Option> { +async fn start_rustfs_or_skip() -> Option> { if env::var("TEST_STORAGE_INTEGRATION").is_err() { eprintln!("Skipping external storages integration tests"); return None; } - match setup_minio_container().await { + match setup_rustfs_container().await { Ok(container) => Some(container), Err(e) if is_docker_pull_rate_limit(&e) => { eprintln!("Skipping test: Docker pull rate limit reached: {e}"); @@ -105,30 +106,30 @@ fn is_retryable(error: &str) -> bool { && !error.contains("failed to initialize a docker client") } -async fn setup_minio_container() -> Result, String> { - let mut delay = MINIO_SETUP_RETRY_DELAY; - let mut last_error = String::from("MinIO container setup was not attempted at all"); +async fn setup_rustfs_container() -> Result, String> { + let mut delay = RUSTFS_SETUP_RETRY_DELAY; + let mut last_error = String::from("RustFS container setup was not attempted at all"); - for attempt in 1..=MINIO_SETUP_ATTEMPTS { + for attempt in 1..=RUSTFS_SETUP_ATTEMPTS { last_error = match tokio::time::timeout( - MINIO_SETUP_TIMEOUT, - try_setup_minio_container(), + RUSTFS_SETUP_TIMEOUT, + try_setup_rustfs_container(), ) .await { Ok(Ok(container)) => return Ok(container), Ok(Err(e)) => e, Err(_) => format!( - "Timed out after {MINIO_SETUP_TIMEOUT:?} while starting the MinIO container" + "Timed out after {RUSTFS_SETUP_TIMEOUT:?} while starting the RustFS container" ), }; - if attempt == MINIO_SETUP_ATTEMPTS || !is_retryable(&last_error) { + if attempt == RUSTFS_SETUP_ATTEMPTS || !is_retryable(&last_error) { break; } eprintln!( - "MinIO container setup failed (attempt {attempt}/{MINIO_SETUP_ATTEMPTS}), \ + "RustFS container setup failed (attempt {attempt}/{RUSTFS_SETUP_ATTEMPTS}), \ retrying in {delay:?}: {last_error}" ); tokio::time::sleep(delay).await; @@ -138,69 +139,76 @@ async fn setup_minio_container() -> Result, String> Err(last_error) } -/// A single attempt at starting and provisioning a MinIO container. +/// A single attempt at starting and provisioning a RustFS container. /// /// The container is removed again if provisioning fails, so that the next /// attempt starts from a clean state. -async fn try_setup_minio_container() -> Result, String> { - let container = start_minio_container().await?; +async fn try_setup_rustfs_container() -> Result, String> { + let container = start_rustfs_container().await?; - match provision_minio_container(&container).await { + match provision_rustfs_container(&container).await { Ok(()) => Ok(container), Err(e) => { if let Err(rm_error) = container.rm().await { - eprintln!("Failed to remove the MinIO container: {rm_error}"); + eprintln!("Failed to remove the RustFS container: {rm_error}"); } Err(e) } } } -async fn start_minio_container() -> Result, String> { - let data_path = - PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../datafusion/core/tests/data"); - - let absolute_data_path = data_path - .canonicalize() - .expect("Failed to get absolute path for test data"); - - minio::MinIO::default() - .with_name(MINIO_IMAGE_NAME) - .with_env_var("MINIO_ROOT_USER", MINIO_ROOT_USER) - .with_env_var("MINIO_ROOT_PASSWORD", MINIO_ROOT_PASSWORD) - .with_mount(Mount::bind_mount( - absolute_data_path.to_str().unwrap(), - "/source", - )) +async fn start_rustfs_container() -> Result, String> { + GenericImage::new(RUSTFS_IMAGE_NAME, RUSTFS_IMAGE_TAG) + .with_exposed_port(9000.tcp()) + .with_env_var("RUSTFS_ACCESS_KEY", RUSTFS_ACCESS_KEY) + .with_env_var("RUSTFS_SECRET_KEY", RUSTFS_SECRET_KEY) + .with_env_var("RUSTFS_CONSOLE_ENABLE", "false") .start() .await .map_err(|e| match e { TestcontainersError::Client(e) => format!( - "Failed to start MinIO container. Ensure Docker is running and accessible: {e}" + "Failed to start RustFS container. Ensure Docker is running and accessible: {e}" ), - e => format!("Failed to start MinIO container: {e}"), + e => format!("Failed to start RustFS container: {e}"), }) } -/// Waits for MinIO to be healthy and uploads the test files. +/// Waits for RustFS to be healthy and uploads the test files. /// -/// This is done via the `mc` CLI shipped in the image to avoid an s3 dependency. -async fn provision_minio_container( - container: &ContainerAsync, +/// Use the image's `curl` to check readiness and create the bucket, then upload +/// the fixtures through `object_store`. +async fn provision_rustfs_container( + container: &ContainerAsync, ) -> Result<(), String> { + let credentials = format!("{RUSTFS_ACCESS_KEY}:{RUSTFS_SECRET_KEY}"); let commands = [ - ExecCommand::new(["/usr/bin/mc", "ready", "local"]), ExecCommand::new([ - "/usr/bin/mc", - "alias", - "set", - "localminio", - "http://localhost:9000", - MINIO_ROOT_USER, - MINIO_ROOT_PASSWORD, + "curl", + "--fail", + "--silent", + "--show-error", + "--retry", + "60", + "--retry-delay", + "1", + "--retry-all-errors", + "--max-time", + "5", + "http://localhost:9000/health/ready", + ]), + ExecCommand::new([ + "curl", + "--fail", + "--silent", + "--show-error", + "--aws-sigv4", + "aws:amz:us-east-1:s3", + "--user", + &credentials, + "--request", + "PUT", + "http://localhost:9000/data", ]), - ExecCommand::new(["/usr/bin/mc", "mb", "localminio/data"]), - ExecCommand::new(["/usr/bin/mc", "cp", "-r", "/source/", "localminio/data/"]), ]; for command in commands { @@ -223,14 +231,45 @@ async fn provision_minio_container( } } + let port = container + .get_host_port_ipv4(9000) + .await + .map_err(|e| e.to_string())?; + let store = AmazonS3Builder::new() + .with_bucket_name("data") + .with_region("us-east-1") + .with_access_key_id(RUSTFS_ACCESS_KEY) + .with_secret_access_key(RUSTFS_SECRET_KEY) + .with_endpoint(format!("http://localhost:{port}")) + .with_allow_http(true) + .build() + .map_err(|e| e.to_string())?; + let data_path = + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../datafusion/core/tests/data"); + let source = + LocalFileSystem::new_with_prefix(data_path).map_err(|e| e.to_string())?; + let mut files = source.list(None); + while let Some(file) = files.try_next().await.map_err(|e| e.to_string())? { + let data = source + .get(&file.location) + .await + .map_err(|e| e.to_string())? + .bytes() + .await + .map_err(|e| e.to_string())?; + store + .put(&file.location, data.into()) + .await + .map_err(|e| e.to_string())?; + } + Ok(()) } -/// CI pre-pulls the MinIO image so that the storage integration tests do not -/// have to pull it themselves. Guard against that pre-pull going stale when -/// `testcontainers-modules` bumps the image it uses. +/// CI pre-pulls the RustFS image so that the storage integration tests do not +/// have to pull it themselves. Keep the CI image and the test image in sync. #[test] -fn minio_image_matches_ci_prepull() { +fn rustfs_image_matches_ci_prepull() { let workflow = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../.github/workflows/rust.yml"); @@ -239,25 +278,12 @@ fn minio_image_matches_ci_prepull() { return; }; - let image = minio::MinIO::default(); + let image_ref = format!("{RUSTFS_IMAGE_NAME}:{RUSTFS_IMAGE_TAG}"); - // The override only redirects the registry, so it would silently stop - // tracking upstream if the crate ever pinned a different image. assert!( - MINIO_IMAGE_NAME.ends_with(&format!("/{}", image.name())), - "`testcontainers-modules` now uses `{}`, which MINIO_IMAGE_NAME \ - (`{MINIO_IMAGE_NAME}`) no longer mirrors.", - image.name() - ); - - // Match the assignment: `minio/minio:` is a substring of the quay - // reference and would pass either way. - let image_ref = format!("{MINIO_IMAGE_NAME}:{}", image.tag()); - - assert!( - contents.contains(&format!("MINIO_IMAGE: {image_ref}")), - "{} does not pre-pull `{image_ref}`. Update MINIO_IMAGE in the \ - `Pre-pull MinIO image` step to match the image used by the tests.", + contents.contains(&format!("RUSTFS_IMAGE: {image_ref}")), + "{} does not pre-pull `{image_ref}`. Update RUSTFS_IMAGE in the \ + `Pre-pull RustFS image` step to match the image used by the tests.", workflow.display() ); } @@ -664,7 +690,7 @@ fn test_cli_wide_result_set_no_crash() { #[tokio::test] async fn test_cli() { - let Some(container) = start_minio_or_skip().await else { + let Some(container) = start_rustfs_or_skip().await else { return; }; @@ -678,8 +704,8 @@ async fn test_cli() { assert_cmd_snapshot!( cli() .env_clear() - .env("AWS_ACCESS_KEY_ID", MINIO_ROOT_USER) - .env("AWS_SECRET_ACCESS_KEY", MINIO_ROOT_PASSWORD) + .env("AWS_ACCESS_KEY_ID", RUSTFS_ACCESS_KEY) + .env("AWS_SECRET_ACCESS_KEY", RUSTFS_SECRET_KEY) .env("AWS_ENDPOINT", format!("http://localhost:{port}")) .env("AWS_ALLOW_HTTP", "true") .pass_stdin(input) @@ -691,7 +717,7 @@ async fn test_cli() { async fn test_aws_options() { // Separate test is needed to pass aws as options in sql and not via env - let Some(container) = start_minio_or_skip().await else { + let Some(container) = start_rustfs_or_skip().await else { return; }; @@ -705,8 +731,8 @@ async fn test_aws_options() { STORED AS CSV LOCATION 's3://data/cars.csv' OPTIONS( - 'aws.access_key_id' '{MINIO_ROOT_USER}', - 'aws.secret_access_key' '{MINIO_ROOT_PASSWORD}', + 'aws.access_key_id' '{RUSTFS_ACCESS_KEY}', + 'aws.secret_access_key' '{RUSTFS_SECRET_KEY}', 'aws.endpoint' 'http://localhost:{port}', 'aws.allow_http' 'true' ); @@ -781,7 +807,7 @@ fn test_backtrace_output(#[case] query: &str) { #[tokio::test] async fn test_s3_url_fallback() { - let Some(container) = start_minio_or_skip().await else { + let Some(container) = start_rustfs_or_skip().await else { return; }; @@ -802,13 +828,13 @@ OPTIONS ( SELECT * FROM partitioned_data ORDER BY column_1, column_2 LIMIT 5; "#; - assert_cmd_snapshot!(cli().with_minio(&container).await.pass_stdin(input)); + assert_cmd_snapshot!(cli().with_rustfs(&container).await.pass_stdin(input)); } /// Validate object store profiling output #[tokio::test] async fn test_object_store_profiling() { - let Some(container) = start_minio_or_skip().await else { + let Some(container) = start_rustfs_or_skip().await else { return; }; @@ -854,27 +880,29 @@ SELECT * from CARS LIMIT 1; SELECT * from CARS LIMIT 1; "#; - assert_cmd_snapshot!(cli().with_minio(&container).await.pass_stdin(input)); + assert_cmd_snapshot!(cli().with_rustfs(&container).await.pass_stdin(input)); } -/// Extension trait to Add the minio connection information to a Command +/// Add the RustFS connection information to a Command. #[async_trait] -trait MinioCommandExt { - async fn with_minio(&mut self, container: &ContainerAsync) - -> &mut Self; +trait RustfsCommandExt { + async fn with_rustfs( + &mut self, + container: &ContainerAsync, + ) -> &mut Self; } #[async_trait] -impl MinioCommandExt for Command { - async fn with_minio( +impl RustfsCommandExt for Command { + async fn with_rustfs( &mut self, - container: &ContainerAsync, + container: &ContainerAsync, ) -> &mut Self { let port = container.get_host_port_ipv4(9000).await.unwrap(); self.env_clear() - .env("AWS_ACCESS_KEY_ID", MINIO_ROOT_USER) - .env("AWS_SECRET_ACCESS_KEY", MINIO_ROOT_PASSWORD) + .env("AWS_ACCESS_KEY_ID", RUSTFS_ACCESS_KEY) + .env("AWS_SECRET_ACCESS_KEY", RUSTFS_SECRET_KEY) .env("AWS_ENDPOINT", format!("http://localhost:{port}")) .env("AWS_ALLOW_HTTP", "true") } From 5e62bc399356b6c50047d8a411a7784ee0688f13 Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 29 Sep 2026 09:09:18 +0800 Subject: [PATCH 16/16] [branch-55] fix: preserve fetch across distribution reoptimization (#24809) (#25821) ## Which issue does this PR close? - Fixes #24807 in the `branch-55` release line. - Part of the 55.2.0 release tracking issue #25758. - Backports the merged fix in #24809. ## Rationale for this change In DataFusion 55.1.0, running `EnsureRequirements` again on a plan with a pushed-down `fetch` can remove the global limit and return too many rows. This is a correctness fix for a bug present in the 55 release line, so it belongs in the 55.2.0 patch release. The incorrect result and expected behavior are described in #24807. ## What changes are included in this PR? - Cherry-picks the merged commit from #24809 onto `branch-55`. - Keeps fetched distribution operators as row-selection boundaries during reoptimization. - Preserves a fetched ordered merge's TopK selection when replacing order-preserving variants. - Carries over the original regression tests and adjusts the backport for the 55 branch. ## What is the testing strategy for this PR? - `cargo fmt --all` - `cargo test -p datafusion --test core_integration physical_optimizer::enforce_distribution` (87 passed) - `cargo clippy --all-targets --all-features -- -D warnings` - `./dev/rust_lint.sh` (passed with the CI-pinned `hawkeye 6.2.0`) - Initialized the repository's test-data submodules, then ran the documented extended workspace test command. One unrelated macOS RSS-threshold test failed: `smj_with_mem_limit_1_runner` measured 154.2 MB against a 152.6 MB allowance; it failed again in isolation at 158.6 MB. This test exercises sort-merge join memory accounting, which this backport does not change. - Ran the extended workspace suite with only that RSS test skipped: all `datafusion` library, integration, and fuzz tests passed. The custom SQLLogic runner does not accept libtest's `--skip` flag, so I ran all other workspace packages separately without a skip; those passed, including all 503 SQLLogic files. - No local benchmark directly exercises reoptimization of a fetched physical plan. ## Are there any user-facing changes? Queries keep their requested global limit when distribution requirements are optimized more than once. This corrects query results without changing public APIs. --- .../enforce_distribution.rs | 325 +++++++++++++++++- .../physical_optimizer/enforce_sorting.rs | 4 +- .../replace_with_order_preserving_variants.rs | 53 +-- .../enforce_distribution.rs | 116 +++---- .../enforce_sorting/mod.rs | 14 +- .../replace_with_order_preserving_variants.rs | 30 +- 6 files changed, 404 insertions(+), 138 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 2dbacf1d898ac..63ee0c8483076 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -26,7 +26,7 @@ use crate::physical_optimizer::test_utils::{ sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; -use arrow::array::{RecordBatch, UInt8Array, UInt64Array}; +use arrow::array::{Int64Array, RecordBatch, UInt8Array, UInt64Array}; use arrow::compute::SortOptions; use arrow_schema::{DataType, Field, Schema, SchemaRef}; use datafusion::config::ConfigOptions; @@ -45,6 +45,7 @@ use datafusion_common::tree_node::{ }; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_datasource::memory::MemorySourceConfig; use datafusion_expr::{JoinType, Operator}; use datafusion_functions_aggregate::count::count_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -71,11 +72,12 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::JoinOn; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, - PlanProperties, ReplaceChildrenOptions, displayable, + PlanProperties, ReplaceChildrenOptions, collect, displayable, }; use insta::Settings; @@ -4444,11 +4446,12 @@ fn test_replace_order_preserving_variants_with_fetch() -> Result<()> { // Apply the function let result = replace_order_preserving_variants(dist_context)?; - // Verify the plan was transformed to CoalescePartitionsExec + // A fetched ordered merge must still select the TopK rows. + let result = check_integrity(result)?; result .plan - .downcast_ref::() - .expect("Expected CoalescePartitionsExec"); + .downcast_ref::() + .expect("Expected a TopK SortExec"); // Verify fetch was preserved assert_eq!( @@ -4460,6 +4463,318 @@ fn test_replace_order_preserving_variants_with_fetch() -> Result<()> { Ok(()) } +#[test] +fn preserve_fetch_when_reoptimizing_ordered_merge() -> Result<()> { + let schema = schema(); + let sort_key: LexOrdering = + [PhysicalSortExpr::new_default(col("c", &schema)?)].into(); + let input = parquet_exec_multiple_sorted(vec![sort_key.clone()]); + let plan: Arc = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + + let optimized = + EnsureRequirements::new().optimize(plan, &test_suite_default_config_options())?; + let plan = displayable(optimized.as_ref()).indent(true).to_string(); + + assert!( + plan.contains("SortPreservingMergeExec: [c@2 ASC], fetch=5"), + "expected the optimizer to preserve fetch:\n{plan}" + ); + + Ok(()) +} + +#[test] +fn preserve_fetch_when_reoptimizing_coalesce_partitions() -> Result<()> { + let input = parquet_exec_multiple(); + let plan: Arc = + Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5))); + + let optimized = + EnsureRequirements::new().optimize(plan, &test_suite_default_config_options())?; + + assert_eq!(optimized.fetch(), Some(5)); + optimized + .downcast_ref::() + .expect("expected CoalescePartitionsExec"); + + Ok(()) +} + +#[tokio::test] +async fn move_fetch_to_replacement_sort() -> Result<()> { + for (options, partitions, expected) in [ + ( + SortOptions::default(), + [ + vec![None, Some(1), Some(1), Some(6)], + vec![None, Some(1), Some(2), Some(7)], + ], + vec![None, None, Some(1), Some(1), Some(1)], + ), + ( + SortOptions { + descending: true, + nulls_first: false, + }, + [vec![Some(7), Some(1), None], vec![Some(6), Some(1), None]], + vec![Some(7), Some(6), Some(1), Some(1), None], + ), + ] { + let (input, sort_key) = sorted_memory_input(partitions, options)?; + let merge: Arc = Arc::new( + SortPreservingMergeExec::new(sort_key.clone(), input).with_fetch(Some(5)), + ); + assert_eq!(fetch_test_values(Arc::clone(&merge)).await?, expected); + let plan = sort_required_exec_with_req(merge, sort_key); + let optimized = ensure_distribution_helper(plan, 10, false)?; + let replacement = Arc::clone(optimized.children()[0]); + let sort = replacement + .downcast_ref::() + .expect("expected a replacement sort"); + assert_eq!(sort.fetch(), Some(5)); + assert_eq!(fetch_test_values(replacement).await?, expected); + } + Ok(()) +} + +#[tokio::test] +async fn preserve_fetch_in_nested_distribution_operators() -> Result<()> { + for outer_fetch in [0, 3, 10] { + let (input, sort_key) = sorted_memory_input( + [0, 1].map(|start| (start..10).step_by(2).map(Some).collect()), + SortOptions::default(), + )?; + let merge: Arc = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + let plan: Arc = + Arc::new(CoalescePartitionsExec::new(merge).with_fetch(Some(outer_fetch))); + let expected = (0..outer_fetch.min(5)) + .map(|value| Some(value as i64)) + .collect::>(); + assert_reoptimized_fetch_values(plan, &expected).await?; + } + Ok(()) +} + +#[tokio::test] +async fn preserve_topk_when_parent_changes_ordering() -> Result<()> { + let (input, sort_key) = sorted_memory_input( + [0, 1].map(|start| (start..10).step_by(2).map(Some).collect()), + SortOptions::default(), + )?; + let descending = [PhysicalSortExpr::new( + col("c", &input.schema())?, + SortOptions { + descending: true, + nulls_first: false, + }, + )] + .into(); + let merge: Arc = + Arc::new(SortPreservingMergeExec::new(sort_key, input).with_fetch(Some(5))); + let plan: Arc = Arc::new(SortExec::new(descending, merge)); + assert_reoptimized_fetch_values(plan, &[Some(4), Some(3), Some(2), Some(1), Some(0)]) + .await +} + +#[tokio::test] +async fn preserve_fetch_when_parallelizing_sort_above_filter() -> Result<()> { + let (input, sort_key) = sorted_memory_input( + [ + vec![Some(-4), Some(-2), Some(2), Some(4), Some(6)], + vec![Some(-3), Some(-1), Some(3), Some(5), Some(7)], + ], + SortOptions::default(), + )?; + let predicate = Arc::new(BinaryExpr::new( + col("c", &input.schema())?, + Operator::Gt, + lit(0_i64), + )); + let coalesce: Arc = + Arc::new(CoalescePartitionsExec::new(input).with_fetch(Some(5))); + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, coalesce)?); + let mut plan: Arc = Arc::new(SortExec::new(sort_key, filter)); + let mut config = test_suite_default_config_options(); + config.optimizer.enable_round_robin_repartition = false; + config.optimizer.repartition_sorts = true; + for iteration in 0..3 { + if iteration > 0 { + plan = EnsureRequirements::new().optimize(plan, &config)?; + } + // Either input batch can arrive first. Both contain three positive + // rows, so keeping the limit below the filter always returns three. + assert_eq!( + fetch_test_values(Arc::clone(&plan)).await?.len(), + 3, + "iteration {iteration}:\n{}", + displayable(plan.as_ref()).indent(true) + ); + } + Ok(()) +} + +fn sorted_memory_input( + partitions: [Vec>; 2], + options: SortOptions, +) -> Result<(Arc, LexOrdering)> { + let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int64, true)])); + let order: LexOrdering = [PhysicalSortExpr::new(col("c", &schema)?, options)].into(); + let partitions = partitions + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + ) + .map(|batch| vec![batch]) + }) + .collect::, _>>()?; + let source = MemorySourceConfig::try_new(&partitions, schema, None)? + .try_with_sort_information(vec![order.clone()])?; + Ok((DataSourceExec::from_data_source(source), order)) +} + +async fn fetch_test_values(plan: Arc) -> Result>> { + let batches = collect(plan, SessionContext::new().task_ctx()).await?; + Ok(batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect()) +} + +async fn assert_reoptimized_fetch_values( + plan: Arc, + expected: &[Option], +) -> Result<()> { + for repartition_sorts in [false, true] { + let mut optimized = Arc::clone(&plan); + let mut config = test_suite_default_config_options(); + config.optimizer.enable_round_robin_repartition = false; + config.optimizer.repartition_sorts = repartition_sorts; + for iteration in 0..3 { + if iteration > 0 { + let distribution = + DistributionContext::new_default(Arc::clone(&optimized)) + .transform_up(|context| ensure_distribution(context, &config))? + .data; + check_integrity(distribution)?; + optimized = EnsureRequirements::new().optimize(optimized, &config)?; + } + assert_eq!( + fetch_test_values(Arc::clone(&optimized)).await?, + expected, + "iteration {iteration}, repartition_sorts={repartition_sorts}:\n{}", + displayable(optimized.as_ref()).indent(true) + ); + } + } + Ok(()) +} + +#[tokio::test] +async fn preserve_fetch_below_filter_when_reoptimizing() -> Result<()> { + check_fetch_below_filter( + Operator::Gt, + [vec![-2, 0, 2, 4], vec![-1, 1, 3, 5]], + &[1, 2], + ) + .await +} + +#[tokio::test] +async fn preserve_fetch_below_filter_with_constant_ordering() -> Result<()> { + check_fetch_below_filter( + Operator::Eq, + [vec![-2, 0, 0, 0], vec![-1, 0, 0, 0]], + &[0, 0, 0], + ) + .await +} + +async fn check_fetch_below_filter( + op: Operator, + partitions: [Vec; 2], + expected: &[i64], +) -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int64, false)])); + let sort_key: LexOrdering = + [PhysicalSortExpr::new_default(col("c", &schema)?)].into(); + let partitions = partitions + .into_iter() + .map(|values| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + ) + .map(|batch| vec![batch]) + }) + .collect::, _>>()?; + let source = MemorySourceConfig::try_new(&partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_key.clone()])?; + let merge: Arc = Arc::new( + SortPreservingMergeExec::new( + sort_key.clone(), + DataSourceExec::from_data_source(source), + ) + .with_fetch(Some(5)), + ); + let predicate = Arc::new(BinaryExpr::new(col("c", &schema)?, op, lit(0_i64))); + let filter: Arc = Arc::new(FilterExec::try_new(predicate, merge)?); + let mut plan = sort_required_exec_with_req(filter, sort_key); + let mut config = test_suite_default_config_options(); + config.optimizer.enable_round_robin_repartition = false; + let task_context = SessionContext::new().task_ctx(); + + // The test operator only declares ordering requirements. Execute its child + // to compare query results before optimization and after repeated passes. + for iteration in 0..3 { + if iteration > 0 { + let distribution = DistributionContext::new_default(Arc::clone(&plan)) + .transform_up(|context| ensure_distribution(context, &config))? + .data; + check_integrity(distribution)?; + plan = EnsureRequirements::new().optimize(plan, &config)?; + } + let input = Arc::clone(plan.children()[0]); + let batches = collect(input, Arc::clone(&task_context)).await?; + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!( + values, + expected, + "iteration {iteration}:\n{}", + displayable(plan.as_ref()).indent(true) + ); + assert!( + plan.children()[0].is::(), + "fetch must stay below the filter:\n{}", + displayable(plan.as_ref()).indent(true) + ); + } + Ok(()) +} + /// When a parent requires SinglePartition and maintains input order, order-preserving /// variants (e.g. SortPreservingMergeExec) should be kept so that ordering can /// propagate to ancestors. Replacing them with CoalescePartitionsExec would destroy diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 76af9b0c29218..87d8ac3b159e1 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -2327,7 +2327,9 @@ async fn test_remove_unnecessary_spm2() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - DataSourceExec: partitions=1, partition_sizes=[0] + LocalLimitExec: fetch=100 + SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) diff --git a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs index 601667ea02c0d..b8adb283c2d31 100644 --- a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs +++ b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs @@ -31,7 +31,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use insta::{allow_duplicates, assert_snapshot}; use datafusion_common::tree_node::{TransformedResult, TreeNode}; -use datafusion_common::{assert_contains, NullEquality, Result}; +use datafusion_common::{NullEquality, Result}; use datafusion_common::config::ConfigOptions; use datafusion_datasource::source::DataSourceExec; use datafusion_execution::TaskContext; @@ -1199,43 +1199,22 @@ fn test_plan_with_order_preserving_variants_preserves_fetch() -> Result<()> { .with_fetch(Some(10)) .unwrap(); - // Test sort's fetch is greater than coalesce fetch, return error because it's not reasonable - let requirements = OrderPreservationContext::new( - coalesced.clone(), - false, - vec![OrderPreservationContext::new( - parquet_exec.clone(), - false, - vec![], - )], - ); - let res = plan_with_order_preserving_variants(requirements, false, true, Some(15)); - assert_contains!( - res.unwrap_err().to_string(), - "CoalescePartitionsExec fetch [10] should be greater than or equal to SortExec fetch [15]" - ); - - // Test sort is without fetch, expected to get the fetch value from the coalesced - let requirements = OrderPreservationContext::new( - coalesced.clone(), - false, - vec![OrderPreservationContext::new( - parquet_exec.clone(), + // Keep the coalesce's row selection independently of an ancestor's fetch. + for sort_fetch in [Some(15), None, Some(5)] { + let requirements = OrderPreservationContext::new( + Arc::clone(&coalesced), false, - vec![], - )], - ); - let res = plan_with_order_preserving_variants(requirements, false, true, None)?; - assert_eq!(res.plan.fetch(), Some(10),); - - // Test sort's fetch is less than coalesces fetch, expected to get the fetch value from the sort - let requirements = OrderPreservationContext::new( - coalesced, - false, - vec![OrderPreservationContext::new(parquet_exec, false, vec![])], - ); - let res = plan_with_order_preserving_variants(requirements, false, true, Some(5))?; - assert_eq!(res.plan.fetch(), Some(5),); + vec![OrderPreservationContext::new( + parquet_exec.clone(), + false, + vec![], + )], + ); + let res = + plan_with_order_preserving_variants(requirements, false, true, sort_fetch)?; + assert!(Arc::ptr_eq(&res.plan, &coalesced)); + assert_eq!(res.plan.fetch(), Some(10)); + } Ok(()) } diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index 07bc98b2db798..d092af4f7887f 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -63,6 +63,7 @@ use datafusion_physical_plan::joins::{ }; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::tree_node::PlanContext; @@ -782,10 +783,7 @@ fn preserving_order_enables_streaming( /// /// Updated node with an execution plan, where the desired single distribution /// requirement is satisfied. -fn add_merge_on_top( - input: DistributionContext, - fetch: Option, -) -> DistributionContext { +fn add_merge_on_top(input: DistributionContext) -> DistributionContext { // Apply only when the partition count is larger than one. if input.plan.output_partitioning().partition_count() > 1 { // When there is an existing ordering, we preserve ordering @@ -794,21 +792,16 @@ fn add_merge_on_top( // - Preserving ordering is not helpful in terms of satisfying ordering requirements // - Usage of order preserving variants is not desirable // (determined by flag `config.optimizer.prefer_existing_sort`) - let new_plan: Arc = if let Some(req) = - input.plan.output_ordering() - { - let mut spm = - SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan)); - if let Some(f) = fetch { - spm = spm.with_fetch(Some(f)); - } - Arc::new(spm) - } else { - // If there is no input order, we can simply coalesce partitions: - Arc::new( - CoalescePartitionsExec::new(Arc::clone(&input.plan)).with_fetch(fetch), - ) - }; + let new_plan: Arc = + if let Some(req) = input.plan.output_ordering() { + Arc::new(SortPreservingMergeExec::new( + req.clone(), + Arc::clone(&input.plan), + )) + } else { + // If there is no input order, we can simply coalesce partitions: + Arc::new(CoalescePartitionsExec::new(Arc::clone(&input.plan))) + }; DistributionContext::new(new_plan, true, vec![input]) } else { @@ -833,41 +826,21 @@ fn add_merge_on_top( /// ```text /// "DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet", /// ``` -/// Returned by [`remove_dist_changing_operators`] to carry the fetch value -/// that may have been on a removed `SortPreservingMergeExec` or `CoalescePartitionsExec`. -struct RemovedDistOps { - context: DistributionContext, - /// The fetch value from the removed SPM/Coalesce, if any. - /// Must be re-applied when distribution operators are re-inserted. - removed_fetch: Option, -} - +/// A distribution operator with a fetch also selects rows. Stop at that +/// boundary so neither its limit nor an ordered merge's TopK selection is +/// moved across another operator. fn remove_dist_changing_operators( mut distribution_context: DistributionContext, -) -> Result { - let mut removed_fetch = None; - while is_repartition(&distribution_context.plan) - || is_coalesce_partitions(&distribution_context.plan) - || is_sort_preserving_merge(&distribution_context.plan) +) -> DistributionContext { + while distribution_context.plan.fetch().is_none() + && (is_repartition(&distribution_context.plan) + || is_coalesce_partitions(&distribution_context.plan) + || is_sort_preserving_merge(&distribution_context.plan)) { - // Preserve fetch from SPM or CoalescePartitions before removing (#14150). - if let Some(fetch) = distribution_context.plan.fetch() { - removed_fetch = Some( - removed_fetch - .map(|existing: usize| existing.min(fetch)) - .unwrap_or(fetch), - ); - } - // All of above operators have a single child. First child is only child. - // Remove any distribution changing operators at the beginning: + // All of the above operators have a single child. distribution_context = distribution_context.children.swap_remove(0); - // Note that they will be re-inserted later on if necessary or helpful. } - - Ok(RemovedDistOps { - context: distribution_context, - removed_fetch, - }) + distribution_context } /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable. @@ -903,11 +876,21 @@ pub fn replace_order_preserving_variants( }) .collect::>>()?; - if is_sort_preserving_merge(&context.plan) { + if let Some(spm) = context.plan.downcast_ref::() { let child_plan = Arc::clone(&context.children[0].plan); - context.plan = Arc::new( - CoalescePartitionsExec::new(child_plan).with_fetch(context.plan.fetch()), - ); + let fetch = spm.fetch(); + if fetch.is_some() { + // A fetched merge selects the first rows in its sort order. Moving + // fetch to an ancestor's sort can cross a filter, or lose the limit + // entirely if the ancestor needs no additional sort. + let ordering = spm.expr().clone(); + context.plan = Arc::new(CoalescePartitionsExec::new(child_plan)); + let sort = Arc::new( + SortExec::new(ordering, Arc::clone(&context.plan)).with_fetch(fetch), + ); + return Ok(DistributionContext::new(sort, false, vec![context])); + } + context.plan = Arc::new(CoalescePartitionsExec::new(child_plan)); return Ok(context); } else if let Some(repartition) = context.plan.downcast_ref::() && repartition.preserve_order() @@ -1166,17 +1149,13 @@ pub fn ensure_distribution( let order_preserving_variants_desirable = unbounded_and_pipeline_friendly || config.optimizer.prefer_existing_sort; - // Remove unnecessary repartition from the physical plan if any. - // Preserve fetch from removed SPM/Coalesce (#14150). - let RemovedDistOps { - context: - DistributionContext { - mut plan, - data, - children, - }, - removed_fetch, - } = remove_dist_changing_operators(dist_context)?; + // Remove distribution-only operators, retaining any fetched operator as + // a row-selection boundary. + let DistributionContext { + mut plan, + data, + children, + } = remove_dist_changing_operators(dist_context); if let Some(exec) = plan.downcast_ref::() { if let Some(updated_window) = get_best_fitting_window( @@ -1331,7 +1310,7 @@ pub fn ensure_distribution( // Satisfy the distribution requirement if it is unmet. match &requirement { Distribution::SinglePartition => { - child = add_merge_on_top(child, removed_fetch); + child = add_merge_on_top(child); } Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) => { @@ -1449,12 +1428,13 @@ pub fn ensure_distribution( // make sure ordering requirements are still satisfied after. if ordering_satisfied { // Make sure to satisfy ordering requirement: + let output_fetch = plan + .downcast_ref::() + .and_then(|output| output.fetch()); context = add_sort_above_with_check( context, sort_req, - plan.downcast_ref::() - .map(|output| output.fetch()) - .unwrap_or(None), + output_fetch, )?; } } diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 90b19ca95bcbf..38708d73da720 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -170,8 +170,10 @@ fn update_coalesce_ctx_children( // Plan has no children, it cannot be a `CoalescePartitionsExec`. false } else if is_coalesce_partitions(&coalesce_context.plan) { - // Initiate a connection: - true + // A fetched coalesce selects rows before the sort. Removing it would + // lose that limit, and moving its fetch onto the sort changes which + // rows are selected. + coalesce_context.plan.fetch().is_none() } else { children.iter().enumerate().any(|(idx, node)| { // Only consider operators that don't require a single partition, @@ -668,11 +670,9 @@ fn remove_bottleneck_in_subplan_impl( Some(Distribution::SinglePartition) ) }; - let remove_from_first_child = requirements - .children - .first() - .is_some_and(|child| is_coalesce_partitions(&child.plan)) - && removable(0); + let remove_from_first_child = requirements.children.first().is_some_and(|child| { + is_coalesce_partitions(&child.plan) && child.plan.fetch().is_none() + }) && removable(0); let children = &mut requirements.children; if remove_from_first_child { // We can safely use the 0th index since we have a `CoalescePartitionsExec`. diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs index 6ab84dc95eab9..2f9d19569b357 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs @@ -25,9 +25,9 @@ use crate::utils::{ is_coalesce_partitions, is_repartition, is_sort, is_sort_preserving_merge, }; +use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::Transformed; -use datafusion_common::{Result, assert_or_internal_err}; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::execution_plan::EmissionType; @@ -64,7 +64,9 @@ pub fn update_order_preservation_ctx_children_data(opc: &mut OrderPreservationCo } let plan_children = plan.children(); - *data = if plan_children.is_empty() { + *data = if plan_children.is_empty() + || (is_coalesce_partitions(plan) && plan.fetch().is_some()) + { false } else if !children[0].data && ((is_repartition(plan) && !maintains_input_order[0]) @@ -102,6 +104,12 @@ pub fn plan_with_order_preserving_variants( is_spm_better: bool, fetch: Option, ) -> Result { + if is_coalesce_partitions(&sort_input.plan) && sort_input.plan.fetch().is_some() { + // A fetched coalesce selects rows in arrival order. An ordered merge + // would select different rows, even with the same fetch value. + sort_input.data = false; + return Ok(sort_input); + } sort_input.children = sort_input .children .into_iter() @@ -137,24 +145,6 @@ pub fn plan_with_order_preserving_variants( } else if is_coalesce_partitions(&sort_input.plan) && is_spm_better { let child = &sort_input.children[0].plan; if let Some(ordering) = child.output_ordering() { - let mut fetch = fetch; - if let Some(coalesce_fetch) = sort_input.plan.fetch() { - fetch = match fetch { - Some(sort_fetch) => { - assert_or_internal_err!( - coalesce_fetch >= sort_fetch, - "CoalescePartitionsExec fetch [{:?}] should be greater than or equal to SortExec fetch [{:?}]", - coalesce_fetch, - sort_fetch - ); - Some(sort_fetch) - } - None => { - // If the sort node does not have a fetch, we need to keep the coalesce node's fetch. - Some(coalesce_fetch) - } - }; - }; // When the input of a `CoalescePartitionsExec` has an ordering, // replace it with a `SortPreservingMergeExec` if appropriate: let spm = SortPreservingMergeExec::new(ordering.clone(), Arc::clone(child))