ueshin commented on code in PR #42272: URL: https://github.com/apache/spark/pull/42272#discussion_r1302130261
########## examples/src/main/python/sql/udtf.py: ########## @@ -0,0 +1,230 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +A simple example demonstrating Python UDTFs in Spark +Run with: + ./bin/spark-submit examples/src/main/python/sql/udtf.py +""" + +# NOTE that this file is imported in the User Guides in PySpark documentation. +# The codes are referred via line numbers. See also `literalinclude` directive in Sphinx. +from pyspark.sql import SparkSession +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version + +# Python UDTFs use Arrow by default. +require_minimum_pandas_version() +require_minimum_pyarrow_version() + + +def python_udtf_simple_example(spark: SparkSession) -> None: + + # Define the UDTF class and implement the required `eval` method. + class SquareNumbers: + def eval(self, start: int, end: int): # type: ignore[no-untyped-def] + for num in range(start, end + 1): + yield (num, num * num) + + from pyspark.sql.functions import lit, udtf + + # Create a UDTF using the class definition and the `udtf` function. + square_num = udtf(SquareNumbers, returnType="num: int, squared: int") + + # Invoke the UDTF in PySpark. + square_num(lit(1), lit(3)).show() # type: ignore + # +---+------+ + # |num|squred| + # +---+------+ + # | 1| 1| + # | 2| 4| + # | 3| 9| + # +---+------+ + + +def python_udtf_decorator_example(spark: SparkSession) -> None: + + from pyspark.sql.functions import lit, udtf + + # Define a UDTF using the `udtf` decorator directly on the class. + @udtf(returnType="num: int, squared: int") # type: ignore + class SquareNumbers: + def eval(self, start: int, end: int): # type: ignore[no-untyped-def] + for num in range(start, end + 1): + yield (num, num * num) + + # Invoke the UDTF in PySpark using the SquareNumbers class directly. + SquareNumbers(lit(1), lit(3)).show() # type: ignore + # +---+------+ + # |num|squred| Review Comment: ditto. ########## examples/src/main/python/sql/udtf.py: ########## @@ -0,0 +1,230 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +A simple example demonstrating Python UDTFs in Spark +Run with: + ./bin/spark-submit examples/src/main/python/sql/udtf.py +""" + +# NOTE that this file is imported in the User Guides in PySpark documentation. +# The codes are referred via line numbers. See also `literalinclude` directive in Sphinx. +from pyspark.sql import SparkSession +from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version + +# Python UDTFs use Arrow by default. +require_minimum_pandas_version() +require_minimum_pyarrow_version() + + +def python_udtf_simple_example(spark: SparkSession) -> None: + + # Define the UDTF class and implement the required `eval` method. + class SquareNumbers: + def eval(self, start: int, end: int): # type: ignore[no-untyped-def] + for num in range(start, end + 1): + yield (num, num * num) + + from pyspark.sql.functions import lit, udtf + + # Create a UDTF using the class definition and the `udtf` function. + square_num = udtf(SquareNumbers, returnType="num: int, squared: int") + + # Invoke the UDTF in PySpark. + square_num(lit(1), lit(3)).show() # type: ignore + # +---+------+ + # |num|squred| Review Comment: The column name seems to be wrong? `squred` -> `squared` ########## python/docs/source/user_guide/sql/python_udtf.rst: ########## @@ -0,0 +1,215 @@ +.. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + +.. http://www.apache.org/licenses/LICENSE-2.0 + +.. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +=========================================== +Python User-defined Table Functions (UDTFs) +=========================================== + +Spark 3.5 introduces Python user-defined table functions (UDTFs), a new type of user-defined function. +Unlike scalar functions that return a single result value for every input, UDTFs is invoked in the `FROM` +clause of a SQL query and returns an entire relation (table) as output. +Each UDTF call can accept zero or more arguments. +These arguments can either be scalar expressions or TABLE arguments that represent entire input relations. + +Implementing a Python UDTF +-------------------------- + +.. currentmodule:: pyspark.sql.functions + +To implement a Python UDTF, you can define a class implementing the methods: + +.. code-block:: python + + class PythonUDTF: + + def __init__(self) -> None: + """ + Initialize the user-defined table function (UDTF). + + This method serves as the default constructor and is called once when the + UDTF is instantiated on the executor side. + + Any class fields assigned in this method will be available for subsequent + calls to the `eval` and `terminate` methods. + + Notes + ----- + - This method does not accept any extra arguments. + - You cannot create or reference the Spark session within the UDTF. Any + attempt to do so will result in a serialization error. + """ + ... + + def eval(self, *args: Any) -> Iterator[Any]: + """ + Evaluate the function using the given input arguments. + + This method is required and must be implemented. + + Argument Mapping: + - Each provided scalar expression maps to exactly one value in the `*args` list. + - Each provided table argument maps to a pyspark.sql.Row object containing the columns + in the order they appear in the provided input relation. + + This method is called on every input row, and can produce zero or more + output rows. Each element in the output tuple corresponds to one column + specified in the return type of the UDTF. + + Parameters + ---------- + *args : Any + Arbitrary positional arguments representing the input to the UDTF. + + Yields + ------ + tuple + A tuple representing a single row in the UDTF result relation. + Yield as many times as needed to produce multiple rows. + + Notes + ----- + - The result of the function must be a tuple representing a single row + in the UDTF result relation. + - UDTFs currently do not accept keyword arguments during the function call. + + Examples + -------- + >>> def eval(self, x: int, y: int) -> Iterator[Any]: + >>> yield x + y, x - y + >>> yield y + x, y - x + """ + ... + + def terminate(self) -> Iterator[Any]: + """ + Called when the UDTF has processed all input rows. + + This method is optional to implement and is useful for performing any + cleanup or finalization operations after the UDTF has finished processing + all rows. It can also be used to yield additional rows if needed. + + Yields + ------ + tuple + A tuple representing a single row in the UDTF result relation. + Yield this if you want to return additional rows during termination. + + Examples + -------- + >>> def terminate(self) -> Iterator[Any]: + >>> yield "done", None + """ + ... + + +The return type of the UDTF defines the schema of the table it outputs. +It must be either a ``StructType`` or a DDL string representing a struct type. + +**Example UDTF Implementation:** + +Here is a simple example of a UDTF implementation: + +.. literalinclude:: ../../../../../examples/src/main/python/sql/udtf.py + :language: python + :lines: 38-42 Review Comment: We need to revisit the line numbers? I guess this is 36-40. And ditto for the following examples. -- 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]
