refactor: Make backwardfill and forwardfill function expr non-anonymo…

…us (#11630)
pola-rs · Oct 10, 2023 · 9a78bfa · 9a78bfa
1 parent 52a2632
commit 9a78bfa
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 10 deletions.
diff --git a/crates/polars-plan/src/dsl/function_expr/dispatch.rs b/crates/polars-plan/src/dsl/function_expr/dispatch.rs
@@ -51,3 +51,11 @@ pub(super) fn value_counts(s: &Series, sort: bool, parallel: bool) -> PolarsResu
 pub(super) fn unique_counts(s: &Series) -> PolarsResult<Series> {
     Ok(s.unique_counts().into_series())
 }
+
+pub(super) fn backward_fill(s: &Series, limit: FillNullLimit) -> PolarsResult<Series> {
+    s.fill_null(FillNullStrategy::Backward(limit))
+}
+
+pub(super) fn forward_fill(s: &Series, limit: FillNullLimit) -> PolarsResult<Series> {
+    s.fill_null(FillNullStrategy::Forward(limit))
+}
diff --git a/crates/polars-plan/src/dsl/function_expr/mod.rs b/crates/polars-plan/src/dsl/function_expr/mod.rs
@@ -250,6 +250,12 @@ pub enum FunctionExpr {
         lib: Arc<str>,
         symbol: Arc<str>,
     },
+    BackwardFill {
+        limit: FillNullLimit,
+    },
+    ForwardFill {
+        limit: FillNullLimit,
+    },
 }
 
 impl Hash for FunctionExpr {
@@ -419,6 +425,8 @@ impl Display for FunctionExpr {
             SetSortedFlag(_) => "set_sorted",
             #[cfg(feature = "ffi_plugin")]
             FfiPlugin { lib, symbol, .. } => return write!(f, "{lib}:{symbol}"),
+            BackwardFill { .. } => "backward_fill",
+            ForwardFill { .. } => "forward_fill",
         };
         write!(f, "{s}")
     }
@@ -736,6 +744,8 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             FfiPlugin { lib, symbol, .. } => unsafe {
                 map_as_slice!(plugin::call_plugin, lib.as_ref(), symbol.as_ref())
             },
+            BackwardFill { limit } => map!(dispatch::backward_fill, limit),
+            ForwardFill { limit } => map!(dispatch::forward_fill, limit),
         }
     }
 }

diff --git a/crates/polars-plan/src/dsl/function_expr/schema.rs b/crates/polars-plan/src/dsl/function_expr/schema.rs
@@ -234,6 +234,8 @@ impl FunctionExpr {
             FfiPlugin { lib, symbol } => unsafe {
                 plugin::plugin_field(fields, lib, &format!("__polars_field_{}", symbol.as_ref()))
             },
+            BackwardFill { .. } => mapper.with_same_dtype(),
+            ForwardFill { .. } => mapper.with_same_dtype(),
         }
     }
 }

diff --git a/crates/polars-plan/src/dsl/mod.rs b/crates/polars-plan/src/dsl/mod.rs
@@ -793,20 +793,12 @@ impl Expr {
 
     /// Fill missing value with next non-null.
     pub fn backward_fill(self, limit: FillNullLimit) -> Self {
-        self.apply(
-            move |s: Series| s.fill_null(FillNullStrategy::Backward(limit)).map(Some),
-            GetOutput::same_type(),
-        )
-        .with_fmt("backward_fill")
+        self.apply_private(FunctionExpr::BackwardFill { limit })
     }
 
     /// Fill missing value with previous non-null.
     pub fn forward_fill(self, limit: FillNullLimit) -> Self {
-        self.apply(
-            move |s: Series| s.fill_null(FillNullStrategy::Forward(limit)).map(Some),
-            GetOutput::same_type(),
-        )
-        .with_fmt("forward_fill")
+        self.apply_private(FunctionExpr::ForwardFill { limit })
     }
 
     /// Round underlying floating point array to given decimal numbers.