Enable casting between Dictionary of DecimalArray and DecimalArray

apache · Nov 29, 2022 · 7cc177d · 7cc177d
1 parent 0b12828
commit 7cc177d
Showing 1 changed file with 89 additions and 2 deletions.
diff --git a/arrow-cast/src/cast.rs b/arrow-cast/src/cast.rs
@@ -692,6 +692,36 @@ pub fn cast_with_options(
                 Float64 => {
                     cast_decimal_to_float!(array, scale, Float64Builder, f64)
                 }
+                Dictionary(index_type, value_type) => match **index_type {
+                    Int8 => {
+                        cast_to_dictionary::<Int8Type>(array, value_type, cast_options)
+                    }
+                    Int16 => {
+                        cast_to_dictionary::<Int16Type>(array, value_type, cast_options)
+                    }
+                    Int32 => {
+                        cast_to_dictionary::<Int32Type>(array, value_type, cast_options)
+                    }
+                    Int64 => {
+                        cast_to_dictionary::<Int64Type>(array, value_type, cast_options)
+                    }
+                    UInt8 => {
+                        cast_to_dictionary::<UInt8Type>(array, value_type, cast_options)
+                    }
+                    UInt16 => {
+                        cast_to_dictionary::<UInt16Type>(array, value_type, cast_options)
+                    }
+                    UInt32 => {
+                        cast_to_dictionary::<UInt32Type>(array, value_type, cast_options)
+                    }
+                    UInt64 => {
+                        cast_to_dictionary::<UInt64Type>(array, value_type, cast_options)
+                    }
+                    _ => Err(ArrowError::CastError(format!(
+                        "Casting from type {:?} to dictionary type {:?} not supported",
+                        from_type, to_type,
+                    ))),
+                },
                 Null => Ok(new_null_array(to_type, array.len())),
                 _ => Err(ArrowError::CastError(format!(
                     "Casting from {:?} to {:?} not supported",
@@ -726,6 +756,36 @@ pub fn cast_with_options(
                     *scale,
                     cast_options,
                 ),
+                Dictionary(index_type, value_type) => match **index_type {
+                    Int8 => {
+                        cast_to_dictionary::<Int8Type>(array, value_type, cast_options)
+                    }
+                    Int16 => {
+                        cast_to_dictionary::<Int16Type>(array, value_type, cast_options)
+                    }
+                    Int32 => {
+                        cast_to_dictionary::<Int32Type>(array, value_type, cast_options)
+                    }
+                    Int64 => {
+                        cast_to_dictionary::<Int64Type>(array, value_type, cast_options)
+                    }
+                    UInt8 => {
+                        cast_to_dictionary::<UInt8Type>(array, value_type, cast_options)
+                    }
+                    UInt16 => {
+                        cast_to_dictionary::<UInt16Type>(array, value_type, cast_options)
+                    }
+                    UInt32 => {
+                        cast_to_dictionary::<UInt32Type>(array, value_type, cast_options)
+                    }
+                    UInt64 => {
+                        cast_to_dictionary::<UInt64Type>(array, value_type, cast_options)
+                    }
+                    _ => Err(ArrowError::CastError(format!(
+                        "Casting from type {:?} to dictionary type {:?} not supported",
+                        from_type, to_type,
+                    ))),
+                },
                 Null => Ok(new_null_array(to_type, array.len())),
                 _ => Err(ArrowError::CastError(format!(
                     "Casting from {:?} to {:?} not supported",
@@ -804,6 +864,20 @@ pub fn cast_with_options(
                     *scale,
                     cast_options,
                 ),
+                Dictionary(index_type, _) => match **index_type {
+                    Int8 => dictionary_cast::<Int8Type>(array, to_type, cast_options),
+                    Int16 => dictionary_cast::<Int16Type>(array, to_type, cast_options),
+                    Int32 => dictionary_cast::<Int32Type>(array, to_type, cast_options),
+                    Int64 => dictionary_cast::<Int64Type>(array, to_type, cast_options),
+                    UInt8 => dictionary_cast::<UInt8Type>(array, to_type, cast_options),
+                    UInt16 => dictionary_cast::<UInt16Type>(array, to_type, cast_options),
+                    UInt32 => dictionary_cast::<UInt32Type>(array, to_type, cast_options),
+                    UInt64 => dictionary_cast::<UInt64Type>(array, to_type, cast_options),
+                    _ => Err(ArrowError::CastError(format!(
+                        "Casting from dictionary type {:?} to {:?} not supported",
+                        from_type, to_type,
+                    ))),
+                },
                 Null => Ok(new_null_array(to_type, array.len())),
                 _ => Err(ArrowError::CastError(format!(
                     "Casting from {:?} to {:?} not supported",
@@ -855,6 +929,20 @@ pub fn cast_with_options(
                     *scale,
                     cast_options,
                 ),
+                Dictionary(index_type, _) => match **index_type {
+                    Int8 => dictionary_cast::<Int8Type>(array, to_type, cast_options),
+                    Int16 => dictionary_cast::<Int16Type>(array, to_type, cast_options),
+                    Int32 => dictionary_cast::<Int32Type>(array, to_type, cast_options),
+                    Int64 => dictionary_cast::<Int64Type>(array, to_type, cast_options),
+                    UInt8 => dictionary_cast::<UInt8Type>(array, to_type, cast_options),
+                    UInt16 => dictionary_cast::<UInt16Type>(array, to_type, cast_options),
+                    UInt32 => dictionary_cast::<UInt32Type>(array, to_type, cast_options),
+                    UInt64 => dictionary_cast::<UInt64Type>(array, to_type, cast_options),
+                    _ => Err(ArrowError::CastError(format!(
+                        "Casting from dictionary type {:?} to {:?} not supported",
+                        from_type, to_type,
+                    ))),
+                },
                 Null => Ok(new_null_array(to_type, array.len())),
                 _ => Err(ArrowError::CastError(format!(
                     "Casting from {:?} to {:?} not supported",
@@ -7028,8 +7116,7 @@ mod tests {
     #[cfg(feature = "chrono-tz")]
     fn make_dictionary_primitive<K: ArrowDictionaryKeyType>() -> ArrayRef {
         // Pick Int32 arbitrarily for dictionary values
-        let mut b: PrimitiveDictionaryBuilder<K, Int32Type> =
-            PrimitiveDictionaryBuilder::new();
+        let mut b: PrimitiveDictionaryBuilder<K, V> = PrimitiveDictionaryBuilder::new();
         b.append(1).unwrap();
         b.append(2).unwrap();
         Arc::new(b.finish())