[spark] branch master updated: [SPARK-28056][PYTHON] add doc for SCALAR_ITER Pandas UDF

[me...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

meng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 1b2448b  [SPARK-28056][PYTHON] add doc for SCALAR_ITER Pandas UDF
1b2448b is described below

commit 1b2448bc10a8ee732d08fa1abae6d64ae25e3a14
Author: Xiangrui Meng <me...@databricks.com>
AuthorDate: Mon Jun 17 20:51:36 2019 -0700

    [SPARK-28056][PYTHON] add doc for SCALAR_ITER Pandas UDF
    
    ## What changes were proposed in this pull request?
    
    Add docs for `SCALAR_ITER` Pandas UDF.
    
    cc: WeichenXu123 HyukjinKwon
    
    ## How was this patch tested?
    
    Tested example code manually.
    
    Closes #24897 from mengxr/SPARK-28056.
    
    Authored-by: Xiangrui Meng <me...@databricks.com>
    Signed-off-by: Xiangrui Meng <me...@databricks.com>
---
 docs/sql-pyspark-pandas-with-arrow.md | 17 +++++++
 examples/src/main/python/sql/arrow.py | 86 +++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+)

diff --git a/docs/sql-pyspark-pandas-with-arrow.md b/docs/sql-pyspark-pandas-with-arrow.md
index 6cf280c..9cab4be 100644
--- a/docs/sql-pyspark-pandas-with-arrow.md
+++ b/docs/sql-pyspark-pandas-with-arrow.md
@@ -86,6 +86,23 @@ The following example shows how to create a scalar Pandas UDF that computes the
 </div>
 </div>
 
+### Scalar Iterator
+
+Scalar iterator (`SCALAR_ITER`) Pandas UDF is the same as scalar Pandas UDF above except that the
+underlying Python function takes an iterator of batches as input instead of a single batch and,
+instead of returning a single output batch, it yields output batches or returns an iterator of
+output batches.
+It is useful when the UDF execution requires initializing some states, e.g., loading an machine
+learning model file to apply inference to every input batch.
+
+The following example shows how to create scalar iterator Pandas UDFs:
+
+<div class="codetabs">
+<div data-lang="python" markdown="1">
+{% include_example scalar_iter_pandas_udf python/sql/arrow.py %}
+</div>
+</div>
+
 ### Grouped Map
 Grouped map Pandas UDFs are used with `groupBy().apply()` which implements the "split-apply-combine" pattern.
 Split-apply-combine consists of three steps:
diff --git a/examples/src/main/python/sql/arrow.py b/examples/src/main/python/sql/arrow.py
index c1e2d29..ede121b 100644
--- a/examples/src/main/python/sql/arrow.py
+++ b/examples/src/main/python/sql/arrow.py
@@ -86,6 +86,92 @@ def scalar_pandas_udf_example(spark):
     # $example off:scalar_pandas_udf$
 
 
+def scalar_iter_pandas_udf_example(spark):
+    # $example on:scalar_iter_pandas_udf$
+    import pandas as pd
+
+    from pyspark.sql.functions import col, pandas_udf, struct, PandasUDFType
+
+    pdf = pd.DataFrame([1, 2, 3], columns=["x"])
+    df = spark.createDataFrame(pdf)
+
+    # When the UDF is called with a single column that is not StructType,
+    # the input to the underlying function is an iterator of pd.Series.
+    @pandas_udf("long", PandasUDFType.SCALAR_ITER)
+    def plus_one(batch_iter):
+        for x in batch_iter:
+            yield x + 1
+
+    df.select(plus_one(col("x"))).show()
+    # +-----------+
+    # |plus_one(x)|
+    # +-----------+
+    # |          2|
+    # |          3|
+    # |          4|
+    # +-----------+
+
+    # When the UDF is called with more than one columns,
+    # the input to the underlying function is an iterator of pd.Series tuple.
+    @pandas_udf("long", PandasUDFType.SCALAR_ITER)
+    def multiply_two_cols(batch_iter):
+        for a, b in batch_iter:
+            yield a * b
+
+    df.select(multiply_two_cols(col("x"), col("x"))).show()
+    # +-----------------------+
+    # |multiply_two_cols(x, x)|
+    # +-----------------------+
+    # |                      1|
+    # |                      4|
+    # |                      9|
+    # +-----------------------+
+
+    # When the UDF is called with a single column that is StructType,
+    # the input to the underlying function is an iterator of pd.DataFrame.
+    @pandas_udf("long", PandasUDFType.SCALAR_ITER)
+    def multiply_two_nested_cols(pdf_iter):
+        for pdf in pdf_iter:
+            yield pdf["a"] * pdf["b"]
+
+    df.select(
+        multiply_two_nested_cols(
+            struct(col("x").alias("a"), col("x").alias("b"))
+        ).alias("y")
+    ).show()
+    # +---+
+    # |  y|
+    # +---+
+    # |  1|
+    # |  4|
+    # |  9|
+    # +---+
+
+    # In the UDF, you can initialize some states before processing batches.
+    # Wrap your code with try/finally or use context managers to ensure
+    # the release of resources at the end.
+    y_bc = spark.sparkContext.broadcast(1)
+
+    @pandas_udf("long", PandasUDFType.SCALAR_ITER)
+    def plus_y(batch_iter):
+        y = y_bc.value  # initialize states
+        try:
+            for x in batch_iter:
+                yield x + y
+        finally:
+            pass  # release resources here, if any
+
+    df.select(plus_y(col("x"))).show()
+    # +---------+
+    # |plus_y(x)|
+    # +---------+
+    # |        2|
+    # |        3|
+    # |        4|
+    # +---------+
+    # $example off:scalar_iter_pandas_udf$
+
+
 def grouped_map_pandas_udf_example(spark):
     # $example on:grouped_map_pandas_udf$
     from pyspark.sql.functions import pandas_udf, PandasUDFType


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org