feat: close upstream coverage gaps for DataFusion 55.1.0

.ai/skills/check-upstream/SKILL.md

@@ -146,6 +146,7 @@ The user may specify an area via `$ARGUMENTS`. If no area is specified or "all"
 - `show_limit` — already covered by `DataFrame.show()`, which provides the same functionality with a simpler API
 - `with_param_values` — already covered by the `param_values` argument on `SessionContext.sql()`, which accomplishes the same thing more robustly
 - `union_by_name_distinct` — already covered by `DataFrame.union_by_name(distinct=True)`, which provides a more Pythonic API
+- `to_string` — `str(df)` is the Pythonic way to get a string and already goes through `__repr__` and the configurable formatter. A separate `to_string()` would either duplicate `str(df)` or render every row through a different path (session `datafusion.format.*` options, no formatter), giving a third text rendering alongside `repr` and `show()`
 
 **How to check:**
 1. Fetch the upstream DataFrame documentation page listing all methods

crates/core/src/dataframe.rs

@@ -844,13 +844,24 @@ impl PyDataFrame {
     }
 
     /// Print the query plan
-    #[pyo3(signature = (verbose=false, analyze=false, format=None))]
+    #[pyo3(signature = (
+        verbose=false,
+        analyze=false,
+        format=None,
+        show_statistics=None,
+        analyze_level=None,
+        analyze_categories=None
+    ))]
+    #[allow(clippy::too_many_arguments)]
     fn explain(
         &self,
         py: Python,
         verbose: bool,
         analyze: bool,
         format: Option<&str>,
+        show_statistics: Option<bool>,
+        analyze_level: Option<&str>,
+        analyze_categories: Option<Vec<String>>,
     ) -> PyDataFusionResult<()> {
         let explain_format = match format {
             Some(f) => f
@@ -860,10 +871,24 @@ impl PyDataFrame {
                 })?,
             None => datafusion::common::format::ExplainFormat::Indent,
         };
+        let analyze_level = analyze_level
+            .map(|l| l.parse::<datafusion::common::format::MetricType>())
+            .transpose()?;
+        let analyze_categories = analyze_categories
+            .map(|cats| {
+                cats.iter()
+                    .map(|c| c.parse::<datafusion::common::format::MetricCategory>())
+                    .collect::<datafusion::common::Result<Vec<_>>>()
+                    .map(datafusion::common::format::ExplainAnalyzeCategories::Only)
+            })
+            .transpose()?;
         let opts = datafusion::logical_expr::ExplainOption::default()
             .with_verbose(verbose)
             .with_analyze(analyze)
-            .with_format(explain_format);
+            .with_format(explain_format)
+            .with_show_statistics(show_statistics)
+            .with_analyze_level(analyze_level)
+            .with_analyze_categories(analyze_categories);
         let df = self.df.as_ref().clone().explain_with_options(opts)?;
         print_dataframe(py, df)
     }
@@ -1320,6 +1345,26 @@ impl PyDataFrame {
         let df = self.df.as_ref().fill_null(&scalar_value.0, &cols)?;
         Ok(Self::new(df))
     }
+
+    /// Fill NaN values with a specified value for specific floating-point columns
+    #[pyo3(signature = (value, columns=None))]
+    fn fill_nan(
+        &self,
+        value: Py<PyAny>,
+        columns: Option<Vec<PyBackedStr>>,
+        py: Python,
+    ) -> PyDataFusionResult<Self> {
+        let scalar_value: PyScalarValue = value.extract(py)?;
+
+        let cols = match columns {
+            Some(col_names) => col_names.iter().map(|c| c.to_string()).collect(),
+            None => Vec::new(), // Empty vector means fill NaN for all columns
+        };
+
+        let cols = cols.iter().map(String::as_str).collect::<Vec<_>>();
+        let df = self.df.as_ref().fill_nan(&scalar_value.0, &cols)?;
+        Ok(Self::new(df))
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]

crates/core/src/expr.rs

@@ -55,7 +55,7 @@ use crate::expr::aggregate_expr::PyAggregateFunction;
 use crate::expr::binary_expr::PyBinaryExpr;
 use crate::expr::column::PyColumn;
 use crate::expr::literal::PyLiteral;
-use crate::functions::add_builder_fns_to_window;
+use crate::functions::{add_builder_fns_to_window, apply_window_options};
 use crate::pyarrow_util::scalar_to_pyarrow;
 use crate::sql::logical::PyLogicalPlan;
 
@@ -625,34 +625,61 @@ impl PyExpr {
     // Expression Function Builder functions
 
     pub fn order_by(&self, order_by: Vec<PySortExpr>) -> PyExprFuncBuilder {
-        self.expr
-            .clone()
+        builder_from_expr(&self.expr)
             .order_by(to_sort_expressions(order_by))
             .into()
     }
 
     pub fn filter(&self, filter: PyExpr) -> PyExprFuncBuilder {
-        self.expr.clone().filter(filter.expr.clone()).into()
+        builder_from_expr(&self.expr)
+            .filter(filter.expr.clone())
+            .into()
     }
 
     pub fn distinct(&self) -> PyExprFuncBuilder {
-        self.expr.clone().distinct().into()
+        // Only aggregates support DISTINCT, including an aggregate run as a window
+        // function. For anything else, upstream's empty builder makes `build()`
+        // raise instead of dropping the option.
+        let supports_distinct = match &self.expr {
+            Expr::AggregateFunction(_) => true,
+            Expr::WindowFunction(window) => {
+                matches!(window.fun, WindowFunctionDefinition::AggregateUDF(_))
+            }
+            _ => false,
+        };
+        if supports_distinct {
+            builder_from_expr(&self.expr).distinct().into()
+        } else {
+            self.expr.clone().distinct().into()
+        }
     }
 
     pub fn null_treatment(&self, null_treatment: NullTreatment) -> PyExprFuncBuilder {
-        self.expr
-            .clone()
+        builder_from_expr(&self.expr)
             .null_treatment(Some(null_treatment.into()))
             .into()
     }
 
     pub fn partition_by(&self, partition_by: Vec<PyExpr>) -> PyExprFuncBuilder {
         let partition_by = partition_by.iter().map(|e| e.expr.clone()).collect();
-        self.expr.clone().partition_by(partition_by).into()
+        // Window-only option: on an aggregate, upstream's empty builder makes
+        // `build()` raise instead of dropping it.
+        if matches!(self.expr, Expr::AggregateFunction(_)) {
+            return self.expr.clone().partition_by(partition_by).into();
+        }
+        builder_from_expr(&self.expr)
+            .partition_by(partition_by)
+            .into()
     }
 
     pub fn window_frame(&self, window_frame: PyWindowFrame) -> PyExprFuncBuilder {
-        self.expr.clone().window_frame(window_frame.into()).into()
+        // Window-only option; see `partition_by`.
+        if matches!(self.expr, Expr::AggregateFunction(_)) {
+            return self.expr.clone().window_frame(window_frame.into()).into();
+        }
+        builder_from_expr(&self.expr)
+            .window_frame(window_frame.into())
+            .into()
     }
 
     #[pyo3(signature = (partition_by=None, window_frame=None, order_by=None, null_treatment=None))]
@@ -678,8 +705,8 @@ impl PyExpr {
                     null_treatment,
                 )
             }
-            Expr::WindowFunction(_) => add_builder_fns_to_window(
-                self.expr.clone(),
+            Expr::WindowFunction(_) => apply_window_options(
+                builder_from_expr(&self.expr),
                 partition_by,
                 window_frame,
                 order_by,
@@ -743,6 +770,69 @@ impl PyExpr {
     }
 }
 
+/// Start an [`ExprFuncBuilder`] that keeps the options already set on `expr`.
+///
+/// Upstream's `ExprFunctionExt` methods on an `Expr` start from an empty
+/// builder, so `build()` would reset every option not set again. The Python
+/// function wrappers already apply their keyword options, so chaining another
+/// builder method onto their result must not discard them.
+///
+/// A built window function always stores a concrete frame, so whether the user
+/// chose it is lost. A frame equal to the default for the current order-by is
+/// treated as unset, which depends only on the expression and so behaves the
+/// same after a copy, pickle, or round trip through protobuf or SQL.
+fn builder_from_expr(expr: &Expr) -> ExprFuncBuilder {
+    match expr {
+        Expr::AggregateFunction(agg) => {
+            let params = &agg.params;
+            let mut builder = expr.clone().null_treatment(params.null_treatment);
+            if !params.order_by.is_empty() {
+                builder = builder.order_by(params.order_by.clone());
+            }
+            if let Some(filter) = &params.filter {
+                builder = builder.filter(filter.as_ref().clone());
+            }
+            if params.distinct {
+                builder = builder.distinct();
+            }
+            builder
+        }
+        Expr::WindowFunction(window) => {
+            let params = &window.params;
+            let mut builder = expr.clone().null_treatment(params.null_treatment);
+            if !params.partition_by.is_empty() {
+                builder = builder.partition_by(params.partition_by.clone());
+            }
+            let has_order_by = !params.order_by.is_empty();
+            if has_order_by {
+                builder = builder.order_by(params.order_by.clone());
+            }
+            // A frame equal to the default `build()` derived from the order-by is
+            // left unset, so it is derived again from the final order-by. An
+            // absent and an empty order-by derive different frames but are both
+            // stored as empty, so either frame counts as the default here.
+            let is_default_frame = if has_order_by {
+                params.window_frame == datafusion::logical_expr::WindowFrame::new(Some(true))
+            } else {
+                params.window_frame == datafusion::logical_expr::WindowFrame::new(None)
+                    || params.window_frame
+                        == datafusion::logical_expr::WindowFrame::new(Some(false))
+            };
+            if !is_default_frame {
+                builder = builder.window_frame(params.window_frame.clone());
+            }
+            if let Some(filter) = &params.filter {
+                builder = builder.filter(filter.as_ref().clone());
+            }
+            if params.distinct {
+                builder = builder.distinct();
+            }
+            builder
+        }
+        _ => expr.clone().null_treatment(None),
+    }
+}
+
 #[pyclass(
     from_py_object,
     frozen,