[SPARK-50133][PYTHON] Support df.argument() for conversion to table argument in Spark Classic

connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -587,6 +587,11 @@ class Dataset[T] private[sql] (
 
   // TODO(SPARK-50134): Support scalar Subquery API in Spark Connect
   // scalastyle:off not.implemented.error.usage
+  /** @inheritdoc */
+  def argument(): Column = {
+    ???
+  }
+
   /** @inheritdoc */
   def scalar(): Column = {
     ???

python/pyspark/sql/classic/dataframe.py

@@ -1786,6 +1786,9 @@ def transpose(self, indexColumn: Optional["ColumnOrName"] = None) -> ParentDataF
         else:
             return DataFrame(self._jdf.transpose(), self.sparkSession)
 
+    def argument(self) -> Column:
+        return Column(self._jdf.argument())
+
     def scalar(self) -> Column:
         return Column(self._jdf.scalar())
 

python/pyspark/sql/connect/dataframe.py

@@ -1784,6 +1784,12 @@ def transpose(self, indexColumn: Optional["ColumnOrName"] = None) -> ParentDataF
             self._session,
         )
 
+    def argument(self) -> Column:
+        raise PySparkNotImplementedError(
+            errorClass="NOT_IMPLEMENTED",
+            messageParameters={"feature": "argument()"},
+        )
+
     def scalar(self) -> Column:
         # TODO(SPARK-50134): Implement this method
         raise PySparkNotImplementedError(

python/pyspark/sql/dataframe.py

@@ -6476,6 +6476,39 @@ def transpose(self, indexColumn: Optional["ColumnOrName"] = None) -> "DataFrame"
         """
         ...
 
+    def argument(self) -> Column:
+        """
+        Converts the DataFrame into a `Column` object for use with table-valued functions (TVFs)
+        or user-defined table functions (UDTFs).
+
+        .. versionadded:: 4.0.0
+
+        Returns
+        -------
+        :class:`Column`
+            A `Column` object representing the DataFrame.
+
+        Examples
+        --------
+        >>> from pyspark.sql import Row
+        >>> from pyspark.sql.functions import udtf
+        >>>
+        >>> @udtf(returnType="a: int")
+        ... class TestUDTF:
+        ...     def eval(self, row: Row):
+        ...         if row[0] > 5:
+        ...             yield row[0],
+        >>> df = spark.range(8)
+        >>> TestUDTF(df.argument()).show()  # doctest: +SKIP
+        +---+
+        |  a|
+        +---+
+        |  6|
+        |  7|
+        +---+
+        """
+        ...
+
     def scalar(self) -> Column:
         """
         Return a `Column` object for a SCALAR Subquery containing exactly one row and one column.

python/pyspark/sql/tests/connect/test_parity_udtf.py

@@ -76,6 +76,10 @@ def test_udtf_with_analyze_using_file(self):
     def test_udtf_access_spark_session(self):
         super().test_udtf_access_spark_session()
 
+    @unittest.skip("Spark Connect does not support df.argument()")
+    def test_df_argument(self):
+        super().test_df_argument()
+
     def _add_pyfile(self, path):
         self.spark.addArtifacts(path, pyfile=True)
 

python/pyspark/sql/tests/test_udtf.py

@@ -1031,6 +1031,19 @@ def eval(self, row: Row):
             [Row(a=6), Row(a=7)],
         )
 
+    def test_df_argument(self):
+        class TestUDTF:
+            def eval(self, row: Row):
+                if row["id"] > 5:
+                    yield row["id"],
+
+        func = udtf(TestUDTF, returnType="a: int")
+        df = self.spark.range(8)
+        self.assertEqual(
+            func(df.argument()).collect(),
+            [Row(a=6), Row(a=7)],
+        )
+
     def test_udtf_with_int_and_table_argument_query(self):
         class TestUDTF:
             def eval(self, i: int, row: Row):

sql/api/src/main/scala/org/apache/spark/sql/api/Dataset.scala

@@ -1699,6 +1699,15 @@ abstract class Dataset[T] extends Serializable {
    */
   def transpose(): Dataset[Row]
 
+  /**
+   * Converts the DataFrame into a `Column` object for use with table-valued functions (TVFs) or
+   * user-defined table functions (UDTFs).
+   *
+   * @group typedrel
+   * @since 4.0.0
+   */
+  def argument(): Column
+
   /**
    * Return a `Column` object for a SCALAR Subquery containing exactly one row and one column.
    *

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -1003,6 +1003,12 @@ class Dataset[T] private[sql](
     )
   }
 
+  /** @inheritdoc */
+  def argument(): Column = {
+    val tableExpr = FunctionTableSubqueryArgumentExpression(logicalPlan)
+    Column(tableExpr)
+  }
+
   /** @inheritdoc */
   def scalar(): Column = {
     Column(ExpressionColumnNode(