complete

Author: alamb

parquet/src/arrow/array_reader/cached.rs

@@ -43,7 +43,6 @@ pub(crate) struct CachedPredicateResultBuilder {
     /// <https://github.com/apache/arrow-rs/issues/6692>
     in_progress_arrays: Vec<Box<dyn InProgressArray>>,
     filters: Vec<BooleanArray>,
-    batch_size: usize,
 }
 
 impl CachedPredicateResultBuilder {
@@ -79,7 +78,6 @@ impl CachedPredicateResultBuilder {
                         Self {
                             in_progress_arrays,
                             filters: vec![],
-                            batch_size,
                         }
                     };
                 }
@@ -106,7 +104,6 @@ impl CachedPredicateResultBuilder {
         Self {
             in_progress_arrays,
             filters: vec![],
-            batch_size,
         }
     }
 
@@ -135,7 +132,6 @@ impl CachedPredicateResultBuilder {
         let Self {
             in_progress_arrays,
             filters,
-            batch_size: _,
         } = self;
 
         let new_selection = RowSelection::from_filters(&filters);
@@ -316,29 +312,36 @@ trait InProgressArray {
 /// results is not possible.
 fn create_in_progress_array(
     in_projection: bool,
-    _data_type: &DataType,
+    data_type: &DataType,
     batch_size: usize,
 ) -> Box<dyn InProgressArray> {
-    if in_projection {
-        Box::new(InProgressArrayImpl::new(batch_size, GenericArrayBuilder::new()))
-        
-    } else {
+    if !in_projection {
         // column is not in the projection, so no need to cache
-        Box::new(NoOpInProgressArray::new())
-        
+        return Box::new(NoOpInProgressArray::new())
+    }
+
+    match data_type {
+        DataType::Utf8View => {
+            Box::new(InProgressArrayImpl::new(batch_size,
+                                              InProgressStringViewBuilder::new_with_capacity(batch_size)))
+        }
+        _ => {
+            // TODO implement more specific types
+            Box::new(InProgressArrayImpl::new(batch_size, GenericArrayBuilder::new()))
+        }
     }
 }
 
 
 /// A builder for creating an InProgressArray. Trait so we can use Dyn dispatch
 trait InProgressArrayBuilder {
     /// Appends all values of the array to the in progress array
-    /// 
+    ///
     /// TODO: potentially pass in filter and unfiltered array to avoid a copy
     fn append(&mut self, array: ArrayRef);
 
     /// Finalizes the in progress array, resetting state and returning the new array.
-    /// 
+    ///
     /// Returns None if there are no rows in progress
     fn try_build(&mut self) -> Result<Option<ArrayRef>>;
 }
@@ -363,7 +366,7 @@ impl <B: InProgressArrayBuilder> InProgressArrayImpl<B> {
             inner,
         }
     }
-    
+
     /// Combines all arrays in `current` into a new array in `finished` and returns the
     /// number of rows in the array added to `self.finished`
     fn finish_current(&mut self) -> Result<usize> {
@@ -372,11 +375,11 @@ impl <B: InProgressArrayBuilder> InProgressArrayImpl<B> {
             return Ok(0);
         }
         let Some(new_array) = self.inner.try_build()? else {
-            // no rows in current 
+            // no rows in current
             self.current_rows = 0;
             return Ok(0);
         };
-        
+
         let num_rows = new_array.len();
         self.finished.push(new_array);
         self.current_rows = 0;
@@ -442,7 +445,7 @@ impl NoOpInProgressArray {
 }
 
 /// Implements a GenericArrayBuilder used for any array type by using buffering and `concat`
-/// 
+///
 /// TODO avoid this by using type specific array builders
 struct GenericArrayBuilder {
     arrays: Vec<ArrayRef>
@@ -480,38 +483,42 @@ impl InProgressArrayBuilder for GenericArrayBuilder {
 }
 
 
-
-
 /// An implementation of InProgressArray for StringViewArray
 /// that knows how to efficiently and incrementally concatenate arrays
-struct StringViewInProgressArray {
+///
+/// TODO move this to StringViewBuilder (basically probably add `append_array` to it)
+struct InProgressStringViewBuilder {
     new_views: Vec<u128>,
     null_buffer_builder: NullBufferBuilder,
     buffers: Vec<Buffer>,
+    initial_capacity: usize,
 }
 
-impl StringViewInProgressArray {
-    fn new_with_capacity(num_rows: usize) -> Self {
-        todo!()
-    }
+impl InProgressStringViewBuilder {
+    fn new_with_capacity(initial_capacity: usize) -> Self {
+     Self {
+            new_views: Vec::with_capacity(initial_capacity),
+            null_buffer_builder: NullBufferBuilder::new(initial_capacity),
+            buffers: Vec::with_capacity(100), // TODO better estimate of number of buffers
+         initial_capacity,
+        }
+     }
 }
 
-fn concat_string_view_arrays(arrays: &[ArrayRef]) -> Result<ArrayRef> {
-    // Special case for StringViewArray inspired by DataFusion:
-    // https://github.com/apache/datafusion/blob/9d2f04996604e709ee440b65f41e7b882f50b788/datafusion/physical-plan/src/coalesce/mod.rs#L222-L221
-
-    let total_rows = arrays.iter().map(|a| a.len()).sum();
-    let mut new_views = Vec::with_capacity(total_rows);
-    let mut null_buffer_builder = NullBufferBuilder::new(total_rows);
-    let mut buffers = Vec::with_capacity(100); // better estimate of buffer sizes
-
-    // copy each input array into the to the output, one at a time
-    for array in arrays.iter() {
+impl InProgressArrayBuilder for InProgressStringViewBuilder {
+    fn append(&mut self, array: ArrayRef) {
+        // Special case for StringViewArray inspired by DataFusion:
+        // https://github.com/apache/datafusion/blob/9d2f04996604e709ee440b65f41e7b882f50b788/datafusion/physical-plan/src/coalesce/mod.rs#L222-L221
         let num_rows = array.len();
         if num_rows == 0 {
-            continue;
+            return; // nothing to do
         }
         let array = array.as_string_view();
+
+        let null_buffer_builder = &mut self.null_buffer_builder;
+        let buffers = &mut self.buffers;
+        let new_views = &mut self.new_views;
+
         // Copy nulls
         if let Some(nulls) = array.nulls() {
             null_buffer_builder.append_buffer(nulls);
@@ -550,7 +557,7 @@ fn concat_string_view_arrays(arrays: &[ArrayRef]) -> Result<ArrayRef> {
             // buffer as well as updating the views
             let mut new_buffer: Vec<u8> = Vec::with_capacity(ideal_buffer_size);
             let new_buffer_index = buffers.len() as u32; // making one new buffer
-                                                         // Update any views that point to the old buffers.
+            // Update any views that point to the old buffers.
             for v in new_views[starting_view..].iter_mut() {
                 let view_len = *v as u32;
                 // if view_len is 12 or less, data is inlined and doesn't need an update
@@ -572,12 +579,18 @@ fn concat_string_view_arrays(arrays: &[ArrayRef]) -> Result<ArrayRef> {
         }
     }
 
-    // Form output array
-    let nulls = null_buffer_builder.finish();
-    // safety: we know what we are doing above
-    let new_array =
-        unsafe { StringViewArray::new_unchecked(ScalarBuffer::from(new_views), buffers, nulls) };
-    Ok(Arc::new(new_array))
+    fn try_build(&mut self) -> Result<Option<ArrayRef>> {
+        // Form output array
+        let nulls = self.null_buffer_builder.finish();
+        let new_views = std::mem::replace(&mut self.new_views, Vec::with_capacity(self.initial_capacity));
+        let buffers= std::mem::replace(&mut self.buffers, Vec::with_capacity(100)); // TODO better buffer estimate
+
+        // safety: we know what we are doing above
+        let new_array =
+            unsafe { StringViewArray::new_unchecked(ScalarBuffer::from(new_views), buffers, nulls) };
+
+        Ok(Some(Arc::new(new_array)))
+    }
 }
 
 /// return the size required for buffers to hold all strings