From ff134c9fd05f9eae0007ff9187cbf71b5e2d4f9b Mon Sep 17 00:00:00 2001 From: naman Date: Wed, 9 Sep 2026 02:08:12 +0530 Subject: [PATCH 1/3] fix: Derive Substrait intersection nullability from every input The Substrait consumer derived all three intersection schemas from the primary input alone, so a field that the intersection makes required stayed nullable in the logical output schema. Narrow an intersection's nullability to `left AND right` per field. The left semi join it compiles to matches nulls with nulls, so a field is nullable in the result only when both inputs make it nullable, which reproduces the spec's rule for the multiset intersections and, because the right side is the union of the secondary inputs, for the primary intersection as well. Closes #25042. --- .../src/logical_plan/consumer/rel/set_rel.rs | 88 +++++++++- .../substrait/tests/cases/logical_plans.rs | 44 +++++ ...tiset_all_mixed_nullability.substrait.json | 160 ++++++++++++++++++ ..._multiset_mixed_nullability.substrait.json | 160 ++++++++++++++++++ ...t_primary_mixed_nullability.substrait.json | 160 ++++++++++++++++++ 5 files changed, 604 insertions(+), 8 deletions(-) create mode 100644 datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json create mode 100644 datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json create mode 100644 datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs index 36bf8dbae4a92..7052be2dd2e36 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -16,8 +16,9 @@ // under the License. use crate::logical_plan::consumer::SubstraitConsumer; -use datafusion::common::{not_impl_err, substrait_err}; -use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder}; +use datafusion::common::{DFSchema, not_impl_err, substrait_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder, Projection}; +use std::sync::Arc; use substrait::proto::set_rel::SetOp; use substrait::proto::{Rel, SetRel}; @@ -31,7 +32,7 @@ pub async fn from_set_rel( match set.op() { SetOp::UnionAll => union_rels(consumer, &set.inputs, true).await, SetOp::UnionDistinct => union_rels(consumer, &set.inputs, false).await, - SetOp::IntersectionPrimary => LogicalPlanBuilder::intersect( + SetOp::IntersectionPrimary => intersect_rel( consumer.consume_rel(&set.inputs[0]).await?, union_rels(consumer, &set.inputs[1..], true).await?, false, @@ -77,16 +78,87 @@ async fn intersect_rels( let mut rel = consumer.consume_rel(&rels[0]).await?; for input in &rels[1..] { - rel = LogicalPlanBuilder::intersect( - rel, - consumer.consume_rel(input).await?, - is_all, - )?; + rel = intersect_rel(rel, consumer.consume_rel(input).await?, is_all)?; } Ok(rel) } +/// Intersects two relations, giving the result the nullability the Substrait +/// [Set Operation rules] prescribe. +/// +/// [`LogicalPlanBuilder::intersect`] compiles an intersection into a left semi +/// join, so on its own the result keeps the left input's nullability. The join +/// matches nulls with nulls, so a left row holding a null in some field only +/// survives when the right input holds a null there too. A field is therefore +/// nullable in the result only when it is nullable in *both* inputs. +/// +/// Applied to each step of a chain, that gives the spec's rule for the multiset +/// intersections - a field is required when any input requires it. For +/// `INTERSECTION_PRIMARY` the right side is the union of the secondary inputs, +/// whose field is nullable exactly when some secondary input makes it nullable, +/// so the same rule yields "nullable in the primary input and in at least one +/// secondary input". +/// +/// [Set Operation rules]: https://substrait.io/relations/logical_relations/#set-operation +fn intersect_rel( + left: LogicalPlan, + right: LogicalPlan, + is_all: bool, +) -> datafusion::common::Result { + let right_nullability: Vec = right + .schema() + .fields() + .iter() + .map(|field| field.is_nullable()) + .collect(); + + let plan = LogicalPlanBuilder::intersect(left, right, is_all)?; + + // `intersect` has already checked that both sides have the same width. + let narrowed: Vec = plan + .schema() + .fields() + .iter() + .zip(&right_nullability) + .map(|(field, right_nullable)| field.is_nullable() && !right_nullable) + .collect(); + + if !narrowed.contains(&true) { + return Ok(plan); + } + + let qualified_fields = plan + .schema() + .iter() + .zip(&narrowed) + .map(|((qualifier, field), narrow)| { + let field = if *narrow { + Arc::new(field.as_ref().clone().with_nullable(false)) + } else { + Arc::clone(field) + }; + (qualifier.cloned(), field) + }) + .collect(); + let schema = Arc::new(DFSchema::new_with_metadata( + qualified_fields, + plan.schema().metadata().clone(), + )?); + + let exprs = plan + .schema() + .columns() + .into_iter() + .map(Expr::Column) + .collect(); + Ok(LogicalPlan::Projection(Projection::try_new_with_schema( + exprs, + Arc::new(plan), + schema, + )?)) +} + async fn except_rels( consumer: &impl SubstraitConsumer, rels: &[Rel], diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 522381de6efdf..eee7fd33b83ab 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -229,6 +229,50 @@ mod tests { Ok(()) } + #[tokio::test] + async fn intersect_nullability() -> Result<()> { + // Substrait's set operation rules derive an intersection's nullability from + // every input, not only the primary one. Each plan below intersects three + // tables carrying the same four columns, with these nullabilities + // (`?` marks a nullable column): + // + // primary a? b? c? d? + // secondary a b c? d? + // secondary a b? c d? + for (file, expected) in [ + // Nullable in the primary input and in at least one secondary input. + ("intersect_primary_mixed_nullability", "a, b?, c?, d?"), + // Required as soon as any input requires it. + ("intersect_multiset_mixed_nullability", "a, b, c, d?"), + ("intersect_multiset_all_mixed_nullability", "a, b, c, d?"), + ] { + let proto_plan = + read_json(&format!("tests/testdata/test_plans/{file}.substrait.json")); + let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + let nullability = plan + .schema() + .fields() + .iter() + .map(|field| { + format!( + "{}{}", + field.name(), + if field.is_nullable() { "?" } else { "" } + ) + }) + .collect::>() + .join(", "); + assert_eq!(nullability, expected, "nullability of {file}"); + + // Trigger execution to ensure plan validity + DataFrame::new(ctx.state(), plan).show().await?; + } + + Ok(()) + } + #[tokio::test] async fn multilayer_aggregate() -> Result<()> { let proto_plan = diff --git a/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json new file mode 100644 index 0000000000000..c2463b04beb5c --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json @@ -0,0 +1,160 @@ +{ + "relations": [ + { + "root": { + "input": { + "set": { + "inputs": [ + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data2" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + } + ], + "op": "SET_OP_INTERSECTION_MULTISET_ALL" + } + }, + "names": [ + "a", + "b", + "c", + "d" + ] + } + } + ], + "version": { + "minorNumber": 54, + "producer": "datafusion-test" + } +} diff --git a/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json new file mode 100644 index 0000000000000..566f40214e276 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json @@ -0,0 +1,160 @@ +{ + "relations": [ + { + "root": { + "input": { + "set": { + "inputs": [ + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data2" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + } + ], + "op": "SET_OP_INTERSECTION_MULTISET" + } + }, + "names": [ + "a", + "b", + "c", + "d" + ] + } + } + ], + "version": { + "minorNumber": 54, + "producer": "datafusion-test" + } +} diff --git a/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json b/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json new file mode 100644 index 0000000000000..f035005f21091 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json @@ -0,0 +1,160 @@ +{ + "relations": [ + { + "root": { + "input": { + "set": { + "inputs": [ + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data2" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + } + ], + "op": "SET_OP_INTERSECTION_PRIMARY" + } + }, + "names": [ + "a", + "b", + "c", + "d" + ] + } + } + ], + "version": { + "minorNumber": 54, + "producer": "datafusion-test" + } +} From 1a420328d6cf01e3c3453bf3993db14599534e8b Mon Sep 17 00:00:00 2001 From: naman Date: Fri, 18 Sep 2026 23:28:59 +0530 Subject: [PATCH 2/3] Derive the narrowed intersection fields in the physical plan too A Projection built with Projection::try_new_with_schema only narrowed the logical schema: the physical ProjectionExec takes field metadata from it but keeps the nullability of its input, so the physical plan and the collected batches still reported the left input's nullability. When the right input requires a field that the left input leaves nullable, build the intersection as an inner join (nulls equal nulls) against the distinct right rows and read that field from the right side. Both planners then derive the field as non-nullable from its source column. A field is only read from the right when its type and metadata match the left's, and the schema metadata matches too, so the result keeps the left input's attributes. The test now gives the tables rows and checks that the logical schema, the physical plan schema and every batch schema agree, and adds columns with NULLABILITY_UNSPECIFIED on a secondary input. --- .../src/logical_plan/consumer/rel/set_rel.rs | 107 +++++++++------- .../substrait/tests/cases/logical_plans.rs | 119 ++++++++++++++++-- ...tiset_all_mixed_nullability.substrait.json | 46 ++++++- ..._multiset_mixed_nullability.substrait.json | 46 ++++++- ...t_primary_mixed_nullability.substrait.json | 46 ++++++- 5 files changed, 293 insertions(+), 71 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs index 7052be2dd2e36..12b26df023db8 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -16,9 +16,10 @@ // under the License. use crate::logical_plan::consumer::SubstraitConsumer; -use datafusion::common::{DFSchema, not_impl_err, substrait_err}; -use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder, Projection}; -use std::sync::Arc; +use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err}; +use datafusion::logical_expr::{ + Expr, LogicalPlan, LogicalPlanBuilder, requalify_sides_if_needed, +}; use substrait::proto::set_rel::SetOp; use substrait::proto::{Rel, SetRel}; @@ -100,63 +101,77 @@ async fn intersect_rels( /// so the same rule yields "nullable in the primary input and in at least one /// secondary input". /// +/// When the right input requires a field the left input leaves nullable, the +/// intersection is built as an inner join against the distinct right rows +/// instead, and that field is read from the right side. Matched rows hold equal +/// values, so the result is unchanged, and the field is non-nullable because +/// its source is: the logical and the physical planner both derive that from +/// the input schema, so the plan, the physical plan and the batches agree. +/// Joining against distinct right rows keeps each left row at most once, as the +/// semi join does. +/// /// [Set Operation rules]: https://substrait.io/relations/logical_relations/#set-operation fn intersect_rel( left: LogicalPlan, right: LogicalPlan, is_all: bool, ) -> datafusion::common::Result { - let right_nullability: Vec = right - .schema() - .fields() - .iter() - .map(|field| field.is_nullable()) - .collect(); - - let plan = LogicalPlanBuilder::intersect(left, right, is_all)?; - - // `intersect` has already checked that both sides have the same width. - let narrowed: Vec = plan - .schema() - .fields() + let left_fields = left.schema().fields(); + let right_fields = right.schema().fields(); + // Only a field that differs from its right counterpart in nullability alone + // is read from the right side, so every other attribute stays the left's. + let from_right: Vec = left_fields .iter() - .zip(&right_nullability) - .map(|(field, right_nullable)| field.is_nullable() && !right_nullable) + .zip(right_fields.iter()) + .map(|(left, right)| { + left.is_nullable() + && !right.is_nullable() + && left.data_type() == right.data_type() + && left.metadata() == right.metadata() + }) .collect(); - if !narrowed.contains(&true) { - return Ok(plan); + // `intersect` also reports inputs of different widths. The join would merge + // the right input's schema metadata into the result, so that must match too. + if left_fields.len() != right_fields.len() + || left.schema().metadata() != right.schema().metadata() + || !from_right.contains(&true) + { + return LogicalPlanBuilder::intersect(left, right, is_all); } - let qualified_fields = plan - .schema() + let (left, right, _) = requalify_sides_if_needed( + LogicalPlanBuilder::from(left), + LogicalPlanBuilder::from(right), + )?; + let left = if is_all { left } else { left.distinct()? }; + let right = right.distinct()?.build()?; + + let left_columns = left.schema().columns(); + let right_columns = right.schema().columns(); + let exprs = left_columns .iter() - .zip(&narrowed) - .map(|((qualifier, field), narrow)| { - let field = if *narrow { - Arc::new(field.as_ref().clone().with_nullable(false)) + .zip(&right_columns) + .zip(&from_right) + .map(|((left, right), from_right)| { + if *from_right { + Expr::Column(right.clone()) + .alias_qualified(left.relation.clone(), &left.name) } else { - Arc::clone(field) - }; - (qualifier.cloned(), field) + Expr::Column(left.clone()) + } }) - .collect(); - let schema = Arc::new(DFSchema::new_with_metadata( - qualified_fields, - plan.schema().metadata().clone(), - )?); - - let exprs = plan - .schema() - .columns() - .into_iter() - .map(Expr::Column) - .collect(); - Ok(LogicalPlan::Projection(Projection::try_new_with_schema( - exprs, - Arc::new(plan), - schema, - )?)) + .collect::>(); + + left.join_detailed( + right, + JoinType::Inner, + (left_columns, right_columns), + None, + NullEquality::NullEqualsNull, + )? + .project(exprs)? + .build() } async fn except_rels( diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index eee7fd33b83ab..4295a7d5ef835 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -21,8 +21,12 @@ mod tests { use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; + use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch}; + use datafusion::assert_batches_sorted_eq; use datafusion::common::test_util::format_batches; + use datafusion::datasource::MemTable; use std::collections::HashSet; + use std::sync::Arc; use datafusion::common::Result; use datafusion::dataframe::DataFrame; @@ -233,22 +237,97 @@ mod tests { async fn intersect_nullability() -> Result<()> { // Substrait's set operation rules derive an intersection's nullability from // every input, not only the primary one. Each plan below intersects three - // tables carrying the same four columns, with these nullabilities - // (`?` marks a nullable column): + // tables carrying the same six columns, with these nullabilities (`?` marks + // a nullable column, `~` a column with unspecified nullability, which the + // consumer reads as nullable): // - // primary a? b? c? d? - // secondary a b c? d? - // secondary a b? c d? - for (file, expected) in [ + // primary a? b? c? d? e? f? + // secondary a b c? d? e~ f~ + // secondary a b? c d? e? f + let rows: [(&str, &[[Option; 6]]); 3] = [ + ( + "data", + &[ + [Some(1), Some(1), Some(1), None, None, Some(1)], + [Some(2), None, Some(2), Some(2), Some(2), Some(2)], + [Some(3), Some(3), None, Some(3), Some(3), Some(3)], + [None, Some(4), Some(4), Some(4), Some(4), Some(4)], + ], + ), + ( + "data2", + &[ + [Some(1), Some(1), Some(1), None, None, Some(1)], + [Some(3), Some(3), None, Some(3), Some(3), Some(3)], + ], + ), + ( + "data3", + &[ + [Some(1), Some(1), Some(1), None, None, Some(1)], + [Some(2), None, Some(2), Some(2), Some(2), Some(2)], + ], + ), + ]; + + for (file, expected_nullability, expected_rows) in [ // Nullable in the primary input and in at least one secondary input. - ("intersect_primary_mixed_nullability", "a, b?, c?, d?"), + ( + "intersect_primary_mixed_nullability", + "a, b?, c?, d?, e?, f?", + &[ + "+---+---+---+---+---+---+", + "| a | b | c | d | e | f |", + "+---+---+---+---+---+---+", + "| 1 | 1 | 1 | | | 1 |", + "| 2 | | 2 | 2 | 2 | 2 |", + "| 3 | 3 | | 3 | 3 | 3 |", + "+---+---+---+---+---+---+", + ][..], + ), // Required as soon as any input requires it. - ("intersect_multiset_mixed_nullability", "a, b, c, d?"), - ("intersect_multiset_all_mixed_nullability", "a, b, c, d?"), + ( + "intersect_multiset_mixed_nullability", + "a, b, c, d?, e?, f", + &[ + "+---+---+---+---+---+---+", + "| a | b | c | d | e | f |", + "+---+---+---+---+---+---+", + "| 1 | 1 | 1 | | | 1 |", + "+---+---+---+---+---+---+", + ][..], + ), + ( + "intersect_multiset_all_mixed_nullability", + "a, b, c, d?, e?, f", + &[ + "+---+---+---+---+---+---+", + "| a | b | c | d | e | f |", + "+---+---+---+---+---+---+", + "| 1 | 1 | 1 | | | 1 |", + "+---+---+---+---+---+---+", + ][..], + ), ] { let proto_plan = read_json(&format!("tests/testdata/test_plans/{file}.substrait.json")); let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + // Give each table rows, so the batch schemas below come from real batches + for (table, rows) in rows { + let schema = ctx.table_provider(table).await?.schema(); + let columns = (0..schema.fields().len()) + .map(|i| { + Arc::new(rows.iter().map(|row| row[i]).collect::()) + as ArrayRef + }) + .collect(); + let batch = RecordBatch::try_new(Arc::clone(&schema), columns)?; + ctx.deregister_table(table)?; + ctx.register_table( + table, + Arc::new(MemTable::try_new(schema, vec![vec![batch]])?), + )?; + } let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; let nullability = plan @@ -264,10 +343,24 @@ mod tests { }) .collect::>() .join(", "); - assert_eq!(nullability, expected, "nullability of {file}"); - - // Trigger execution to ensure plan validity - DataFrame::new(ctx.state(), plan).show().await?; + assert_eq!(nullability, expected_nullability, "nullability of {file}"); + + // The physical plan and the batches it produces must carry the same + // schema as the logical plan, not the left input's nullability + let logical_schema = Arc::clone(plan.schema().inner()); + let df = DataFrame::new(ctx.state(), plan); + let physical_plan = df.clone().create_physical_plan().await?; + assert_eq!( + physical_plan.schema(), + logical_schema, + "physical schema of {file}" + ); + let batches = df.collect().await?; + assert!(!batches.is_empty(), "no batches for {file}"); + for batch in &batches { + assert_eq!(batch.schema(), logical_schema, "batch schema of {file}"); + } + assert_batches_sorted_eq!(expected_rows, &batches); } Ok(()) diff --git a/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json index c2463b04beb5c..a99e7dd5720c6 100644 --- a/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json @@ -15,7 +15,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -34,6 +36,16 @@ "nullability": "NULLABILITY_NULLABLE" } }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, { "i64": { "nullability": "NULLABILITY_NULLABLE" @@ -60,7 +72,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -83,6 +97,16 @@ "i64": { "nullability": "NULLABILITY_NULLABLE" } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } } ], "nullability": "NULLABILITY_REQUIRED" @@ -105,7 +129,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -128,6 +154,16 @@ "i64": { "nullability": "NULLABILITY_NULLABLE" } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } } ], "nullability": "NULLABILITY_REQUIRED" @@ -148,7 +184,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ] } } diff --git a/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json index 566f40214e276..e896015c8a4e8 100644 --- a/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json @@ -15,7 +15,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -34,6 +36,16 @@ "nullability": "NULLABILITY_NULLABLE" } }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, { "i64": { "nullability": "NULLABILITY_NULLABLE" @@ -60,7 +72,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -83,6 +97,16 @@ "i64": { "nullability": "NULLABILITY_NULLABLE" } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } } ], "nullability": "NULLABILITY_REQUIRED" @@ -105,7 +129,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -128,6 +154,16 @@ "i64": { "nullability": "NULLABILITY_NULLABLE" } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } } ], "nullability": "NULLABILITY_REQUIRED" @@ -148,7 +184,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ] } } diff --git a/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json b/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json index f035005f21091..77ac9d50da576 100644 --- a/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json @@ -15,7 +15,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -34,6 +36,16 @@ "nullability": "NULLABILITY_NULLABLE" } }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, { "i64": { "nullability": "NULLABILITY_NULLABLE" @@ -60,7 +72,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -83,6 +97,16 @@ "i64": { "nullability": "NULLABILITY_NULLABLE" } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } } ], "nullability": "NULLABILITY_REQUIRED" @@ -105,7 +129,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ], "struct": { "types": [ @@ -128,6 +154,16 @@ "i64": { "nullability": "NULLABILITY_NULLABLE" } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } } ], "nullability": "NULLABILITY_REQUIRED" @@ -148,7 +184,9 @@ "a", "b", "c", - "d" + "d", + "e", + "f" ] } } From 650a5de1b928fb48df5b206a19d6a94964d7ca65 Mon Sep 17 00:00:00 2001 From: naman Date: Mon, 21 Sep 2026 17:43:12 +0530 Subject: [PATCH 3/3] Keep the intersection nullability fix when input metadata differs The narrowing path fell back to `LogicalPlanBuilder::intersect` when the inputs' schema metadata differed, and skipped a field whose metadata differed, so the nullable left-side schema came back for inputs that `ensure_schema_compatibility` accepts. Take the narrowing path regardless of metadata. A column read from the right is aliased with the left field's metadata, and the inner join already lets the left input's schema metadata win, so the left's metadata wins wherever the inputs disagree. `intersect_nullability` now also runs every plan over tables with differing schema and field metadata, and checks that the logical, the physical and the collected batch schemas agree. Co-Authored-By: Claude Sonnet 5 --- .../src/logical_plan/consumer/rel/set_rel.rs | 38 ++++++---- .../substrait/tests/cases/logical_plans.rs | 72 +++++++++++++++++-- 2 files changed, 93 insertions(+), 17 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs index 12b26df023db8..e27ade18cf144 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs @@ -16,6 +16,7 @@ // under the License. use crate::logical_plan::consumer::SubstraitConsumer; +use datafusion::common::metadata::FieldMetadata; use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err}; use datafusion::logical_expr::{ Expr, LogicalPlan, LogicalPlanBuilder, requalify_sides_if_needed, @@ -110,6 +111,15 @@ async fn intersect_rels( /// Joining against distinct right rows keeps each left row at most once, as the /// semi join does. /// +/// Differing metadata does not change which path is taken, as it is no part of +/// nullability. The result should describe the left input, as +/// [`LogicalPlanBuilder::intersect`] does, so on conflicting keys the left's +/// metadata wins: a column read from the right is aliased with the left field's +/// metadata, and an inner join's schema already lets the left input's schema +/// metadata win, in the logical and in the physical plan alike. Keys only the +/// right input carries are merged in, and the physical plan and the batches +/// carry the same metadata as the logical plan. +/// /// [Set Operation rules]: https://substrait.io/relations/logical_relations/#set-operation fn intersect_rel( left: LogicalPlan, @@ -118,8 +128,9 @@ fn intersect_rel( ) -> datafusion::common::Result { let left_fields = left.schema().fields(); let right_fields = right.schema().fields(); - // Only a field that differs from its right counterpart in nullability alone - // is read from the right side, so every other attribute stays the left's. + // A field is read from the right side when the left leaves it nullable and + // the right requires it. Its metadata does not matter here: it is read from + // the right with the left field's metadata layered over it. let from_right: Vec = left_fields .iter() .zip(right_fields.iter()) @@ -127,16 +138,11 @@ fn intersect_rel( left.is_nullable() && !right.is_nullable() && left.data_type() == right.data_type() - && left.metadata() == right.metadata() }) .collect(); - // `intersect` also reports inputs of different widths. The join would merge - // the right input's schema metadata into the result, so that must match too. - if left_fields.len() != right_fields.len() - || left.schema().metadata() != right.schema().metadata() - || !from_right.contains(&true) - { + // `intersect` also reports inputs of different widths. + if left_fields.len() != right_fields.len() || !from_right.contains(&true) { return LogicalPlanBuilder::intersect(left, right, is_all); } @@ -149,14 +155,20 @@ fn intersect_rel( let left_columns = left.schema().columns(); let right_columns = right.schema().columns(); - let exprs = left_columns + let exprs = left + .schema() + .fields() .iter() + .zip(&left_columns) .zip(&right_columns) .zip(&from_right) - .map(|((left, right), from_right)| { + .map(|(((field, left), right), from_right)| { if *from_right { - Expr::Column(right.clone()) - .alias_qualified(left.relation.clone(), &left.name) + Expr::Column(right.clone()).alias_qualified_with_metadata( + left.relation.clone(), + &left.name, + Some(FieldMetadata::from(field.metadata().clone())), + ) } else { Expr::Column(left.clone()) } diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 4295a7d5ef835..6219b2b4ec47f 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -22,10 +22,11 @@ mod tests { use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch}; + use datafusion::arrow::datatypes::{Field, Schema, SchemaRef}; use datafusion::assert_batches_sorted_eq; use datafusion::common::test_util::format_batches; use datafusion::datasource::MemTable; - use std::collections::HashSet; + use std::collections::{HashMap, HashSet}; use std::sync::Arc; use datafusion::common::Result; @@ -270,7 +271,39 @@ mod tests { ), ]; - for (file, expected_nullability, expected_rows) in [ + // Schema and field metadata are no part of an input's nullability, so the + // result must not depend on whether the tables carry any. With metadata, + // every table describes itself and the secondary tables add keys the + // primary one lacks: the result has to keep the primary table's metadata + // where they disagree, including on the columns read from a secondary + // input, and the physical plan and the batches have to agree with it. The + // secondary tables share their metadata, as `INTERSECTION_PRIMARY` unions + // them, and a union of inputs with differing field metadata reports + // different metadata in its logical and in its physical schema. + let with_metadata = |schema: &SchemaRef, table: &str| -> SchemaRef { + let role = if table == "data" { "data" } else { "secondary" }; + let tag = |mut metadata: HashMap| { + if role == "secondary" { + metadata.insert("only_in_secondary".to_string(), "yes".to_string()); + } + metadata + }; + let fields: Vec = schema + .fields() + .iter() + .map(|field| { + let metadata = HashMap::from([( + "column".to_string(), + format!("{role}.{}", field.name()), + )]); + field.as_ref().clone().with_metadata(tag(metadata)) + }) + .collect(); + let metadata = HashMap::from([("table".to_string(), role.to_string())]); + Arc::new(Schema::new_with_metadata(fields, tag(metadata))) + }; + + let cases = [ // Nullable in the primary input and in at least one secondary input. ( "intersect_primary_mixed_nullability", @@ -308,13 +341,23 @@ mod tests { "+---+---+---+---+---+---+", ][..], ), - ] { + ]; + + for ((file, expected_nullability, expected_rows), tagged) in cases + .into_iter() + .flat_map(|case| [(case, false), (case, true)]) + { let proto_plan = read_json(&format!("tests/testdata/test_plans/{file}.substrait.json")); let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; // Give each table rows, so the batch schemas below come from real batches for (table, rows) in rows { let schema = ctx.table_provider(table).await?.schema(); + let schema = if tagged { + with_metadata(&schema, table) + } else { + schema + }; let columns = (0..schema.fields().len()) .map(|i| { Arc::new(rows.iter().map(|row| row[i]).collect::()) @@ -343,7 +386,28 @@ mod tests { }) .collect::>() .join(", "); - assert_eq!(nullability, expected_nullability, "nullability of {file}"); + assert_eq!( + nullability, expected_nullability, + "nullability of {file} (tagged: {tagged})" + ); + + if tagged { + // Where the inputs disagree, the primary table's metadata wins, + // also on the columns read from a secondary input + assert_eq!( + plan.schema().metadata().get("table").map(String::as_str), + Some("data"), + "schema metadata of {file}" + ); + for field in plan.schema().fields() { + assert_eq!( + field.metadata().get("column"), + Some(&format!("data.{}", field.name())), + "metadata of column {} of {file}", + field.name() + ); + } + } // The physical plan and the batches it produces must carry the same // schema as the logical plan, not the left input's nullability