Binary / row helpers

arrow-row/src/lib.rs

@@ -132,7 +132,7 @@ use std::sync::Arc;
 use arrow_array::cast::*;
 use arrow_array::types::ArrowDictionaryKeyType;
 use arrow_array::*;
-use arrow_buffer::ArrowNativeType;
+use arrow_buffer::{ArrowNativeType, OffsetBuffer, ScalarBuffer};
 use arrow_data::ArrayDataBuilder;
 use arrow_schema::*;
 use variable::{decode_binary_view, decode_string_view};
@@ -738,6 +738,23 @@ impl RowConverter {
         }
     }
 
+    /// Create a new [Rows] instance from the given binary data.
+    pub fn from_binary(&self, array: BinaryArray) -> Rows {
+        assert_eq!(
+            array.null_count(),
+            0,
+            "can't construct Rows instance from array with nulls"
+        );
+        Rows {
+            buffer: array.values().to_vec(),
+            offsets: array.offsets().into_iter().map(|&i| i.as_usize()).collect(),
+            config: RowConfig {
+                fields: Arc::clone(&self.fields),
+                validate_utf8: true,
+            },
+        }
+    }
+
     /// Convert raw bytes into [`ArrayRef`]
     ///
     /// # Safety
@@ -868,6 +885,24 @@ impl Rows {
             + self.buffer.len()
             + self.offsets.len() * std::mem::size_of::<usize>()
     }
+
+    /// Create a [BinaryArray] from the [Rows] data without reallocating the
+    /// underlying bytes.
+    pub fn into_binary(self) -> BinaryArray {
+        assert!(
+            self.buffer.len() <= i32::MAX as usize,
+            "rows buffer too large"
+        );
+        let offsets_scalar = ScalarBuffer::from_iter(self.offsets.into_iter().map(i32::usize_as));
+        // SAFETY: offsets buffer is nonempty, monotonically increasing, and all represent valid indexes into buffer.
+        unsafe {
+            BinaryArray::new_unchecked(
+                OffsetBuffer::new_unchecked(offsets_scalar),
+                self.buffer.into(),
+                None,
+            )
+        }
+    }
 }
 
 impl<'a> IntoIterator for &'a Rows {
@@ -948,6 +983,11 @@ impl<'a> Row<'a> {
             config: self.config.clone(),
         }
     }
+
+    /// The row's bytes, with the lifetime of the underlying data.
+    pub fn data(&self) -> &'a [u8] {
+        self.data
+    }
 }
 
 // Manually derive these as don't wish to include `fields`
@@ -2271,7 +2311,7 @@ mod tests {
 
             let comparator = LexicographicalComparator::try_new(&sort_columns).unwrap();
 
-            let columns = options
+            let columns: Vec<SortField> = options
                 .into_iter()
                 .zip(&arrays)
                 .map(|(o, a)| SortField::new_with_options(a.data_type().clone(), o))
@@ -2304,6 +2344,24 @@ mod tests {
                 actual.to_data().validate_full().unwrap();
                 dictionary_eq(actual, expected)
             }
+
+            // Check that we can convert
+            let rows = rows.into_binary();
+            let parser = converter.parser();
+            let back = converter
+                .convert_rows(rows.iter().map(|b| parser.parse(b.expect("valid bytes"))))
+                .unwrap();
+            for (actual, expected) in back.iter().zip(&arrays) {
+                actual.to_data().validate_full().unwrap();
+                dictionary_eq(actual, expected)
+            }
+
+            let rows = converter.from_binary(rows);
+            let back = converter.convert_rows(&rows).unwrap();
+            for (actual, expected) in back.iter().zip(&arrays) {
+                actual.to_data().validate_full().unwrap();
+                dictionary_eq(actual, expected)
+            }
         }
     }