From 28aad96a6a97a1b0136244356b410c0b38c93573 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:26:50 +0800 Subject: [PATCH 1/3] fix(substrait): support Dictionary literals in producer (#24) Encode ScalarValue::Dictionary as its inner value in to_substrait_literal. Dictionary type info is carried by the type layer as a Map with DICTIONARY_MAP_TYPE_VARIATION_REF, so the literal layer only needs the inner value; null dictionaries are handled by the is_null() check at the top. GreptimeDB flow queries against dictionary-encoded PK string columns (e.g. metric tables) fail with 'Failed to encode DataFusion plan: NotImplemented("Unsupported literal: Dictionary(UInt32, Utf8(...))")'. Fixing the producer unblocks all substrait encode paths (flow, dist plan, TQL) at once. Also add unit tests for non-null and null dictionary literals. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/logical_plan/producer/expr/literal.rs | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index 8882c992dca1c..a103ac530b037 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs @@ -360,6 +360,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, @@ -558,4 +572,47 @@ mod tests { assert_eq!(scalar, roundtrip_scalar); Ok(()) } + + #[test] + fn test_dictionary_literal() -> Result<()> { + let state = SessionContext::default().state(); + let mut producer = DefaultSubstraitProducer::new(&state); + + // Dictionary(UInt32, Utf8("a")) is encoded as its inner Utf8 literal. + let dict = ScalarValue::Dictionary( + Box::new(DataType::UInt32), + Box::new(ScalarValue::Utf8(Some("a".to_string()))), + ); + let literal = to_substrait_literal(&mut producer, &dict)?; + assert_eq!( + literal.literal_type, + Some(LiteralType::String("a".to_string())) + ); + assert!(!literal.nullable); + assert_eq!( + literal.type_variation_reference, + DEFAULT_CONTAINER_TYPE_VARIATION_REF + ); + + // The encoded literal decodes back to the unwrapped Utf8 value. + let roundtrip = from_substrait_literal_without_names(&test_consumer(), &literal)?; + assert_eq!(roundtrip, ScalarValue::Utf8(Some("a".to_string()))); + 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(()) + } } From 139cfb62ef31ee6c7cd180a1a2f24ad5f096b7e2 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:17:16 +0800 Subject: [PATCH 2/3] fix(substrait): preserve dictionary type when encoding dictionary literals (#25) Address review feedback: encoding ScalarValue::Dictionary as its bare inner literal loses the dictionary type in positions without a separate type layer (projection, VALUES), silently changing the output schema. Wrap the inner literal in a cast to the dictionary type in to_substrait_literal_expr so the consumer reconstructs a Dictionary-typed value. Also add a full expression round-trip test verifying the type is preserved. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/logical_plan/producer/expr/literal.rs | 91 ++++++++++++++++--- 1 file changed, 77 insertions(+), 14 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index a103ac530b037..f35e643eec4a8 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)), @@ -421,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; @@ -573,30 +602,64 @@ mod tests { Ok(()) } - #[test] - fn test_dictionary_literal() -> Result<()> { + #[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 its inner Utf8 literal. + // 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 literal = to_substrait_literal(&mut producer, &dict)?; - assert_eq!( - literal.literal_type, - Some(LiteralType::String("a".to_string())) - ); - assert!(!literal.nullable); + 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!( - literal.type_variation_reference, - DEFAULT_CONTAINER_TYPE_VARIATION_REF + from_substrait_type_without_names(&test_consumer(), cast_type)?, + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)) ); - // The encoded literal decodes back to the unwrapped Utf8 value. - let roundtrip = from_substrait_literal_without_names(&test_consumer(), &literal)?; - assert_eq!(roundtrip, ScalarValue::Utf8(Some("a".to_string()))); + // 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.field.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(()) } From dd591397177a3855d91485058691fa6f6f7b6235 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:56:59 +0800 Subject: [PATCH 3/3] fix(substrait): preserve dictionary type in VALUES literals (#26 follow-up) Route VALUES rows containing non-null dictionary literals through the expression format (VirtualTable.expressions) instead of the deprecated literal-row path, which encoded only the inner scalar and lost the dictionary type. The consumer decodes the resulting cast back to a Dictionary-typed expression, so physical planning no longer hits RecordBatch::try_new_with_options schema mismatch. Also fix a compile error in the dictionary literal test from #26 (cast.field.data_type() -> cast.data_type). Adds roundtrip_values_with_dictionary end-to-end test. Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/logical_plan/producer/expr/literal.rs | 4 +- .../src/logical_plan/producer/rel/read_rel.rs | 38 ++++++++-- .../tests/cases/roundtrip_logical_plan.rs | 69 +++++++++++++++++++ 3 files changed, 105 insertions(+), 6 deletions(-) diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index f35e643eec4a8..bc8ee89d046cf 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs @@ -645,8 +645,8 @@ mod tests { match from_substrait_rex(&test_consumer(), &expr, &DFSchema::empty()).await? { Expr::Cast(cast) => { assert_eq!( - cast.field.data_type(), - &DataType::Dictionary( + cast.data_type, + DataType::Dictionary( Box::new(DataType::UInt32), Box::new(DataType::Utf8) ) 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?;