From 647fddb7a6e35f185f8256bc99e8a24f9f78d026 Mon Sep 17 00:00:00 2001
From: sundy-li <543950155@qq.com>
Date: Sat, 7 Aug 2021 12:38:37 +0800
Subject: [PATCH] [cast] make CastOptions private and introduce
 wrapping_dictionary_to_dictionary_keys
 wrapping_dictionary_to_dictionary_values

---
 src/compute/cast/dictionary_to.rs | 38 ++++++++++++++++++++++++++++---
 src/compute/cast/mod.rs           |  4 ++--
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/compute/cast/dictionary_to.rs b/src/compute/cast/dictionary_to.rs
index 08243e7e586..9e053c9d9a0 100644
--- a/src/compute/cast/dictionary_to.rs
+++ b/src/compute/cast/dictionary_to.rs
@@ -1,4 +1,4 @@
-use super::{primitive_to_primitive, CastOptions};
+use super::{primitive_as_primitive, primitive_to_primitive, CastOptions};
 use crate::{
     array::{Array, DictionaryArray, DictionaryKey, PrimitiveArray},
     compute::{cast::cast_with_options, take::take},
@@ -28,12 +28,24 @@ macro_rules! key_cast {
 pub fn dictionary_to_dictionary_values<K: DictionaryKey>(
     from: &DictionaryArray<K>,
     values_type: &DataType,
-    options: CastOptions,
 ) -> Result<DictionaryArray<K>> {
     let keys = from.keys();
     let values = from.values();
 
-    let values = cast_with_options(values.as_ref(), values_type, options)?.into();
+    let values = cast_with_options(values.as_ref(), values_type, CastOptions::default())?.into();
+    Ok(DictionaryArray::from_data(keys.clone(), values))
+}
+
+/// Similar to dictionary_to_dictionary_values, but overflowing cast is wrapped
+pub fn wrapping_dictionary_to_dictionary_values<K: DictionaryKey>(
+    from: &DictionaryArray<K>,
+    values_type: &DataType,
+) -> Result<DictionaryArray<K>> {
+    let keys = from.keys();
+    let values = from.values();
+
+    let values =
+        cast_with_options(values.as_ref(), values_type, CastOptions { wrapped: true })?.into();
     Ok(DictionaryArray::from_data(keys.clone(), values))
 }
 
@@ -61,6 +73,26 @@ where
     }
 }
 
+/// Similar to dictionary_to_dictionary_keys, but overflowing cast is wrapped
+pub fn wrapping_dictionary_to_dictionary_keys<K1, K2>(
+    from: &DictionaryArray<K1>,
+) -> Result<DictionaryArray<K2>>
+where
+    K1: DictionaryKey + num::traits::AsPrimitive<K2>,
+    K2: DictionaryKey,
+{
+    let keys = from.keys();
+    let values = from.values();
+
+    let casted_keys = primitive_as_primitive::<K1, K2>(keys, &K2::DATA_TYPE);
+
+    if casted_keys.null_count() > keys.null_count() {
+        Err(ArrowError::KeyOverflowError)
+    } else {
+        Ok(DictionaryArray::from_data(casted_keys, values.clone()))
+    }
+}
+
 pub(super) fn dictionary_cast_dyn<K: DictionaryKey>(
     array: &dyn Array,
     to_type: &DataType,
diff --git a/src/compute/cast/mod.rs b/src/compute/cast/mod.rs
index 096ddf1b359..2a99c899c31 100644
--- a/src/compute/cast/mod.rs
+++ b/src/compute/cast/mod.rs
@@ -38,7 +38,7 @@ pub use utf8_to::*;
 
 /// options defining how Cast kernels behave
 #[derive(Clone, Copy, Debug)]
-pub struct CastOptions {
+struct CastOptions {
     /// default to false
     /// whether an overflowing cast should be converted to `None` (default), or be wrapped (i.e. `256i16 as u8 = 0` vectorized).
     /// Settings this to `true` is 5-6x faster for numeric types.
@@ -341,7 +341,7 @@ pub fn wrapping_cast(array: &dyn Array, to_type: &DataType) -> Result<Box<dyn Ar
 }
 
 #[inline]
-pub fn cast_with_options(
+fn cast_with_options(
     array: &dyn Array,
     to_type: &DataType,
     options: CastOptions,