Conversation
There was a problem hiding this comment.
Pull request overview
This PR extends the Python DataFrame API to expose additional upstream DataFusion DataFrame functionality (set operations, explain formatting, window expressions, sort convenience, and configurable unnesting) and adds corresponding unit tests.
Changes:
- Added
ExplainFormatand extendedDataFrame.explain()to support configurable output formats. - Exposed additional DataFrame operations to Python:
window,sort_by,distinct_on,except_distinct,intersect_distinct,union_by_name, andunion_by_name_distinct. - Extended
unnest_column(s)to accept optionalrecursionsand added tests for the new APIs.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
python/tests/test_dataframe.py |
Adds unit tests covering newly exposed DataFrame methods and explain formatting. |
python/datafusion/dataframe.py |
Adds ExplainFormat and new/extended DataFrame methods (window, set ops, distinct_on, sort_by, unnest recursion options). |
python/datafusion/__init__.py |
Re-exports ExplainFormat from the top-level package API. |
crates/core/src/dataframe.rs |
Adds/extends PyO3 bindings for the new DataFrame methods and explain formatting + unnest recursion options. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Expose upstream DataFusion DataFrame methods that were not yet available in the Python API. Closes apache#1455. Set operations: - except_distinct: set difference with deduplication - intersect_distinct: set intersection with deduplication - union_by_name: union matching columns by name instead of position - union_by_name_distinct: union by name with deduplication Query: - distinct_on: deduplicate rows based on specific columns - sort_by: sort by expressions with ascending order and nulls last Note: show_limit is already covered by the existing show(num) method. explain_with_options and with_param_values are deferred as they require exposing additional types (ExplainOption, ParamValues). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Extend the existing explain() method with an optional format parameter instead of adding a separate explain_with_options() method. This keeps the API simple while exposing all upstream ExplainOption functionality. Available formats: indent (default), tree, pgjson, graphviz. The ExplainFormat enum is exported from the top-level datafusion module. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Expose remaining DataFrame methods from upstream DataFusion. Closes apache#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>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Clarify except_distinct/intersect_distinct docstrings, add deterministic sort to test_window, add sort_by ascending verification test, and add smoke tests for PGJSON and GRAPHVIZ explain formats. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Combine set operation tests (except_distinct, intersect_distinct, union_by_name, union_by_name_distinct) into a single parametrized test_set_operations_distinct. Merge sort_by tests and convert explain format tests to parametrized form. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
3ab9a6c to
bb18bf9
Compare
Add >>> style usage examples for window, explain, except_distinct, intersect_distinct, union_by_name, union_by_name_distinct, distinct_on, sort_by, and unnest_columns to match existing docstring conventions. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Provide actionable error message for invalid explain format strings - Remove recursions param from deprecated unnest_column (use unnest_columns) - Add null-handling test case for sort_by to verify nulls-last behavior - Add format-specific assertions to explain tests (TREE, PGJSON, GRAPHVIZ) - Add deep recursion test for unnest_columns with depth > 1 - Add multi-expression window test to verify variadic *exprs Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Combine test_window and test_window_multiple_expressions into a single parametrized test. Merge unnest recursion tests into one parametrized test covering basic, explicit depth 1, and deep recursion cases. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
ntjohnson1
left a comment
There was a problem hiding this comment.
Related to #1340
This doesn't cover find_qualified_cols from apache/datafusion#19549
(which should probably have a nicer name/alias in python) in the docs https://docs.rs/datafusion/53.0.0/datafusion/dataframe/struct.DataFrame.html#method.find_qualified_columns
Which is needed to resolve something like #1305 (comment)
| .parse::<datafusion::common::format::ExplainFormat>() | ||
| .map_err(|_| { | ||
| PyDataFusionError::Common(format!( | ||
| "Invalid explain format: '{}'. Valid options: indent, tree, pgjson, graphviz", |
There was a problem hiding this comment.
NIT: nice if the options weren't hard coded for easier maintainability
| """ | ||
| return DataFrame(self.df.union_by_name(other.df)) | ||
|
|
||
| def union_by_name_distinct(self, other: DataFrame) -> DataFrame: |
There was a problem hiding this comment.
NIT: I don't know if the style throughout is to try to keep clearer 1 to 1 mapping with rust but I would find
def union_by_name(self, other: DataFrame, distinct=False) to be more intuitive/pythonic
| Returns: | ||
| A DataFrame with the columns expanded. | ||
|
|
||
| Examples: |
There was a problem hiding this comment.
NIT: Doesn't do an example with preserve nulls but that is pretty clear
Which issue does this PR close?
Closes #1455
Closes #1456
Rationale for this change
These features exist in the upstream repository but have not been exposed in Python.
What changes are included in this PR?
Exposes functions to python
Adds unit tests
Are there any user-facing changes?
Addition only