fix(substrait): preserve dictionary type in VALUES literals

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<Expression> {
+    // 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<T: OffsetSizeTrait>(
 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(())
+    }
 }

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![])

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?;