fix: Derive Substrait intersection nullability from every input

datafusion/substrait/src/logical_plan/consumer/rel/set_rel.rs

@@ -16,8 +16,10 @@
 // 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::{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};
 
@@ -31,7 +33,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 +79,101 @@ 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".
+///
+/// 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<LogicalPlan> {
+    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<bool> = left_fields
+        .iter()
+        .zip(right_fields.iter())
+        .map(|(left, right)| {
+            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)
+    {
+        return LogicalPlanBuilder::intersect(left, right, is_all);
+    }
+
+    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(&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 {
+                Expr::Column(left.clone())
+            }
+        })
+        .collect::<Vec<_>>();
+
+    left.join_detailed(
+        right,
+        JoinType::Inner,
+        (left_columns, right_columns),
+        None,
+        NullEquality::NullEqualsNull,
+    )?
+    .project(exprs)?
+    .build()
+}
+
 async fn except_rels(
     consumer: &impl SubstraitConsumer,
     rels: &[Rel],

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;
@@ -229,6 +233,139 @@ 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 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? e? f?
+        //   secondary   a  b  c? d? e~ f~
+        //   secondary   a  b? c  d? e? f
+        let rows: [(&str, &[[Option<i64>; 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?, 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?, 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::<Int64Array>())
+                            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
+                .schema()
+                .fields()
+                .iter()
+                .map(|field| {
+                    format!(
+                        "{}{}",
+                        field.name(),
+                        if field.is_nullable() { "?" } else { "" }
+                    )
+                })
+                .collect::<Vec<_>>()
+                .join(", ");
+            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(())
+    }
+
     #[tokio::test]
     async fn multilayer_aggregate() -> Result<()> {
         let proto_plan =