feat: add Expr.interpolate_by (#16313)

crates/polars-lazy/Cargo.toml

@@ -125,6 +125,7 @@ range = ["polars-plan/range"]
 mode = ["polars-plan/mode"]
 cum_agg = ["polars-plan/cum_agg"]
 interpolate = ["polars-plan/interpolate"]
+interpolate_by = ["polars-plan/interpolate_by"]
 rolling_window = [
   "polars-plan/rolling_window",
 ]
@@ -270,6 +271,7 @@ features = [
   "futures",
   "hist",
   "interpolate",
+  "interpolate_by",
   "ipc",
   "is_first_distinct",
   "is_in",

crates/polars-ops/Cargo.toml

@@ -91,6 +91,7 @@ string_encoding = ["base64", "hex"]
 # ops
 to_dummies = []
 interpolate = []
+interpolate_by = []
 list_to_struct = ["polars-core/dtype-struct"]
 array_to_struct = ["polars-core/dtype-array", "polars-core/dtype-struct"]
 list_count = []

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

@@ -3,8 +3,6 @@ pub mod array;
 mod binary;
 #[cfg(feature = "timezones")]
 pub mod datetime;
-#[cfg(feature = "interpolate")]
-mod interpolate;
 pub mod list;
 #[cfg(feature = "propagate_nans")]
 pub mod nan_propagating_aggregate;
@@ -36,8 +34,6 @@ pub use datetime::*;
 pub use gather::*;
 #[cfg(feature = "hist")]
 pub use hist::*;
-#[cfg(feature = "interpolate")]
-pub use interpolate::*;
 pub use list::*;
 #[allow(unused_imports)]
 use polars_core::prelude::*;

crates/polars-ops/src/chunked_array/interpolate.rs → crates/polars-ops/src/series/ops/interpolation/interpolate.rs

@@ -8,26 +8,7 @@ use polars_core::prelude::*;
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
-fn linear_itp<T>(low: T, step: T, slope: T) -> T
-where
-    T: Sub<Output = T> + Mul<Output = T> + Add<Output = T> + Div<Output = T>,
-{
-    low + step * slope
-}
-
-fn nearest_itp<T>(low: T, step: T, diff: T, steps_n: T) -> T
-where
-    T: Sub<Output = T> + Mul<Output = T> + Add<Output = T> + Div<Output = T> + PartialOrd + Copy,
-{
-    // 5 - 1 = 5 -> low
-    // 5 - 2 = 3 -> low
-    // 5 - 3 = 2 -> high
-    if (steps_n - step) > step {
-        low
-    } else {
-        low + diff
-    }
-}
+use super::{linear_itp, nearest_itp};
 
 fn near_interp<T>(low: T, high: T, steps: IdxSize, steps_n: T, out: &mut Vec<T>)
 where
@@ -75,11 +56,11 @@ where
     let first = chunked_arr.first_non_null().unwrap();
     let last = chunked_arr.last_non_null().unwrap() + 1;
 
-    // Fill out with first.
+    // Fill out with `first` nulls.
     let mut out = Vec::with_capacity(chunked_arr.len());
     let mut iter = chunked_arr.iter().skip(first);
     for _ in 0..first {
-        out.push(Zero::zero())
+        out.push(Zero::zero());
     }
 
     // The next element of `iter` is definitely `Some(Some(v))`, because we skipped the first
@@ -109,11 +90,11 @@ where
         validity.extend_constant(chunked_arr.len(), true);
 
         for i in 0..first {
-            validity.set(i, false);
+            unsafe { validity.set_unchecked(i, false) };
         }
 
         for i in last..chunked_arr.len() {
-            validity.set(i, false);
+            unsafe { validity.set_unchecked(i, false) };
             out.push(Zero::zero())
         }