Optimized writing of byte array to parquet (apache#1764)

tustvold · Jul 29, 2022 · 43dcd97 · 43dcd97
1 parent 6ce4c4e
commit 43dcd97
Show file tree

Hide file tree

Showing 7 changed files with 731 additions and 218 deletions.
diff --git a/arrow/src/array/data.rs b/arrow/src/array/data.rs
@@ -23,7 +23,7 @@ use crate::datatypes::{
     UnionMode,
 };
 use crate::error::{ArrowError, Result};
-use crate::util::bit_iterator::BitSliceIterator;
+use crate::util::bit_iterator::{BitIndexIterator, BitSliceIterator};
 use crate::{bitmap::Bitmap, datatypes::ArrowNativeType};
 use crate::{
     buffer::{Buffer, MutableBuffer},
@@ -1277,6 +1277,26 @@ impl ArrayData {
             .all(|(a, b)| a.ptr_eq(b))
     }
 
+    /// Returns an iterator over the valid indexes of this array
+    pub fn iter_valid_idx(&self) -> Option<BitIndexIterator<'_>> {
+        let nulls = self.null_bitmap()?;
+        Some(BitIndexIterator::new(
+            nulls.buffer_ref().as_slice(),
+            self.offset,
+            self.len,
+        ))
+    }
+
+    /// Returns an iterator over the valid slices of this array
+    pub fn iter_valid_slices(&self) -> Option<BitSliceIterator<'_>> {
+        let nulls = self.null_bitmap()?;
+        Some(BitSliceIterator::new(
+            nulls.buffer_ref().as_slice(),
+            self.offset,
+            self.len,
+        ))
+    }
+
     /// Converts this [`ArrayData`] into an [`ArrayDataBuilder`]
     pub fn into_builder(self) -> ArrayDataBuilder {
         self.into()