Use new API for appending values to avoid slices

Author: alamb

arrow-select/src/coalesce.rs

@@ -226,6 +226,15 @@ impl BatchCoalescer {
             return Ok(());
         }
 
+        // setup input rows
+        assert_eq!(arrays.len(), self.in_progress_arrays.len());
+        self.in_progress_arrays
+            .iter_mut()
+            .zip(arrays)
+            .for_each(|(in_progress, array)| {
+                in_progress.set_source(Some(array));
+            });
+
         // If pushing this batch would exceed the target batch size,
         // finish the current batch and start a new one
         let mut offset = 0;
@@ -234,10 +243,8 @@ impl BatchCoalescer {
             debug_assert!(remaining_rows > 0);
 
             // Copy remaining_rows from each array
-            for (idx, array) in arrays.iter().enumerate() {
-                let array = array.slice(offset, remaining_rows);
-                // TODO move slicing into InProgressArray
-                self.in_progress_arrays[idx].push_array(array);
+            for in_progress in self.in_progress_arrays.iter_mut() {
+                in_progress.copy_rows(offset, remaining_rows)?;
             }
 
             self.buffered_rows += remaining_rows;
@@ -250,18 +257,21 @@ impl BatchCoalescer {
         // Add any the remaining rows to the buffer
         self.buffered_rows += num_rows;
         if num_rows > 0 {
-            for (idx, array) in arrays.into_iter().enumerate() {
-                // Push the array to the in-progress array
-                let array = array.slice(offset, num_rows);
-                // TODO move slicing into InProgressArray
-                self.in_progress_arrays[idx].push_array(array);
+            for in_progress in self.in_progress_arrays.iter_mut() {
+                in_progress.copy_rows(offset, num_rows)?;
             }
         }
 
         // If we have reached the target batch size, finalize the buffered batch
         if self.buffered_rows >= self.batch_size {
             self.finish_buffered_batch()?;
         }
+
+        // clear in progress sources (to allow the memory to be freed)
+        for in_progress in self.in_progress_arrays.iter_mut() {
+            in_progress.set_source(None);
+        }
+
         Ok(())
     }
 
@@ -335,10 +345,17 @@ fn create_in_progress_array(data_type: &DataType, batch_size: usize) -> Box<dyn
 /// [`StringViewArray`]: arrow_array::StringViewArray
 /// [`UInt32Array`]: arrow_array::UInt32Array
 trait InProgressArray: std::fmt::Debug + Send + Sync {
-    /// Push a new array to the in-progress array
-    fn push_array(&mut self, array: ArrayRef);
+    /// Set the source array.
+    ///
+    /// Subsequent calls to "copy" copy rows from this array into the in-progress array
+    fn set_source(&mut self, source: Option<ArrayRef>);
+
+    /// copy rows from the source array into the in-progress array
+    ///
+    /// Return an error if the source is not set
+    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError>;
 
-    /// Finish the currently in-progress array and clear state for the next
+    /// Finish the currently in-progress array and clear state
     fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
 }
 
@@ -552,7 +569,6 @@ mod tests {
         let array = col_as_string_view("c0", &batch);
         let gc_array = col_as_string_view("c0", output_batches.first().unwrap());
         assert_eq!(array.data_buffers().len(), 5);
-        assert_eq!(gc_array.data_buffers().len(), 1); // compacted into a single buffer
 
         expect_buffer_layout(
             gc_array,
@@ -624,7 +640,7 @@ mod tests {
 
     #[test]
     fn test_string_view_many_small_compact() {
-        // The strings are 37 bytes long, so each batch has 200 * 28 = 5600 bytes
+        // The strings are 28 long, so each batch has 400 * 28 = 5600 bytes
         let batch = stringview_batch_repeated(
             400,
             [Some("This string is 28 bytes long"), Some("small string")],

arrow-select/src/coalesce/byte_view.rs

@@ -27,6 +27,8 @@ use std::sync::Arc;
 
 /// InProgressArray for StringViewArray and BinaryViewArray
 pub(crate) struct InProgressByteViewArray<B: ByteViewType> {
+    /// The source array
+    source: Option<Source>,
     /// the target batch size (and thus size for views allocation)
     batch_size: usize,
     /// The in progress vies
@@ -44,6 +46,15 @@ pub(crate) struct InProgressByteViewArray<B: ByteViewType> {
     _phantom: PhantomData<B>,
 }
 
+struct Source {
+    /// The array to copy form
+    array: ArrayRef,
+    /// Should the strings from the source array be copied into new buffers?
+    need_gc: bool,
+    /// How many bytes were actually used in the source array's buffers?
+    ideal_buffer_size: usize,
+}
+
 // manually implement Debug because ByteViewType doesn't implement Debug
 impl<B: ByteViewType> std::fmt::Debug for InProgressByteViewArray<B> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -63,6 +74,7 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
 
         Self {
             batch_size,
+            source: None,
             views: Vec::new(),                         // allocate in push
             nulls: NullBufferBuilder::new(batch_size), // no allocation
             current: None,
@@ -80,15 +92,6 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
         self.views.reserve(self.batch_size);
     }
 
-    /// Update self.nulls with the nulls from the StringViewArray
-    fn push_nulls(&mut self, s: &GenericByteViewArray<B>) {
-        if let Some(nulls) = s.nulls().as_ref() {
-            self.nulls.append_buffer(nulls);
-        } else {
-            self.nulls.append_n_non_nulls(s.len());
-        }
-    }
-
     /// Finishes in progress block, if any
     fn finish_current(&mut self) {
         let Some(next_buffer) = self.current.take() else {
@@ -263,38 +266,67 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
 }
 
 impl<B: ByteViewType> InProgressArray for InProgressByteViewArray<B> {
-    fn push_array(&mut self, array: ArrayRef) {
-        // If creating StringViewArray output, ensure input was valid utf8 too
+    fn set_source(&mut self, source: Option<ArrayRef>) {
+        self.source = source.map(|array| {
+            let s = array.as_byte_view::<B>();
+
+            let (need_gc, ideal_buffer_size) = if s.data_buffers().is_empty() {
+                (false, 0)
+            } else {
+                let ideal_buffer_size = s.total_buffer_bytes_used();
+                let actual_buffer_size = s.get_buffer_memory_size();
+                // copying strings is expensive, so only do it if the array is
+                // sparse (uses at least 2x the memory it needs)
+                let need_gc =
+                    ideal_buffer_size != 0 && actual_buffer_size > (ideal_buffer_size * 2);
+                (need_gc, ideal_buffer_size)
+            };
+
+            Source {
+                array,
+                need_gc,
+                ideal_buffer_size,
+            }
+        })
+    }
+
+    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
         self.ensure_capacity();
-        let s = array.as_byte_view::<B>();
+        let source = self.source.take().ok_or_else(|| {
+            ArrowError::InvalidArgumentError("InProgressByteViewArray: source not set".to_string())
+        })?;
 
-        // add any nulls, as necessary
-        self.push_nulls(s);
+        // If creating StringViewArray output, ensure input was valid utf8 too
+        let s = source.array.as_byte_view::<B>();
 
-        // If there are no data buffers in s (all inlined views), can append the
-        // views/nulls and done
-        if s.data_buffers().is_empty() {
-            self.views.extend_from_slice(s.views().as_ref());
-            return;
-        }
+        // add any nulls, as necessary
+        if let Some(nulls) = s.nulls().as_ref() {
+            let nulls = nulls.slice(offset, len);
+            self.nulls.append_buffer(&nulls);
+        } else {
+            self.nulls.append_n_non_nulls(len);
+        };
 
-        let ideal_buffer_size = s.total_buffer_bytes_used();
-        let actual_buffer_size = s.get_buffer_memory_size();
         let buffers = s.data_buffers();
+        let views = &s.views().as_ref()[offset..offset + len];
 
-        // None of the views references the buffers (e.g. sliced)
-        if ideal_buffer_size == 0 {
-            self.views.extend_from_slice(s.views().as_ref());
-            return;
+        // If there are no data buffers in s (all inlined views), can append the
+        // views/nulls and done
+        if source.ideal_buffer_size == 0 {
+            self.views.extend_from_slice(views);
+            self.source = Some(source);
+            return Ok(());
         }
 
         // Copying the strings into a buffer can be time-consuming so
         // only do it if the array is sparse
-        if actual_buffer_size > (ideal_buffer_size * 2) {
-            self.append_views_and_copy_strings(s.views(), ideal_buffer_size, buffers);
+        if source.need_gc {
+            self.append_views_and_copy_strings(views, source.ideal_buffer_size, buffers);
         } else {
-            self.append_views_and_update_buffer_index(s.views(), buffers);
+            self.append_views_and_update_buffer_index(views, buffers);
         }
+        self.source = Some(source);
+        Ok(())
     }
 
     fn finish(&mut self) -> Result<ArrayRef, ArrowError> {

arrow-select/src/coalesce/generic.rs

@@ -27,21 +27,33 @@ use arrow_schema::ArrowError;
 /// [`concat`]: crate::concat::concat
 #[derive(Debug)]
 pub(crate) struct GenericInProgressArray {
-    /// The buffered arrays
+    /// The current source
+    source: Option<ArrayRef>,
+    /// The buffered array slices
     buffered_arrays: Vec<ArrayRef>,
 }
 
 impl GenericInProgressArray {
     /// Create a new `GenericInProgressArray`
     pub(crate) fn new() -> Self {
         Self {
+            source: None,
             buffered_arrays: vec![],
         }
     }
 }
 impl InProgressArray for GenericInProgressArray {
-    fn push_array(&mut self, array: ArrayRef) {
+    fn set_source(&mut self, source: Option<ArrayRef>) {
+        self.source = source
+    }
+
+    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
+        let source = self.source.as_ref().ok_or_else(|| {
+            ArrowError::InvalidArgumentError("GenericInProgressArray: source not set".to_string())
+        })?;
+        let array = source.slice(offset, len);
         self.buffered_arrays.push(array);
+        Ok(())
     }
 
     fn finish(&mut self) -> Result<ArrayRef, ArrowError> {