Add support for interpolate_by (#100)

Author: ozhilin

ext/polars/Cargo.toml

@@ -63,6 +63,7 @@ features = [
     "gcp",
     "http",
     "interpolate",
+    "interpolate_by",
     "ipc",
     "ipc_streaming",
     "is_between",

ext/polars/src/expr/general.rs

@@ -671,6 +671,10 @@ impl RbExpr {
         self.inner.clone().interpolate(method.0).into()
     }
 
+    pub fn interpolate_by(&self, by: &Self) -> Self {
+        self.inner.clone().interpolate_by(by.inner.clone()).into()
+    }
+
     pub fn lower_bound(&self) -> Self {
         self.inner.clone().lower_bound().into()
     }

ext/polars/src/lib.rs

@@ -430,6 +430,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     class.define_method("mode", method!(RbExpr::mode, 0))?;
     class.define_method("exclude", method!(RbExpr::exclude, 1))?;
     class.define_method("interpolate", method!(RbExpr::interpolate, 1))?;
+    class.define_method("interpolate_by", method!(RbExpr::interpolate_by, 1))?;
     class.define_method("rolling_sum", method!(RbExpr::rolling_sum, 4))?;
     class.define_method("rolling_sum_by", method!(RbExpr::rolling_sum_by, 4))?;
     class.define_method("rolling_min", method!(RbExpr::rolling_min, 4))?;

lib/polars/expr.rb

@@ -3994,6 +3994,37 @@ def interpolate(method: "linear")
       _from_rbexpr(_rbexpr.interpolate(method))
     end
 
+    # Fill null values using interpolation based on another column.
+    #
+    # @return [Expr]
+    #
+    # @param by [Expr] Column to interpolate values based on.
+    #
+    # @example Fill null values using linear interpolation.
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, nil, nil, 3],
+    #       "b" => [1, 2, 7, 8]
+    #     }
+    #   )
+    #   df.with_columns(a_interpolated: Polars.col("a").interpolate_by("b"))
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌──────┬─────┬────────────────┐
+    #   # │ a    ┆ b   ┆ a_interpolated │
+    #   # │ ---  ┆ --- ┆ ---            │
+    #   # │ i64  ┆ i64 ┆ f64            │
+    #   # ╞══════╪═════╪════════════════╡
+    #   # │ 1    ┆ 1   ┆ 1.0            │
+    #   # │ null ┆ 2   ┆ 1.285714       │
+    #   # │ null ┆ 7   ┆ 2.714286       │
+    #   # │ 3    ┆ 8   ┆ 3.0            │
+    #   # └──────┴─────┴────────────────┘
+    def interpolate_by(by)
+      by = Utils.parse_into_expression(by)
+      _from_rbexpr(_rbexpr.interpolate_by(by))
+    end
+
     # Apply a rolling min based on another column.
     #
     # @param by [String]

lib/polars/series.rb

@@ -3815,6 +3815,30 @@ def interpolate(method: "linear")
       super
     end
 
+    # Fill null values using interpolation based on another column.
+    #
+    # @param by [Expr]
+    #   Column to interpolate values based on.
+    #
+    # @return [Series]
+    #
+    # @example Fill null values using linear interpolation.
+    #   s = Polars::Series.new("a", [1, nil, nil, 3])
+    #   by = Polars::Series.new("b", [1, 2, 7, 8])
+    #   s.interpolate_by(by)
+    #   # =>
+    #   # shape: (4,)
+    #   # Series: 'a' [f64]
+    #   # [
+    #   #         1.0
+    #   #         1.285714
+    #   #         2.714286
+    #   #         3.0
+    #   # ]
+    def interpolate_by(by)
+      super
+    end
+
     # Compute absolute values.
     #
     # @return [Series]

test/series_test.rb

@@ -734,6 +734,13 @@ def test_interpolate
     assert_series [1, 2, 3, 4, 5], s.interpolate
   end
 
+  def test_interpolate_by
+    s = Polars::Series.new("a", [1, nil, nil, 3])
+    by = Polars::Series.new("b", [1, 2, 7, 8])
+    s.interpolate_by(by)
+    assert_series [1.0, 1.2857142857142856, 2.7142857142857144, 3.0], s.interpolate_by(by)
+  end
+
   def test_shrink_to_fit
     s = Polars::Series.new([1, 2, 3])
     s.shrink_to_fit