refactor: add gc_stringview_arrays util and use that

Author: Samyak2

datafusion/physical-plan/src/common.rs

@@ -25,8 +25,9 @@ use super::SendableRecordBatchStream;
 use crate::stream::RecordBatchReceiverStream;
 use crate::{ColumnStatistics, Statistics};
 
-use arrow::array::Array;
+use arrow::array::{Array, StringViewArray};
 use arrow::datatypes::Schema;
+use arrow::error::ArrowError;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::stats::Precision;
 use datafusion_common::{Result, plan_err};
@@ -190,7 +191,7 @@ pub fn can_project(
                 .max()
                 .is_some_and(|&i| i >= schema.fields().len())
             {
-                Err(arrow::error::ArrowError::SchemaError(format!(
+                Err(ArrowError::SchemaError(format!(
                     "project index {} out of bounds, max field {}",
                     columns.iter().max().unwrap(),
                     schema.fields().len()
@@ -204,6 +205,31 @@ pub fn can_project(
     }
 }
 
+/// Return a new `RecordBatch` with [`StringViewArray::gc`] called on such columns (if any).
+pub(crate) fn gc_stringview_arrays(
+    batch: RecordBatch,
+) -> Result<RecordBatch, ArrowError> {
+    let mut new_columns: Vec<Arc<dyn Array>> = Vec::with_capacity(batch.num_columns());
+
+    let mut arr_mutated = false;
+    for array in batch.columns() {
+        if let Some(string_view_array) = array.as_any().downcast_ref::<StringViewArray>()
+        {
+            let new_array = string_view_array.gc();
+            new_columns.push(Arc::new(new_array));
+            arr_mutated = true;
+        } else {
+            new_columns.push(Arc::clone(array));
+        }
+    }
+
+    if arr_mutated {
+        RecordBatch::try_new(batch.schema(), new_columns)
+    } else {
+        Ok(batch)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

datafusion/physical-plan/src/repartition/mod.rs

@@ -31,6 +31,7 @@ use super::{
     DisplayAs, ExecutionPlanProperties, RecordBatchStream, SendableRecordBatchStream,
 };
 use crate::coalesce::LimitedBatchCoalescer;
+use crate::common::gc_stringview_arrays;
 use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType};
 use crate::hash_utils::create_hashes;
 use crate::metrics::{BaselineMetrics, SpillMetrics};
@@ -611,27 +612,20 @@ impl BatchPartitioner {
                             let columns =
                                 take_arrays(batch.columns(), &indices_array, None)?;
 
-                            let new_columns = columns
-                                .into_iter()
-                                .map(|col| {
-                                    if let Some(sv) =
-                                        col.as_any().downcast_ref::<StringViewArray>()
-                                    {
-                                        Arc::new(sv.gc())
-                                    } else {
-                                        col
-                                    }
-                                })
-                                .collect();
-
                             let mut options = RecordBatchOptions::new();
                             options = options.with_row_count(Some(indices_array.len()));
                             let batch = RecordBatch::try_new_with_options(
                                 batch.schema(),
-                                new_columns,
+                                columns,
                                 &options,
-                            )
-                            .unwrap();
+                            )?;
+
+                            // When `StringViewArray`s are present, the `take_arrays` call above
+                            // re-uses data buffers from the original array. This causes the memory
+                            // pool to count the same data buffers multiple times, once for each
+                            // consumer of the repartition.
+                            // So we gc the output arrays, which creates new data buffers.
+                            let batch = gc_stringview_arrays(batch)?;
 
                             partitioned_batches.push(Ok((partition, batch)));
 

datafusion/physical-plan/src/sorts/sort.rs

@@ -26,7 +26,7 @@ use std::sync::Arc;
 
 use parking_lot::RwLock;
 
-use crate::common::spawn_buffered;
+use crate::common::{gc_stringview_arrays, spawn_buffered};
 use crate::execution_plan::{
     Boundedness, CardinalityEffect, EmissionType, has_same_children_properties,
 };
@@ -483,33 +483,10 @@ impl ExternalSorter {
     fn organize_stringview_arrays(
         globally_sorted_batches: &mut Vec<RecordBatch>,
     ) -> Result<()> {
-        let mut organized_batches = Vec::with_capacity(globally_sorted_batches.len());
-
-        for batch in globally_sorted_batches.drain(..) {
-            let mut new_columns: Vec<Arc<dyn Array>> =
-                Vec::with_capacity(batch.num_columns());
-
-            let mut arr_mutated = false;
-            for array in batch.columns() {
-                if let Some(string_view_array) =
-                    array.as_any().downcast_ref::<StringViewArray>()
-                {
-                    let new_array = string_view_array.gc();
-                    new_columns.push(Arc::new(new_array));
-                    arr_mutated = true;
-                } else {
-                    new_columns.push(Arc::clone(array));
-                }
-            }
-
-            let organized_batch = if arr_mutated {
-                RecordBatch::try_new(batch.schema(), new_columns)?
-            } else {
-                batch
-            };
-
-            organized_batches.push(organized_batch);
-        }
+        let organized_batches = globally_sorted_batches
+            .drain(..)
+            .map(gc_stringview_arrays)
+            .collect::<Result<_, _>>()?;
 
         *globally_sorted_batches = organized_batches;