From 6b14d7e4b49a8de34decc792ddd41d66d1995776 Mon Sep 17 00:00:00 2001
From: Orson Peters <orsonpeters@gmail.com>
Date: Fri, 24 May 2024 18:01:29 +0200
Subject: [PATCH] perf: improved numeric fill_(forward/backward)

---
 crates/polars-arrow/src/legacy/utils.rs       | 30 ++++----
 crates/polars-arrow/src/trusted_len.rs        |  2 +-
 .../src/chunked_array/ops/fill_null.rs        | 68 ++++++++++++++++++-
 3 files changed, 85 insertions(+), 15 deletions(-)

diff --git a/crates/polars-arrow/src/legacy/utils.rs b/crates/polars-arrow/src/legacy/utils.rs
index 5670fed471ce..a9a6b7d0ed78 100644
--- a/crates/polars-arrow/src/legacy/utils.rs
+++ b/crates/polars-arrow/src/legacy/utils.rs
@@ -113,21 +113,27 @@ impl<T: NativeType> FromTrustedLenIterator<T> for PrimitiveArray<T> {
     }
 }
 
-impl<T: NativeType> FromIteratorReversed<T> for PrimitiveArray<T> {
+impl<T> FromIteratorReversed<T> for Vec<T> {
     fn from_trusted_len_iter_rev<I: TrustedLen<Item = T>>(iter: I) -> Self {
-        let size = iter.size_hint().1.unwrap();
-
-        let mut vals: Vec<T> = Vec::with_capacity(size);
         unsafe {
-            // Set to end of buffer.
-            let mut ptr = vals.as_mut_ptr().add(size);
-
-            iter.for_each(|item| {
-                ptr = ptr.sub(1);
-                std::ptr::write(ptr, item);
-            });
-            vals.set_len(size)
+            let len = iter.size_hint().1.unwrap();
+            let mut out: Vec<T> = Vec::with_capacity(len);
+            let mut idx = len;
+            for x in iter {
+                debug_assert!(idx > 0);
+                idx -= 1;
+                out.as_mut_ptr().add(idx).write(x);
+            }
+            debug_assert!(idx == 0);
+            out.set_len(len);
+            out
         }
+    }
+}
+
+impl<T: NativeType> FromIteratorReversed<T> for PrimitiveArray<T> {
+    fn from_trusted_len_iter_rev<I: TrustedLen<Item = T>>(iter: I) -> Self {
+        let vals: Vec<T> = iter.collect_reversed();
         PrimitiveArray::new(ArrowDataType::from(T::PRIMITIVE), vals.into(), None)
     }
 }
diff --git a/crates/polars-arrow/src/trusted_len.rs b/crates/polars-arrow/src/trusted_len.rs
index 4bdce32e4990..3237ba83cbb2 100644
--- a/crates/polars-arrow/src/trusted_len.rs
+++ b/crates/polars-arrow/src/trusted_len.rs
@@ -71,7 +71,7 @@ unsafe impl<A: TrustedLen> TrustedLen for std::iter::StepBy<A> {}
 unsafe impl<I, St, F, B> TrustedLen for Scan<I, St, F>
 where
     F: FnMut(&mut St, I::Item) -> Option<B>,
-    I: TrustedLen + Iterator<Item = B>,
+    I: TrustedLen,
 {
 }
 
diff --git a/crates/polars-core/src/chunked_array/ops/fill_null.rs b/crates/polars-core/src/chunked_array/ops/fill_null.rs
index e52cef3c296e..88dafe4edd58 100644
--- a/crates/polars-core/src/chunked_array/ops/fill_null.rs
+++ b/crates/polars-core/src/chunked_array/ops/fill_null.rs
@@ -1,6 +1,8 @@
+use arrow::bitmap::MutableBitmap;
 use arrow::legacy::kernels::set::set_at_nulls;
 use arrow::legacy::trusted_len::FromIteratorReversed;
 use arrow::legacy::utils::FromTrustedLenIterator;
+use bytemuck::Zeroable;
 use num_traits::{Bounded, NumCast, One, Zero};
 
 use crate::prelude::*;
@@ -280,6 +282,68 @@ macro_rules! impl_fill_backward_limit {
     }};
 }
 
+fn fill_forward_numeric<'a, T, I>(ca: &'a ChunkedArray<T>) -> ChunkedArray<T>
+where
+    T: PolarsDataType,
+    &'a ChunkedArray<T>: IntoIterator<IntoIter = I>,
+    I: TrustedLen + Iterator<Item = Option<T::Physical<'a>>>,
+    T::ZeroablePhysical<'a>: LocalCopy,
+{
+    // Compute values.
+    let values: Vec<T::ZeroablePhysical<'a>> = ca
+        .into_iter()
+        .scan(T::ZeroablePhysical::zeroed(), |prev, v| {
+            *prev = v.map(|v| v.into()).unwrap_or(prev.cheap_clone());
+            Some(prev.cheap_clone())
+        })
+        .collect_trusted();
+
+    // Compute bitmask.
+    let num_start_nulls = ca.first_non_null().unwrap_or(0);
+    let mut bm = MutableBitmap::with_capacity(ca.len());
+    bm.extend_constant(num_start_nulls, false);
+    bm.extend_constant(ca.len() - num_start_nulls, true);
+    ChunkedArray::from_chunk_iter_like(
+        ca,
+        [
+            T::Array::from_zeroable_vec(values, ca.dtype().to_arrow(true))
+                .with_validity_typed(Some(bm.into())),
+        ],
+    )
+}
+
+fn fill_backward_numeric<'a, T, I>(ca: &'a ChunkedArray<T>) -> ChunkedArray<T>
+where
+    T: PolarsDataType,
+    &'a ChunkedArray<T>: IntoIterator<IntoIter = I>,
+    I: TrustedLen + Iterator<Item = Option<T::Physical<'a>>> + DoubleEndedIterator,
+    T::ZeroablePhysical<'a>: LocalCopy,
+{
+    // Compute values.
+    let values: Vec<T::ZeroablePhysical<'a>> = ca
+        .into_iter()
+        .rev()
+        .scan(T::ZeroablePhysical::zeroed(), |prev, v| {
+            *prev = v.map(|v| v.into()).unwrap_or(prev.cheap_clone());
+            Some(prev.cheap_clone())
+        })
+        .collect_reversed();
+
+    // Compute bitmask.
+    let last_idx = ca.len().saturating_sub(1);
+    let num_end_nulls = last_idx - ca.last_non_null().unwrap_or(last_idx);
+    let mut bm = MutableBitmap::with_capacity(ca.len());
+    bm.extend_constant(ca.len() - num_end_nulls, true);
+    bm.extend_constant(num_end_nulls, false);
+    ChunkedArray::from_chunk_iter_like(
+        ca,
+        [
+            T::Array::from_zeroable_vec(values, ca.dtype().to_arrow(true))
+                .with_validity_typed(Some(bm.into())),
+        ],
+    )
+}
+
 fn fill_null_numeric<T>(
     ca: &ChunkedArray<T>,
     strategy: FillNullStrategy,
@@ -293,9 +357,9 @@ where
         return Ok(ca.clone());
     }
     let mut out = match strategy {
-        FillNullStrategy::Forward(None) => fill_forward(ca),
+        FillNullStrategy::Forward(None) => fill_forward_numeric(ca),
         FillNullStrategy::Forward(Some(limit)) => fill_forward_limit(ca, limit),
-        FillNullStrategy::Backward(None) => fill_backward(ca),
+        FillNullStrategy::Backward(None) => fill_backward_numeric(ca),
         FillNullStrategy::Backward(Some(limit)) => fill_backward_limit(ca, limit),
         FillNullStrategy::Min => {
             ca.fill_null_with_values(ChunkAgg::min(ca).ok_or_else(err_fill_null)?)?