fix: Derive Substrait intersection nullability from every input

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

@@ -16,8 +16,12 @@
 // 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::arrow::datatypes::Field;
+use datafusion::common::{JoinType, NullEquality, not_impl_err, substrait_err};
+use datafusion::logical_expr::{
+    Cast, Expr, LogicalPlan, LogicalPlanBuilder, requalify_sides_if_needed,
+};
+use std::sync::Arc;
 use substrait::proto::set_rel::SetOp;
 use substrait::proto::{Rel, SetRel};
 
@@ -31,7 +35,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 +81,129 @@ 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.
+///
+/// Differing metadata does not change which path is taken, as it is no part of
+/// nullability. The result should describe the left input, exactly as
+/// [`LogicalPlanBuilder::intersect`] does: a column read from the right is cast
+/// to an explicit target field carrying the left field's metadata, and an
+/// explicit-field cast target replaces the source's metadata outright rather
+/// than merging into it, so a key only the right input carries is dropped
+/// rather than surviving into the result. A column read from the left is
+/// already exactly the left input's own field, in the logical and in the
+/// physical plan alike, so it needs no such treatment. 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,
+    right: LogicalPlan,
+    is_all: bool,
+) -> datafusion::common::Result<LogicalPlan> {
+    let left_fields = left.schema().fields();
+    let right_fields = right.schema().fields();
+    // 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 substituted for its own.
+    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()
+        })
+        .collect();
+
+    // `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);
+    }
+
+    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
+        .schema()
+        .fields()
+        .iter()
+        .zip(&left_columns)
+        .zip(&right_columns)
+        .zip(&from_right)
+        .map(|(((field, left), right), from_right)| {
+            if *from_right {
+                // `alias_qualified_with_metadata` would not do: `Expr::Alias`'s
+                // field derivation extends the aliased expression's own
+                // metadata with the alias's, so a key only the right column
+                // carries would survive alongside the left field's metadata.
+                // An explicit-field `Cast` target's metadata is instead used
+                // exactly as given, in both the logical and the physical
+                // plan, so casting to the left field's type (already proven
+                // equal to the right's) and metadata drops the right's own
+                // metadata outright. The qualifier and name still need
+                // `alias_qualified` on top, since a `Cast`'s own field is not
+                // renamed to its target field's name.
+                let target_field = Arc::new(
+                    Field::new(&left.name, field.data_type().clone(), false)
+                        .with_metadata(field.metadata().clone()),
+                );
+                Expr::Cast(Cast::new_from_field(
+                    Box::new(Expr::Column(right.clone())),
+                    target_field,
+                ))
+                .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,13 @@
 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 std::collections::HashSet;
+    use datafusion::datasource::MemTable;
+    use std::collections::{HashMap, HashSet};
+    use std::sync::Arc;
 
     use datafusion::common::Result;
     use datafusion::dataframe::DataFrame;
@@ -229,6 +234,210 @@ 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)],
+                ],
+            ),
+        ];
+
+        // 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<String, String>| {
+                if role == "secondary" {
+                    metadata.insert("only_in_secondary".to_string(), "yes".to_string());
+                }
+                metadata
+            };
+            let fields: Vec<Field> = 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",
+                "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 |",
+                    "+---+---+---+---+---+---+",
+                ][..],
+            ),
+        ];
+
+        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::<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} (tagged: {tagged})"
+            );
+
+            if tagged {
+                // Where the inputs disagree, the primary table's metadata wins,
+                // also on the columns read from a secondary input. Comparing
+                // the complete map, not just the "column" key, catches a
+                // secondary-only key ("only_in_secondary") leaking through:
+                // the primary table never carries that key, so its presence
+                // would fail this exact-equality check.
+                assert_eq!(
+                    plan.schema().metadata().get("table").map(String::as_str),
+                    Some("data"),
+                    "schema metadata of {file}"
+                );
+                for field in plan.schema().fields() {
+                    let expected_metadata = HashMap::from([(
+                        "column".to_string(),
+                        format!("data.{}", field.name()),
+                    )]);
+                    assert_eq!(
+                        field.metadata(),
+                        &expected_metadata,
+                        "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
+            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 =