Skip to content

[SPARK-58500][PYTHON] Accept tuple in DataFrame.describe and selectExpr #57702

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Spenserrrr wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[SPARK-58500][PYTHON] Accept tuple in DataFrame.describe and selectExpr #57702

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 13 additions & 7 deletions python/pyspark/sql/classic/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1033,9 +1033,15 @@ def _jcols_ordinal(
_cols.append(c) # type: ignore[arg-type]
return self._jseq(_cols, _to_java_column)

def describe(self, *cols: Union[str, List[str]]) -> ParentDataFrame:
if len(cols) == 1 and isinstance(cols[0], list):
cols = cols[0] # type: ignore[assignment]
@overload
def describe(self, *cols: str) -> ParentDataFrame: ...

@overload
def describe(self, __cols: Sequence[str]) -> ParentDataFrame: ...

def describe(self, *cols: Union[str, Sequence[str]]) -> ParentDataFrame:
if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence):
cols = tuple(cols[0])
jdf = self._jdf.describe(self._jseq(cols))
return DataFrame(jdf, self.sparkSession)

Expand Down Expand Up @@ -1116,11 +1122,11 @@ def select(self, *cols: Union[Sequence["ColumnOrName"], "ColumnOrName"]) -> Pare
def selectExpr(self, *expr: str) -> ParentDataFrame: ...

@overload
def selectExpr(self, *expr: List[str]) -> ParentDataFrame: ...
def selectExpr(self, __expr: Sequence[str]) -> ParentDataFrame: ...

def selectExpr(self, *expr: Union[str, List[str]]) -> ParentDataFrame:
if len(expr) == 1 and isinstance(expr[0], list):
expr = expr[0] # type: ignore[assignment]
def selectExpr(self, *expr: Union[str, Sequence[str]]) -> ParentDataFrame:
if len(expr) == 1 and not isinstance(expr[0], str) and isinstance(expr[0], Sequence):
expr = tuple(expr[0])
jdf = self._jdf.selectExpr(self._jseq(expr))
return DataFrame(jdf, self.sparkSession)

Expand Down
24 changes: 18 additions & 6 deletions python/pyspark/sql/connect/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -300,10 +300,16 @@ def select(self, *cols: Union[Sequence["ColumnOrName"], "ColumnOrName"]) -> Pare
session=self._session,
)

def selectExpr(self, *expr: Union[str, List[str]]) -> ParentDataFrame:
@overload
def selectExpr(self, *expr: str) -> ParentDataFrame: ...

@overload
def selectExpr(self, __expr: Sequence[str]) -> ParentDataFrame: ...

def selectExpr(self, *expr: Union[str, Sequence[str]]) -> ParentDataFrame:
sql_expr = []
if len(expr) == 1 and isinstance(expr[0], list):
expr = expr[0] # type: ignore[assignment]
if len(expr) == 1 and not isinstance(expr[0], str) and isinstance(expr[0], Sequence):
expr = tuple(expr[0])
for element in expr:
if isinstance(element, str):
sql_expr.append(F.expr(element))
Expand Down Expand Up @@ -1601,9 +1607,15 @@ def summary(self, *statistics: str) -> ParentDataFrame:
session=self._session,
)

def describe(self, *cols: Union[str, List[str]]) -> ParentDataFrame:
if len(cols) == 1 and isinstance(cols[0], list):
cols = cols[0] # type: ignore[assignment]
@overload
def describe(self, *cols: str) -> ParentDataFrame: ...

@overload
def describe(self, __cols: Sequence[str]) -> ParentDataFrame: ...

def describe(self, *cols: Union[str, Sequence[str]]) -> ParentDataFrame:
if len(cols) == 1 and not isinstance(cols[0], str) and isinstance(cols[0], Sequence):
cols = tuple(cols[0])

_cols = []
for column in cols:
Expand Down
12 changes: 9 additions & 3 deletions python/pyspark/sql/dataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3360,8 +3360,14 @@ def _get_col(

orderBy = sort

@overload
def describe(self, *cols: str) -> "DataFrame": ...

@overload
def describe(self, __cols: Sequence[str]) -> "DataFrame": ...

@dispatch_df_method
def describe(self, *cols: Union[str, List[str]]) -> "DataFrame":
def describe(self, *cols: Union[str, Sequence[str]]) -> "DataFrame":
"""Computes basic statistics for numeric and string columns.

.. versionadded:: 1.3.1
Expand Down Expand Up @@ -3777,10 +3783,10 @@ def select(self, *cols: Union[Sequence["ColumnOrName"], "ColumnOrName"]) -> "Dat
def selectExpr(self, *expr: str) -> "DataFrame": ...

@overload
def selectExpr(self, *expr: List[str]) -> "DataFrame": ...
def selectExpr(self, __expr: Sequence[str]) -> "DataFrame": ...

@dispatch_df_method
def selectExpr(self, *expr: Union[str, List[str]]) -> "DataFrame":
def selectExpr(self, *expr: Union[str, Sequence[str]]) -> "DataFrame":
"""Projects a set of SQL expressions and returns a new :class:`DataFrame`.

This is a variant of :func:`select` that accepts SQL expressions.
Expand Down
9 changes: 9 additions & 0 deletions python/pyspark/sql/tests/connect/test_connect_basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -842,6 +842,15 @@ def test_select_expr(self):
.toPandas(),
)

self.assert_eq(
self.connect.read.table(self.tbl_name)
.selectExpr(("id * 2", "cast(name as long) as name"))

@HyukjinKwon HyukjinKwon Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I feel like we shouldn't allow this case explicitly. I know the support was already there but adding the corresponding type hints means that we will officially support this.

I would rather just issue a deprecated warning or sth.

@HyukjinKwon HyukjinKwon Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are there a lot of API that supports a tuple? or only selectExpr and describe are all?

@HyukjinKwon HyukjinKwon Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's your take @gaogaotiantian ?

@Spenserrrr Spenserrrr Aug 3, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you so much for taking a look, @HyukjinKwon!

This is one of the three related PRs on cleaning up the pattern that varargs also accepts a single sequence, so I think this question applies to all three as well. The other two PR: #57693, #57726.

Also, these PRs are suggested by Tian's comment on #57611, and he suggested handling this pattern at the repo level, and he pointed me to his #54764 as a reference. He widened List to Sequence and added the two overloads. So that's the pattern I followed here to make these consistent with what is already in dataframe.py.

For how widespread this is, on the current master, the single-collection form is already accepted at around a dozen places. For example, partitionBy/clusterBy on the writers and stream writers (classic and connect), and struct/create_map/array/map_concat, which takes a list or a set (and in connect a tuple too, so a tuple works in connect but not classic for the same public function today). And in dataframe.py, select/groupBy/rollup/cube/sort/sortWithinPartitions already carry the Sequence[...] annotation and the two overloads from Tian's edit.

Also, a quick summary of my PRs: #57693 only corrects annotations for the writer methods, which already accept list and tuple; this PR and #57726 widen list-only sites, which do add tuple support; and I plan to create another PR for struct/create_map/array/map_concat because a set isn't a Sequence and classic/connect disagree there.

.toPandas(),
self.spark.read.table(self.tbl_name)
.selectExpr(("id * 2", "cast(name as long) as name"))
.toPandas(),
)

def test_select_star(self):
data = [Row(a=1, b=Row(c=2, d=Row(e=3)))]

Expand Down
4 changes: 4 additions & 0 deletions python/pyspark/sql/tests/connect/test_connect_stat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,10 @@ def test_describe(self):
self.connect.read.table(self.tbl_name).describe(["id", "name"]).toPandas(),
self.spark.read.table(self.tbl_name).describe(["id", "name"]).toPandas(),
)
self.assert_eq(
self.connect.read.table(self.tbl_name).describe(("id", "name")).toPandas(),
self.spark.read.table(self.tbl_name).describe(("id", "name")).toPandas(),
)

def test_stat_cov(self):
# SPARK-41067: Test the stat.cov method
Expand Down