[SPARK-50392][PYTHON] DataFrame conversion to table argument in Spark Classic

[xinrong-meng]

What changes were proposed in this pull request?

Support DataFrame conversion to table arguments in Spark Classic, and enable UDTFs to accept table arguments in both PySpark and Scala.

Spark Connect support will be a follow-up, with the goal of completing it by the end of this month.

Why are the changes needed?

Part of SPARK-50391.
Table-Valued Functions (TVFs) and User-Defined Table Functions (UDTFs) are widely used in Spark workflows. These functions often require a table argument, which Spark internally represents as a Catalyst expression. While Spark SQL supports constructs like TABLE() for this purpose, there is no direct API in PySpark or Scala to convert a DataFrame into a table argument. So we propose to support DataFrame conversion to table arguments (in Spark Classic first), and enable UDTFs to accept table arguments in both PySpark and Scala..

Does this PR introduce any user-facing change?

Yes DataFrame conversion to table argument is supported in Spark Classic, and UDTFs accept table arguments in both PySpark and Scala.

>>> from pyspark.sql.functions import udtf
>>> from pyspark.sql import Row
>>> 
>>> @udtf(returnType="a: int")
... class TestUDTF:
...     def eval(self, row: Row):
...         if row[0] > 5:
...             yield row[0],
... 
>>> df = spark.range(8)
>>> 
>>> TestUDTF(df.asTable()).show()
+---+                                                                           
|  a|
+---+
|  6|
|  7|
+---+

>>> TestUDTF(df.asTable().partitionBy(df.id)).show()
+---+
|  a|
+---+
|  6|
|  7|
+---+

>>> TestUDTF(df.asTable().partitionBy(df.id).orderBy(df.id)).show()
+---+
|  a|
+---+
|  6|
|  7|
+---+

>>> TestUDTF(df.asTable().withSinglePartition()).show()
+---+
|  a|
+---+
|  6|
|  7|
+---+

>>> TestUDTF(df.asTable().partitionBy(df.id).withSinglePartition()).show()
Traceback (most recent call last):
...
pyspark.errors.exceptions.captured.IllegalArgumentException: Cannot call withSinglePartition() after partitionBy() has been called.

How was this patch tested?

Unit tests.

Was this patch authored or co-authored using generative AI tooling?

No.

[HyukjinKwon]

I reviewed the design, and LGTM

[ueshin]

How about TVFArgument, TableFunctionArgument or TableValuedFunctionArgument?
In SQL, this is not only for UDTF but TVFs in general, although currently there is no builtin tvf that supports table arguments.
cc @dtenedor @allisonwang-db

[xinrong-meng]

Good point! UDTFs are a type of TVFs, how about TableValuedFunctionArgument?

[allisonwang-db]

TableValuedFunctionArgument sounds good. This way we don't need to limit it to user-defined table functions.

[ueshin]

Wha't the behavior of this? In SQL, order by only is not allowed, IIRC. cc @dtenedor

[ueshin]

Similarly, what happens if, e.g., func(df.asTable().partitionBy(df.key).orderBy(df.value)).partitionBy()?

[xinrong-meng]

It is allowed here

>>> TestUDTF(df.asTable().partitionBy("id").orderBy("id").partitionBy()).show()
+---+
|  a|
+---+
|  6|
|  7|
+---+

[ueshin]

@dtenedor Wha't the behavior of this? I don't think we should allow this?

[xinrong-meng]

order by only is not supported in SQL as

[PARSE_SYNTAX_ERROR] Syntax error at or near 'ORDER'. SQLSTATE: 42601 (line 1, pos 59)

== SQL ==
SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)) ORDER BY id)
-----------------------------------------------------------^^^

Adjusted

[xinrong-meng]

Multiple partition bys are not supported as

[PARSE_SYNTAX_ERROR] Syntax error at or near 'PARTITION'. SQLSTATE: 42601 (line 1, pos 87)

== SQL ==

SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)) PARTITION BY id ORDER BY id PARTITION BY id)
---------------------------------------------------------------------------------------^^^

or

== SQL ==
SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)) PARTITION BY id PARTITION BY id)
---------------------------------------------------------------------------^^^

Adjusted.

[ueshin]

Could you add tests with named arguments?

[xinrong-meng]

def partitionBy(self, *cols: "ColumnOrName") -> "TableArg":
does not accept keyword arguments, would you clarify what we are expecting here?

[ueshin]

func(row = df.asTable() ...)

[xinrong-meng]

Got it, included named arguments

[xinrong-meng]

I created SPARK-50392 as a follow-up to support named arguments.
It requires a change in PythonSQLUtils.namedArgumentExpression, which depends on the TableArg class in Spark Connect.

[ueshin]

Otherwise, LGTM, pending #49055 (comment) and tests.

[ueshin]

Do we need this?

[xinrong-meng]

I think so, I added the line due to a test failure hint.
Let me verify here.

[xinrong-meng]

Otherwise, LGTM, pending #49055 (comment) and tests.

I appreciate your detailed review. That's very helpful!

[dongjoon-hyun]

Hi, @xinrong-meng , @HyukjinKwon , @ueshin , @allisonwang-db .

Newly added table_arg.py seems to break Spark Connect Python-only CI. Could you take a look at the failure, please?

• https://github.com/apache/spark/blob/master/.github/workflows/build_python_connect.yml
  • https://github.com/apache/spark/actions/runs/12657646176/job/35285622746

  File "/opt/hostedtoolcache/Python/3.11.11/x64/lib/python3.11/site-packages/pyspark/sql/table_arg.py", line 20, in <module>
    from pyspark.sql.classic.column import _to_java_column, _to_seq
ModuleNotFoundError: No module named 'pyspark.sql.classic'

[ueshin]

@dongjoon-hyun I submitted a follow-up PR #49472 to fix it. Thanks. cc @xinrong-meng

[dongjoon-hyun]

Thank you!