perf: improve string/binary reverse performance (#14016)

pola-rs · Jan 26, 2024 · b4d8594 · b4d8594
1 parent 3651ecd
commit b4d8594
Showing 1 changed file with 35 additions and 2 deletions.
diff --git a/crates/polars-core/src/chunked_array/ops/reverse.rs b/crates/polars-core/src/chunked_array/ops/reverse.rs
@@ -40,11 +40,44 @@ macro_rules! impl_reverse {
 }
 
 impl_reverse!(BooleanType, BooleanChunked);
-impl_reverse!(StringType, StringChunked);
-impl_reverse!(BinaryType, BinaryChunked);
 impl_reverse!(BinaryOffsetType, BinaryOffsetChunked);
 impl_reverse!(ListType, ListChunked);
 
+impl ChunkReverse for BinaryChunked {
+    fn reverse(&self) -> Self {
+        if self.chunks.len() == 1 {
+            let arr = self.downcast_iter().next().unwrap();
+            let views = arr.views().iter().copied().rev().collect::<Vec<_>>();
+
+            unsafe {
+                let arr = BinaryViewArray::new_unchecked(
+                    arr.data_type().clone(),
+                    views.into(),
+                    arr.data_buffers().clone(),
+                    arr.validity().map(|bitmap| bitmap.iter().rev().collect()),
+                    arr.total_bytes_len(),
+                    arr.total_buffer_len(),
+                )
+                .boxed();
+                BinaryChunked::from_chunks_and_dtype_unchecked(
+                    self.name(),
+                    vec![arr],
+                    self.dtype().clone(),
+                )
+            }
+        } else {
+            let ca = IdxCa::from_vec("", (0..self.len() as IdxSize).collect());
+            unsafe { self.take_unchecked(&ca) }
+        }
+    }
+}
+
+impl ChunkReverse for StringChunked {
+    fn reverse(&self) -> Self {
+        unsafe { self.as_binary().reverse().to_string() }
+    }
+}
+
 #[cfg(feature = "dtype-array")]
 impl ChunkReverse for ArrayChunked {
     fn reverse(&self) -> Self {