From 2bba020e91a77a074f75ab472e91017856063116 Mon Sep 17 00:00:00 2001 From: dario curreri Date: Wed, 5 Aug 2026 11:58:53 +0200 Subject: [PATCH 1/3] physical-plan: coerce UNION/INTERLEAVE schema mismatches at plan time Follow-up to #23861 (issue #15394). Moves the schema re-stamping for nullability-mismatched UNION ALL / INTERLEAVE inputs out of `execute()` and into plan construction, via a new `CoerceSchemaExec` node inserted by `UnionExec::try_new`/`InterleaveExec::try_new` whenever a child's own schema disagrees with the computed union schema. The node is now visible in `EXPLAIN` output, is transparent for statistics/pushdown/proto purposes, and adds no measurable overhead versus inline re-stamping. --- datafusion/core/tests/dataframe/mod.rs | 3 +- datafusion/physical-plan/src/union.rs | 298 ++++++++++++++++-- .../sqllogictest/test_files/array_agg.slt | 23 +- datafusion/sqllogictest/test_files/union.slt | 55 ++-- 4 files changed, 323 insertions(+), 56 deletions(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 73a9177ab738a..3480850dde2ce 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -7075,7 +7075,8 @@ async fn test_copy_to_preserves_order() -> Result<()> { DataSinkExec: sink=CsvSink(file_groups=[]) SortExec: expr=[column1@0 DESC], preserve_partitioning=[false] DataSourceExec: partitions=1, partition_sizes=[1] - DataSourceExec: partitions=1, partition_sizes=[1] + CoerceSchemaExec + DataSourceExec: partitions=1, partition_sizes=[1] " ); Ok(()) diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 8d77556509b9e..0cfb1bc5f755b 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -67,20 +67,14 @@ use tokio::macros::support::thread_rng_n; /// Wraps a child stream so that every batch it yields is re-stamped with /// `schema` instead of the child's own schema. /// -/// This is used by both [`UnionExec`] and [`InterleaveExec`] when a child's -/// output schema disagrees with the operator's declared output schema -- -/// in practice this only happens for nullability (the declared schema is -/// nullable wherever *any* input's field is, but casts are only inserted -/// between inputs when the *type* differs, not when only nullability -/// does). For [`UnionExec`], [`UnionExec::try_new`] guarantees this: it -/// calls `calculate_union`, which rejects any input whose field data types -/// don't match the computed union schema. [`InterleaveExec::try_new`] does -/// not repeat that check -- its inputs are only ever produced by the -/// optimizer rewriting an already-validated `UnionExec`, whose children's -/// types are therefore already known to agree -- but if this wrapper ever -/// did see a genuine data type mismatch (e.g. from a hand-built -/// `InterleaveExec`), `RecordBatch::try_new_with_options` below reports it -/// as an error rather than silently yielding a corrupt batch. +/// Used by [`CoerceSchemaExec`], which [`UnionExec::try_new`] and +/// [`InterleaveExec::try_new`] insert above any child whose output schema +/// disagrees with the operator's declared output schema -- in practice this +/// only happens for nullability (the declared schema is nullable wherever +/// *any* input's field is, but casts are only inserted between inputs when +/// the *type* differs, not when only nullability does). If this wrapper +/// ever did see a genuine data type mismatch, `RecordBatch::try_new_with_options` +/// below reports it as an error rather than silently yielding a corrupt batch. struct SchemaConformingStream { schema: SchemaRef, inner: SendableRecordBatchStream, @@ -141,6 +135,179 @@ fn conform_stream_schema( } } +/// Coerces a single child's declared output schema to `schema`, re-stamping +/// every batch it produces to match. [`UnionExec::try_new`] and +/// [`InterleaveExec::try_new`] insert this above any child whose own schema +/// disagrees with the computed union schema, so that the coercion is visible +/// in the plan tree (e.g. in `EXPLAIN`) instead of happening invisibly inside +/// the union operator's own `execute()`. +/// +/// A genuine data type mismatch (as opposed to a nullability-only one) is +/// rejected eagerly, at construction time, by `EquivalenceProperties:: +/// with_new_schema` below -- unlike the old purely-runtime approach, a +/// hand-built union/interleave with mismatched child types now fails in +/// `try_new` rather than in `execute()`. +/// +/// This node is a strict 1:1, order-preserving passthrough of `input` (it +/// only ever changes a batch's declared schema, never its rows), so every +/// `ExecutionPlan` method below that isn't about the schema itself just +/// delegates straight to `input`. +/// +/// See . +#[derive(Debug)] +struct CoerceSchemaExec { + input: Arc, + cache: Arc, + metrics: ExecutionPlanMetricsSet, +} + +impl CoerceSchemaExec { + /// Wraps `input` in a [`CoerceSchemaExec`] targeting `schema` if its own + /// schema disagrees with `schema`, otherwise returns it unchanged. + fn wrap_if_needed( + input: Arc, + schema: &SchemaRef, + ) -> Result> { + if &input.schema() == schema { + Ok(input) + } else { + Ok(Arc::new(Self::new(input, schema)?)) + } + } + + fn new(input: Arc, schema: &SchemaRef) -> Result { + let eq_properties = input + .equivalence_properties() + .clone() + .with_new_schema(Arc::clone(schema))?; + let output_partitioning = input.output_partitioning().clone(); + let cache = PlanProperties::new( + eq_properties, + output_partitioning, + emission_type_from_children(std::iter::once(&input)), + boundedness_from_children(std::iter::once(&input)), + ); + Ok(Self { + input, + cache: Arc::new(cache), + metrics: ExecutionPlanMetricsSet::new(), + }) + } +} + +impl DisplayAs for CoerceSchemaExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "CoerceSchemaExec") + } + DisplayFormatType::TreeRender => Ok(()), + } + } +} + +impl ExecutionPlan for CoerceSchemaExec { + fn name(&self) -> &'static str { + "CoerceSchemaExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + // A 1:1 passthrough never combines partitions, so re-deriving the cache + // (rather than collapsing back to the raw child via `wrap_if_needed`) is + // always safe here -- it just keeps this node from vanishing and + // changing arity out from under a caller mid-rewrite. + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + assert_or_internal_err!( + children.len() == 1, + "CoerceSchemaExec expects exactly one child" + ); + Ok(Arc::new(Self::new(children.remove(0), &self.schema())?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + let stream = self.input.execute(partition, context)?; + let stream = conform_stream_schema(self.schema(), stream); + Ok(Box::pin(ObservedStream::new( + stream, + baseline_metrics, + None, + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec>, + _config: &ConfigOptions, + ) -> Result { + FilterDescription::from_children(parent_filters, &self.children()) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + // No dedicated protobuf variant: this node is fully determined by its + // child and the parent union/interleave's declared schema, so it's + // serialized as if it weren't there. `UnionExec`/`InterleaveExec:: + // try_from_proto` both go through `try_new`, which re-inserts the + // wrapper via `wrap_if_needed` on decode. + Ok(Some(ctx.encode_child(&self.input)?)) + } +} + /// `UnionExec`: `UNION ALL` execution plan. /// /// `UnionExec` combines multiple inputs with the same schema by @@ -208,6 +375,10 @@ impl UnionExec { // The schema of the inputs and the union schema is consistent when: // - They have the same number of fields, and // - Their fields have same types at the same indices. + let inputs = inputs + .into_iter() + .map(|input| CoerceSchemaExec::wrap_if_needed(input, &schema)) + .collect::>>()?; let cache = Self::compute_properties(&inputs, schema)?; Ok(Arc::new(UnionExec { inputs, @@ -372,7 +543,6 @@ impl ExecutionPlan for UnionExec { if partition < input.output_partitioning().partition_count() { let stream = input.execute(partition, context)?; debug!("Found a Union partition to execute"); - let stream = conform_stream_schema(self.schema(), stream); return Ok(Box::pin(ObservedStream::new( stream, baseline_metrics, @@ -642,7 +812,12 @@ impl InterleaveExec { can_interleave(inputs.iter()), "Not all InterleaveExec children have a consistent hash or range partitioning" ); - let cache = Self::compute_properties(&inputs)?; + let schema = union_schema(&inputs)?; + let inputs = inputs + .into_iter() + .map(|input| CoerceSchemaExec::wrap_if_needed(input, &schema)) + .collect::>>()?; + let cache = Self::compute_properties(&inputs, schema)?; Ok(InterleaveExec { inputs, metrics: ExecutionPlanMetricsSet::new(), @@ -656,8 +831,10 @@ impl InterleaveExec { } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties(inputs: &[Arc]) -> Result { - let schema = union_schema(inputs)?; + fn compute_properties( + inputs: &[Arc], + schema: SchemaRef, + ) -> Result { let eq_properties = EquivalenceProperties::new(schema); // Get output partitioning: let output_partitioning = inputs[0].output_partitioning().clone(); @@ -748,7 +925,7 @@ impl ExecutionPlan for InterleaveExec { for input in self.inputs.iter() { if partition < input.output_partitioning().partition_count() { let stream = input.execute(partition, Arc::clone(&context))?; - input_stream_vec.push(conform_stream_schema(self.schema(), stream)); + input_stream_vec.push(stream); } else { // Do not find a partition to execute break; @@ -1262,6 +1439,89 @@ mod tests { Ok(()) } + #[test] + fn test_union_partition_statistics_with_mismatched_nullability() -> Result<()> { + // Regression test for the `CoerceSchemaExec` wrapper `UnionExec::try_new` + // inserts above the non-nullable leg here: before it forwarded + // `child_stats_requests`/`statistics_from_inputs` to its child, it + // reported `Statistics::new_unknown`, poisoning the merge into all-`Absent` + // even though both legs have exact statistics. + let (_, left, right, expected) = stats_merge_inputs(); + + let non_nullable_schema = + Schema::new(vec![Field::new("a", DataType::UInt32, false)]); + let nullable_schema = Schema::new(vec![Field::new("a", DataType::UInt32, true)]); + + let left: Arc = + Arc::new(StatisticsExec::new(left, non_nullable_schema)); + let right: Arc = + Arc::new(StatisticsExec::new(right, nullable_schema)); + + let union = UnionExec::try_new(vec![left, right])?; + let stats = + StatisticsContext::new().compute(union.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(stats.as_ref(), &expected); + Ok(()) + } + + #[tokio::test] + async fn test_coerce_schema_exec_execution_plan_methods() -> Result<()> { + // Most sqllogictest coverage only observes `CoerceSchemaExec` through + // `EXPLAIN`; exercise its `ExecutionPlan` methods directly here. + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch_not_null = RecordBatch::try_new( + Arc::clone(&schema_not_null), + vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2]))], + )?; + let input: Arc = TestMemoryExec::try_new_exec( + &[vec![batch_not_null]], + Arc::clone(&schema_not_null), + None, + )?; + + // Matching schema: `wrap_if_needed` is a no-op. + let unwrapped = + CoerceSchemaExec::wrap_if_needed(Arc::clone(&input), &schema_not_null)?; + assert!(Arc::ptr_eq(&unwrapped, &input)); + + // Mismatched nullability: the input gets wrapped. + let nullable_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let wrapped = + CoerceSchemaExec::wrap_if_needed(Arc::clone(&input), &nullable_schema)?; + assert_eq!(wrapped.name(), "CoerceSchemaExec"); + assert_eq!(&wrapped.schema(), &nullable_schema); + assert_eq!(wrapped.maintains_input_order(), vec![true]); + assert_eq!(wrapped.benefits_from_input_partitioning(), vec![false]); + assert!(wrapped.supports_limit_pushdown()); + assert!(matches!( + wrapped.cardinality_effect(), + CardinalityEffect::Equal + )); + assert!(wrapped.metrics().is_some()); + + wrapped.gather_filters_for_pushdown( + FilterPushdownPhase::Pre, + vec![], + &ConfigOptions::new(), + )?; + + // Re-deriving via `with_new_children` keeps targeting the same schema. + let new_child: Arc = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_not_null), None)?; + let rewrapped = Arc::clone(&wrapped).with_new_children(vec![new_child])?; + assert_eq!(&rewrapped.schema(), &nullable_schema); + + let task_ctx = Arc::new(TaskContext::default()); + let batches = collect(wrapped, task_ctx).await?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema(), nullable_schema); + + Ok(()) + } + #[test] fn test_interleave_partition_statistics_uses_shared_statistics_merge() -> Result<()> { let (schema, left, right, expected) = stats_merge_inputs(); diff --git a/datafusion/sqllogictest/test_files/array_agg.slt b/datafusion/sqllogictest/test_files/array_agg.slt index f44e7f7d02e9c..36ad13d2491bc 100644 --- a/datafusion/sqllogictest/test_files/array_agg.slt +++ b/datafusion/sqllogictest/test_files/array_agg.slt @@ -534,16 +534,19 @@ physical_plan 03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5 04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted 05)--------UnionExec -06)----------ProjectionExec: expr=[1 as id, 2 as foo] -07)------------PlaceholderRowExec -08)----------ProjectionExec: expr=[1 as id, NULL as foo] -09)------------PlaceholderRowExec -10)----------ProjectionExec: expr=[1 as id, NULL as foo] -11)------------PlaceholderRowExec -12)----------ProjectionExec: expr=[1 as id, 3 as foo] -13)------------PlaceholderRowExec -14)----------ProjectionExec: expr=[1 as id, 2 as foo] -15)------------PlaceholderRowExec +06)----------CoerceSchemaExec +07)------------ProjectionExec: expr=[1 as id, 2 as foo] +08)--------------PlaceholderRowExec +09)----------ProjectionExec: expr=[1 as id, NULL as foo] +10)------------PlaceholderRowExec +11)----------ProjectionExec: expr=[1 as id, NULL as foo] +12)------------PlaceholderRowExec +13)----------CoerceSchemaExec +14)------------ProjectionExec: expr=[1 as id, 3 as foo] +15)--------------PlaceholderRowExec +16)----------CoerceSchemaExec +17)------------ProjectionExec: expr=[1 as id, 2 as foo] +18)--------------PlaceholderRowExec ####### # Unsupported syntax diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index cb5a06f7296fd..f5dd6db505ea2 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -572,25 +572,27 @@ logical_plan physical_plan 01)CoalescePartitionsExec: fetch=3 02)--UnionExec -03)----ProjectionExec: expr=[count(Int64(1))@0 as cnt] -04)------GlobalLimitExec: skip=0, fetch=3 -05)--------AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] -06)----------CoalescePartitionsExec -07)------------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] -08)--------------ProjectionExec: expr=[] -09)----------------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[] -10)------------------RepartitionExec: partitioning=Hash([c1@0], 4), input_partitions=4 -11)--------------------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[] -12)----------------------FilterExec: c13@1 != C2GT5KVyOPZpgKVl110TyZO0NcJ434, projection=[c1@0] -13)------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -14)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c13], file_type=csv, has_header=true -15)----ProjectionExec: expr=[1 as cnt] -16)------PlaceholderRowExec -17)----ProjectionExec: expr=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as cnt] -18)------GlobalLimitExec: skip=0, fetch=3 -19)--------BoundedWindowAggExec: wdw=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Field { "lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING], mode=[Sorted] -20)----------ProjectionExec: expr=[1 as c1] -21)------------PlaceholderRowExec +03)----CoerceSchemaExec +04)------ProjectionExec: expr=[count(Int64(1))@0 as cnt] +05)--------GlobalLimitExec: skip=0, fetch=3 +06)----------AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] +07)------------CoalescePartitionsExec +08)--------------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] +09)----------------ProjectionExec: expr=[] +10)------------------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[] +11)--------------------RepartitionExec: partitioning=Hash([c1@0], 4), input_partitions=4 +12)----------------------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[] +13)------------------------FilterExec: c13@1 != C2GT5KVyOPZpgKVl110TyZO0NcJ434, projection=[c1@0] +14)--------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +15)----------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c13], file_type=csv, has_header=true +16)----CoerceSchemaExec +17)------ProjectionExec: expr=[1 as cnt] +18)--------PlaceholderRowExec +19)----ProjectionExec: expr=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as cnt] +20)------GlobalLimitExec: skip=0, fetch=3 +21)--------BoundedWindowAggExec: wdw=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Field { "lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING], mode=[Sorted] +22)----------ProjectionExec: expr=[1 as c1] +23)------------PlaceholderRowExec ######## @@ -721,13 +723,14 @@ logical_plan 11)----------EmptyRelation: rows=1 physical_plan 01)UnionExec -02)--ProjectionExec: expr=[count(Int64(1))@1 as count, n@0 as n] -03)----AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))], ordering_mode=Sorted -04)------ProjectionExec: expr=[5 as n] -05)--------PlaceholderRowExec -06)--ProjectionExec: expr=[1 as count, max(Int64(10))@0 as n] -07)----AggregateExec: mode=Single, gby=[], aggr=[max(Int64(10))] -08)------PlaceholderRowExec +02)--CoerceSchemaExec +03)----ProjectionExec: expr=[count(Int64(1))@1 as count, n@0 as n] +04)------AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))], ordering_mode=Sorted +05)--------ProjectionExec: expr=[5 as n] +06)----------PlaceholderRowExec +07)--ProjectionExec: expr=[1 as count, max(Int64(10))@0 as n] +08)----AggregateExec: mode=Single, gby=[], aggr=[max(Int64(10))] +09)------PlaceholderRowExec # Test issue: https://github.com/apache/datafusion/issues/11409 From 578f6047268d60931f644d02ba1c08d9e5cc1cde Mon Sep 17 00:00:00 2001 From: dario curreri Date: Wed, 5 Aug 2026 13:08:31 +0200 Subject: [PATCH 2/3] test(proto): roundtrip UNION/INTERLEAVE with mismatched nullability Addresses review feedback on #24094 (kosiew): verify that `CoerceSchemaExec::try_to_proto`'s wrapper-erasure trick is correctly undone by `UnionExec`/`InterleaveExec::try_from_proto` re-inserting the wrapper via `try_new`, and that the decoded plan's emitted batches actually carry the coerced nullable schema, not just its EXPLAIN string. --- .../tests/cases/roundtrip_physical_plan.rs | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index b22cfd7764a21..f42296380c977 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2330,6 +2330,54 @@ fn roundtrip_union() -> Result<()> { roundtrip_test(union) } +/// `CoerceSchemaExec::try_to_proto` intentionally leaves the wrapper out of +/// the serialized plan, relying on `UnionExec::try_from_proto` to rebuild it +/// via `try_new`. Verify the decoded plan still contains the coercion and +/// that its emitted batches expose the union's nullable schema. +#[tokio::test] +async fn roundtrip_union_with_mismatched_nullability_executes() -> Result<()> { + let literal_leg = |value: ScalarValue| -> Result> { + Ok(Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: lit(value), + alias: "a".to_string(), + }], + Arc::new(PlaceholderRowExec::new(Arc::new(Schema::empty()))), + )?)) + }; + let non_nullable_leg = literal_leg(ScalarValue::Int64(Some(1)))?; + let nullable_leg = literal_leg(ScalarValue::Int64(None))?; + + let union: Arc = + UnionExec::try_new(vec![non_nullable_leg, nullable_leg])?; + assert!(union.schema().field(0).is_nullable()); + assert!( + format!("{union:?}").contains("CoerceSchemaExec"), + "expected CoerceSchemaExec in plan:\n{union:?}" + ); + + let ctx = SessionContext::new(); + let bytes = datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&union))?; + let roundtripped = datafusion_proto::bytes::physical_plan_from_bytes( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + )?; + assert!(roundtripped.schema().field(0).is_nullable()); + assert!( + format!("{roundtripped:?}").contains("CoerceSchemaExec"), + "expected CoerceSchemaExec after roundtrip:\n{roundtripped:?}" + ); + + let batches = + datafusion::physical_plan::collect(roundtripped, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + for batch in &batches { + assert!(batch.schema().field(0).is_nullable()); + } + + Ok(()) +} + #[test] fn roundtrip_repartition_preserve_order() -> Result<()> { let field_a = Field::new("a", DataType::Int64, false); @@ -2476,6 +2524,59 @@ fn roundtrip_interleave() -> Result<()> { roundtrip_test(Arc::new(interleave)) } +/// See [`roundtrip_union_with_mismatched_nullability_executes`]: the same +/// wrapper-reinsertion behavior applies to `InterleaveExec::try_from_proto`. +#[tokio::test] +async fn roundtrip_interleave_with_mismatched_nullability_executes() -> Result<()> { + let partition = Partitioning::Hash(vec![], 3); + let literal_leg = |value: ScalarValue| -> Result> { + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: lit(value), + alias: "a".to_string(), + }], + Arc::new(PlaceholderRowExec::new(Arc::new(Schema::empty()))), + )?; + Ok(Arc::new(RepartitionExec::try_new( + Arc::new(projection), + partition.clone(), + )?)) + }; + let non_nullable_leg = literal_leg(ScalarValue::Int64(Some(1)))?; + let nullable_leg = literal_leg(ScalarValue::Int64(None))?; + + let interleave: Arc = Arc::new(InterleaveExec::try_new(vec![ + non_nullable_leg, + nullable_leg, + ])?); + assert!(interleave.schema().field(0).is_nullable()); + assert!( + format!("{interleave:?}").contains("CoerceSchemaExec"), + "expected CoerceSchemaExec in plan:\n{interleave:?}" + ); + + let ctx = SessionContext::new(); + let bytes = datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&interleave))?; + let roundtripped = datafusion_proto::bytes::physical_plan_from_bytes( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + )?; + assert!(roundtripped.schema().field(0).is_nullable()); + assert!( + format!("{roundtripped:?}").contains("CoerceSchemaExec"), + "expected CoerceSchemaExec after roundtrip:\n{roundtripped:?}" + ); + + let batches = + datafusion::physical_plan::collect(roundtripped, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + for batch in &batches { + assert!(batch.schema().field(0).is_nullable()); + } + + Ok(()) +} + #[test] fn roundtrip_unnest() -> Result<()> { let fa = Field::new("a", DataType::Int64, true); From bc2f35f84c711d23d9168bda122ef0ca141761b8 Mon Sep 17 00:00:00 2001 From: dario curreri Date: Wed, 5 Aug 2026 22:55:10 +0200 Subject: [PATCH 3/3] physical-plan: replace CoerceSchemaExec with ProjectionExec + CastExpr Per alamb's suggestion on #24094: CastExpr::new_with_target_field already lets a cast carry an explicit target Field (not just a DataType), and the cast kernel has a same-type fast path (Arc::clone, no data copy), so a nullability-only coercion is just a same-type cast. UnionExec/InterleaveExec now build a ProjectionExec with a CastExpr (or a plain Column when a leg's field already matches exactly) instead of a hand-rolled ExecutionPlan. This deletes the hand-rolled node's ~150 lines of trait boilerplate (statistics/pushdown/proto plumbing) and, since ProjectionExec has an ordinary protobuf message, removes the wrapper-erasure trick entirely -- there's no more invisible reinsertion on proto decode to reason about. It also gets the existing projection-collapsing optimizer pass for free: when a coerced leg's own top node is already a ProjectionExec, the two fuse into one instead of stacking. Also fixes two narrow gaps this surfaced in ProjectionExec's statistics propagation (datafusion-physical-expr's project_column_statistics_through_expr): a CastExpr whose source values are already of the target DataType is a value-preserving relabeling, so unlike a real type-changing cast, sum_value and byte_size should carry over unchanged rather than degrading to Absent. --- datafusion/core/tests/dataframe/mod.rs | 2 +- datafusion/physical-expr/src/projection.rs | 45 ++ datafusion/physical-plan/src/union.rs | 395 ++++++------------ .../tests/cases/roundtrip_physical_plan.rs | 27 +- .../sqllogictest/test_files/array_agg.slt | 23 +- datafusion/sqllogictest/test_files/union.slt | 55 ++- 6 files changed, 214 insertions(+), 333 deletions(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 3480850dde2ce..a51c752c5cc40 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -7075,7 +7075,7 @@ async fn test_copy_to_preserves_order() -> Result<()> { DataSinkExec: sink=CsvSink(file_groups=[]) SortExec: expr=[column1@0 DESC], preserve_partitioning=[false] DataSourceExec: partitions=1, partition_sizes=[1] - CoerceSchemaExec + ProjectionExec: expr=[CAST(column1@0 AS UInt64) as count] DataSourceExec: partitions=1, partition_sizes=[1] " ); diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index e3fd6ddf744a9..f8f4cf51faa63 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -848,6 +848,22 @@ fn project_column_statistics_through_expr( let inner_stats = project_column_statistics_through_expr(cast_expr.expr.as_ref(), column_stats); let target_type = cast_expr.cast_type(); + + // A cast whose source values are already of the target `DataType` never + // changes any value -- see `cast_array_by_name`'s same-type fast path in + // `ColumnarValue::cast_to`. In that case every statistic, not just + // min/max, carries over unchanged (this is what a cast that only + // re-stamps a column's nullability, as `UnionExec`/`InterleaveExec` + // insert, looks like here). + let already_target_type = matches!( + (inner_stats.min_value.get_value(), inner_stats.max_value.get_value()), + (Some(min), Some(max)) + if min.data_type() == *target_type && max.data_type() == *target_type + ); + if already_target_type { + return inner_stats; + } + ColumnStatistics { min_value: inner_stats .min_value @@ -2934,6 +2950,35 @@ pub(crate) mod tests { Ok(()) } + #[test] + fn test_project_statistics_with_same_type_cast_is_exact_passthrough() -> Result<()> { + // A cast to the column's own `DataType` (e.g. one that only re-stamps + // nullability via `CastExpr::new_with_target_field`, as `UnionExec`/ + // `InterleaveExec` insert) never changes any value, so every + // statistic -- not just min/max -- should carry over unchanged. + let input_stats = get_stats(); + let col0_stats = input_stats.column_statistics[0].clone(); + let input_schema = get_schema(); + + let projection = ProjectionExprs::new(vec![ProjectionExpr { + expr: Arc::new(CastExpr::new( + Arc::new(Column::new("col0", 0)), + DataType::Int64, + None, + )), + alias: "casted".to_string(), + }]); + + let output_stats = projection.project_statistics( + input_stats, + &projection.project_schema(&input_schema)?, + )?; + + assert_eq!(output_stats.column_statistics[0], col0_stats); + + Ok(()) + } + #[test] fn test_project_statistics_with_cast() -> Result<()> { let input_stats = get_stats(); diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 0cfb1bc5f755b..646f8db39112f 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -42,270 +42,91 @@ use crate::filter_pushdown::{ FilterPushdownPropagation, PushedDown, }; use crate::metrics::BaselineMetrics; -use crate::projection::{ProjectionExec, make_with_child}; +use crate::projection::{ProjectionExec, ProjectionExpr, make_with_child}; use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::ObservedStream; -use arrow::array::RecordBatchOptions; use arrow::datatypes::{Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::NdvFallback; use datafusion_common::{ - Result, assert_or_internal_err, exec_err, internal_datafusion_err, + Result, assert_or_internal_err, exec_err, internal_datafusion_err, plan_err, }; use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::{CastExpr, Column}; use datafusion_physical_expr::{ EquivalenceProperties, PhysicalExpr, calculate_union, conjunction, }; -use futures::{Stream, StreamExt}; +use futures::Stream; use itertools::Itertools; use log::{debug, trace, warn}; use tokio::macros::support::thread_rng_n; -/// Wraps a child stream so that every batch it yields is re-stamped with -/// `schema` instead of the child's own schema. +/// Coerces `input`'s output schema to exactly `schema` via a `ProjectionExec` +/// that re-stamps each column with the union's merged field (same +/// `DataType`, but the union's merged nullability/name/metadata), or returns +/// `input` unchanged if its schema already matches. [`UnionExec::try_new`] +/// and [`InterleaveExec::try_new`] call this on every child, so the coercion +/// is visible in the plan tree (e.g. in `EXPLAIN`) instead of happening +/// invisibly inside the union operator's own `execute()`. /// -/// Used by [`CoerceSchemaExec`], which [`UnionExec::try_new`] and -/// [`InterleaveExec::try_new`] insert above any child whose output schema -/// disagrees with the operator's declared output schema -- in practice this -/// only happens for nullability (the declared schema is nullable wherever -/// *any* input's field is, but casts are only inserted between inputs when -/// the *type* differs, not when only nullability does). If this wrapper -/// ever did see a genuine data type mismatch, `RecordBatch::try_new_with_options` -/// below reports it as an error rather than silently yielding a corrupt batch. -struct SchemaConformingStream { - schema: SchemaRef, - inner: SendableRecordBatchStream, -} - -impl SchemaConformingStream { - fn new(schema: SchemaRef, inner: SendableRecordBatchStream) -> Self { - Self { schema, inner } - } -} - -impl RecordBatchStream for SchemaConformingStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} - -impl Stream for SchemaConformingStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - self.inner.poll_next_unpin(cx).map(|opt| { - opt.map(|batch_result| { - batch_result.and_then(|batch| { - let options = - RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); - RecordBatch::try_new_with_options( - Arc::clone(&self.schema), - batch.columns().to_vec(), - &options, - ) - .map_err(Into::into) - }) - }) - }) - } - - fn size_hint(&self) -> (usize, Option) { - self.inner.size_hint() - } -} - -/// Wraps `stream` in a [`SchemaConformingStream`] if its schema disagrees -/// with `schema`, otherwise returns it unchanged. See -/// [`SchemaConformingStream`] and -/// . -fn conform_stream_schema( - schema: SchemaRef, - stream: SendableRecordBatchStream, -) -> SendableRecordBatchStream { - if stream.schema() == schema { - stream - } else { - Box::pin(SchemaConformingStream::new(schema, stream)) - } -} - -/// Coerces a single child's declared output schema to `schema`, re-stamping -/// every batch it produces to match. [`UnionExec::try_new`] and -/// [`InterleaveExec::try_new`] insert this above any child whose own schema -/// disagrees with the computed union schema, so that the coercion is visible -/// in the plan tree (e.g. in `EXPLAIN`) instead of happening invisibly inside -/// the union operator's own `execute()`. -/// -/// A genuine data type mismatch (as opposed to a nullability-only one) is -/// rejected eagerly, at construction time, by `EquivalenceProperties:: -/// with_new_schema` below -- unlike the old purely-runtime approach, a -/// hand-built union/interleave with mismatched child types now fails in -/// `try_new` rather than in `execute()`. +/// A column whose `DataType` doesn't already match the union's is a genuine +/// data type mismatch (as opposed to a nullability/name/metadata-only one), +/// and is rejected eagerly here rather than silently cast or deferred to a +/// runtime failure -- this only ever changes a column's declared schema, +/// never its values. /// -/// This node is a strict 1:1, order-preserving passthrough of `input` (it -/// only ever changes a batch's declared schema, never its rows), so every -/// `ExecutionPlan` method below that isn't about the schema itself just -/// delegates straight to `input`. +/// Casting a column to its own `DataType` (only the `Field`'s nullability, +/// name, or metadata changes) is a zero-copy relabeling: the cast kernel's +/// same-type fast path (`cast_array_by_name`) just clones the `Arc`, so this carries no runtime overhead over the schema it replaces. /// /// See . -#[derive(Debug)] -struct CoerceSchemaExec { +fn coerce_schema( input: Arc, - cache: Arc, - metrics: ExecutionPlanMetricsSet, -} - -impl CoerceSchemaExec { - /// Wraps `input` in a [`CoerceSchemaExec`] targeting `schema` if its own - /// schema disagrees with `schema`, otherwise returns it unchanged. - fn wrap_if_needed( - input: Arc, - schema: &SchemaRef, - ) -> Result> { - if &input.schema() == schema { - Ok(input) - } else { - Ok(Arc::new(Self::new(input, schema)?)) - } + schema: &SchemaRef, +) -> Result> { + let input_schema = input.schema(); + if &input_schema == schema { + return Ok(input); } - fn new(input: Arc, schema: &SchemaRef) -> Result { - let eq_properties = input - .equivalence_properties() - .clone() - .with_new_schema(Arc::clone(schema))?; - let output_partitioning = input.output_partitioning().clone(); - let cache = PlanProperties::new( - eq_properties, - output_partitioning, - emission_type_from_children(std::iter::once(&input)), - boundedness_from_children(std::iter::once(&input)), - ); - Ok(Self { - input, - cache: Arc::new(cache), - metrics: ExecutionPlanMetricsSet::new(), - }) - } -} - -impl DisplayAs for CoerceSchemaExec { - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "CoerceSchemaExec") + let exprs = input_schema + .fields() + .iter() + .zip(schema.fields()) + .enumerate() + .map(|(i, (input_field, target_field))| { + if input_field.data_type() != target_field.data_type() { + return plan_err!( + "UnionExec/InterleaveExec requires all inputs to have the same \ + data type per column; column {i} has type {} in one input, but \ + the union schema expects {}", + input_field.data_type(), + target_field.data_type() + ); } - DisplayFormatType::TreeRender => Ok(()), - } - } -} - -impl ExecutionPlan for CoerceSchemaExec { - fn name(&self) -> &'static str { - "CoerceSchemaExec" - } - - fn properties(&self) -> &Arc { - &self.cache - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn maintains_input_order(&self) -> Vec { - vec![true] - } - - // A 1:1 passthrough never combines partitions, so re-deriving the cache - // (rather than collapsing back to the raw child via `wrap_if_needed`) is - // always safe here -- it just keeps this node from vanishing and - // changing arity out from under a caller mid-rewrite. - fn with_new_children( - self: Arc, - mut children: Vec>, - ) -> Result> { - assert_or_internal_err!( - children.len() == 1, - "CoerceSchemaExec expects exactly one child" - ); - Ok(Arc::new(Self::new(children.remove(0), &self.schema())?)) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); - let stream = self.input.execute(partition, context)?; - let stream = conform_stream_schema(self.schema(), stream); - Ok(Box::pin(ObservedStream::new( - stream, - baseline_metrics, - None, - ))) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn benefits_from_input_partitioning(&self) -> Vec { - vec![false] - } - - fn supports_limit_pushdown(&self) -> bool { - true - } - - fn cardinality_effect(&self) -> CardinalityEffect { - CardinalityEffect::Equal - } - - fn child_stats_requests(&self, partition: Option) -> Vec { - vec![ChildStats::At(partition)] - } - - fn statistics_from_inputs( - &self, - input_stats: &[Arc], - _args: &StatisticsArgs, - ) -> Result> { - Ok(Arc::clone(&input_stats[0])) - } - - fn gather_filters_for_pushdown( - &self, - _phase: FilterPushdownPhase, - parent_filters: Vec>, - _config: &ConfigOptions, - ) -> Result { - FilterDescription::from_children(parent_filters, &self.children()) - } + let column: Arc = + Arc::new(Column::new(input_field.name(), i)); + let expr = if input_field == target_field { + column + } else { + Arc::new(CastExpr::new_with_target_field( + column, + Arc::clone(target_field), + None, + )) as Arc + }; + Ok(ProjectionExpr { + expr, + alias: target_field.name().clone(), + }) + }) + .collect::>>()?; - #[cfg(feature = "proto")] - fn try_to_proto( - &self, - ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, - ) -> Result> { - // No dedicated protobuf variant: this node is fully determined by its - // child and the parent union/interleave's declared schema, so it's - // serialized as if it weren't there. `UnionExec`/`InterleaveExec:: - // try_from_proto` both go through `try_new`, which re-inserts the - // wrapper via `wrap_if_needed` on decode. - Ok(Some(ctx.encode_child(&self.input)?)) - } + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) } /// `UnionExec`: `UNION ALL` execution plan. @@ -377,7 +198,7 @@ impl UnionExec { // - Their fields have same types at the same indices. let inputs = inputs .into_iter() - .map(|input| CoerceSchemaExec::wrap_if_needed(input, &schema)) + .map(|input| coerce_schema(input, &schema)) .collect::>>()?; let cache = Self::compute_properties(&inputs, schema)?; Ok(Arc::new(UnionExec { @@ -815,7 +636,7 @@ impl InterleaveExec { let schema = union_schema(&inputs)?; let inputs = inputs .into_iter() - .map(|input| CoerceSchemaExec::wrap_if_needed(input, &schema)) + .map(|input| coerce_schema(input, &schema)) .collect::>>()?; let cache = Self::compute_properties(&inputs, schema)?; Ok(InterleaveExec { @@ -1441,13 +1262,23 @@ mod tests { #[test] fn test_union_partition_statistics_with_mismatched_nullability() -> Result<()> { - // Regression test for the `CoerceSchemaExec` wrapper `UnionExec::try_new` - // inserts above the non-nullable leg here: before it forwarded - // `child_stats_requests`/`statistics_from_inputs` to its child, it - // reported `Statistics::new_unknown`, poisoning the merge into all-`Absent` - // even though both legs have exact statistics. + // Regression test for the `ProjectionExec` wrapper `UnionExec::try_new` + // inserts above the non-nullable leg here (via `coerce_schema`): + // exact column statistics (min/max/null/distinct/sum/byte_size) must + // still make it through the wrapper's same-type `CastExpr`, not get + // poisoned into `Absent` the way a generic (type-changing) cast's + // statistics would be. let (_, left, right, expected) = stats_merge_inputs(); + // `total_byte_size` differs from the plain-merge fixture (52): the + // wrapper is a `ProjectionExec`, whose `statistics_from_inputs` + // recomputes `total_byte_size` from the (unchanged) schema's row + // width times row count, rather than trusting the wrapped leg's own + // self-reported total -- still `Exact`, just derived differently. + // left: 5 rows * 4 bytes (UInt32) = 20 (was 23); right is untouched + // (already nullable, so `coerce_schema` doesn't wrap it): 20 + 29 = 49. + let expected = expected.with_total_byte_size(Precision::Exact(49)); + let non_nullable_schema = Schema::new(vec![Field::new("a", DataType::UInt32, false)]); let nullable_schema = Schema::new(vec![Field::new("a", DataType::UInt32, true)]); @@ -1466,9 +1297,23 @@ mod tests { } #[tokio::test] - async fn test_coerce_schema_exec_execution_plan_methods() -> Result<()> { - // Most sqllogictest coverage only observes `CoerceSchemaExec` through - // `EXPLAIN`; exercise its `ExecutionPlan` methods directly here. + async fn test_coerce_schema_no_op_when_already_matching() -> Result<()> { + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input: Arc = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_not_null), None)?; + + let coerced = coerce_schema(Arc::clone(&input), &schema_not_null)?; + assert!(Arc::ptr_eq(&coerced, &input)); + + Ok(()) + } + + #[tokio::test] + async fn test_coerce_schema_casts_only_nullability() -> Result<()> { + // Mismatched nullability: the input gets wrapped in a `ProjectionExec` + // whose `CastExpr` re-stamps the column with the target's `Field` + // (same `DataType`, so this is a zero-copy relabeling, not a real cast). let schema_not_null = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); let batch_not_null = RecordBatch::try_new( @@ -1481,47 +1326,41 @@ mod tests { None, )?; - // Matching schema: `wrap_if_needed` is a no-op. - let unwrapped = - CoerceSchemaExec::wrap_if_needed(Arc::clone(&input), &schema_not_null)?; - assert!(Arc::ptr_eq(&unwrapped, &input)); - - // Mismatched nullability: the input gets wrapped. let nullable_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); - let wrapped = - CoerceSchemaExec::wrap_if_needed(Arc::clone(&input), &nullable_schema)?; - assert_eq!(wrapped.name(), "CoerceSchemaExec"); - assert_eq!(&wrapped.schema(), &nullable_schema); - assert_eq!(wrapped.maintains_input_order(), vec![true]); - assert_eq!(wrapped.benefits_from_input_partitioning(), vec![false]); - assert!(wrapped.supports_limit_pushdown()); - assert!(matches!( - wrapped.cardinality_effect(), - CardinalityEffect::Equal - )); - assert!(wrapped.metrics().is_some()); - - wrapped.gather_filters_for_pushdown( - FilterPushdownPhase::Pre, - vec![], - &ConfigOptions::new(), - )?; - - // Re-deriving via `with_new_children` keeps targeting the same schema. - let new_child: Arc = - TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_not_null), None)?; - let rewrapped = Arc::clone(&wrapped).with_new_children(vec![new_child])?; - assert_eq!(&rewrapped.schema(), &nullable_schema); + let coerced = coerce_schema(Arc::clone(&input), &nullable_schema)?; + assert_eq!(&coerced.schema(), &nullable_schema); + let plan_str = crate::displayable(coerced.as_ref()) + .indent(true) + .to_string(); + assert!( + plan_str.contains("CAST"), + "expected a CAST in the coerced plan:\n{plan_str}" + ); let task_ctx = Arc::new(TaskContext::default()); - let batches = collect(wrapped, task_ctx).await?; + let batches = collect(coerced, task_ctx).await?; assert_eq!(batches.len(), 1); assert_eq!(batches[0].schema(), nullable_schema); Ok(()) } + #[test] + fn test_coerce_schema_rejects_genuine_type_mismatch() -> Result<()> { + let schema_int = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input: Arc = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_int), None)?; + + let schema_utf8 = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)])); + let err = coerce_schema(input, &schema_utf8).unwrap_err(); + assert!(err.to_string().contains("same data type per column")); + + Ok(()) + } + #[test] fn test_interleave_partition_statistics_uses_shared_statistics_merge() -> Result<()> { let (schema, left, right, expected) = stats_merge_inputs(); diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index f42296380c977..a8c71a181811c 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2330,10 +2330,13 @@ fn roundtrip_union() -> Result<()> { roundtrip_test(union) } -/// `CoerceSchemaExec::try_to_proto` intentionally leaves the wrapper out of -/// the serialized plan, relying on `UnionExec::try_from_proto` to rebuild it -/// via `try_new`. Verify the decoded plan still contains the coercion and -/// that its emitted batches expose the union's nullable schema. +/// `UnionExec::try_new` coerces a nullability-mismatched leg by wrapping it +/// in a `ProjectionExec` with a same-type `CastExpr` (see `coerce_schema` in +/// `datafusion-physical-plan`'s `union` module) -- a zero-copy relabeling, +/// not a real cast. `ProjectionExec` has an ordinary protobuf message, so +/// unlike the node this replaced, there's no wrapper-erasure trick to verify; +/// just that the decoded plan still contains the coercion and that its +/// emitted batches expose the union's nullable schema. #[tokio::test] async fn roundtrip_union_with_mismatched_nullability_executes() -> Result<()> { let literal_leg = |value: ScalarValue| -> Result> { @@ -2352,8 +2355,8 @@ async fn roundtrip_union_with_mismatched_nullability_executes() -> Result<()> { UnionExec::try_new(vec![non_nullable_leg, nullable_leg])?; assert!(union.schema().field(0).is_nullable()); assert!( - format!("{union:?}").contains("CoerceSchemaExec"), - "expected CoerceSchemaExec in plan:\n{union:?}" + format!("{union:?}").contains("CastExpr"), + "expected a coercing CastExpr in plan:\n{union:?}" ); let ctx = SessionContext::new(); @@ -2364,8 +2367,8 @@ async fn roundtrip_union_with_mismatched_nullability_executes() -> Result<()> { )?; assert!(roundtripped.schema().field(0).is_nullable()); assert!( - format!("{roundtripped:?}").contains("CoerceSchemaExec"), - "expected CoerceSchemaExec after roundtrip:\n{roundtripped:?}" + format!("{roundtripped:?}").contains("CastExpr"), + "expected a coercing CastExpr after roundtrip:\n{roundtripped:?}" ); let batches = @@ -2551,8 +2554,8 @@ async fn roundtrip_interleave_with_mismatched_nullability_executes() -> Result<( ])?); assert!(interleave.schema().field(0).is_nullable()); assert!( - format!("{interleave:?}").contains("CoerceSchemaExec"), - "expected CoerceSchemaExec in plan:\n{interleave:?}" + format!("{interleave:?}").contains("CastExpr"), + "expected a coercing CastExpr in plan:\n{interleave:?}" ); let ctx = SessionContext::new(); @@ -2563,8 +2566,8 @@ async fn roundtrip_interleave_with_mismatched_nullability_executes() -> Result<( )?; assert!(roundtripped.schema().field(0).is_nullable()); assert!( - format!("{roundtripped:?}").contains("CoerceSchemaExec"), - "expected CoerceSchemaExec after roundtrip:\n{roundtripped:?}" + format!("{roundtripped:?}").contains("CastExpr"), + "expected a coercing CastExpr after roundtrip:\n{roundtripped:?}" ); let batches = diff --git a/datafusion/sqllogictest/test_files/array_agg.slt b/datafusion/sqllogictest/test_files/array_agg.slt index 36ad13d2491bc..d5aaf8cab17c1 100644 --- a/datafusion/sqllogictest/test_files/array_agg.slt +++ b/datafusion/sqllogictest/test_files/array_agg.slt @@ -534,19 +534,16 @@ physical_plan 03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5 04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted 05)--------UnionExec -06)----------CoerceSchemaExec -07)------------ProjectionExec: expr=[1 as id, 2 as foo] -08)--------------PlaceholderRowExec -09)----------ProjectionExec: expr=[1 as id, NULL as foo] -10)------------PlaceholderRowExec -11)----------ProjectionExec: expr=[1 as id, NULL as foo] -12)------------PlaceholderRowExec -13)----------CoerceSchemaExec -14)------------ProjectionExec: expr=[1 as id, 3 as foo] -15)--------------PlaceholderRowExec -16)----------CoerceSchemaExec -17)------------ProjectionExec: expr=[1 as id, 2 as foo] -18)--------------PlaceholderRowExec +06)----------ProjectionExec: expr=[1 as id, CAST(2 AS Int64) as foo] +07)------------PlaceholderRowExec +08)----------ProjectionExec: expr=[1 as id, NULL as foo] +09)------------PlaceholderRowExec +10)----------ProjectionExec: expr=[1 as id, NULL as foo] +11)------------PlaceholderRowExec +12)----------ProjectionExec: expr=[1 as id, CAST(3 AS Int64) as foo] +13)------------PlaceholderRowExec +14)----------ProjectionExec: expr=[1 as id, CAST(2 AS Int64) as foo] +15)------------PlaceholderRowExec ####### # Unsupported syntax diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index f5dd6db505ea2..d4776dd0c0ddb 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -572,27 +572,25 @@ logical_plan physical_plan 01)CoalescePartitionsExec: fetch=3 02)--UnionExec -03)----CoerceSchemaExec -04)------ProjectionExec: expr=[count(Int64(1))@0 as cnt] -05)--------GlobalLimitExec: skip=0, fetch=3 -06)----------AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] -07)------------CoalescePartitionsExec -08)--------------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] -09)----------------ProjectionExec: expr=[] -10)------------------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[] -11)--------------------RepartitionExec: partitioning=Hash([c1@0], 4), input_partitions=4 -12)----------------------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[] -13)------------------------FilterExec: c13@1 != C2GT5KVyOPZpgKVl110TyZO0NcJ434, projection=[c1@0] -14)--------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -15)----------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c13], file_type=csv, has_header=true -16)----CoerceSchemaExec -17)------ProjectionExec: expr=[1 as cnt] -18)--------PlaceholderRowExec -19)----ProjectionExec: expr=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as cnt] -20)------GlobalLimitExec: skip=0, fetch=3 -21)--------BoundedWindowAggExec: wdw=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Field { "lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING], mode=[Sorted] -22)----------ProjectionExec: expr=[1 as c1] -23)------------PlaceholderRowExec +03)----ProjectionExec: expr=[CAST(count(Int64(1))@0 AS Int64) as cnt] +04)------GlobalLimitExec: skip=0, fetch=3 +05)--------AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] +06)----------CoalescePartitionsExec +07)------------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] +08)--------------ProjectionExec: expr=[] +09)----------------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[] +10)------------------RepartitionExec: partitioning=Hash([c1@0], 4), input_partitions=4 +11)--------------------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[] +12)----------------------FilterExec: c13@1 != C2GT5KVyOPZpgKVl110TyZO0NcJ434, projection=[c1@0] +13)------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +14)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c13], file_type=csv, has_header=true +15)----ProjectionExec: expr=[CAST(1 AS Int64) as cnt] +16)------PlaceholderRowExec +17)----ProjectionExec: expr=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as cnt] +18)------GlobalLimitExec: skip=0, fetch=3 +19)--------BoundedWindowAggExec: wdw=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Field { "lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING], mode=[Sorted] +20)----------ProjectionExec: expr=[1 as c1] +21)------------PlaceholderRowExec ######## @@ -723,14 +721,13 @@ logical_plan 11)----------EmptyRelation: rows=1 physical_plan 01)UnionExec -02)--CoerceSchemaExec -03)----ProjectionExec: expr=[count(Int64(1))@1 as count, n@0 as n] -04)------AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))], ordering_mode=Sorted -05)--------ProjectionExec: expr=[5 as n] -06)----------PlaceholderRowExec -07)--ProjectionExec: expr=[1 as count, max(Int64(10))@0 as n] -08)----AggregateExec: mode=Single, gby=[], aggr=[max(Int64(10))] -09)------PlaceholderRowExec +02)--ProjectionExec: expr=[count(Int64(1))@1 as count, CAST(n@0 AS Int64) as n] +03)----AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))], ordering_mode=Sorted +04)------ProjectionExec: expr=[5 as n] +05)--------PlaceholderRowExec +06)--ProjectionExec: expr=[1 as count, max(Int64(10))@0 as n] +07)----AggregateExec: mode=Single, gby=[], aggr=[max(Int64(10))] +08)------PlaceholderRowExec # Test issue: https://github.com/apache/datafusion/issues/11409