Skip to content

Commit

Permalink
feat: Support Map type in Substrait conversions
Browse files Browse the repository at this point in the history
This only supports Map _type_, not Map literals,
since Map isn't yet implemented in ScalarValue
(apache#11128)
  • Loading branch information
Blizzara committed Jun 26, 2024
1 parent 7adc940 commit 1fa6596
Show file tree
Hide file tree
Showing 2 changed files with 64 additions and 0 deletions.
30 changes: 30 additions & 0 deletions datafusion/substrait/src/logical_plan/consumer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1449,6 +1449,36 @@ fn from_substrait_type(
)?,
}
}
r#type::Kind::Map(map) => {
let key_type = map.key.as_ref().ok_or_else(|| {
substrait_datafusion_err!("Map type must have key type")
})?;
let value_type = map.value.as_ref().ok_or_else(|| {
substrait_datafusion_err!("Map type must have value type")
})?;
let key_field = Arc::new(Field::new(
"key",
from_substrait_type(key_type, dfs_names, name_idx)?,
false,
));
let value_field = Arc::new(Field::new(
"value",
from_substrait_type(value_type, dfs_names, name_idx)?,
true,
));
match map.type_variation_reference {
DEFAULT_CONTAINER_TYPE_VARIATION_REF => {
Ok(DataType::Map(Arc::new(Field::new_struct(
"entries",
[key_field, value_field],
false, // The inner map field is always non-nullable (Arrow #1697),
)), false))
},
v => not_impl_err!(
"Unsupported Substrait type variation {v} of type {s_kind:?}"
)?,
}
}
r#type::Kind::Decimal(d) => match d.type_variation_reference {
DECIMAL_128_TYPE_VARIATION_REF => {
Ok(DataType::Decimal128(d.precision as u8, d.scale as i8))
Expand Down
34 changes: 34 additions & 0 deletions datafusion/substrait/src/logical_plan/producer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1619,6 +1619,27 @@ fn to_substrait_type(dt: &DataType, nullable: bool) -> Result<substrait::proto::
}))),
})
}
DataType::Map(inner, _) => match inner.data_type() {
DataType::Struct(key_and_value) if key_and_value.len() == 2 => {
let key_type = to_substrait_type(
&key_and_value[0].data_type(),
key_and_value[0].is_nullable(),
)?;
let value_type = to_substrait_type(
&key_and_value[1].data_type(),
key_and_value[1].is_nullable(),
)?;
Ok(substrait::proto::Type {
kind: Some(r#type::Kind::Map(Box::new(r#type::Map {
key: Some(Box::new(key_type)),
value: Some(Box::new(value_type)),
type_variation_reference: DEFAULT_CONTAINER_TYPE_VARIATION_REF,
nullability,
}))),
})
}
_ => plan_err!("Map fields must contain a Struct with exactly 2 fields"),
},
DataType::Struct(fields) => {
let field_types = fields
.iter()
Expand Down Expand Up @@ -2316,6 +2337,19 @@ mod test {
Field::new_list_field(DataType::Int32, true).into(),
))?;

round_trip_type(DataType::Map(
Field::new_struct(
"entries",
[
Field::new("key", DataType::Utf8, false).into(),
Field::new("value", DataType::Int32, true).into(),
],
false,
)
.into(),
false,
))?;

round_trip_type(DataType::Struct(
vec![
Field::new("c0", DataType::Int32, true),
Expand Down

0 comments on commit 1fa6596

Please sign in to comment.