Pandas UDF and Python Type Hint in Apache Spark 3.0Hyukjin Kwon Databricks Software Engineer
Hyukjin Kwon
▪ Apache Spark PMC / Committer
▪ Major Koalas contributor
▪ Databricks Software Engineer
▪ @HyukjinKwon in Github
Agenda
▪ Pandas UDFs
▪ Python Type Hints
▪ Proliferation of Pandas UDF Types
▪ New Pandas APIs with Python Type Hints ▪ Pandas UDFs ▪ Pandas Function APIs
Pandas UDFs
Pandas UDFs
▪ Apache Arrow, to exchange data between JVM and Python driver/executors with near-zero (de)serialization cost
▪ Vectorization
▪ Rich APIs in Pandas and NumPy
Pandas UDFs
from pyspark.sql.functions import pandas_udf, PandasUDFType
@pandas_udf('double', PandasUDFType.SCALAR) def pandas_plus_one(v): # `v` is a pandas Series return v.add(1) # outputs a pandas Series
spark.range(10).select(pandas_plus_one("id")).show()
Scalar Pandas UDF example that adds one
Spark�DataFrame
Spark�Columns
ValuePandas�Series�in�Pandas�UDF
Pandas UDFs
Spark Executor Spark Executor Spark Executor
Python Worker Python Worker Python Worker
Pandas UDFs
Partition
Partition
Partition
Partition
Partition
Partition
Spark DataFrame
Pandas UDFs
Partition
Partition
Partition
Partition
Partition
Partition
Arrow Batch Arrow Batch Arrow Batch
Near-zero (de)serialization
Pandas UDFs
Partition
Partition
Partition
Partition
Partition
Partition
Arrow Batch
Pandas Series
Arrow Batch Arrow Batch
Pandas Series Pandas Series
def pandas_plus_one(v): # `v` is a pandas Series return v.add(1) # outputs a pandas Series
Vectorized execution
Performance
Performance: Pandas UDF vs regular UDF
Python Type Hints
Python Type Hints
def greeting(name): return 'Hello ' + name
Typical Python codes
def greeting(name: str) -> str: return 'Hello ' + name
Python codes with type hints
Python Type Hints
▪ PEP 484 ▪ Standard syntax for type annotations in Python 3 ▪ Optional
▪ Static analysis ▪ IDE can automatically detects and reports the type mismatch ▪ Static analysis such as mypy ▪ Easier to refactor codes
▪ Runtime type checking and code generation ▪ Infer the type of codes to run ▪ Runtime type checking
IDE Support
def merge( self, right: "DataFrame", how: str = "inner", ...
Python type hint support in IDE
Static Analysis and Documentation
databricks/koalas/frame.py: note: In member "join" of class "DataFrame": databricks/koalas/frame.py:7546: error: Argument "how" to "merge" of "DataFrame" has incompatible type "int"; expected "str" Found 1 error in 1 file (checked 65 source files)
mypy static analysis
Auto-documentation
Python Type Hints
▪ Early but still growing ▪ Arguably still premature ▪ Type hinting APIs are still being changed and under development.
▪ Started being used in production ▪ Type hinting is being encouraged, and being used in production
▪ PySpark type hints support, pyspark-stubs ▪ Third-party, optional PySpark type hinting support.
Proliferation of Pandas UDF Types
Pandas UDFs in Apache Spark 2.4
▪ Scalar Pandas UDF ▪ Transforms Pandas Series to Pandas Series and returns a Spark Column ▪ The same length of the input and output
▪ Grouped Map Pandas UDF ▪ Splits each group as a Pandas DataFrame, applies a function on each, and combines as a Spark DataFrame ▪ The function takes a Pandas DataFrame and returns a Pandas DataFrame
▪ Grouped Aggregate Pandas UDF ▪ Splits each group as a Pandas Series, applies a function on each, and combines as a Spark Column ▪ The function takes a Pandas Series and returns single aggregated scalar value
Pandas UDFs proposed in Apache Spark 3.0
▪ Scalar Iterator Pandas UDF ▪ Transforms an iterator of Pandas Series to an iterator Pandas Series and returns a Spark Column
▪ Map Pandas UDF ▪ Transforms an iterator of Pandas DataFrame to an iterator of Pandas DataFrame in a Spark DataFrame
▪ Cogrouped Map Pandas UDF ▪ Splits each cogroup as a Pandas DataFrame, applies a function on each, and combines as a Spark DataFrame ▪ The function takes and returns a Pandas DataFrame
Complexity and Confusion @pandas_udf("long", PandasUDFType.SCALAR) def pandas_plus_one(v):
return v + 1
spark.range(3).select(pandas_plus_one("id").alias("id")).show()
@pandas_udf("long", PandasUDFType.SCALAR_ITER) def pandas_plus_one(vv):
return map(lambda v: v + 1, vv)
spark.range(3).select(pandas_plus_one("id").alias("id")).show()
@pandas_udf("id long", PandasUDFType.GROUPED_MAP) def pandas_plus_one(v):
return v + 1
spark.range(3).groupby("id").apply(pandas_plus_one).show()
+---+ | id| +---+ | 1| | 2| | 3| +---+
Same output
Adds one
Complexity and Confusion @pandas_udf("long", PandasUDFType.SCALAR) def pandas_plus_one(v): # `v` is a pandas Series return v + 1 # outputs a pandas Series
spark.range(3).select(pandas_plus_one("id").alias("id")).show()
@pandas_udf("long", PandasUDFType.SCALAR_ITER) def pandas_plus_one(vv): # `vv` is an iterator of pandas Series. # outputs an iterator of pandas Series. return map(lambda v: v + 1, vv)
spark.range(3).select(pandas_plus_one("id").alias("id")).show()
@pandas_udf("id long", PandasUDFType.GROUPED_MAP) def pandas_plus_one(v): # `v` is a pandas DataFrame return v + 1 # outputs a pandas DataFrame
spark.range(3).groupby("id").apply(pandas_plus_one).show()
▪ What types are expected in the function?
▪ How does each UDF work?
▪ Why should I specify the UDF type?
Adds one
Complexity and Confusion @pandas_udf("long", PandasUDFType.SCALAR) def pandas_plus_one(v): return v + 1
df = spark.range(3) df.select(pandas_plus_one("id") + cos("id")).show()
@pandas_udf("id long", PandasUDFType.GROUPED_MAP) def pandas_plus_one(v): return v + 1
df = spark.range(3)
df.groupby("id").apply(pandas_plus_one("id") + col(“id")).show()
Adds one and cosine
Adds one and cosine(?)
Traceback (most recent call last): File "<stdin>", line 1, in <module> File "...", line 70, in apply ... ValueError: Invalid udf: the udf argument must be a pandas_udf of type GROUPED_MAP.
+-------------------------------+ |(pandas_plus_one(id) + COS(id))| +-------------------------------+ | 2.0| | 2.5403023058681398| | 2.5838531634528574| +-------------------------------+
Complexity and Confusion @pandas_udf("long", PandasUDFType.SCALAR) def pandas_plus_one(v): return v + 1
df = spark.range(3) df.select(pandas_plus_one("id") + cos("id")).show()
@pandas_udf("id long", PandasUDFType.GROUPED_MAP) def pandas_plus_one(v): return v + 1
df = spark.range(3) # `pandas_plus_one` can _only_ be used with `groupby(...).apply(...)` df.groupby("id").apply(pandas_plus_one("id") + col("id")).show()
Adds one and cosine
Adds one and cosine(?)
▪ Expression
▪ Query execution plan
New Pandas APIs with Python Type Hints
Python Type Hints@pandas_udf("long") def pandas_plus_one(v: pd.Series) -> pd.Series: return v + 1
spark.range(3).select(pandas_plus_one("id").alias("id")).show()
@pandas_udf("long") def pandas_plus_one(vv: Iterator[pd.Series]) -> Iterator[pd.Series]: return map(lambda v: v + 1, vv)
spark.range(3).select(pandas_plus_one("id").alias("id")).show()
@pandas_udf("id long") def pandas_plus_one(v: pd.DataFrame) -> pd.DataFrame: return v + 1
spark.range(3).groupby("id").apply(pandas_plus_one).show()
▪ Self-descriptive ▪ Describe what the pandas UDF is supposed to take and
return. ▪ Shows the relationship between input and output.
▪ Static analysis ▪ IDE detects if non-pandas instances are used mistakenly. ▪ Other tools such as mypy can be integrated for a better
code quality in the pandas UDFs.
▪ Auto-documentation ▪ Type hints in the pandas UDF automatically documents the
input and output.
▪ Pandas UDFs ▪ Works as a function, internally an expression ▪ Consistent with Scala UDFs and regular Python UDFs ▪ Returns a regular PySpark column
▪ Pandas Function APIs ▪ Works as an API in DataFrame, query plan internally ▪ Consistent with APIs such as map, mapGroups, etc.
API Separation
@pandas_udf("long") def pandas_plus_one(v: pd.Series) -> pd.Series: return v + 1
df = spark.range(3) df.select(pandas_plus_one("id") + cos("id")).show()
def pandas_plus_one(v: pd.DataFrame) -> pd.DataFrame: return v + 1
df = spark.range(3) df.groupby("id").applyInPandas(pandas_plus_one).show()
▪ Series to Series ▪ A Pandas UDF ▪ pandas.Series, ... -> pandas.Series ▪ Length of each input series and output series should be the same ▪ StructType in input and output is represented via pandas.DataFrame
New Pandas UDFs
import pandas as pd from pyspark.sql.functions import pandas_udf
@pandas_udf('long') def pandas_plus_one(s: pd.Series) -> pd.Series: return s + 1
spark.range(10).select(pandas_plus_one("id")).show()
New Style
from pyspark.sql.functions import pandas_udf, PandasUDFType
@pandas_udf('long', PandasUDFType.SCALAR) def pandas_plus_one(v): return v + 1
spark.range(10).select(pandas_plus_one("id")).show()
Old Style (Scalar Pandas UDF)
New Pandas UDFs
▪ Iterator of Series to Iterator of Series ▪ A Pandas UDF ▪ Iterator[pd.Series] -> Iterator[pd.Series] ▪ Length of the whole input iterator and output iterator should be the same ▪ StructType in input and output is represented via pandas.DataFrame
from typing import Iterator import pandas as pd from pyspark.sql.functions import pandas_udf
@pandas_udf('long') def pandas_plus_one(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]: return map(lambda s: s + 1, iterator)
spark.range(10).select(pandas_plus_one("id")).show()
New Style
from pyspark.sql.functions import pandas_udf, PandasUDFType
@pandas_udf('long', PandasUDFType.SCALAR_ITER) def pandas_plus_one(iterator): return map(lambda s: s + 1, iterator)
spark.range(10).select(pandas_plus_one("id")).show()
Old Style (Scalar Iterator Pandas UDF)
New Pandas UDFs
▪ Iterator of Multiple Series to Iterator of Series ▪ A Pandas UDF ▪ Iterator[Tuple[pandas.Series, ...]] -> Iterator[pandas.Series] ▪ Length of the whole input iterator and output iterator should be the same ▪ StructType in input and output is represented via pandas.DataFrame
from typing import Iterator, Tuple import pandas as pd from pyspark.sql.functions import pandas_udf
@pandas_udf("long") def multiply_two( iterator: Iterator[Tuple[pd.Series, pd.Series]]) -> Iterator[pd.Series]: return (a * b for a, b in iterator)
spark.range(10).select(multiply_two("id", "id")).show()
New Style
from pyspark.sql.functions import pandas_udf, PandasUDFType
@pandas_udf('long', PandasUDFType.SCALAR_ITER) def multiply_two(iterator): return (a * b for a, b in iterator)
spark.range(10).select(multiply_two("id", "id")).show()
Old Style (Scalar Iterator Pandas UDF)
New Pandas UDFs
▪ Iterator of Series to Iterator of Series ▪ Iterator of Multiple Series to Iterator of Series
▪ Useful when it requires to execute to calculate one expensive state to share ▪ Prefetch the data within the iterator
@pandas_udf("long") def calculate(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]: # Do some expensive initialization with a state state = very_expensive_initialization() for x in iterator: # Use that state for the whole iterator. yield calculate_with_state(x, state)
df.select(calculate("value")).show()
Initializing a expensive state
@pandas_udf("long") def calculate(iterator: Iterator[pd.Series]) -> Iterator[pd.Series]:
# Pre-fetch the iterator threading.Thread(consume, args=(iterator, queue)) for s in queue: yield func(s)
df.select(calculate("value")).show()
Pre-fetching input iterator
New Pandas UDFs
▪ Series to Scalar ▪ A Pandas UDF ▪ pandas.Series, ... -> Any (any scalar value) ▪ Should output a scalar value a Python primitive type such as int, or NumPy data type such as numpy.int64.
Any should ideally be a specific scalar type accordingly ▪ StructType in input is represented via pandas.DataFrame ▪ Typically assumes an aggregation
import pandas as pd from pyspark.sql.functions import pandas_udf
df = spark.createDataFrame( [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
@pandas_udf("double") def pandas_mean(v: pd.Series) -> float: return v.sum()
df.select(pandas_mean(df['v'])).show()
New Style
import pandas as pd from pyspark.sql.functions import pandas_udf, PandasUDFType
df = spark.createDataFrame( [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
@pandas_udf("double", PandasUDFType.GROUPED_AGG) def pandas_mean(v): return v.sum()
df.select(pandas_mean(df['v'])).show()
Old Style (Grouped Aggregate Pandas UDF)
Pandas Function APIs: Grouped Map
▪ Grouped Map ▪ A Pandas Function API that applies a function on each group ▪ Optional Python type hints currently in Spark 3.0 ▪ Length of output can be arbitrary ▪ StructType is unsupported
import pandas as pd
df = spark.createDataFrame( [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
def subtract_mean(pdf: pd.DataFrame) -> pd.DataFrame: v = pdf.v return pdf.assign(v=v - v.mean())
df.groupby(“id").applyInPandas(subtract_mean, df.schema).show()
New Style
import pandas as pd from pyspark.sql.functions import pandas_udf, PandasUDFType
df = spark.createDataFrame( [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "v"))
@pandas_udf(df.schema, PandasUDFType.GROUPED_MAP) def subtract_mean(pdf): v = pdf.v return pdf.assign(v=v - v.mean())
df.groupby("id").apply(subtract_mean).show()
Old Style (Grouped Map Pandas UDF)
Pandas Function APIs: Grouped Map
▪ Map ▪ A Pandas Function API that applies a function on the Spark DataFrame ▪ Similar characteristics with the iterator support of Python UDF ▪ Optional Python type hints currently in Spark 3.0 ▪ Length of output can be arbitrary ▪ StructType is unsupported
from typing import Iterator import pandas as pd
df = spark.createDataFrame([(1, 21), (2, 30)], ("id", "age"))
def pandas_filter(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]: for pdf in iterator: yield pdf[pdf.id == 1]
df.mapInPandas(pandas_filter, df.schema).show()
Pandas Function APIs: Grouped Map
▪ Co-grouped Map ▪ A Pandas Function API that applies a function on each co-group ▪ Requires two grouped Spark DataFrames ▪ Optional Python type hints currently in Spark 3.0 ▪ Length of output can be arbitrary ▪ StructType is unsupported
import pandas as pd
df1 = spark.createDataFrame( [(1201, 1, 1.0), (1201, 2, 2.0), (1202, 1, 3.0), (1202, 2, 4.0)], ("time", "id", "v1")) df2 = spark.createDataFrame( [(1201, 1, "x"), (1201, 2, "y")], ("time", "id", "v2"))
def asof_join(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame: return pd.merge_asof(left, right, on="time", by="id")
df1.groupby("id").cogroup( df2.groupby("id") ).applyInPandas(asof_join, "time int, id int, v1 double, v2 string").show()
Re-cap
▪ Pandas APIs leverage Python type hints for static analysis, auto-documentation and self-descriptive UDF
▪ Old Pandas UDFs separation to Pandas UDF and Pandas Function API
▪ New APIs ▪ Iterator support in Pandas UDF ▪ Cogrouped-map and map Pandas Function APIs
Questions?