From 391d0a2dd75fe8984d2754049c71571024d4b856 Mon Sep 17 00:00:00 2001 From: Huy Mac Date: Fri, 7 Aug 2026 15:06:50 +0900 Subject: [PATCH 1/2] fix(substrait): preserve nested nullability Closes #22095 --- .../src/logical_plan/consumer/expr/literal.rs | 386 ++++++++++++++---- .../src/logical_plan/consumer/rel/read_rel.rs | 107 +++-- .../src/logical_plan/consumer/types.rs | 8 +- .../src/logical_plan/producer/expr/literal.rs | 96 ++++- .../src/logical_plan/producer/types.rs | 26 +- .../substrait/tests/cases/logical_plans.rs | 9 +- .../tests/cases/roundtrip_logical_plan.rs | 25 +- 7 files changed, 525 insertions(+), 132 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs index d7d7a69581f05..3431a74debcd4 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs @@ -30,9 +30,13 @@ use crate::variation_const::{ TIMESTAMP_SECOND_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF, }; -use datafusion::arrow::array::{AsArray, MapArray, new_empty_array}; +use datafusion::arrow::array::{ + Array, AsArray, LargeListArray, ListArray, MapArray, new_empty_array, +}; use datafusion::arrow::buffer::OffsetBuffer; -use datafusion::arrow::datatypes::{Field, IntervalDayTime, IntervalMonthDayNano}; +use datafusion::arrow::datatypes::{ + DataType, Field, FieldRef, IntervalDayTime, IntervalMonthDayNano, +}; use datafusion::arrow::temporal_conversions::NANOSECONDS; use datafusion::common::scalar::ScalarStructBuilder; use datafusion::common::{ @@ -69,6 +73,24 @@ pub(crate) fn from_substrait_literal( lit: &Literal, dfs_names: &Vec, name_idx: &mut usize, +) -> datafusion::common::Result { + from_substrait_literal_impl(consumer, lit, dfs_names, name_idx, None) +} + +pub(crate) fn from_substrait_literal_with_expected_field( + consumer: &impl SubstraitConsumer, + lit: &Literal, + expected_field: &Field, +) -> datafusion::common::Result { + from_substrait_literal_impl(consumer, lit, &vec![], &mut 0, Some(expected_field)) +} + +fn from_substrait_literal_impl( + consumer: &impl SubstraitConsumer, + lit: &Literal, + dfs_names: &Vec, + name_idx: &mut usize, + expected_field: Option<&Field>, ) -> datafusion::common::Result { let scalar_value = match &lit.literal_type { Some(LiteralType::Boolean(b)) => ScalarValue::Boolean(Some(*b)), @@ -234,6 +256,8 @@ pub(crate) fn from_substrait_literal( ScalarValue::Decimal128(Some(i128::from_le_bytes(value)), p, s) } Some(LiteralType::List(l)) => { + let expected_item = + expected_list_item(expected_field, lit.type_variation_reference)?; // Each element should start the name index from the same value, then we increase it // once at the end let mut element_name_idx = *name_idx; @@ -242,7 +266,13 @@ pub(crate) fn from_substrait_literal( .iter() .map(|el| { element_name_idx = *name_idx; - from_substrait_literal(consumer, el, dfs_names, &mut element_name_idx) + from_substrait_literal_impl( + consumer, + el, + dfs_names, + &mut element_name_idx, + expected_item.map(FieldRef::as_ref), + ) }) .collect::>>()?; *name_idx = element_name_idx; @@ -251,39 +281,57 @@ pub(crate) fn from_substrait_literal( "Empty list must be encoded as EmptyList literal type, not List" ); } - let element_type = elements[0].data_type(); - match lit.type_variation_reference { - DEFAULT_CONTAINER_TYPE_VARIATION_REF => ScalarValue::List( - ScalarValue::new_list_nullable(elements.as_slice(), &element_type), - ), - LARGE_CONTAINER_TYPE_VARIATION_REF => ScalarValue::LargeList( - ScalarValue::new_large_list(elements.as_slice(), &element_type), - ), - others => { - return substrait_err!("Unknown type variation reference {others}"); + if let Some(expected_item) = expected_item { + make_list_scalar(elements, expected_item, lit.type_variation_reference)? + } else { + let element_type = elements[0].data_type(); + match lit.type_variation_reference { + DEFAULT_CONTAINER_TYPE_VARIATION_REF => { + ScalarValue::List(ScalarValue::new_list_nullable( + elements.as_slice(), + &element_type, + )) + } + LARGE_CONTAINER_TYPE_VARIATION_REF => ScalarValue::LargeList( + ScalarValue::new_large_list(elements.as_slice(), &element_type), + ), + others => { + return substrait_err!( + "Unknown type variation reference {others}" + ); + } } } } Some(LiteralType::EmptyList(l)) => { - let element_type = from_substrait_type( - consumer, - l.r#type.clone().unwrap().as_ref(), - dfs_names, - name_idx, - )?; - match lit.type_variation_reference { - DEFAULT_CONTAINER_TYPE_VARIATION_REF => { - ScalarValue::List(ScalarValue::new_list_nullable(&[], &element_type)) - } - LARGE_CONTAINER_TYPE_VARIATION_REF => ScalarValue::LargeList( - ScalarValue::new_large_list(&[], &element_type), - ), - others => { - return substrait_err!("Unknown type variation reference {others}"); + let expected_item = + expected_list_item(expected_field, lit.type_variation_reference)?; + if let Some(expected_item) = expected_item { + make_list_scalar(vec![], expected_item, lit.type_variation_reference)? + } else { + let element_type = from_substrait_type( + consumer, + l.r#type.clone().unwrap().as_ref(), + dfs_names, + name_idx, + )?; + match lit.type_variation_reference { + DEFAULT_CONTAINER_TYPE_VARIATION_REF => ScalarValue::List( + ScalarValue::new_list_nullable(&[], &element_type), + ), + LARGE_CONTAINER_TYPE_VARIATION_REF => ScalarValue::LargeList( + ScalarValue::new_large_list(&[], &element_type), + ), + others => { + return substrait_err!( + "Unknown type variation reference {others}" + ); + } } } } Some(LiteralType::Map(m)) => { + let expected_map = expected_map_fields(expected_field)?; // Each entry should start the name index from the same value, then we increase it // once at the end let mut entry_name_idx = *name_idx; @@ -292,25 +340,37 @@ pub(crate) fn from_substrait_literal( .iter() .map(|kv| { entry_name_idx = *name_idx; - let key_sv = from_substrait_literal( + let key_sv = from_substrait_literal_impl( consumer, kv.key.as_ref().unwrap(), dfs_names, &mut entry_name_idx, + expected_map.map(|fields| fields.key.as_ref()), )?; - let value_sv = from_substrait_literal( + let value_sv = from_substrait_literal_impl( consumer, kv.value.as_ref().unwrap(), dfs_names, &mut entry_name_idx, + expected_map.map(|fields| fields.value.as_ref()), )?; - ScalarStructBuilder::new() - .with_scalar(Field::new("key", key_sv.data_type(), false), key_sv) - .with_scalar( - Field::new("value", value_sv.data_type(), true), - value_sv, - ) - .build() + if let Some(expected_map) = expected_map { + ScalarStructBuilder::new() + .with_scalar(expected_map.key, key_sv) + .with_scalar(expected_map.value, value_sv) + .build() + } else { + ScalarStructBuilder::new() + .with_scalar( + Field::new("key", key_sv.data_type(), false), + key_sv, + ) + .with_scalar( + Field::new("value", value_sv.data_type(), true), + value_sv, + ) + .build() + } }) .collect::>>()?; *name_idx = entry_name_idx; @@ -321,60 +381,117 @@ pub(crate) fn from_substrait_literal( ); } + let entries_array = + ScalarValue::iter_to_array(entries)?.as_struct().to_owned(); + let (entries_field, keys_sorted) = if let Some(expected_map) = expected_map { + (Arc::clone(expected_map.entries), expected_map.keys_sorted) + } else { + ( + Arc::new(Field::new( + "entries", + entries_array.data_type().clone(), + false, + )), + false, + ) + }; ScalarValue::Map(Arc::new(MapArray::new( - Arc::new(Field::new("entries", entries[0].data_type(), false)), - OffsetBuffer::new(vec![0, entries.len() as i32].into()), - ScalarValue::iter_to_array(entries)?.as_struct().to_owned(), + entries_field, + OffsetBuffer::new(vec![0, entries_array.len() as i32].into()), + entries_array, None, - false, + keys_sorted, ))) } Some(LiteralType::EmptyMap(m)) => { - let key = match &m.key { - Some(k) => Ok(k), - _ => plan_err!("Missing key type for empty map"), - }?; - let value = match &m.value { - Some(v) => Ok(v), - _ => plan_err!("Missing value type for empty map"), - }?; - let key_type = from_substrait_type(consumer, key, dfs_names, name_idx)?; - let value_type = from_substrait_type(consumer, value, dfs_names, name_idx)?; + let expected_map = expected_map_fields(expected_field)?; + if let Some(expected_map) = expected_map { + let struct_array = new_empty_array(expected_map.entries.data_type()) + .as_struct() + .to_owned(); + ScalarValue::Map(Arc::new(MapArray::new( + Arc::clone(expected_map.entries), + OffsetBuffer::new(vec![0, 0].into()), + struct_array, + None, + expected_map.keys_sorted, + ))) + } else { + let key = match &m.key { + Some(k) => Ok(k), + _ => plan_err!("Missing key type for empty map"), + }?; + let value = match &m.value { + Some(v) => Ok(v), + _ => plan_err!("Missing value type for empty map"), + }?; + let key_type = from_substrait_type(consumer, key, dfs_names, name_idx)?; + let value_type = + from_substrait_type(consumer, value, dfs_names, name_idx)?; - // new_empty_array on a MapType creates a too empty array - // We want it to contain an empty struct array to align with an empty MapBuilder one - let entries = Field::new_struct( - "entries", - vec![ - Field::new("key", key_type, false), - Field::new("value", value_type, true), - ], - false, - ); - let struct_array = - new_empty_array(entries.data_type()).as_struct().to_owned(); - ScalarValue::Map(Arc::new(MapArray::new( - Arc::new(entries), - OffsetBuffer::new(vec![0, 0].into()), - struct_array, - None, - false, - ))) + // new_empty_array on a MapType creates a too empty array + // We want it to contain an empty struct array to align with an empty MapBuilder one + let entries = Field::new_struct( + "entries", + vec![ + Field::new("key", key_type, false), + Field::new("value", value_type, true), + ], + false, + ); + let struct_array = + new_empty_array(entries.data_type()).as_struct().to_owned(); + ScalarValue::Map(Arc::new(MapArray::new( + Arc::new(entries), + OffsetBuffer::new(vec![0, 0].into()), + struct_array, + None, + false, + ))) + } } Some(LiteralType::Struct(s)) => { + let expected_fields = expected_struct_fields(expected_field)?; + if let Some(expected_fields) = expected_fields + && s.fields.len() != expected_fields.len() + { + return substrait_err!( + "Struct literal field count mismatch: expected {} fields but found {}", + expected_fields.len(), + s.fields.len() + ); + } let mut builder = ScalarStructBuilder::new(); for (i, field) in s.fields.iter().enumerate() { - let name = next_struct_field_name(i, dfs_names, name_idx)?; - let sv = from_substrait_literal(consumer, field, dfs_names, name_idx)?; - // We assume everything to be nullable, since Arrow's strict about things matching - // and it's hard to match otherwise. - builder = builder.with_scalar(Field::new(name, sv.data_type(), true), sv); + if let Some(expected_field) = + expected_fields.map(|fields| fields[i].as_ref()) + { + let sv = from_substrait_literal_impl( + consumer, + field, + dfs_names, + name_idx, + Some(expected_field), + )?; + builder = builder.with_scalar(expected_field.clone(), sv); + } else { + let name = next_struct_field_name(i, dfs_names, name_idx)?; + let sv = from_substrait_literal_impl( + consumer, field, dfs_names, name_idx, None, + )?; + // Schema-less literals retain the existing nullable-field behavior. + builder = + builder.with_scalar(Field::new(name, sv.data_type(), true), sv); + } } builder.build()? } Some(LiteralType::Null(null_type)) => { - let data_type = - from_substrait_type(consumer, null_type, dfs_names, name_idx)?; + let data_type = if let Some(expected_field) = expected_field { + expected_field.data_type().clone() + } else { + from_substrait_type(consumer, null_type, dfs_names, name_idx)? + }; ScalarValue::try_from(&data_type)? } Some(LiteralType::IntervalDayToSecond(IntervalDayToSecond { @@ -583,9 +700,124 @@ pub(crate) fn from_substrait_literal( _ => return not_impl_err!("Unsupported literal_type: {:?}", lit.literal_type), }; + if let Some(expected_field) = expected_field + && scalar_value.data_type() != *expected_field.data_type() + { + return substrait_err!( + "Literal type mismatch: expected {:?} but found {:?}", + expected_field.data_type(), + scalar_value.data_type() + ); + } + Ok(scalar_value) } +fn expected_list_item( + expected_field: Option<&Field>, + type_variation_reference: u32, +) -> datafusion::common::Result> { + let Some(expected_field) = expected_field else { + return Ok(None); + }; + + match (expected_field.data_type(), type_variation_reference) { + (DataType::List(item), DEFAULT_CONTAINER_TYPE_VARIATION_REF) + | (DataType::LargeList(item), LARGE_CONTAINER_TYPE_VARIATION_REF) => { + Ok(Some(item)) + } + (expected, _) => substrait_err!( + "Expected List literal to have List or LargeList type, found {expected:?}" + ), + } +} + +fn make_list_scalar( + elements: Vec, + expected_item: &FieldRef, + type_variation_reference: u32, +) -> datafusion::common::Result { + let values = if elements.is_empty() { + new_empty_array(expected_item.data_type()) + } else { + ScalarValue::iter_to_array(elements)? + }; + + match type_variation_reference { + DEFAULT_CONTAINER_TYPE_VARIATION_REF => { + Ok(ScalarValue::List(Arc::new(ListArray::new( + Arc::clone(expected_item), + OffsetBuffer::new(vec![0, values.len() as i32].into()), + values, + None, + )))) + } + LARGE_CONTAINER_TYPE_VARIATION_REF => { + Ok(ScalarValue::LargeList(Arc::new(LargeListArray::new( + Arc::clone(expected_item), + OffsetBuffer::new(vec![0, values.len() as i64].into()), + values, + None, + )))) + } + others => substrait_err!("Unknown type variation reference {others}"), + } +} + +fn expected_struct_fields( + expected_field: Option<&Field>, +) -> datafusion::common::Result> { + let Some(expected_field) = expected_field else { + return Ok(None); + }; + + match expected_field.data_type() { + DataType::Struct(fields) => Ok(Some(fields)), + expected => substrait_err!( + "Expected Struct literal to have Struct type, found {expected:?}" + ), + } +} + +#[derive(Clone, Copy)] +struct ExpectedMapFields<'a> { + entries: &'a FieldRef, + key: &'a FieldRef, + value: &'a FieldRef, + keys_sorted: bool, +} + +fn expected_map_fields( + expected_field: Option<&Field>, +) -> datafusion::common::Result>> { + let Some(expected_field) = expected_field else { + return Ok(None); + }; + + let DataType::Map(entries, keys_sorted) = expected_field.data_type() else { + return substrait_err!( + "Expected Map literal to have Map type, found {:?}", + expected_field.data_type() + ); + }; + let DataType::Struct(fields) = entries.data_type() else { + return substrait_err!("Expected Map entries field to contain a Struct"); + }; + if fields.len() != 2 { + return substrait_err!( + "Expected Map entries Struct to have 2 fields, found {}", + fields.len() + ); + } + + Ok(Some(ExpectedMapFields { + entries, + key: &fields[0], + value: &fields[1], + keys_sorted: *keys_sorted, + })) +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index 78951a3aff549..dbb78f54806fc 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs @@ -16,7 +16,7 @@ // under the License. use crate::logical_plan::consumer::SubstraitConsumer; -use crate::logical_plan::consumer::from_substrait_literal; +use crate::logical_plan::consumer::from_substrait_literal_with_expected_field; use crate::logical_plan::consumer::from_substrait_named_struct; use crate::logical_plan::consumer::utils::ensure_schema_compatibility; use datafusion::common::{ @@ -157,21 +157,16 @@ pub async fn from_read_rel( } let mut row_exprs = vec![]; - let mut name_idx = 0; - for expression in &row.fields { - // Top-level names are provided through schema - // Each expression consumes at least one name, and Literals may consume additional names. - name_idx += 1; + for (expression, expected_field) in + row.fields.iter().zip(substrait_schema.fields()) + { let expr = match expression.rex_type.as_ref() { Some(substrait::proto::expression::RexType::Literal(lit)) => { - // Values literals need 'named_struct.names' so nested struct fields keep their names from the ReadRel base schema. - // This is important for nested struct fields to retain their names. Expr::Literal( - from_substrait_literal( + from_substrait_literal_with_expected_field( consumer, lit, - &named_struct.names, - &mut name_idx, + expected_field, )?, None, ) @@ -184,19 +179,11 @@ pub async fn from_read_rel( }; row_exprs.push(expr); } - - if name_idx != named_struct.names.len() { - return substrait_err!( - "Names list must match exactly to nested schema, but found {} uses for {} names", - name_idx, - named_struct.names.len() - ); - } exprs.push(row_exprs); } exprs } else { - convert_literal_rows(consumer, vt, named_struct)? + convert_literal_rows(consumer, vt, &substrait_schema)? }; Ok(LogicalPlan::Values(Values { @@ -254,38 +241,35 @@ pub async fn from_read_rel( /// Converts Substrait literal rows from a VirtualTable into DataFusion expressions. /// /// This function processes the deprecated `values` field of VirtualTable, converting -/// each literal value into a `Expr::Literal` while tracking and validating the name -/// indices against the provided named struct schema. +/// each literal value into a `Expr::Literal` using the matching field from the schema. fn convert_literal_rows( consumer: &impl SubstraitConsumer, vt: &substrait::proto::read_rel::VirtualTable, - named_struct: &substrait::proto::NamedStruct, + schema: &DFSchema, ) -> datafusion::common::Result>> { #[expect(deprecated)] vt.values .iter() .map(|row| { - let mut name_idx = 0; + if row.fields.len() != schema.fields().len() { + return substrait_err!( + "Field count mismatch: expected {} fields but found {} in virtual table row", + schema.fields().len(), + row.fields.len() + ); + } let lits = row .fields .iter() - .map(|lit| { - name_idx += 1; // top-level names are provided through schema - Ok(Expr::Literal(from_substrait_literal( + .zip(schema.fields()) + .map(|(lit, expected_field)| { + Ok(Expr::Literal(from_substrait_literal_with_expected_field( consumer, lit, - &named_struct.names, - &mut name_idx, + expected_field, )?, None)) }) .collect::>()?; - if name_idx != named_struct.names.len() { - return substrait_err!( - "Names list must match exactly to nested schema, but found {} uses for {} names", - name_idx, - named_struct.names.len() - ); - } Ok(lits) }) .collect::>() @@ -365,3 +349,54 @@ fn apply_projection( _ => plan_err!("DataFrame passed to apply_projection must be a TableScan"), } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::logical_plan::consumer::utils::tests::test_consumer; + use datafusion::arrow::array::Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::ScalarValue; + use substrait::proto::expression::Literal; + use substrait::proto::expression::literal::{ + List, LiteralType, Struct as LiteralStruct, + }; + use substrait::proto::read_rel::VirtualTable; + + #[test] + fn deprecated_literal_rows_use_expected_fields() -> datafusion::common::Result<()> { + let list_type = + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))); + let schema = DFSchema::try_from(Schema::new(vec![Field::new( + "list", + list_type.clone(), + false, + )]))?; + let list_literal = Literal { + nullable: false, + type_variation_reference: 0, + literal_type: Some(LiteralType::List(List { + values: vec![Literal { + nullable: false, + type_variation_reference: 0, + literal_type: Some(LiteralType::I32(1)), + }], + })), + }; + #[expect(deprecated)] + let virtual_table = VirtualTable { + values: vec![LiteralStruct { + fields: vec![list_literal], + }], + ..Default::default() + }; + + let rows = convert_literal_rows(&test_consumer(), &virtual_table, &schema)?; + let Expr::Literal(ScalarValue::List(list), _) = &rows[0][0] else { + panic!("expected list literal") + }; + assert_eq!(list.data_type(), &list_type); + + Ok(()) + } +} diff --git a/datafusion/substrait/src/logical_plan/consumer/types.rs b/datafusion/substrait/src/logical_plan/consumer/types.rs index 6ca9ced2fcb06..cae051baff082 100644 --- a/datafusion/substrait/src/logical_plan/consumer/types.rs +++ b/datafusion/substrait/src/logical_plan/consumer/types.rs @@ -179,9 +179,7 @@ pub fn from_substrait_type( })?; let field = Arc::new(Field::new_list_field( from_substrait_type(consumer, inner_type, dfs_names, name_idx)?, - // We ignore Substrait's nullability here to match to_substrait_literal - // which always creates nullable lists - true, + type_is_nullable(inner_type)?, )); match list.type_variation_reference { DEFAULT_CONTAINER_TYPE_VARIATION_REF => Ok(DataType::List(field)), @@ -198,6 +196,7 @@ pub fn from_substrait_type( let value_type = map.value.as_ref().ok_or_else(|| { substrait_datafusion_err!("Map type must have value type") })?; + let value_nullable = type_is_nullable(value_type)?; let key_type = from_substrait_type(consumer, key_type, dfs_names, name_idx)?; let value_type = @@ -206,7 +205,8 @@ pub fn from_substrait_type( match map.type_variation_reference { DEFAULT_MAP_TYPE_VARIATION_REF => { let key_field = Arc::new(Field::new("key", key_type, false)); - let value_field = Arc::new(Field::new("value", value_type, true)); + let value_field = + Arc::new(Field::new("value", value_type, value_nullable)); Ok(DataType::Map( Arc::new(Field::new_struct( "entries", diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index f26984ad4a394..58127c8a67ca0 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs @@ -406,10 +406,15 @@ fn convert_array_to_literal_list( #[cfg(test)] mod tests { use super::*; - use crate::logical_plan::consumer::from_substrait_literal_without_names; use crate::logical_plan::consumer::tests::test_consumer; + use crate::logical_plan::consumer::{ + from_substrait_literal_with_expected_field, from_substrait_literal_without_names, + }; use crate::logical_plan::producer::DefaultSubstraitProducer; - use datafusion::arrow::array::{Int64Builder, MapBuilder, StringBuilder}; + use datafusion::arrow::array::{ + AsArray, Int64Builder, MapArray, MapBuilder, StringBuilder, new_empty_array, + }; + use datafusion::arrow::buffer::OffsetBuffer; use datafusion::arrow::datatypes::{ DataType, Field, IntervalDayTime, IntervalMonthDayNano, }; @@ -548,6 +553,79 @@ mod tests { Ok(()) } + #[test] + fn round_trip_literals_with_expected_field() -> Result<()> { + let required_struct = ScalarStructBuilder::new() + .with_scalar( + Field::new("required", DataType::Int32, false), + ScalarValue::Int32(Some(1)), + ) + .build()?; + let state = SessionContext::default().state(); + let mut producer = DefaultSubstraitProducer::new(&state); + let substrait_literal = to_substrait_literal(&mut producer, &required_struct)?; + let schema_less = + from_substrait_literal_without_names(&test_consumer(), &substrait_literal)?; + let DataType::Struct(fields) = schema_less.data_type() else { + panic!("expected struct literal") + }; + assert!(fields[0].is_nullable()); + + round_trip_literal_with_expected_field(required_struct.clone())?; + + let required_list = ScalarValue::List(ScalarValue::new_list( + std::slice::from_ref(&required_struct), + &required_struct.data_type(), + false, + )); + round_trip_literal_with_expected_field(required_list)?; + + let empty_required_list = ScalarValue::List(ScalarValue::new_list( + &[], + &required_struct.data_type(), + false, + )); + round_trip_literal_with_expected_field(empty_required_list)?; + + let required_entry = ScalarStructBuilder::new() + .with_scalar( + Field::new("key", DataType::Utf8, false), + ScalarValue::Utf8(Some("key".to_string())), + ) + .with_scalar( + Field::new("value", DataType::Int64, false), + ScalarValue::Int64(Some(1)), + ) + .build()?; + let entries = ScalarValue::iter_to_array([required_entry])? + .as_struct() + .to_owned(); + let entries_field = + Arc::new(Field::new("entries", entries.data_type().clone(), false)); + let required_map = ScalarValue::Map(Arc::new(MapArray::new( + Arc::clone(&entries_field), + OffsetBuffer::new(vec![0, 1].into()), + entries, + None, + false, + ))); + round_trip_literal_with_expected_field(required_map)?; + + let empty_entries = new_empty_array(entries_field.data_type()) + .as_struct() + .to_owned(); + let empty_required_map = ScalarValue::Map(Arc::new(MapArray::new( + entries_field, + OffsetBuffer::new(vec![0, 0].into()), + empty_entries, + None, + false, + ))); + round_trip_literal_with_expected_field(empty_required_map)?; + + Ok(()) + } + fn round_trip_literal(scalar: ScalarValue) -> Result<()> { println!("Checking round trip of {scalar:?}"); let state = SessionContext::default().state(); @@ -558,4 +636,18 @@ mod tests { assert_eq!(scalar, roundtrip_scalar); Ok(()) } + + fn round_trip_literal_with_expected_field(scalar: ScalarValue) -> Result<()> { + let state = SessionContext::default().state(); + let mut producer = DefaultSubstraitProducer::new(&state); + let substrait_literal = to_substrait_literal(&mut producer, &scalar)?; + let expected_field = Field::new("expected", scalar.data_type(), false); + let roundtrip_scalar = from_substrait_literal_with_expected_field( + &test_consumer(), + &substrait_literal, + &expected_field, + )?; + assert_eq!(scalar, roundtrip_scalar); + Ok(()) + } } diff --git a/datafusion/substrait/src/logical_plan/producer/types.rs b/datafusion/substrait/src/logical_plan/producer/types.rs index ebc7de5a291dd..1c6ac501764f7 100644 --- a/datafusion/substrait/src/logical_plan/producer/types.rs +++ b/datafusion/substrait/src/logical_plan/producer/types.rs @@ -437,12 +437,14 @@ mod tests { round_trip_type(DataType::Decimal128(10, 2))?; round_trip_type(DataType::Decimal256(30, 2))?; - round_trip_type(DataType::List( - Field::new_list_field(DataType::Int32, true).into(), - ))?; - round_trip_type(DataType::LargeList( - Field::new_list_field(DataType::Int32, true).into(), - ))?; + for nullable in [true, false] { + round_trip_type(DataType::List( + Field::new_list_field(DataType::Int32, nullable).into(), + ))?; + round_trip_type(DataType::LargeList( + Field::new_list_field(DataType::Int32, nullable).into(), + ))?; + } round_trip_type(DataType::Map( Field::new_struct( @@ -456,6 +458,18 @@ mod tests { .into(), false, ))?; + round_trip_type(DataType::Map( + Field::new_struct( + "entries", + [ + Field::new("key", DataType::Utf8, false).into(), + Field::new("value", DataType::Int32, false).into(), + ], + false, + ) + .into(), + false, + ))?; round_trip_type(DataType::Dictionary( Box::new(DataType::Utf8), Box::new(DataType::Int32), diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 522381de6efdf..32001b7e18740 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -21,6 +21,7 @@ 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::datatypes::DataType; use datafusion::common::test_util::format_batches; use std::collections::HashSet; @@ -210,14 +211,16 @@ mod tests { #[tokio::test] async fn non_nullable_lists() -> Result<()> { - // DataFusion's Substrait consumer treats all lists as nullable, even if the Substrait plan specifies them as non-nullable. - // That's because implementing the non-nullability consistently is non-trivial. - // This test confirms that reading a plan with non-nullable lists works as expected. let proto_plan = read_json("tests/testdata/test_plans/non_nullable_lists.substrait.json"); let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + let DataType::List(item) = plan.schema().field(0).data_type() else { + panic!("expected list field") + }; + assert!(!item.is_nullable()); + assert_snapshot!( &plan, @"Values: (List([1, 2]))" diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index f084d3170edcc..9db6964c00c03 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -1598,11 +1598,8 @@ async fn roundtrip_values_duplicate_column_join() -> Result<()> { async fn roundtrip_preserves_field_nullability() -> Result<()> { use datafusion::arrow::datatypes::Fields; - // Verify that required and nullable fields, including nested struct fields, + // Verify that required and nullable fields, including nested fields, // preserve their nullability through a Substrait round-trip. - // - // List child nullability is intentionally omitted because it is not - // preserved today. let ctx = create_context().await?; let df_schema = DFSchema::try_from(Schema::new(vec![ Field::new("required_int", DataType::Int32, false), @@ -1615,6 +1612,26 @@ async fn roundtrip_preserves_field_nullability() -> Result<()> { ])), false, ), + Field::new( + "required_list_item", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))), + false, + ), + Field::new( + "required_map_value", + DataType::Map( + Arc::new(Field::new_struct( + "entries", + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, false), + ], + false, + )), + false, + ), + false, + ), ]))?; let plan = LogicalPlan::EmptyRelation(EmptyRelation { produce_one_row: false, From 22cf4bc0aaf1b25244c9a3de48c278a693326aa5 Mon Sep 17 00:00:00 2001 From: Huy Mac Date: Fri, 7 Aug 2026 17:23:07 +0900 Subject: [PATCH 2/2] test(substrait): cover literal validation --- .../src/logical_plan/consumer/expr/literal.rs | 136 ++++++++++++++++++ .../src/logical_plan/consumer/rel/read_rel.rs | 98 +++++++++++-- 2 files changed, 221 insertions(+), 13 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs index 3431a74debcd4..7bf698bd1384c 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/literal.rs @@ -822,6 +822,22 @@ fn expected_map_fields( mod tests { use super::*; use crate::logical_plan::consumer::utils::tests::test_consumer; + use substrait::proto::expression::literal::map::KeyValue; + + fn literal(literal_type: Option) -> Literal { + Literal { + nullable: false, + type_variation_reference: DEFAULT_TYPE_VARIATION_REF, + literal_type, + } + } + + fn assert_expected_field_error(lit: &Literal, field: &Field, expected: &str) { + let err = + from_substrait_literal_with_expected_field(&test_consumer(), lit, field) + .unwrap_err(); + assert!(err.to_string().contains(expected), "got: {err}"); + } #[test] fn interval_compound_different_precision() -> datafusion::common::Result<()> { @@ -858,4 +874,124 @@ mod tests { Ok(()) } + + #[test] + fn expected_field_validation_errors() { + let i32_literal = literal(Some(LiteralType::I32(1))); + assert_expected_field_error( + &i32_literal, + &Field::new("expected", DataType::Utf8, false), + "Literal type mismatch", + ); + + let list_literal = + literal(Some(LiteralType::List(proto::expression::literal::List { + values: vec![i32_literal.clone()], + }))); + assert_expected_field_error( + &list_literal, + &Field::new("expected", DataType::Int32, false), + "Expected List literal", + ); + + let struct_literal = literal(Some(LiteralType::Struct( + proto::expression::literal::Struct { + fields: vec![i32_literal.clone()], + }, + ))); + assert_expected_field_error( + &struct_literal, + &Field::new("expected", DataType::Int32, false), + "Expected Struct literal", + ); + assert_expected_field_error( + &struct_literal, + &Field::new_struct("expected", Vec::::new(), false), + "Struct literal field count mismatch", + ); + + let map_literal = + literal(Some(LiteralType::Map(proto::expression::literal::Map { + key_values: vec![KeyValue { + key: Some(i32_literal.clone()), + value: Some(i32_literal.clone()), + }], + }))); + assert_expected_field_error( + &map_literal, + &Field::new("expected", DataType::Int32, false), + "Expected Map literal", + ); + + let non_struct_entries = Field::new("entries", DataType::Int32, false); + assert_expected_field_error( + &map_literal, + &Field::new( + "expected", + DataType::Map(Arc::new(non_struct_entries), false), + false, + ), + "Expected Map entries field to contain a Struct", + ); + + let one_field_entries = Field::new_struct( + "entries", + vec![Field::new("key", DataType::Int32, false)], + false, + ); + assert_expected_field_error( + &map_literal, + &Field::new( + "expected", + DataType::Map(Arc::new(one_field_entries), false), + false, + ), + "Expected Map entries Struct to have 2 fields", + ); + + let mismatched_key_entries = Field::new_struct( + "entries", + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, false), + ], + false, + ); + assert_expected_field_error( + &map_literal, + &Field::new( + "expected", + DataType::Map(Arc::new(mismatched_key_entries), false), + false, + ), + "Literal type mismatch", + ); + + let unsupported_literal = literal(None); + assert_expected_field_error( + &unsupported_literal, + &Field::new("expected", DataType::Int32, false), + "Unsupported literal_type", + ); + + let invalid_list_variation = Literal { + type_variation_reference: u32::MAX, + ..list_literal + }; + let err = from_substrait_literal_without_names( + &test_consumer(), + &invalid_list_variation, + ) + .unwrap_err(); + assert!( + err.to_string().contains("Unknown type variation reference"), + "got: {err}" + ); + + let empty_map = + literal(Some(LiteralType::EmptyMap(proto::r#type::Map::default()))); + let err = from_substrait_literal_without_names(&test_consumer(), &empty_map) + .unwrap_err(); + assert!(err.to_string().contains("Missing key type"), "got: {err}"); + } } diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index dbb78f54806fc..cb2eb82b3f304 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs @@ -354,25 +354,23 @@ fn apply_projection( mod tests { use super::*; use crate::logical_plan::consumer::utils::tests::test_consumer; + use crate::logical_plan::producer::{ + DefaultSubstraitProducer, to_substrait_named_struct, + }; use datafusion::arrow::array::Array; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::ScalarValue; + use datafusion::prelude::SessionContext; use substrait::proto::expression::Literal; + use substrait::proto::expression::RexType; use substrait::proto::expression::literal::{ List, LiteralType, Struct as LiteralStruct, }; + use substrait::proto::expression::nested::Struct as ExpressionStruct; use substrait::proto::read_rel::VirtualTable; - #[test] - fn deprecated_literal_rows_use_expected_fields() -> datafusion::common::Result<()> { - let list_type = - DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))); - let schema = DFSchema::try_from(Schema::new(vec![Field::new( - "list", - list_type.clone(), - false, - )]))?; - let list_literal = Literal { + fn list_literal() -> Literal { + Literal { nullable: false, type_variation_reference: 0, literal_type: Some(LiteralType::List(List { @@ -382,11 +380,23 @@ mod tests { literal_type: Some(LiteralType::I32(1)), }], })), - }; + } + } + + fn list_schema() -> datafusion::common::Result { + let list_type = + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))); + DFSchema::try_from(Schema::new(vec![Field::new("list", list_type, false)])) + } + + #[test] + fn deprecated_literal_rows_use_expected_fields() -> datafusion::common::Result<()> { + let schema = list_schema()?; + let list_type = schema.field(0).data_type(); #[expect(deprecated)] let virtual_table = VirtualTable { values: vec![LiteralStruct { - fields: vec![list_literal], + fields: vec![list_literal()], }], ..Default::default() }; @@ -395,7 +405,69 @@ mod tests { let Expr::Literal(ScalarValue::List(list), _) = &rows[0][0] else { panic!("expected list literal") }; - assert_eq!(list.data_type(), &list_type); + assert_eq!(list.data_type(), list_type); + + #[expect(deprecated)] + let mismatched_table = VirtualTable { + values: vec![LiteralStruct { fields: vec![] }], + ..Default::default() + }; + let err = convert_literal_rows(&test_consumer(), &mismatched_table, &schema) + .unwrap_err(); + assert!( + err.to_string().contains("Field count mismatch"), + "got: {err}" + ); + + Ok(()) + } + + #[tokio::test] + async fn expression_rows_use_expected_fields() -> datafusion::common::Result<()> { + let schema = list_schema()?; + let state = SessionContext::new().state(); + let mut producer = DefaultSubstraitProducer::new(&state); + let base_schema = + to_substrait_named_struct(&mut producer, &DFSchemaRef::new(schema.clone()))?; + let expression = Expression { + rex_type: Some(RexType::Literal(list_literal())), + }; + let virtual_table = VirtualTable { + expressions: vec![ExpressionStruct { + fields: vec![expression], + }], + ..Default::default() + }; + let read = ReadRel { + base_schema: Some(base_schema.clone()), + read_type: Some(ReadType::VirtualTable(virtual_table)), + ..Default::default() + }; + + let plan = from_read_rel(&test_consumer(), &read).await?; + let LogicalPlan::Values(values) = plan else { + panic!("expected Values plan") + }; + let Expr::Literal(ScalarValue::List(list), _) = &values.values[0][0] else { + panic!("expected list literal") + }; + assert_eq!(list.data_type(), schema.field(0).data_type()); + + let mismatched_read = ReadRel { + base_schema: Some(base_schema), + read_type: Some(ReadType::VirtualTable(VirtualTable { + expressions: vec![ExpressionStruct { fields: vec![] }], + ..Default::default() + })), + ..Default::default() + }; + let err = from_read_rel(&test_consumer(), &mismatched_read) + .await + .unwrap_err(); + assert!( + err.to_string().contains("Field count mismatch"), + "got: {err}" + ); Ok(()) }