Add DataFrame.window() and unnest recursion options

Expose remaining DataFrame methods from upstream DataFusion.
Closes #1456.

- window(*exprs): apply window function expressions and append results
  as new columns
- unnest_column/unnest_columns: add optional recursions parameter for
  controlling unnest depth via (input_column, output_column, depth)
  tuples

Note: drop_columns is already exposed as the existing drop() method.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timsaucer

crates/core/src/dataframe.rs

@@ -582,6 +582,14 @@ impl PyDataFrame {
         Ok(Self::new(df))
     }
 
+    /// Apply window function expressions to the DataFrame
+    #[pyo3(signature = (*exprs))]
+    fn window(&self, exprs: Vec<PyExpr>) -> PyDataFusionResult<Self> {
+        let window_exprs = exprs.into_iter().map(|e| e.into()).collect();
+        let df = self.df.as_ref().clone().window(window_exprs)?;
+        Ok(Self::new(df))
+    }
+
     fn filter(&self, predicate: PyExpr) -> PyDataFusionResult<Self> {
         let df = self.df.as_ref().clone().filter(predicate.into())?;
         Ok(Self::new(df))
@@ -891,11 +899,14 @@ impl PyDataFrame {
         Ok(Self::new(new_df))
     }
 
-    #[pyo3(signature = (column, preserve_nulls=true))]
-    fn unnest_column(&self, column: &str, preserve_nulls: bool) -> PyDataFusionResult<Self> {
-        // TODO: expose RecursionUnnestOptions
-        // REF: https://github.com/apache/datafusion/pull/11577
-        let unnest_options = UnnestOptions::default().with_preserve_nulls(preserve_nulls);
+    #[pyo3(signature = (column, preserve_nulls=true, recursions=None))]
+    fn unnest_column(
+        &self,
+        column: &str,
+        preserve_nulls: bool,
+        recursions: Option<Vec<(String, String, usize)>>,
+    ) -> PyDataFusionResult<Self> {
+        let unnest_options = build_unnest_options(preserve_nulls, recursions);
         let df = self
             .df
             .as_ref()
@@ -904,15 +915,14 @@ impl PyDataFrame {
         Ok(Self::new(df))
     }
 
-    #[pyo3(signature = (columns, preserve_nulls=true))]
+    #[pyo3(signature = (columns, preserve_nulls=true, recursions=None))]
     fn unnest_columns(
         &self,
         columns: Vec<String>,
         preserve_nulls: bool,
+        recursions: Option<Vec<(String, String, usize)>>,
     ) -> PyDataFusionResult<Self> {
-        // TODO: expose RecursionUnnestOptions
-        // REF: https://github.com/apache/datafusion/pull/11577
-        let unnest_options = UnnestOptions::default().with_preserve_nulls(preserve_nulls);
+        let unnest_options = build_unnest_options(preserve_nulls, recursions);
         let cols = columns.iter().map(|s| s.as_ref()).collect::<Vec<&str>>();
         let df = self
             .df
@@ -1376,6 +1386,26 @@ impl PyDataFrameWriteOptions {
     }
 }
 
+fn build_unnest_options(
+    preserve_nulls: bool,
+    recursions: Option<Vec<(String, String, usize)>>,
+) -> UnnestOptions {
+    let mut opts = UnnestOptions::default().with_preserve_nulls(preserve_nulls);
+    if let Some(recs) = recursions {
+        opts.recursions = recs
+            .into_iter()
+            .map(
+                |(input, output, depth)| datafusion::common::RecursionUnnestOption {
+                    input_column: datafusion::common::Column::from(input.as_str()),
+                    output_column: datafusion::common::Column::from(output.as_str()),
+                    depth,
+                },
+            )
+            .collect();
+    }
+    opts
+}
+
 /// Print DataFrame
 fn print_dataframe(py: Python, df: DataFrame) -> PyDataFusionResult<()> {
     // Get string representation of record batches

python/datafusion/dataframe.py

@@ -487,6 +487,21 @@ def drop(self, *columns: str) -> DataFrame:
         """
         return DataFrame(self.df.drop(*columns))
 
+    def window(self, *exprs: Expr) -> DataFrame:
+        """Add window function columns to the DataFrame.
+
+        Applies the given window function expressions and appends the results
+        as new columns.
+
+        Args:
+            exprs: Window function expressions to evaluate.
+
+        Returns:
+            DataFrame with new window function columns appended.
+        """
+        raw = expr_list_to_raw_expr_list(exprs)
+        return DataFrame(self.df.window(*raw))
+
     def filter(self, *predicates: Expr | str) -> DataFrame:
         """Return a DataFrame for which ``predicate`` evaluates to ``True``.
 
@@ -1426,23 +1441,44 @@ def count(self) -> int:
         return self.df.count()
 
     @deprecated("Use :py:func:`unnest_columns` instead.")
-    def unnest_column(self, column: str, preserve_nulls: bool = True) -> DataFrame:
+    def unnest_column(
+        self,
+        column: str,
+        preserve_nulls: bool = True,
+        recursions: list[tuple[str, str, int]] | None = None,
+    ) -> DataFrame:
         """See :py:func:`unnest_columns`."""
-        return DataFrame(self.df.unnest_column(column, preserve_nulls=preserve_nulls))
+        return DataFrame(
+            self.df.unnest_column(
+                column, preserve_nulls=preserve_nulls, recursions=recursions
+            )
+        )
 
-    def unnest_columns(self, *columns: str, preserve_nulls: bool = True) -> DataFrame:
+    def unnest_columns(
+        self,
+        *columns: str,
+        preserve_nulls: bool = True,
+        recursions: list[tuple[str, str, int]] | None = None,
+    ) -> DataFrame:
         """Expand columns of arrays into a single row per array element.
 
         Args:
             columns: Column names to perform unnest operation on.
             preserve_nulls: If False, rows with null entries will not be
                 returned.
+            recursions: Optional list of ``(input_column, output_column, depth)``
+                tuples that control how deeply nested columns are unnested. Any
+                column not mentioned here is unnested with depth 1.
 
         Returns:
             A DataFrame with the columns expanded.
         """
         columns = list(columns)
-        return DataFrame(self.df.unnest_columns(columns, preserve_nulls=preserve_nulls))
+        return DataFrame(
+            self.df.unnest_columns(
+                columns, preserve_nulls=preserve_nulls, recursions=recursions
+            )
+        )
 
     def __arrow_c_stream__(self, requested_schema: object | None = None) -> object:
         """Export the DataFrame as an Arrow C Stream.

python/tests/test_dataframe.py

@@ -3654,3 +3654,24 @@ def test_explain_with_format(capsys):
     df.explain(verbose=True, analyze=True, format=ExplainFormat.INDENT)
     captured = capsys.readouterr()
     assert "plan_type" in captured.out
+
+
+def test_window():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": [1, 2, 3], "b": ["x", "x", "y"]})
+    result = df.window(
+        f.row_number(partition_by=[column("b")], order_by=[column("a")]).alias("rn")
+    ).collect()[0]
+    assert "rn" in result.schema.names
+    assert result.column(result.schema.get_field_index("rn")).to_pylist() == [1, 2, 1]
+
+
+def test_unnest_columns_with_recursions():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": [[1, 2], [3]], "b": ["x", "y"]})
+    # Basic unnest still works
+    result = df.unnest_columns("a").collect()[0]
+    assert result.column(0).to_pylist() == [1, 2, 3]
+    # With explicit recursion options
+    result = df.unnest_columns("a", recursions=[("a", "a", 1)]).collect()[0]
+    assert result.column(0).to_pylist() == [1, 2, 3]