only attempt merge on supported types

Author: davidhewitt

arrow-select/src/concat.rs

@@ -30,7 +30,9 @@
 //! assert_eq!(arr.len(), 3);
 //! ```
 
-use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values};
+use crate::dictionary::{
+    merge_dictionary_values, should_merge_dictionary_values, ShouldMergeValues,
+};
 use arrow_array::builder::{
     BooleanBuilder, GenericByteBuilder, PrimitiveBuilder, PrimitiveDictionaryBuilder,
 };
@@ -96,9 +98,14 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
         .map(|x| x.as_dictionary::<K>())
         .inspect(|d| output_len += d.len())
         .collect();
-    if !should_merge_dictionary_values::<K>(&dictionaries, output_len) {
-        return concat_fallback(arrays, Capacities::Array(output_len));
-    }
+
+    let is_overflow = match should_merge_dictionary_values::<K>(&dictionaries, output_len) {
+        ShouldMergeValues::ConcatWillOverflow => true,
+        ShouldMergeValues::Yes => false,
+        ShouldMergeValues::No => {
+            return concat_fallback(arrays, Capacities::Array(output_len));
+        }
+    };
 
     macro_rules! primitive_dict_helper {
         ($t:ty) => {
@@ -111,8 +118,12 @@ fn concat_dictionaries<K: ArrowDictionaryKeyType>(
         DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
             merge_concat_byte_dictionaries(&dictionaries, output_len)
         },
+        // merge not yet implemented for this type and it's not going to overflow, so fall back
+        // to concatenating values
+        _ if !is_overflow => concat_fallback(arrays, Capacities::Array(output_len)),
         other => Err(ArrowError::NotYetImplemented(format!(
-            "interleave does not yet support merging dictionaries with value type {other:?}"
+            "concat of dictionaries would overflow key type {key_type:?} with value type {other:?}",
+            key_type = K::DATA_TYPE,
         )))
     }
 }

arrow-select/src/dictionary.rs

@@ -101,6 +101,17 @@ fn bytes_ptr_eq<T: ByteArrayType>(a: &dyn Array, b: &dyn Array) -> bool {
     }
 }
 
+/// Whether selection kernels should attempt to merge dictionary values
+pub enum ShouldMergeValues {
+    /// Concatenation of the dictionary values will lead to overflowing
+    /// the key space; it's necessary to attempt to merge
+    ConcatWillOverflow,
+    /// The heuristic suggests that merging will be beneficial
+    Yes,
+    /// The heuristic suggests that merging is not necessary
+    No,
+}
+
 /// A type-erased function that compares two array for pointer equality
 type PtrEq = fn(&dyn Array, &dyn Array) -> bool;
 
@@ -112,7 +123,7 @@ type PtrEq = fn(&dyn Array, &dyn Array) -> bool;
 pub fn should_merge_dictionary_values<K: ArrowDictionaryKeyType>(
     dictionaries: &[&DictionaryArray<K>],
     len: usize,
-) -> bool {
+) -> ShouldMergeValues {
     use DataType::*;
     let first_values = dictionaries[0].values().as_ref();
     let ptr_eq: PtrEq = match first_values.data_type() {
@@ -136,7 +147,15 @@ pub fn should_merge_dictionary_values<K: ArrowDictionaryKeyType>(
     let overflow = K::Native::from_usize(total_values).is_none();
     let values_exceed_length = total_values >= len;
 
-    !single_dictionary && (overflow || values_exceed_length)
+    if single_dictionary {
+        ShouldMergeValues::No
+    } else if overflow {
+        ShouldMergeValues::ConcatWillOverflow
+    } else if values_exceed_length {
+        ShouldMergeValues::Yes
+    } else {
+        ShouldMergeValues::No
+    }
 }
 
 /// Given an array of dictionaries and an optional key mask compute a values array

arrow-select/src/interleave.rs

@@ -17,7 +17,9 @@
 
 //! Interleave elements from multiple arrays
 
-use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values};
+use crate::dictionary::{
+    merge_dictionary_values, should_merge_dictionary_values, ShouldMergeValues,
+};
 use arrow_array::builder::{
     BooleanBufferBuilder, BufferBuilder, PrimitiveBuilder, PrimitiveDictionaryBuilder,
 };
@@ -198,9 +200,13 @@ fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
     indices: &[(usize, usize)],
 ) -> Result<ArrayRef, ArrowError> {
     let dictionaries: Vec<_> = arrays.iter().map(|x| x.as_dictionary::<K>()).collect();
-    if !should_merge_dictionary_values::<K>(&dictionaries, indices.len()) {
-        return interleave_fallback(arrays, indices);
-    }
+    let is_overflow = match should_merge_dictionary_values::<K>(&dictionaries, indices.len()) {
+        ShouldMergeValues::ConcatWillOverflow => true,
+        ShouldMergeValues::Yes => false,
+        ShouldMergeValues::No => {
+            return interleave_fallback(arrays, indices);
+        }
+    };
 
     macro_rules! primitive_dict_helper {
         ($t:ty) => {
@@ -213,8 +219,12 @@ fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
         DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
             merge_interleave_byte_dictionaries(&dictionaries, indices)
         },
+        // merge not yet implemented for this type and it's not going to overflow, so fall back
+        // to concatenating values
+        _ if !is_overflow => interleave_fallback(arrays, indices),
         other => Err(ArrowError::NotYetImplemented(format!(
-            "interleave does not yet support merging dictionaries with value type {other:?}"
+            "interleave of dictionaries would overflow key type {key_type:?} with value type {other:?}",
+            key_type = K::DATA_TYPE,
         )))
     }
 }