perf: Reduce conversion cost in chunked string gather

crates/polars-compute/src/gather/binview.rs

@@ -1,18 +1,18 @@
-use arrow::array::BinaryViewArray;
+use arrow::array::{BinaryViewArrayGeneric, ViewType};
 
 use self::primitive::take_values_and_validity_unchecked;
 use super::*;
 
 /// # Safety
 /// No bound checks
-pub(super) unsafe fn take_binview_unchecked(
-    arr: &BinaryViewArray,
+pub(super) unsafe fn take_binview_unchecked<V: ViewType + ?Sized>(
+    arr: &BinaryViewArrayGeneric<V>,
     indices: &IdxArr,
-) -> BinaryViewArray {
+) -> BinaryViewArrayGeneric<V> {
     let (views, validity) =
         take_values_and_validity_unchecked(arr.views(), arr.validity(), indices);
 
-    BinaryViewArray::new_unchecked_unknown_md(
+    BinaryViewArrayGeneric::new_unchecked_unknown_md(
         arr.dtype().clone(),
         views.into(),
         arr.data_buffers().clone(),

crates/polars-compute/src/gather/mod.rs

@@ -18,8 +18,8 @@
 //! Defines take kernel for [`Array`]
 
 use arrow::array::{
-    self, new_empty_array, Array, ArrayCollectIterExt, ArrayFromIterDtype, NullArray, StaticArray,
-    Utf8ViewArray,
+    self, new_empty_array, Array, ArrayCollectIterExt, ArrayFromIterDtype, BinaryViewArray,
+    NullArray, StaticArray, Utf8ViewArray,
 };
 use arrow::datatypes::{ArrowDataType, IdxArr};
 use arrow::types::Index;
@@ -76,13 +76,12 @@ pub unsafe fn take_unchecked(values: &dyn Array, indices: &IdxArr) -> Box<dyn Ar
             fixed_size_list::take_unchecked(array, indices)
         },
         BinaryView => {
-            take_binview_unchecked(values.as_any().downcast_ref().unwrap(), indices).boxed()
+            let array: &BinaryViewArray = values.as_any().downcast_ref().unwrap();
+            take_binview_unchecked(array, indices).boxed()
         },
         Utf8View => {
-            let arr: &Utf8ViewArray = values.as_any().downcast_ref().unwrap();
-            take_binview_unchecked(&arr.to_binview(), indices)
-                .to_utf8view_unchecked()
-                .boxed()
+            let array: &Utf8ViewArray = values.as_any().downcast_ref().unwrap();
+            take_binview_unchecked(array, indices).boxed()
         },
         t => unimplemented!("Take not supported for data type {:?}", t),
     }

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

@@ -255,9 +255,20 @@ impl ChunkTakeUnchecked<IdxCa> for BinaryChunked {
 
 impl ChunkTakeUnchecked<IdxCa> for StringChunked {
     unsafe fn take_unchecked(&self, indices: &IdxCa) -> Self {
-        self.as_binary()
-            .take_unchecked(indices)
-            .to_string_unchecked()
+        let rechunked = self.rechunk();
+        let indices = indices.rechunk();
+        let indices_arr = indices.downcast_iter().next().unwrap();
+        let chunks = rechunked
+            .chunks()
+            .iter()
+            .map(|arr| take_unchecked(arr.as_ref(), indices_arr))
+            .collect::<Vec<_>>();
+
+        let mut out = ChunkedArray::from_chunks(self.name().clone(), chunks);
+        let sorted_flag =
+            _update_gather_sorted_flag(self.is_sorted_flag(), indices.is_sorted_flag());
+        out.set_sorted_flag(sorted_flag);
+        out
     }
 }
 
@@ -272,9 +283,8 @@ impl<I: AsRef<[IdxSize]> + ?Sized> ChunkTakeUnchecked<I> for BinaryChunked {
 impl<I: AsRef<[IdxSize]> + ?Sized> ChunkTakeUnchecked<I> for StringChunked {
     /// Gather values from ChunkedArray by index.
     unsafe fn take_unchecked(&self, indices: &I) -> Self {
-        self.as_binary()
-            .take_unchecked(indices)
-            .to_string_unchecked()
+        let indices = IdxCa::mmap_slice(PlSmallStr::EMPTY, indices.as_ref());
+        self.take_unchecked(&indices)
     }
 }