Implement arrow-row encoding/decoding for view types

arrow-array/src/array/byte_view_array.rs

@@ -16,19 +16,22 @@
 // under the License.
 
 use crate::array::print_long_array;
-use crate::builder::GenericByteViewBuilder;
+use crate::builder::{ArrayBuilder, GenericByteViewBuilder};
 use crate::iterator::ArrayIter;
 use crate::types::bytes::ByteArrayNativeType;
 use crate::types::{BinaryViewType, ByteViewType, StringViewType};
-use crate::{Array, ArrayAccessor, ArrayRef, Scalar};
-use arrow_buffer::{Buffer, NullBuffer, ScalarBuffer};
+use crate::{Array, ArrayAccessor, ArrayRef, GenericByteArray, OffsetSizeTrait, Scalar};
+use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, ScalarBuffer};
 use arrow_data::{ArrayData, ArrayDataBuilder, ByteView};
 use arrow_schema::{ArrowError, DataType};
+use num::ToPrimitive;
 use std::any::Any;
 use std::fmt::Debug;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
+use super::ByteArrayType;
+
 /// [Variable-size Binary View Layout]: An array of variable length bytes view arrays.
 ///
 /// Different than [`crate::GenericByteArray`] as it stores both an offset and length
@@ -429,6 +432,49 @@ impl<T: ByteViewType + ?Sized> From<ArrayData> for GenericByteViewArray<T> {
     }
 }
 
+impl<FROM, V> From<&GenericByteArray<FROM>> for GenericByteViewArray<V>
+where
+    FROM: ByteArrayType,
+    FROM::Offset: OffsetSizeTrait + ToPrimitive,
+    V: ByteViewType<Native = FROM::Native>,
+{
+    fn from(byte_array: &GenericByteArray<FROM>) -> Self {
+        let offsets = byte_array.offsets();
+
+        let can_reuse_buffer = match offsets.last() {
+            Some(offset) => offset.as_usize() < u32::MAX as usize,
+            None => true,
+        };
+
+        if can_reuse_buffer {
+            let len = byte_array.len();
+            let mut views_builder = GenericByteViewBuilder::<V>::with_capacity(len);
+            let str_values_buf = byte_array.values().clone();
+            let block = views_builder.append_block(str_values_buf);
+            for (i, w) in offsets.windows(2).enumerate() {
+                let offset = w[0].as_usize();
+                let end = w[1].as_usize();
+                let length = end - offset;
+
+                if byte_array.is_null(i) {
+                    views_builder.append_null();
+                } else {
+                    // Safety: the input was a valid array so it valid UTF8 (if string). And
+                    // all offsets were valid and we created the views correctly
+                    unsafe {
+                        views_builder.append_view_unchecked(block, offset as u32, length as u32)
+                    }
+                }
+            }
+            assert_eq!(views_builder.len(), len);
+            views_builder.finish()
+        } else {
+            // TODO: the first u32::MAX can still be reused
+            GenericByteViewArray::<V>::from_iter(byte_array.iter())
+        }
+    }
+}
+
 impl<T: ByteViewType + ?Sized> From<GenericByteViewArray<T>> for ArrayData {
     fn from(mut array: GenericByteViewArray<T>) -> Self {
         let len = array.len();

arrow-cast/src/cast/mod.rs

@@ -2339,31 +2339,11 @@ fn cast_byte_to_view<FROM, V>(array: &dyn Array) -> Result<ArrayRef, ArrowError>
 where
     FROM: ByteArrayType,
     FROM::Offset: OffsetSizeTrait + ToPrimitive,
-    V: ByteViewType,
+    V: ByteViewType<Native = FROM::Native>,
 {
     let byte_array: &GenericByteArray<FROM> = array.as_bytes();
-    let len = array.len();
-    let str_values_buf = byte_array.values().clone();
-    let offsets = byte_array.offsets();
-
-    let mut views_builder = GenericByteViewBuilder::<V>::with_capacity(len);
-    let block = views_builder.append_block(str_values_buf);
-    for (i, w) in offsets.windows(2).enumerate() {
-        let offset = w[0].to_u32().unwrap();
-        let end = w[1].to_u32().unwrap();
-        let length = end - offset;
-
-        if byte_array.is_null(i) {
-            views_builder.append_null();
-        } else {
-            // Safety: the input was a valid array so it valid UTF8 (if string). And
-            // all offsets were valid and we created the views correctly
-            unsafe { views_builder.append_view_unchecked(block, offset, length) }
-        }
-    }
-
-    assert_eq!(views_builder.len(), len);
-    Ok(Arc::new(views_builder.finish()))
+    let byte_view_array = GenericByteViewArray::<V>::from(byte_array);
+    Ok(Arc::new(byte_view_array))
 }
 
 /// Helper function to cast from one `ByteViewType` array to `ByteArrayType` array.

arrow-ord/src/cmp.rs

@@ -540,98 +540,32 @@ impl<'a, T: ByteArrayType> ArrayOrd for &'a GenericByteArray<T> {
     }
 }
 
-/// Comparing two ByteView types are non-trivial.
-/// It takes a bit of patience to understand why we don't just compare two &[u8] directly.
-///
-/// ByteView types give us the following two advantages, and we need to be careful not to lose them:
-/// (1) For string/byte smaller than 12 bytes, the entire data is inlined in the view.
-///     Meaning that reading one array element requires only one memory access
-///     (two memory access required for StringArray, one for offset buffer, the other for value buffer).
-///
-/// (2) For string/byte larger than 12 bytes, we can still be faster than (for certain operations) StringArray/ByteArray,
-///     thanks to the inlined 4 bytes.
-///     Consider equality check:
-///     If the first four bytes of the two strings are different, we can return false immediately (with just one memory access).
-///     If we are unlucky and the first four bytes are the same, we need to fallback to compare two full strings.
 impl<'a, T: ByteViewType> ArrayOrd for &'a GenericByteViewArray<T> {
-    /// Item.0 is the array, Item.1 is the index into the array.
-    /// Why don't we just store Item.0[Item.1] as the item?
-    ///  - Because if we do so, we materialize the entire string (i.e., make multiple memory accesses), which might be unnecessary.
-    ///  - Most of the time (eq, ord), we only need to look at the first 4 bytes to know the answer,
-    ///     e.g., if the inlined 4 bytes are different, we can directly return unequal without looking at the full string.
+    /// This is the item type for the GenericByteViewArray::compare
+    /// Item.0 is the array, Item.1 is the index
     type Item = (&'a GenericByteViewArray<T>, usize);
 
-    /// # Equality check flow
-    /// (1) if both string are smaller than 12 bytes, we can directly compare the data inlined to the view.
-    /// (2) if any of the string is larger than 12 bytes, we need to compare the full string.
-    ///     (2.1) if the inlined 4 bytes are different, we can return false immediately.
-    ///     (2.2) o.w., we need to compare the full string.
-    ///
-    /// # Safety
-    /// (1) Indexing. The Self::Item.1 encodes the index value, which is already checked in `value` function,
-    ///               so it is safe to index into the views.
-    /// (2) Slice data from view. We know the bytes 4-8 are inlined data (per spec), so it is safe to slice from the view.
     fn is_eq(l: Self::Item, r: Self::Item) -> bool {
+        // # Safety
+        // The index is within bounds as it is checked in value()
         let l_view = unsafe { l.0.views().get_unchecked(l.1) };
         let l_len = *l_view as u32;
 
         let r_view = unsafe { r.0.views().get_unchecked(r.1) };
         let r_len = *r_view as u32;
-
+        // This is a fast path for equality check.
+        // We don't need to look at the actual bytes to determine if they are equal.
         if l_len != r_len {
             return false;
         }
 
-        if l_len <= 12 {
-            let l_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, l_len as usize) };
-            let r_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, r_len as usize) };
-            l_data == r_data
-        } else {
-            let l_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, 4) };
-            let r_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, 4) };
-            if l_inlined_data != r_inlined_data {
-                return false;
-            }
-
-            let l_full_data: &[u8] = unsafe { l.0.value_unchecked(l.1).as_ref() };
-            let r_full_data: &[u8] = unsafe { r.0.value_unchecked(r.1).as_ref() };
-            l_full_data == r_full_data
-        }
+        unsafe { compare_byte_view_unchecked(l.0, l.1, r.0, r.1).is_eq() }
     }
 
-    /// # Ordering check flow
-    /// (1) if both string are smaller than 12 bytes, we can directly compare the data inlined to the view.
-    /// (2) if any of the string is larger than 12 bytes, we need to compare the full string.
-    ///     (2.1) if the inlined 4 bytes are different, we can return the result immediately.
-    ///     (2.2) o.w., we need to compare the full string.
-    ///
-    /// # Safety
-    /// (1) Indexing. The Self::Item.1 encodes the index value, which is already checked in `value` function,
-    ///              so it is safe to index into the views.
-    /// (2) Slice data from view. We know the bytes 4-8 are inlined data (per spec), so it is safe to slice from the view.
     fn is_lt(l: Self::Item, r: Self::Item) -> bool {
-        let l_view = l.0.views().get(l.1).unwrap();
-        let l_len = *l_view as u32;
-
-        let r_view = r.0.views().get(r.1).unwrap();
-        let r_len = *r_view as u32;
-
-        if l_len <= 12 && r_len <= 12 {
-            let l_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, l_len as usize) };
-            let r_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, r_len as usize) };
-            return l_data < r_data;
-        }
-        // one of the string is larger than 12 bytes,
-        // we then try to compare the inlined data first
-        let l_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, 4) };
-        let r_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, 4) };
-        if r_inlined_data != l_inlined_data {
-            return l_inlined_data < r_inlined_data;
-        }
-        // unfortunately, we need to compare the full data
-        let l_full_data: &[u8] = unsafe { l.0.value_unchecked(l.1).as_ref() };
-        let r_full_data: &[u8] = unsafe { r.0.value_unchecked(r.1).as_ref() };
-        l_full_data < r_full_data
+        // # Safety
+        // The index is within bounds as it is checked in value()
+        unsafe { compare_byte_view_unchecked(l.0, l.1, r.0, r.1).is_lt() }
     }
 
     fn len(&self) -> usize {
@@ -663,6 +597,78 @@ impl<'a> ArrayOrd for &'a FixedSizeBinaryArray {
     }
 }
 
+/// Comparing two [`GenericByteViewArray`] at index `left_idx` and `right_idx`
+pub fn compare_byte_view<T: ByteViewType>(
+    left: &GenericByteViewArray<T>,
+    left_idx: usize,
+    right: &GenericByteViewArray<T>,
+    right_idx: usize,
+) -> std::cmp::Ordering {
+    assert!(left_idx < left.len());
+    assert!(right_idx < right.len());
+    unsafe { compare_byte_view_unchecked(left, left_idx, right, right_idx) }
+}
+
+/// Comparing two [`GenericByteViewArray`] at index `left_idx` and `right_idx`
+///
+/// Comparing two ByteView types are non-trivial.
+/// It takes a bit of patience to understand why we don't just compare two &[u8] directly.
+///
+/// ByteView types give us the following two advantages, and we need to be careful not to lose them:
+/// (1) For string/byte smaller than 12 bytes, the entire data is inlined in the view.
+///     Meaning that reading one array element requires only one memory access
+///     (two memory access required for StringArray, one for offset buffer, the other for value buffer).
+///
+/// (2) For string/byte larger than 12 bytes, we can still be faster than (for certain operations) StringArray/ByteArray,
+///     thanks to the inlined 4 bytes.
+///     Consider equality check:
+///     If the first four bytes of the two strings are different, we can return false immediately (with just one memory access).
+///
+/// If we directly compare two &[u8], we materialize the entire string (i.e., make multiple memory accesses), which might be unnecessary.
+/// - Most of the time (eq, ord), we only need to look at the first 4 bytes to know the answer,
+///   e.g., if the inlined 4 bytes are different, we can directly return unequal without looking at the full string.
+///
+/// # Order check flow
+/// (1) if both string are smaller than 12 bytes, we can directly compare the data inlined to the view.
+/// (2) if any of the string is larger than 12 bytes, we need to compare the full string.
+///     (2.1) if the inlined 4 bytes are different, we can return the result immediately.
+///     (2.2) o.w., we need to compare the full string.
+///
+/// # Safety
+/// The left/right_idx must within range of each array
+pub unsafe fn compare_byte_view_unchecked<T: ByteViewType>(
+    left: &GenericByteViewArray<T>,
+    left_idx: usize,
+    right: &GenericByteViewArray<T>,
+    right_idx: usize,
+) -> std::cmp::Ordering {
+    let l_view = left.views().get_unchecked(left_idx);
+    let l_len = *l_view as u32;
+
+    let r_view = right.views().get_unchecked(right_idx);
+    let r_len = *r_view as u32;
+
+    if l_len <= 12 && r_len <= 12 {
+        let l_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, l_len as usize) };
+        let r_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, r_len as usize) };
+        return l_data.cmp(r_data);
+    }
+
+    // one of the string is larger than 12 bytes,
+    // we then try to compare the inlined data first
+    let l_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(l_view, 4) };
+    let r_inlined_data = unsafe { GenericByteViewArray::<T>::inline_value(r_view, 4) };
+    if r_inlined_data != l_inlined_data {
+        return l_inlined_data.cmp(r_inlined_data);
+    }
+
+    // unfortunately, we need to compare the full data
+    let l_full_data: &[u8] = unsafe { left.value_unchecked(left_idx).as_ref() };
+    let r_full_data: &[u8] = unsafe { right.value_unchecked(right_idx).as_ref() };
+
+    l_full_data.cmp(r_full_data)
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

arrow-ord/src/ord.rs

@@ -135,6 +135,21 @@ fn compare_bytes<T: ByteArrayType>(
     })
 }
 
+fn compare_byte_view<T: ByteViewType>(
+    left: &dyn Array,
+    right: &dyn Array,
+    opts: SortOptions,
+) -> DynComparator {
+    let left = left.as_byte_view::<T>();
+    let right = right.as_byte_view::<T>();
+
+    let l = left.clone();
+    let r = right.clone();
+    compare(left, right, opts, move |i, j| {
+        crate::cmp::compare_byte_view(&l, i, &r, j)
+    })
+}
+
 fn compare_dict<K: ArrowDictionaryKeyType>(
     left: &dyn Array,
     right: &dyn Array,
@@ -342,8 +357,10 @@ pub fn make_comparator(
         (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)),
         (Utf8, Utf8) => Ok(compare_bytes::<Utf8Type>(left, right, opts)),
         (LargeUtf8, LargeUtf8) => Ok(compare_bytes::<LargeUtf8Type>(left, right, opts)),
+        (Utf8View, Utf8View) => Ok(compare_byte_view::<StringViewType>(left, right, opts)),
         (Binary, Binary) => Ok(compare_bytes::<BinaryType>(left, right, opts)),
         (LargeBinary, LargeBinary) => Ok(compare_bytes::<LargeBinaryType>(left, right, opts)),
+        (BinaryView, BinaryView) => Ok(compare_byte_view::<BinaryViewType>(left, right, opts)),
         (FixedSizeBinary(_), FixedSizeBinary(_)) => {
             let left = left.as_fixed_size_binary();
             let right = right.as_fixed_size_binary();