Spenserrrr opened a new pull request, #57726:
URL: https://github.com/apache/spark/pull/57726
### What changes were proposed in this pull request?
`Window.partitionBy`/`orderBy` (on both `Window` and `WindowSpec`) and
`TableArg.partitionBy`/`orderBy` accept their columns as varargs, and also
allow a single sequence to be passed in place of the varargs (e.g.
`Window.partitionBy(["a", "b"])`). That single-sequence form was unwrapped only
when it was a list, so a tuple was not unwrapped and was instead treated as a
single, invalid column argument.
This PR:
1. Widens the runtime unwrap check from `isinstance(cols[0], list)` to any
non-`str` `Sequence`, so a single tuple is unpacked the same way a list already
is. (`Column` is not a `Sequence`, so it is excluded by the check itself.)
2. Widens the implementation signatures to `Union[ColumnOrName,
Sequence[ColumnOrName]]` across the base, classic and connect layers, and adds
the standard two-overload public signature (a spread of columns, or a single
sequence of columns) which neither `Window` nor `TableArg` had.
3. Removes the three `# type: ignore[assignment]` comments on the
single-sequence unwrap in the `_to_cols` and `_to_java_cols` helpers, by
assigning `tuple(cols[0])` rather than reusing the tuple-typed variable to hold
a list.
Two `# type: ignore[arg-type]` comments are added in the connect `Window`
static methods. They forward their varargs into
`WindowSpec.partitionBy`/`orderBy`, which are now overloaded, and a spread of
the wide union matches neither overload. This follows the existing precedent
for the same pattern in `DataFrame.agg`, which forwards to
`self.groupBy().agg(*exprs) # type: ignore[arg-type]`. Classic does not need
this because it forwards to the plain `_to_java_cols` helper instead.
### Why are the changes needed?
The annotations were inaccurate today, in opposite directions:
- `Window.partitionBy`/`orderBy` were annotated `*cols: Union[ColumnOrName,
Sequence[ColumnOrName]]`, so passing a tuple type-checked even though the
runtime did not unwrap it — the annotation promised more than the
implementation delivered.
- `TableArg.partitionBy`/`orderBy` were annotated `*cols: ColumnOrName`,
which admits no sequence at all, even though the runtime unwrapped a single
list and the docstrings document `str, Column, or list`. Because the declared
element type had no sequence member, the list-handling branch was dead code as
far as the type checker was concerned (confirmed with `mypy
--warn-unreachable`), so the annotation forbade a call that is supported,
documented, and already covered by an existing test.
This also makes these methods consistent with the other column-collecting
varargs methods on `DataFrame` (`select`, `groupBy`, `rollup`, `cube`), which
already accept any non-`str` `Sequence` and declare it with a two-overload
signature, and with `describe`/`selectExpr`, which were widened the same way in
SPARK-58500.
### Does this PR introduce _any_ user-facing change?
Yes, and it is backward compatible. Calls that work today keep working; a
single tuple (or any other sequence) is now accepted in addition to a list.
Before:
```python
>>> Window.partitionBy(("dept", "team")) # tuple not unwrapped -> invalid
column argument
>>> df.asTable().partitionBy(("key", "number"))
```
After, these behave the same as the already-supported list form:
```python
>>> Window.partitionBy(["dept", "team"]) # unchanged
>>> Window.partitionBy(("dept", "team")) # now equivalent
```
Docstrings are left as-is (`str`, `Column` or `list`), matching the
convention followed in SPARK-58488 and SPARK-58500.
### How was this patch tested?
Added tuple-form cases alongside the existing list-form cases, in mixins
that run under both classic and Spark Connect:
- `test_udtf.py` `test_df_asTable_chaining_methods` (`BaseUDTFTestsMixin`):
`partitionBy(("key", "number"))` and `orderBy(("number", "value"))`.
- `test_functions.py` `test_window_partitionBy_orderBy_with_sequence`
(`FunctionsTestsMixin`): asserts the varargs, list and tuple forms produce
identical results.
Also verified `mypy --namespace-packages --config-file python/mypy.ini
python/pyspark` is clean over the full scope (1290 files), and `ruff format
--check` / `ruff check` pass on all changed files. No typing `.yml` cases
reference these methods, so none needed updating.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
