diff --git a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs index f26984ad4a394..f81333e630fbd 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/literal.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/literal.rs @@ -286,19 +286,20 @@ pub(crate) fn to_substrait_literal( }?; LiteralType::EmptyMap(mt) } else { - let keys = (0..m.keys().len()) + let entries = m.value(0); + let keys = (0..entries.len()) .map(|i| { to_substrait_literal( producer, - &ScalarValue::try_from_array(&m.keys(), i)?, + &ScalarValue::try_from_array(entries.column(0), i)?, ) }) .collect::>>()?; - let values = (0..m.values().len()) + let values = (0..entries.len()) .map(|i| { to_substrait_literal( producer, - &ScalarValue::try_from_array(&m.values(), i)?, + &ScalarValue::try_from_array(entries.column(1), i)?, ) }) .collect::>>()?; @@ -548,6 +549,32 @@ mod tests { Ok(()) } + #[test] + fn round_trip_sliced_map_literals() -> Result<()> { + let mut builder = + MapBuilder::new(None, StringBuilder::new(), Int64Builder::new()); + builder.keys().append_value("prefix"); + builder.values().append_value(10); + builder.append(true)?; + builder.keys().append_value("selected"); + builder.values().append_value(20); + builder.keys().append_value("null_value"); + builder.values().append_null(); + builder.append(true)?; + builder.append(true)?; + builder.append(false)?; + builder.keys().append_value("suffix"); + builder.values().append_value(30); + builder.append(true)?; + let maps = builder.finish(); + + // Scalar extraction retains the other rows' entries in the backing arrays. + for row in 0..maps.len() { + round_trip_literal(ScalarValue::try_from_array(&maps, row)?)?; + } + Ok(()) + } + fn round_trip_literal(scalar: ScalarValue) -> Result<()> { println!("Checking round trip of {scalar:?}"); let state = SessionContext::default().state();