ab0890bdb1
### What changes were proposed in this pull request? This PR proposes to redesign pandas UDFs as described in [the proposal](https://docs.google.com/document/d/1-kV0FS_LF2zvaRh_GhkV32Uqksm_Sq8SvnBBmRyxm30/edit?usp=sharing). ```python from pyspark.sql.functions import pandas_udf import pandas as pd pandas_udf("long") def plug_one(s: pd.Series) -> pd.Series: return s + 1 spark.range(10).select(plug_one("id")).show() ``` ``` +------------+ |plug_one(id)| +------------+ | 1| | 2| | 3| | 4| | 5| | 6| | 7| | 8| | 9| | 10| +------------+ ``` Note that, this PR address one of the future improvements described [here](https://docs.google.com/document/d/1-kV0FS_LF2zvaRh_GhkV32Uqksm_Sq8SvnBBmRyxm30/edit#heading=h.h3ncjpk6ujqu), "A couple of less-intuitive pandas UDF types" (by zero323) together. In short, - Adds new way with type hints as an alternative and experimental way. ```python pandas_udf(schema='...') def func(c1: Series, c2: Series) -> DataFrame: pass ``` - Replace and/or add an alias for three types below from UDF, and make them as separate standalone APIs. So, `pandas_udf` is now consistent with regular `udf`s and other expressions. `df.mapInPandas(udf)` -replace-> `df.mapInPandas(f, schema)` `df.groupby.apply(udf)` -alias-> `df.groupby.applyInPandas(f, schema)` `df.groupby.cogroup.apply(udf)` -replace-> `df.groupby.cogroup.applyInPandas(f, schema)` *`df.groupby.apply` was added from 2.3 while the other were added in the master only. - No deprecation for the existing ways for now. ```python pandas_udf(schema='...', functionType=PandasUDFType.SCALAR) def func(c1, c2): pass ``` If users are happy with this, I plan to deprecate the existing way and declare using type hints is not experimental anymore. One design goal in this PR was that, avoid touching the internal (since we didn't deprecate the old ways for now), but supports type hints with a minimised changes only at the interface. - Once we deprecate or remove the old ways, I think it requires another refactoring for the internal in the future. At the very least, we should rename internal pandas evaluation types. - If users find this experimental type hints isn't quite helpful, we should simply revert the changes at the interface level. ### Why are the changes needed? In order to address old design issues. Please see [the proposal](https://docs.google.com/document/d/1-kV0FS_LF2zvaRh_GhkV32Uqksm_Sq8SvnBBmRyxm30/edit?usp=sharing). ### Does this PR introduce any user-facing change? For behaviour changes, No. It adds new ways to use pandas UDFs by using type hints. See below. **SCALAR**: ```python pandas_udf(schema='...') def func(c1: Series, c2: DataFrame) -> Series: pass # DataFrame represents a struct column ``` **SCALAR_ITER**: ```python pandas_udf(schema='...') def func(iter: Iterator[Tuple[Series, DataFrame, ...]]) -> Iterator[Series]: pass # Same as SCALAR but wrapped by Iterator ``` **GROUPED_AGG**: ```python pandas_udf(schema='...') def func(c1: Series, c2: DataFrame) -> int: pass # DataFrame represents a struct column ``` **GROUPED_MAP**: This was added in Spark 2.3 as of SPARK-20396. As described above, it keeps the existing behaviour. Additionally, we now have a new alias `groupby.applyInPandas` for `groupby.apply`. See the example below: ```python def func(pdf): return pdf df.groupby("...").applyInPandas(func, schema=df.schema) ``` **MAP_ITER**: this is not a pandas UDF anymore This was added in Spark 3.0 as of SPARK-28198; and this PR replaces the usages. See the example below: ```python def func(iter): for df in iter: yield df df.mapInPandas(func, df.schema) ``` **COGROUPED_MAP**: this is not a pandas UDF anymore This was added in Spark 3.0 as of SPARK-27463; and this PR replaces the usages. See the example below: ```python def asof_join(left, right): return pd.merge_asof(left, right, on="...", by="...") df1.groupby("...").cogroup(df2.groupby("...")).applyInPandas(asof_join, schema="...") ``` ### How was this patch tested? Unittests added and tested against Python 2.7, 3.6 and 3.7. Closes #27165 from HyukjinKwon/revisit-pandas. Authored-by: HyukjinKwon <gurwls223@apache.org> Signed-off-by: HyukjinKwon <gurwls223@apache.org>
273 lines
11 KiB
Python
273 lines
11 KiB
Python
#
|
|
# 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.
|
|
#
|
|
import sys
|
|
import warnings
|
|
|
|
from pyspark import since
|
|
from pyspark.rdd import PythonEvalType
|
|
from pyspark.sql.column import Column
|
|
from pyspark.sql.dataframe import DataFrame
|
|
|
|
|
|
class PandasGroupedOpsMixin(object):
|
|
"""
|
|
Min-in for pandas grouped operations. Currently, only :class:`GroupedData`
|
|
can use this class.
|
|
"""
|
|
|
|
@since(2.3)
|
|
def apply(self, udf):
|
|
"""
|
|
It is an alias of :meth:`pyspark.sql.GroupedData.applyInPandas`; however, it takes a
|
|
:meth:`pyspark.sql.functions.pandas_udf` whereas
|
|
:meth:`pyspark.sql.GroupedData.applyInPandas` takes a Python native function.
|
|
|
|
.. note:: It is preferred to use :meth:`pyspark.sql.GroupedData.applyInPandas` over this
|
|
API. This API will be deprecated in the future releases.
|
|
|
|
:param udf: a grouped map user-defined function returned by
|
|
:func:`pyspark.sql.functions.pandas_udf`.
|
|
|
|
>>> 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("id long, v double", PandasUDFType.GROUPED_MAP) # doctest: +SKIP
|
|
... def normalize(pdf):
|
|
... v = pdf.v
|
|
... return pdf.assign(v=(v - v.mean()) / v.std())
|
|
>>> df.groupby("id").apply(normalize).show() # doctest: +SKIP
|
|
+---+-------------------+
|
|
| id| v|
|
|
+---+-------------------+
|
|
| 1|-0.7071067811865475|
|
|
| 1| 0.7071067811865475|
|
|
| 2|-0.8320502943378437|
|
|
| 2|-0.2773500981126146|
|
|
| 2| 1.1094003924504583|
|
|
+---+-------------------+
|
|
|
|
.. seealso:: :meth:`pyspark.sql.functions.pandas_udf`
|
|
|
|
"""
|
|
# Columns are special because hasattr always return True
|
|
if isinstance(udf, Column) or not hasattr(udf, 'func') \
|
|
or udf.evalType != PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF:
|
|
raise ValueError("Invalid udf: the udf argument must be a pandas_udf of type "
|
|
"GROUPED_MAP.")
|
|
|
|
warnings.warn(
|
|
"It is preferred to use 'applyInPandas' over this "
|
|
"API. This API will be deprecated in the future releases. See SPARK-28264 for "
|
|
"more details.", UserWarning)
|
|
|
|
return self.applyInPandas(udf.func, schema=udf.returnType)
|
|
|
|
@since(3.0)
|
|
def applyInPandas(self, func, schema):
|
|
"""
|
|
Maps each group of the current :class:`DataFrame` using a pandas udf and returns the result
|
|
as a `DataFrame`.
|
|
|
|
The function should take a `pandas.DataFrame` and return another
|
|
`pandas.DataFrame`. For each group, all columns are passed together as a `pandas.DataFrame`
|
|
to the user-function and the returned `pandas.DataFrame` are combined as a
|
|
:class:`DataFrame`.
|
|
|
|
The returned `pandas.DataFrame` can be of arbitrary length and its schema must match the
|
|
returnType of the pandas udf.
|
|
|
|
.. note:: This function requires a full shuffle. All the data of a group will be loaded
|
|
into memory, so the user should be aware of the potential OOM risk if data is skewed
|
|
and certain groups are too large to fit in memory.
|
|
|
|
:param func: a Python native function that takes a `pandas.DataFrame`, and outputs a
|
|
`pandas.DataFrame`.
|
|
:param schema: the return type of the `func` in PySpark. The value can be either a
|
|
:class:`pyspark.sql.types.DataType` object or a DDL-formatted type string.
|
|
|
|
.. note:: Experimental
|
|
|
|
>>> 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"))
|
|
>>> def normalize(pdf):
|
|
... v = pdf.v
|
|
... return pdf.assign(v=(v - v.mean()) / v.std())
|
|
>>> df.groupby("id").applyInPandas(normalize, schema="id long, v double").show()
|
|
... # doctest: +SKIP
|
|
+---+-------------------+
|
|
| id| v|
|
|
+---+-------------------+
|
|
| 1|-0.7071067811865475|
|
|
| 1| 0.7071067811865475|
|
|
| 2|-0.8320502943378437|
|
|
| 2|-0.2773500981126146|
|
|
| 2| 1.1094003924504583|
|
|
+---+-------------------+
|
|
|
|
.. seealso:: :meth:`pyspark.sql.functions.pandas_udf`
|
|
|
|
"""
|
|
from pyspark.sql import GroupedData
|
|
from pyspark.sql.functions import pandas_udf, PandasUDFType
|
|
|
|
assert isinstance(self, GroupedData)
|
|
|
|
udf = pandas_udf(
|
|
func, returnType=schema, functionType=PandasUDFType.GROUPED_MAP)
|
|
df = self._df
|
|
udf_column = udf(*[df[col] for col in df.columns])
|
|
jdf = self._jgd.flatMapGroupsInPandas(udf_column._jc.expr())
|
|
return DataFrame(jdf, self.sql_ctx)
|
|
|
|
@since(3.0)
|
|
def cogroup(self, other):
|
|
"""
|
|
Cogroups this group with another group so that we can run cogrouped operations.
|
|
|
|
See :class:`CoGroupedData` for the operations that can be run.
|
|
"""
|
|
from pyspark.sql import GroupedData
|
|
|
|
assert isinstance(self, GroupedData)
|
|
|
|
return PandasCogroupedOps(self, other)
|
|
|
|
|
|
class PandasCogroupedOps(object):
|
|
"""
|
|
A logical grouping of two :class:`GroupedData`,
|
|
created by :func:`GroupedData.cogroup`.
|
|
|
|
.. note:: Experimental
|
|
|
|
.. versionadded:: 3.0
|
|
"""
|
|
|
|
def __init__(self, gd1, gd2):
|
|
self._gd1 = gd1
|
|
self._gd2 = gd2
|
|
self.sql_ctx = gd1.sql_ctx
|
|
|
|
@since(3.0)
|
|
def applyInPandas(self, func, schema):
|
|
"""
|
|
Applies a function to each cogroup using pandas and returns the result
|
|
as a `DataFrame`.
|
|
|
|
The function should take two `pandas.DataFrame`\\s and return another
|
|
`pandas.DataFrame`. For each side of the cogroup, all columns are passed together as a
|
|
`pandas.DataFrame` to the user-function and the returned `pandas.DataFrame` are combined as
|
|
a :class:`DataFrame`.
|
|
|
|
The returned `pandas.DataFrame` can be of arbitrary length and its schema must match the
|
|
returnType of the pandas udf.
|
|
|
|
.. note:: This function requires a full shuffle. All the data of a cogroup will be loaded
|
|
into memory, so the user should be aware of the potential OOM risk if data is skewed
|
|
and certain groups are too large to fit in memory.
|
|
|
|
.. note:: Experimental
|
|
|
|
:param func: a Python native function that takes two `pandas.DataFrame`\\s, and
|
|
outputs a `pandas.DataFrame`, or that takes one tuple (grouping keys) and two
|
|
pandas ``DataFrame``s, and outputs a pandas ``DataFrame``.
|
|
:param schema: the return type of the `func` in PySpark. The value can be either a
|
|
:class:`pyspark.sql.types.DataType` object or a DDL-formatted type string.
|
|
|
|
>>> from pyspark.sql.functions import pandas_udf, PandasUDFType
|
|
>>> df1 = spark.createDataFrame(
|
|
... [(20000101, 1, 1.0), (20000101, 2, 2.0), (20000102, 1, 3.0), (20000102, 2, 4.0)],
|
|
... ("time", "id", "v1"))
|
|
>>> df2 = spark.createDataFrame(
|
|
... [(20000101, 1, "x"), (20000101, 2, "y")],
|
|
... ("time", "id", "v2"))
|
|
>>> def asof_join(l, r):
|
|
... return pd.merge_asof(l, r, on="time", by="id")
|
|
>>> df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas(
|
|
... asof_join, schema="time int, id int, v1 double, v2 string"
|
|
... ).show() # doctest: +SKIP
|
|
+--------+---+---+---+
|
|
| time| id| v1| v2|
|
|
+--------+---+---+---+
|
|
|20000101| 1|1.0| x|
|
|
|20000102| 1|3.0| x|
|
|
|20000101| 2|2.0| y|
|
|
|20000102| 2|4.0| y|
|
|
+--------+---+---+---+
|
|
|
|
Alternatively, the user can define a function that takes three arguments. In this case,
|
|
the grouping key(s) will be passed as the first argument and the data will be passed as the
|
|
second and third arguments. The grouping key(s) will be passed as a tuple of numpy data
|
|
types, e.g., `numpy.int32` and `numpy.float64`. The data will still be passed in as two
|
|
`pandas.DataFrame` containing all columns from the original Spark DataFrames.
|
|
|
|
>>> def asof_join(k, l, r):
|
|
... if k == (1,):
|
|
... return pd.merge_asof(l, r, on="time", by="id")
|
|
... else:
|
|
... return pd.DataFrame(columns=['time', 'id', 'v1', 'v2'])
|
|
>>> df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas(
|
|
... asof_join, "time int, id int, v1 double, v2 string").show() # doctest: +SKIP
|
|
+--------+---+---+---+
|
|
| time| id| v1| v2|
|
|
+--------+---+---+---+
|
|
|20000101| 1|1.0| x|
|
|
|20000102| 1|3.0| x|
|
|
+--------+---+---+---+
|
|
|
|
.. seealso:: :meth:`pyspark.sql.functions.pandas_udf`
|
|
|
|
"""
|
|
from pyspark.sql.pandas.functions import pandas_udf
|
|
|
|
udf = pandas_udf(
|
|
func, returnType=schema, functionType=PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF)
|
|
all_cols = self._extract_cols(self._gd1) + self._extract_cols(self._gd2)
|
|
udf_column = udf(*all_cols)
|
|
jdf = self._gd1._jgd.flatMapCoGroupsInPandas(self._gd2._jgd, udf_column._jc.expr())
|
|
return DataFrame(jdf, self.sql_ctx)
|
|
|
|
@staticmethod
|
|
def _extract_cols(gd):
|
|
df = gd._df
|
|
return [df[col] for col in df.columns]
|
|
|
|
|
|
def _test():
|
|
import doctest
|
|
from pyspark.sql import SparkSession
|
|
import pyspark.sql.pandas.group_ops
|
|
globs = pyspark.sql.pandas.group_ops.__dict__.copy()
|
|
spark = SparkSession.builder\
|
|
.master("local[4]")\
|
|
.appName("sql.pandas.group tests")\
|
|
.getOrCreate()
|
|
globs['spark'] = spark
|
|
(failure_count, test_count) = doctest.testmod(
|
|
pyspark.sql.pandas.group_ops, globs=globs,
|
|
optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE | doctest.REPORT_NDIFF)
|
|
spark.stop()
|
|
if failure_count:
|
|
sys.exit(-1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
_test()
|