Add support for Arrow Time types in Substrait

This commit adds support for the Arrow Dictionary type in Substrait
plans.

Resolves #16273

Author: jkosh44

datafusion/substrait/src/logical_plan/consumer/types.rs

@@ -22,13 +22,13 @@ use crate::variation_const::{
     DATE_32_TYPE_VARIATION_REF, DATE_64_TYPE_VARIATION_REF,
     DECIMAL_128_TYPE_VARIATION_REF, DECIMAL_256_TYPE_VARIATION_REF,
     DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_INTERVAL_DAY_TYPE_VARIATION_REF,
-    DEFAULT_TYPE_VARIATION_REF, DURATION_INTERVAL_DAY_TYPE_VARIATION_REF,
-    INTERVAL_DAY_TIME_TYPE_REF, INTERVAL_MONTH_DAY_NANO_TYPE_NAME,
-    INTERVAL_MONTH_DAY_NANO_TYPE_REF, INTERVAL_YEAR_MONTH_TYPE_REF,
-    LARGE_CONTAINER_TYPE_VARIATION_REF, TIMESTAMP_MICRO_TYPE_VARIATION_REF,
-    TIMESTAMP_MILLI_TYPE_VARIATION_REF, TIMESTAMP_NANO_TYPE_VARIATION_REF,
-    TIMESTAMP_SECOND_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
-    VIEW_CONTAINER_TYPE_VARIATION_REF,
+    DEFAULT_TYPE_VARIATION_REF, DICTIONARY_CONTAINER_TYPE_VARIATION_REF,
+    DURATION_INTERVAL_DAY_TYPE_VARIATION_REF, INTERVAL_DAY_TIME_TYPE_REF,
+    INTERVAL_MONTH_DAY_NANO_TYPE_NAME, INTERVAL_MONTH_DAY_NANO_TYPE_REF,
+    INTERVAL_YEAR_MONTH_TYPE_REF, LARGE_CONTAINER_TYPE_VARIATION_REF,
+    TIMESTAMP_MICRO_TYPE_VARIATION_REF, TIMESTAMP_MILLI_TYPE_VARIATION_REF,
+    TIMESTAMP_NANO_TYPE_VARIATION_REF, TIMESTAMP_SECOND_TYPE_VARIATION_REF,
+    UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF,
 };
 use datafusion::arrow::datatypes::{
     DataType, Field, Fields, IntervalUnit, Schema, TimeUnit,
@@ -181,24 +181,32 @@ pub fn from_substrait_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(consumer, key_type, dfs_names, name_idx)?,
-                    false,
-                ));
-                let value_field = Arc::new(Field::new(
-                    "value",
-                    from_substrait_type(consumer, value_type, dfs_names, name_idx)?,
-                    true,
-                ));
-                Ok(DataType::Map(
-                    Arc::new(Field::new_struct(
-                        "entries",
-                        [key_field, value_field],
-                        false, // The inner map field is always non-nullable (Arrow #1697),
+                let key_type =
+                    from_substrait_type(consumer, key_type, dfs_names, name_idx)?;
+                let value_type =
+                    from_substrait_type(consumer, value_type, dfs_names, name_idx)?;
+
+                match map.type_variation_reference {
+                    DEFAULT_CONTAINER_TYPE_VARIATION_REF => {
+                        let key_field = Arc::new(Field::new("key", key_type, false));
+                        let value_field = Arc::new(Field::new("value", value_type, true));
+                        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, // whether keys are sorted
+                        ))
+                    }
+                    DICTIONARY_CONTAINER_TYPE_VARIATION_REF => Ok(DataType::Dictionary(
+                        Box::new(key_type),
+                        Box::new(value_type),
                     )),
-                    false, // whether keys are sorted
-                ))
+                    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 => {

datafusion/substrait/src/logical_plan/producer/types.rs

@@ -20,9 +20,9 @@ use crate::variation_const::{
     DATE_32_TYPE_VARIATION_REF, DATE_64_TYPE_VARIATION_REF,
     DECIMAL_128_TYPE_VARIATION_REF, DECIMAL_256_TYPE_VARIATION_REF,
     DEFAULT_CONTAINER_TYPE_VARIATION_REF, DEFAULT_INTERVAL_DAY_TYPE_VARIATION_REF,
-    DEFAULT_TYPE_VARIATION_REF, DURATION_INTERVAL_DAY_TYPE_VARIATION_REF,
-    LARGE_CONTAINER_TYPE_VARIATION_REF, UNSIGNED_INTEGER_TYPE_VARIATION_REF,
-    VIEW_CONTAINER_TYPE_VARIATION_REF,
+    DEFAULT_TYPE_VARIATION_REF, DICTIONARY_CONTAINER_TYPE_VARIATION_REF,
+    DURATION_INTERVAL_DAY_TYPE_VARIATION_REF, LARGE_CONTAINER_TYPE_VARIATION_REF,
+    UNSIGNED_INTEGER_TYPE_VARIATION_REF, VIEW_CONTAINER_TYPE_VARIATION_REF,
 };
 use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
 use datafusion::common::{internal_err, not_impl_err, plan_err, DFSchemaRef};
@@ -271,6 +271,18 @@ pub(crate) fn to_substrait_type(
             }
             _ => plan_err!("Map fields must contain a Struct with exactly 2 fields"),
         },
+        DataType::Dictionary(key_type, value_type) => {
+            let key_type = to_substrait_type(key_type, nullable)?;
+            let value_type = to_substrait_type(value_type, 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: DICTIONARY_CONTAINER_TYPE_VARIATION_REF,
+                    nullability,
+                }))),
+            })
+        }
         DataType::Struct(fields) => {
             let field_types = fields
                 .iter()
@@ -391,6 +403,10 @@ mod tests {
             .into(),
             false,
         ))?;
+        round_trip_type(DataType::Dictionary(
+            Box::new(DataType::Utf8),
+            Box::new(DataType::Int32),
+        ))?;
 
         round_trip_type(DataType::Struct(
             vec![

datafusion/substrait/src/variation_const.rs

@@ -53,6 +53,10 @@ pub const DATE_64_TYPE_VARIATION_REF: u32 = 1;
 pub const DEFAULT_CONTAINER_TYPE_VARIATION_REF: u32 = 0;
 pub const LARGE_CONTAINER_TYPE_VARIATION_REF: u32 = 1;
 pub const VIEW_CONTAINER_TYPE_VARIATION_REF: u32 = 2;
+/// Used for the arrow type [`DataType::Map`].
+///
+/// [`DataType::Map`]: datafusion::arrow::datatypes::DataType::Map
+pub const DICTIONARY_CONTAINER_TYPE_VARIATION_REF: u32 = 3;
 pub const DECIMAL_128_TYPE_VARIATION_REF: u32 = 0;
 pub const DECIMAL_256_TYPE_VARIATION_REF: u32 = 1;
 /// Used for the arrow type [`DataType::Interval`] with [`IntervalUnit::DayTime`].