Add support StringView / BinaryView in interleave kernel

arrow-select/src/interleave.rs

@@ -24,7 +24,9 @@ use arrow_array::types::*;
 use arrow_array::*;
 use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer, NullBufferBuilder, OffsetBuffer};
 use arrow_data::transform::MutableArrayData;
+use arrow_data::ByteView;
 use arrow_schema::{ArrowError, DataType};
+use std::collections::HashMap;
 use std::sync::Arc;
 
 macro_rules! primitive_helper {
@@ -97,6 +99,8 @@ pub fn interleave(
         DataType::LargeUtf8 => interleave_bytes::<LargeUtf8Type>(values, indices),
         DataType::Binary => interleave_bytes::<BinaryType>(values, indices),
         DataType::LargeBinary => interleave_bytes::<LargeBinaryType>(values, indices),
+        DataType::BinaryView => interleave_views::<BinaryViewType>(values, indices),
+        DataType::Utf8View => interleave_views::<StringViewType>(values, indices),
         DataType::Dictionary(k, _) => downcast_integer! {
             k.as_ref() => (dict_helper, values, indices),
             _ => unreachable!("illegal dictionary key type {k}")
@@ -231,6 +235,39 @@ fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
     Ok(Arc::new(array))
 }
 
+fn interleave_views<T: ByteViewType>(
+    values: &[&dyn Array],
+    indices: &[(usize, usize)],
+) -> Result<ArrayRef, ArrowError> {
+    let interleaved = Interleave::<'_, GenericByteViewArray<T>>::new(values, indices);
+    let mut views_builder = BufferBuilder::new(indices.len());
+    let mut buffers = Vec::with_capacity(values[0].len());
+
+    let mut buffer_lookup = HashMap::new();
+    for (array_idx, value_idx) in indices {
+        let array = interleaved.arrays[*array_idx];
+        let raw_view = array.views().get(*value_idx).unwrap();
+        let view = ByteView::from(*raw_view);
+
+        if view.length <= 12 {
+            views_builder.append(*raw_view);
+            continue;
+        }
+        // value is big enough to be in a variadic buffer
+        let new_buffer_idx: &mut u32 = buffer_lookup
+            .entry((*array_idx, view.buffer_index))
+            .or_insert_with(|| {
+                buffers.push(array.data_buffers()[view.buffer_index as usize].clone());
+                (buffers.len() - 1) as u32
+            });
+        views_builder.append(view.with_buffer_index(*new_buffer_idx).into());
+    }
+
+    let array =
+        GenericByteViewArray::<T>::try_new(views_builder.into(), buffers, interleaved.nulls)?;
+    Ok(Arc::new(array))
+}
+
 /// Fallback implementation of interleave using [`MutableArrayData`]
 fn interleave_fallback(
     values: &[&dyn Array],
@@ -461,4 +498,125 @@ mod tests {
             DictionaryArray::<Int32Type>::from_iter(vec![Some("0"), Some("1"), Some("2"), None]);
         assert_eq!(array.as_ref(), &expected)
     }
+
+    #[test]
+    fn test_interleave_views() {
+        let values = StringArray::from_iter_values([
+            "hello",
+            "world_long_string_not_inlined",
+            "foo",
+            "bar",
+            "baz",
+        ]);
+        let view_a = StringViewArray::from(&values);
+
+        let values = StringArray::from_iter_values([
+            "test",
+            "data",
+            "more_long_string_not_inlined",
+            "views",
+            "here",
+        ]);
+        let view_b = StringViewArray::from(&values);
+
+        let indices = &[
+            (0, 2), // "foo"
+            (1, 0), // "test"
+            (0, 4), // "baz"
+            (1, 3), // "views"
+            (0, 1), // "world_long_string_not_inlined"
+        ];
+
+        // Test specialized implementation
+        let values = interleave(&[&view_a, &view_b], indices).unwrap();
+        let result = values.as_string_view();
+        assert_eq!(result.data_buffers().len(), 1);
+
+        // Test fallback implementation
+        let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
+        let fallback_result = fallback.as_string_view();
+        // as of commit 97055631, assertion below, commented out to not block future improvements, passes:
+        // note that fallback_result has 2 buffers, but only one long enough string to warrant a buffer
+        // assert_eq!(fallback_result.data_buffers().len(), 2);
+
+        // Convert to strings for easier assertion
+        let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
+
+        let fallback_collected: Vec<_> = fallback_result
+            .iter()
+            .map(|x| x.map(|s| s.to_string()))
+            .collect();
+
+        assert_eq!(&collected, &fallback_collected);
+
+        assert_eq!(
+            &collected,
+            &[
+                Some("foo".to_string()),
+                Some("test".to_string()),
+                Some("baz".to_string()),
+                Some("views".to_string()),
+                Some("world_long_string_not_inlined".to_string()),
+            ]
+        );
+    }
+
+    #[test]
+    fn test_interleave_views_with_nulls() {
+        let values = StringArray::from_iter([
+            Some("hello"),
+            None,
+            Some("foo_long_string_not_inlined"),
+            Some("bar"),
+            None,
+        ]);
+        let view_a = StringViewArray::from(&values);
+
+        let values = StringArray::from_iter([
+            Some("test"),
+            Some("data_long_string_not_inlined"),
+            None,
+            None,
+            Some("here"),
+        ]);
+        let view_b = StringViewArray::from(&values);
+
+        let indices = &[
+            (0, 1), // null
+            (1, 2), // null
+            (0, 2), // "foo_long_string_not_inlined"
+            (1, 3), // null
+            (0, 4), // null
+        ];
+
+        // Test specialized implementation
+        let values = interleave(&[&view_a, &view_b], indices).unwrap();
+        let result = values.as_string_view();
+        assert_eq!(result.data_buffers().len(), 1);
+
+        // Test fallback implementation
+        let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
+        let fallback_result = fallback.as_string_view();
+
+        // Convert to strings for easier assertion
+        let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
+
+        let fallback_collected: Vec<_> = fallback_result
+            .iter()
+            .map(|x| x.map(|s| s.to_string()))
+            .collect();
+
+        assert_eq!(&collected, &fallback_collected);
+
+        assert_eq!(
+            &collected,
+            &[
+                None,
+                None,
+                Some("foo_long_string_not_inlined".to_string()),
+                None,
+                None,
+            ]
+        );
+    }
 }