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..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,8 +16,11 @@ // 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::metadata::FieldMetadata; +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 +34,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 +80,112 @@ 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, 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, + right: LogicalPlan, + is_all: bool, +) -> datafusion::common::Result { + 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 layered over it. + let from_right: Vec = 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 { + Expr::Column(right.clone()).alias_qualified_with_metadata( + left.relation.clone(), + &left.name, + Some(FieldMetadata::from(field.metadata().clone())), + ) + } else { + Expr::Column(left.clone()) + } + }) + .collect::>(); + + 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], diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 522381de6efdf..6219b2b4ec47f 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/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,202 @@ 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; 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| { + 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", + "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::()) + 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::>() + .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 + 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 + 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 = 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..a99e7dd5720c6 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_all_mixed_nullability.substrait.json @@ -0,0 +1,198 @@ +{ + "relations": [ + { + "root": { + "input": { + "set": { + "inputs": [ + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "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", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data2" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + } + ], + "op": "SET_OP_INTERSECTION_MULTISET_ALL" + } + }, + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ] + } + } + ], + "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..e896015c8a4e8 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_multiset_mixed_nullability.substrait.json @@ -0,0 +1,198 @@ +{ + "relations": [ + { + "root": { + "input": { + "set": { + "inputs": [ + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "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", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data2" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + } + ], + "op": "SET_OP_INTERSECTION_MULTISET" + } + }, + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ] + } + } + ], + "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..77ac9d50da576 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/intersect_primary_mixed_nullability.substrait.json @@ -0,0 +1,198 @@ +{ + "relations": [ + { + "root": { + "input": { + "set": { + "inputs": [ + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "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", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_UNSPECIFIED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data2" + ] + } + } + }, + { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + } + ], + "op": "SET_OP_INTERSECTION_PRIMARY" + } + }, + "names": [ + "a", + "b", + "c", + "d", + "e", + "f" + ] + } + } + ], + "version": { + "minorNumber": 54, + "producer": "datafusion-test" + } +}