perf: Use Cow as output for rechunk and add rechunk_mut (#21116)

Author: orlp

crates/polars-core/src/chunked_array/cast.rs

@@ -601,7 +601,7 @@ fn cast_list(
     // We still rechunk because we must bubble up a single data-type
     // TODO!: consider a version that works on chunks and merges the data-types and arrays.
     let ca = ca.rechunk();
-    let arr = ca.downcast_iter().next().unwrap();
+    let arr = ca.downcast_as_array();
     // SAFETY: inner dtype is passed correctly
     let s = unsafe {
         Series::from_chunks_and_dtype_unchecked(
@@ -630,7 +630,7 @@ fn cast_list(
 unsafe fn cast_list_unchecked(ca: &ListChunked, child_type: &DataType) -> PolarsResult<Series> {
     // TODO! add chunked, but this must correct for list offsets.
     let ca = ca.rechunk();
-    let arr = ca.downcast_iter().next().unwrap();
+    let arr = ca.downcast_as_array();
     // SAFETY: inner dtype is passed correctly
     let s = unsafe {
         Series::from_chunks_and_dtype_unchecked(
@@ -666,7 +666,7 @@ fn cast_fixed_size_list(
     options: CastOptions,
 ) -> PolarsResult<(ArrayRef, DataType)> {
     let ca = ca.rechunk();
-    let arr = ca.downcast_iter().next().unwrap();
+    let arr = ca.downcast_as_array();
     // SAFETY: inner dtype is passed correctly
     let s = unsafe {
         Series::from_chunks_and_dtype_unchecked(

crates/polars-core/src/chunked_array/iterator/par/list.rs

@@ -30,7 +30,7 @@ impl ListChunked {
     // Get an indexed parallel iterator over the [`Series`] in this [`ListChunked`].
     // Also might be faster as it doesn't use `flat_map`.
     pub fn par_iter_indexed(&mut self) -> impl IndexedParallelIterator<Item = Option<Series>> + '_ {
-        *self = self.rechunk();
+        self.rechunk_mut();
         let arr = self.downcast_iter().next().unwrap();
 
         let dtype = self.inner_dtype();

crates/polars-core/src/chunked_array/list/mod.rs

@@ -135,7 +135,7 @@ impl ListChunked {
     ) -> PolarsResult<ListChunked> {
         // generated Series will have wrong length otherwise.
         let ca = self.rechunk();
-        let arr = ca.downcast_iter().next().unwrap();
+        let arr = ca.downcast_as_array();
 
         // SAFETY:
         // Inner dtype is passed correctly

crates/polars-core/src/chunked_array/logical/categorical/from.rs

@@ -27,7 +27,7 @@ impl CategoricalChunked {
             ArrowDataType::LargeUtf8
         };
         let keys = self.physical().rechunk();
-        let keys = keys.downcast_iter().next().unwrap();
+        let keys = keys.downcast_as_array();
         let map = &**self.get_rev_map();
         let dtype = ArrowDataType::Dictionary(IntegerType::UInt32, Box::new(values_dtype), false);
         match map {
@@ -40,7 +40,7 @@ impl CategoricalChunked {
             },
             RevMapping::Global(reverse_map, values, _uuid) => {
                 let iter = keys
-                    .into_iter()
+                    .iter()
                     .map(|opt_k| opt_k.map(|k| *reverse_map.get(k).unwrap()));
                 let keys = PrimitiveArray::from_trusted_len_iter(iter);
 
@@ -60,7 +60,7 @@ impl CategoricalChunked {
             ArrowDataType::LargeUtf8
         };
         let keys = self.physical().rechunk();
-        let keys = keys.downcast_iter().next().unwrap();
+        let keys = keys.downcast_as_array();
         let map = &**self.get_rev_map();
         let dtype = ArrowDataType::Dictionary(IntegerType::Int64, Box::new(values_dtype), false);
         match map {
@@ -85,7 +85,7 @@ impl CategoricalChunked {
             },
             RevMapping::Global(reverse_map, values, _uuid) => {
                 let iter = keys
-                    .into_iter()
+                    .iter()
                     .map(|opt_k| opt_k.map(|k| *reverse_map.get(k).unwrap() as i64));
                 let keys = PrimitiveArray::from_trusted_len_iter(iter);
 

crates/polars-core/src/chunked_array/mod.rs

@@ -48,7 +48,6 @@ pub mod temporal;
 mod to_vec;
 mod trusted_len;
 
-use std::mem;
 use std::slice::Iter;
 
 use arrow::legacy::prelude::*;
@@ -154,13 +153,12 @@ impl<T: PolarsDataType> ChunkedArray<T> {
         self.chunks.len() > 1 && self.chunks.len() > self.len() / 3
     }
 
-    fn optional_rechunk(self) -> Self {
+    fn optional_rechunk(mut self) -> Self {
         // Rechunk if we have many small chunks.
         if self.should_rechunk() {
-            self.rechunk()
-        } else {
-            self
+            self.rechunk_mut()
         }
+        self
     }
 
     pub(crate) fn as_any(&self) -> &dyn std::any::Any {
@@ -671,8 +669,7 @@ where
     T: PolarsNumericType,
 {
     fn as_single_ptr(&mut self) -> PolarsResult<usize> {
-        let mut ca = self.rechunk();
-        mem::swap(&mut ca, self);
+        self.rechunk_mut();
         let a = self.data_views().next().unwrap();
         let ptr = a.as_ptr();
         Ok(ptr as usize)

crates/polars-core/src/chunked_array/ops/chunkops.rs

@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::cell::Cell;
 
 use arrow::bitmap::{Bitmap, BitmapBuilder};
@@ -158,27 +159,43 @@ impl<T: PolarsDataType> ChunkedArray<T> {
             .sum::<usize>();
     }
 
-    pub fn rechunk(&self) -> Self {
+    /// Rechunks this ChunkedArray, returning a new Cow::Owned ChunkedArray if it was
+    /// rechunked or simply a Cow::Borrowed of itself if it was already a single chunk.
+    pub fn rechunk(&self) -> Cow<'_, Self> {
         match self.dtype() {
             #[cfg(feature = "object")]
             DataType::Object(_, _) => {
                 panic!("implementation error")
             },
             _ => {
                 if self.chunks.len() == 1 {
-                    self.clone()
+                    Cow::Borrowed(self)
                 } else {
                     let chunks = vec![concatenate_unchecked(&self.chunks).unwrap()];
 
                     let mut ca = unsafe { self.copy_with_chunks(chunks) };
                     use StatisticsFlags as F;
                     ca.retain_flags_from(self, F::IS_SORTED_ANY | F::CAN_FAST_EXPLODE_LIST);
-                    ca
+                    Cow::Owned(ca)
                 }
             },
         }
     }
 
+    /// Rechunks this ChunkedArray in-place.
+    pub fn rechunk_mut(&mut self) {
+        if self.chunks.len() > 1 {
+            let rechunked = concatenate_unchecked(&self.chunks).unwrap();
+            if self.chunks.capacity() <= 8 {
+                // Reuse chunk allocation if not excessive.
+                self.chunks.clear();
+                self.chunks.push(rechunked);
+            } else {
+                self.chunks = vec![rechunked];
+            }
+        }
+    }
+
     pub fn rechunk_validity(&self) -> Option<Bitmap> {
         if self.chunks.len() == 1 {
             return self.chunks[0].validity().cloned();

crates/polars-core/src/chunked_array/ops/downcast.rs

@@ -108,9 +108,9 @@ impl<T: PolarsDataType> ChunkedArray<T> {
     }
 
     #[inline]
-    pub fn downcast_into_array(self) -> T::Array {
+    pub fn downcast_as_array(&self) -> &T::Array {
         assert_eq!(self.chunks.len(), 1);
-        self.downcast_get(0).unwrap().clone()
+        self.downcast_get(0).unwrap()
     }
 
     #[inline]

crates/polars-core/src/chunked_array/ops/explode_and_offsets.rs

@@ -51,7 +51,6 @@ impl ChunkExplode for ListChunked {
         let ca = self.rechunk();
         let listarr: &LargeListArray = ca.downcast_iter().next().unwrap();
         let offsets = listarr.offsets().clone();
-
         Ok(offsets)
     }
 

crates/polars-core/src/chunked_array/ops/extend.rs

@@ -41,7 +41,7 @@ where
         // all to a single chunk
         if self.chunks.len() > 1 {
             self.append(other)?;
-            *self = self.rechunk();
+            self.rechunk_mut();
             return Ok(());
         }
         // Depending on the state of the underlying arrow array we
@@ -119,7 +119,7 @@ impl BooleanChunked {
         // make sure that we are a single chunk already
         if self.chunks.len() > 1 {
             self.append(other)?;
-            *self = self.rechunk();
+            self.rechunk_mut();
             return Ok(());
         }
         let arr = self.downcast_iter().next().unwrap();

crates/polars-core/src/chunked_array/ops/gather.rs

@@ -327,8 +327,8 @@ impl IdxCa {
 impl ChunkTakeUnchecked<IdxCa> for ArrayChunked {
     unsafe fn take_unchecked(&self, indices: &IdxCa) -> Self {
         let chunks = vec![take_unchecked(
-            &self.rechunk().downcast_into_array(),
-            &indices.rechunk().downcast_into_array(),
+            self.rechunk().downcast_as_array(),
+            indices.rechunk().downcast_as_array(),
         )];
         self.copy_with_chunks(chunks)
     }
@@ -345,8 +345,8 @@ impl<I: AsRef<[IdxSize]> + ?Sized> ChunkTakeUnchecked<I> for ArrayChunked {
 impl ChunkTakeUnchecked<IdxCa> for ListChunked {
     unsafe fn take_unchecked(&self, indices: &IdxCa) -> Self {
         let chunks = vec![take_unchecked(
-            &self.rechunk().downcast_into_array(),
-            &indices.rechunk().downcast_into_array(),
+            self.rechunk().downcast_as_array(),
+            indices.rechunk().downcast_as_array(),
         )];
         self.copy_with_chunks(chunks)
     }

crates/polars-core/src/chunked_array/ops/reverse.rs

@@ -88,7 +88,7 @@ impl ChunkReverse for ArrayChunked {
             todo!("reverse for FixedSizeList with non-numeric dtypes not yet supported")
         }
         let ca = self.rechunk();
-        let arr = ca.downcast_iter().next().unwrap();
+        let arr = ca.downcast_as_array();
         let values = arr.values().as_ref();
 
         let mut builder =

crates/polars-core/src/chunked_array/ops/rolling_window.rs

@@ -95,7 +95,7 @@ mod inner_mod {
             options.window_size = std::cmp::min(self.len(), options.window_size);
 
             let len = self.len();
-            let arr = ca.downcast_iter().next().unwrap();
+            let arr = ca.downcast_as_array();
             let mut ca = ChunkedArray::<T>::from_slice(PlSmallStr::EMPTY, &[T::Native::zero()]);
             let ptr = ca.chunks[0].as_mut() as *mut dyn Array as *mut PrimitiveArray<T::Native>;
             let mut series_container = ca.into_series();
@@ -215,7 +215,7 @@ mod inner_mod {
                 return Ok(Self::full_null(self.name().clone(), self.len()));
             }
             let ca = self.rechunk();
-            let arr = ca.downcast_iter().next().unwrap();
+            let arr = ca.downcast_as_array();
 
             // We create a temporary dummy ChunkedArray. This will be a
             // container where we swap the window contents every iteration doing

crates/polars-core/src/chunked_array/ops/sort/mod.rs

@@ -408,7 +408,7 @@ impl ChunkSort<BinaryType> for BinaryChunked {
         // We will sort by the views and reconstruct with sorted views. We leave the buffers as is.
         // We must rechunk to ensure that all views point into the proper buffers.
         let ca = self.rechunk();
-        let arr = ca.downcast_into_array();
+        let arr = ca.downcast_as_array().clone();
 
         let (views, buffers, validity, total_bytes_len, total_buffer_len) = arr.into_inner();
         let mut views = views.make_mut();
@@ -590,7 +590,7 @@ impl ChunkSort<BinaryOffsetType> for BinaryOffsetChunked {
     fn arg_sort(&self, mut options: SortOptions) -> IdxCa {
         options.multithreaded &= POOL.current_num_threads() > 1;
         let ca = self.rechunk();
-        let arr = ca.downcast_into_array();
+        let arr = ca.downcast_as_array();
         let mut idx = (0..(arr.len() as IdxSize)).collect::<Vec<_>>();
 
         let argsort = |args| {

crates/polars-core/src/chunked_array/struct_/mod.rs

@@ -411,7 +411,7 @@ impl StructChunked {
                 .zip(other.chunks())
                 .any(|(a, b)| a.len() != b.len())
         {
-            *self = self.rechunk();
+            self.rechunk_mut();
             let other = other.rechunk();
             return self.zip_outer_validity(&other);
         }

crates/polars-core/src/frame/group_by/aggregations/agg_list.rs

@@ -287,7 +287,7 @@ impl AggList for StructChunked {
             let out = ca.into_series().take_unchecked(&gather);
             out.struct_().unwrap().clone()
         } else {
-            ca.rechunk()
+            ca.rechunk().into_owned()
         };
 
         let arr = gathered.chunks()[0].clone();

crates/polars-core/src/frame/group_by/aggregations/boolean.rs

@@ -1,3 +1,5 @@
+use std::borrow::Cow;
+
 use super::*;
 use crate::chunked_array::cast::CastOptions;
 
@@ -24,10 +26,11 @@ unsafe fn bitwise_agg(
     f: fn(&BooleanChunked) -> Option<bool>,
 ) -> Series {
     // Prevent a rechunk for every individual group.
+
     let s = if groups.len() > 1 {
         ca.rechunk()
     } else {
-        ca.clone()
+        Cow::Borrowed(ca)
     };
 
     match groups {

crates/polars-core/src/frame/group_by/aggregations/mod.rs

@@ -3,6 +3,7 @@ mod boolean;
 mod dispatch;
 mod string;
 
+use std::borrow::Cow;
 use std::cmp::Ordering;
 
 pub use agg_list::*;
@@ -456,10 +457,11 @@ where
         ChunkTakeUnchecked<[IdxSize]> + ChunkBitwiseReduce<Physical = T::Native> + IntoSeries,
 {
     // Prevent a rechunk for every individual group.
+
     let s = if groups.len() > 1 {
         ca.rechunk()
     } else {
-        ca.clone()
+        Cow::Borrowed(ca)
     };
 
     match groups {

crates/polars-core/src/frame/group_by/aggregations/string.rs

@@ -33,7 +33,7 @@ impl BinaryChunked {
         match groups {
             GroupsType::Idx(groups) => {
                 let ca_self = self.rechunk();
-                let arr = ca_self.downcast_iter().next().unwrap();
+                let arr = ca_self.downcast_as_array();
                 let no_nulls = arr.null_count() == 0;
                 _agg_helper_idx_bin(groups, |(first, idx)| {
                     debug_assert!(idx.len() <= ca_self.len());
@@ -95,7 +95,7 @@ impl BinaryChunked {
         match groups {
             GroupsType::Idx(groups) => {
                 let ca_self = self.rechunk();
-                let arr = ca_self.downcast_iter().next().unwrap();
+                let arr = ca_self.downcast_as_array();
                 let no_nulls = arr.null_count() == 0;
                 _agg_helper_idx_bin(groups, |(first, idx)| {
                     debug_assert!(idx.len() <= self.len());

crates/polars-core/src/frame/mod.rs

@@ -3109,7 +3109,7 @@ impl DataFrame {
         for ca in iter {
             acc_ca.append(&ca)?;
         }
-        Ok(acc_ca.rechunk())
+        Ok(acc_ca.rechunk().into_owned())
     }
 
     /// Get the supertype of the columns in this DataFrame

crates/polars-core/src/series/implementations/array.rs

@@ -164,7 +164,7 @@ impl SeriesTrait for SeriesWrap<ArrayChunked> {
     }
 
     fn rechunk(&self) -> Series {
-        self.0.rechunk().into_series()
+        self.0.rechunk().into_owned().into_series()
     }
 
     fn new_from_index(&self, index: usize, length: usize) -> Series {

crates/polars-core/src/series/implementations/binary.rs

@@ -167,7 +167,7 @@ impl SeriesTrait for SeriesWrap<BinaryChunked> {
     }
 
     fn rechunk(&self) -> Series {
-        self.0.rechunk().into_series()
+        self.0.rechunk().into_owned().into_series()
     }
 
     fn new_from_index(&self, index: usize, length: usize) -> Series {

crates/polars-core/src/series/implementations/binary_offset.rs

@@ -139,7 +139,7 @@ impl SeriesTrait for SeriesWrap<BinaryOffsetChunked> {
     }
 
     fn rechunk(&self) -> Series {
-        self.0.rechunk().into_series()
+        self.0.rechunk().into_owned().into_series()
     }
 
     fn new_from_index(&self, index: usize, length: usize) -> Series {

crates/polars-core/src/series/implementations/boolean.rs

@@ -193,7 +193,7 @@ impl SeriesTrait for SeriesWrap<BooleanChunked> {
     }
 
     fn rechunk(&self) -> Series {
-        self.0.rechunk().into_series()
+        self.0.rechunk().into_owned().into_series()
     }
 
     fn new_from_index(&self, index: usize, length: usize) -> Series {