diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index 8882c992dca1c..bc8ee89d046cf 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs @@ -24,6 +24,7 @@ use crate::variation_const::{ VIEW_CONTAINER_TYPE_VARIATION_REF, }; use datafusion::arrow::array::{Array, GenericListArray, OffsetSizeTrait}; +use datafusion::arrow::datatypes::DataType; use datafusion::arrow::temporal_conversions::NANOSECONDS; use datafusion::common::{ScalarValue, exec_err, not_impl_err}; use substrait::proto::expression::literal::interval_day_to_second::PrecisionMode; @@ -46,6 +47,30 @@ pub(crate) fn to_substrait_literal_expr( producer: &mut impl SubstraitProducer, value: &ScalarValue, ) -> datafusion::common::Result { + // Preserve the dictionary type when encoding a dictionary literal as a + // standalone expression (e.g. projection or VALUES): there is no separate + // type layer carrying the dictionary metadata in those positions, so we + // wrap the inner literal in a cast to the dictionary type. This lets the + // consumer reconstruct a `Dictionary` typed value instead of silently + // degrading to the inner (e.g. Utf8) type. + if let ScalarValue::Dictionary(key_type, inner) = value { + let inner_expr = to_substrait_literal_expr(producer, inner)?; + let cast_type = to_substrait_type( + producer, + &DataType::Dictionary(key_type.clone(), Box::new(inner.data_type())), + value.is_null(), + )?; + return Ok(Expression { + rex_type: Some(RexType::Cast(Box::new( + substrait::proto::expression::Cast { + r#type: Some(cast_type), + input: Some(Box::new(inner_expr)), + failure_behavior: substrait::proto::expression::cast::FailureBehavior::ThrowException + .into(), + }, + ))), + }); + } let literal = to_substrait_literal(producer, value)?; Ok(Expression { rex_type: Some(RexType::Literal(literal)), @@ -360,6 +385,20 @@ pub(crate) fn to_substrait_literal( }), DEFAULT_TYPE_VARIATION_REF, ), + // Dictionary literals are encoded as their inner value. The dictionary + // type information is carried separately by the type layer (as a Map + // with DICTIONARY_MAP_TYPE_VARIATION_REF), and the inner value is + // guaranteed to be non-null here since null dictionaries are handled by + // the `value.is_null()` check at the top of this function. + ScalarValue::Dictionary(_, value) => { + let literal = to_substrait_literal(producer, value)?; + ( + literal + .literal_type + .expect("to_substrait_literal always sets a literal type"), + literal.type_variation_reference, + ) + } _ => ( not_impl_err!("Unsupported literal: {value:?}")?, DEFAULT_TYPE_VARIATION_REF, @@ -407,14 +446,18 @@ fn convert_array_to_literal_list( mod tests { use super::*; use crate::logical_plan::consumer::from_substrait_literal_without_names; + use crate::logical_plan::consumer::from_substrait_rex; + use crate::logical_plan::consumer::from_substrait_type_without_names; use crate::logical_plan::consumer::tests::test_consumer; use crate::logical_plan::producer::DefaultSubstraitProducer; use datafusion::arrow::array::{Int64Builder, MapBuilder, StringBuilder}; use datafusion::arrow::datatypes::{ DataType, Field, IntervalDayTime, IntervalMonthDayNano, }; + use datafusion::common::DFSchema; use datafusion::common::Result; use datafusion::common::scalar::ScalarStructBuilder; + use datafusion::logical_expr::Expr; use datafusion::prelude::SessionContext; use std::sync::Arc; @@ -558,4 +601,81 @@ mod tests { assert_eq!(scalar, roundtrip_scalar); Ok(()) } + + #[tokio::test] + async fn test_dictionary_literal() -> Result<()> { + let state = SessionContext::default().state(); + let mut producer = DefaultSubstraitProducer::new(&state); + + // Dictionary(UInt32, Utf8("a")) is encoded as a cast to the dictionary + // type wrapping the inner Utf8 literal, so the dictionary type is + // preserved in positions without a separate type layer (projection, + // VALUES, ...) instead of silently degrading to the inner type. + let dict = ScalarValue::Dictionary( + Box::new(DataType::UInt32), + Box::new(ScalarValue::Utf8(Some("a".to_string()))), + ); + let expr = to_substrait_literal_expr(&mut producer, &dict)?; + let cast = match expr.rex_type.as_ref() { + Some(RexType::Cast(cast)) => cast, + other => panic!("expected Cast rex type, got {other:?}"), + }; + + // The cast output type is the full dictionary type. + let cast_type = cast.r#type.as_ref().expect("cast must have an output type"); + assert_eq!( + from_substrait_type_without_names(&test_consumer(), cast_type)?, + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)) + ); + + // The cast input is the inner Utf8 literal. + match cast.input.as_deref() { + Some(Expression { + rex_type: Some(RexType::Literal(lit)), + .. + }) => { + assert_eq!(lit.literal_type, Some(LiteralType::String("a".to_string()))); + assert!(!lit.nullable); + } + other => panic!("expected inner literal, got {other:?}"), + } + + // A full expression round-trip through the consumer preserves the + // dictionary type. + match from_substrait_rex(&test_consumer(), &expr, &DFSchema::empty()).await? { + Expr::Cast(cast) => { + assert_eq!( + cast.data_type, + DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(DataType::Utf8) + ) + ); + match cast.expr.as_ref() { + Expr::Literal(ScalarValue::Utf8(Some(s)), _) => { + assert_eq!(s, "a"); + } + other => panic!("expected inner Utf8 literal, got {other:?}"), + } + } + other => panic!("expected Cast expr, got {other:?}"), + } + Ok(()) + } + + #[test] + fn test_dictionary_null_literal() -> Result<()> { + let state = SessionContext::default().state(); + let mut producer = DefaultSubstraitProducer::new(&state); + + // A dictionary with a null inner value is encoded as a Null literal. + let dict = ScalarValue::Dictionary( + Box::new(DataType::UInt32), + Box::new(ScalarValue::Utf8(None)), + ); + let literal = to_substrait_literal(&mut producer, &dict)?; + assert!(literal.nullable); + assert!(matches!(literal.literal_type, Some(LiteralType::Null(_)))); + Ok(()) + } } diff --git a/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs index 8dfbb36d3767d..b46b901577690 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs @@ -31,6 +31,17 @@ use substrait::proto::read_rel::{NamedTable, ReadType, VirtualTable}; use substrait::proto::rel::RelType; use substrait::proto::{ReadRel, Rel}; +/// Returns true if the scalar is a non-null Dictionary literal. +/// +/// Non-null dictionary literals cannot be losslessly encoded in the plain +/// `VirtualTable.values` (literal) format because `to_substrait_literal` only +/// encodes the inner value, dropping the dictionary type. Null dictionaries are +/// unaffected: they are encoded as typed Null literals carrying the full +/// dictionary type. +fn is_non_null_dictionary(sv: &ScalarValue) -> bool { + matches!(sv, ScalarValue::Dictionary(_, inner) if !inner.is_null()) +} + /// Converts rows of literal expressions into Substrait literal structs. /// /// Each row is expected to contain only `Expr::Literal` or `Expr::Alias` wrapping literals. @@ -213,14 +224,33 @@ pub fn from_values( let schema_len = v.schema.fields().len(); let empty_schema = Arc::new(DFSchema::empty()); - let use_literals = v.values.iter().all(|row| { - row.iter().all(|expr| match expr { - Expr::Literal(_, _) => true, - Expr::Alias(alias) => matches!(alias.expr.as_ref(), Expr::Literal(_, _)), + // A non-null `ScalarValue::Dictionary` literal cannot be encoded as a plain + // Substrait literal: `to_substrait_literal` only encodes the inner value and + // drops the dictionary type. (Null dictionaries are fine - they are encoded + // as typed Null literals.) When any row contains such a literal, fall back to + // the expression format, where `to_substrait_literal_expr` preserves the + // dictionary type by wrapping the inner literal in a cast to the dictionary + // type, which the consumer decodes back into a `Dictionary` typed value. + let contains_non_null_dictionary = v.values.iter().any(|row| { + row.iter().any(|expr| match expr { + Expr::Literal(sv, _) => is_non_null_dictionary(sv), + Expr::Alias(alias) => match alias.expr.as_ref() { + Expr::Literal(sv, _) => is_non_null_dictionary(sv), + _ => false, + }, _ => false, }) }); + let use_literals = !contains_non_null_dictionary + && v.values.iter().all(|row| { + row.iter().all(|expr| match expr { + Expr::Literal(_, _) => true, + Expr::Alias(alias) => matches!(alias.expr.as_ref(), Expr::Literal(_, _)), + _ => false, + }) + }); + let (values, expressions) = if use_literals { let values = convert_literal_rows(producer, &v.values)?; (values, vec![]) diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 5dd4aa4e2be91..546c4661d17a4 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -1507,6 +1507,75 @@ async fn roundtrip_values_no_columns() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_values_with_dictionary() -> Result<()> { + let ctx = create_context().await?; + // A non-null Dictionary literal in a VALUES clause must keep its Dictionary + // type through the Substrait round-trip. It is encoded in the expression + // format (a cast to the dictionary type wrapping the inner literal) because + // the plain literal format would silently degrade it to the inner (Utf8) + // type, which then fails to materialize against the Dictionary-typed schema + // during physical execution (RecordBatch::try_new_with_options rejects the + // schema vs array mismatch). + let dict = ScalarValue::Dictionary( + Box::new(DataType::UInt32), + Box::new(ScalarValue::Utf8(Some("a".to_string()))), + ); + let plan = + LogicalPlanBuilder::values(vec![vec![Expr::Literal(dict, None)]])?.build()?; + + let proto = to_substrait_plan(&plan, &ctx.state())?; + let plan2 = from_substrait_plan(&ctx.state(), &proto).await?; + + // The decoded Values row contains a cast to the dictionary type wrapping the + // inner Utf8 literal, not a bare inner literal (which would drop the type). + match &plan2 { + LogicalPlan::Values(Values { values, .. }) => match &values[0][0] { + Expr::Cast(cast) => { + assert_eq!( + cast.data_type, + DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(DataType::Utf8) + ) + ); + match cast.expr.as_ref() { + Expr::Literal(ScalarValue::Utf8(Some(s)), _) => { + assert_eq!(s, "a"); + } + other => panic!("expected inner Utf8 literal, got {other:?}"), + } + } + other => panic!("expected Cast expr, got {other:?}"), + }, + other => panic!("expected Values plan, got {other:?}"), + } + + let plan2 = ctx.state().optimize(&plan2)?; + + // The dictionary type must survive the round-trip. + assert_eq!(plan.schema(), plan2.schema()); + assert_eq!( + plan2.schema().field(0).data_type(), + &DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)) + ); + + // Executing the round-tripped plan must succeed and produce a + // Dictionary-typed array matching the schema. + let batches = DataFrame::new(ctx.state(), plan2).collect().await?; + let batch = &batches[0]; + assert_eq!( + batch.schema().field(0).data_type(), + &DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)) + ); + assert_eq!(batch.num_rows(), 1); + assert_eq!( + batch.column(0).data_type(), + &DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)) + ); + Ok(()) +} + #[tokio::test] async fn roundtrip_values_with_scalar_function() -> Result<()> { let ctx = create_context().await?;