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" + } +}