From a0457690fb0f2cc939891c6d27cdd13bbbf3251a Mon Sep 17 00:00:00 2001 From: naman Date: Wed, 9 Sep 2026 02:08:12 +0530 Subject: [PATCH] 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" + } +}