Question

我有这个在pandas数据帧中本地运行的python代码：

df_result = pd.DataFrame(df
                          .groupby('A')
                          .apply(lambda x: myFunction(zip(x.B, x.C), x.name))

我想在PySpark中运行它，但在处理pyspark.sql.group.GroupedData对象时遇到问题。

我尝试过以下方法：

sparkDF
 .groupby('A')
 .agg(myFunction(zip('B', 'C'), 'A'))

返回

KeyError: 'A'

我认为因为＆＃39; A＆＃39;不再是列，我找不到x.name。

的等价物

然后

sparkDF
 .groupby('A')
 .map(lambda row: Row(myFunction(zip('B', 'C'), 'A'))) 
 .toDF()

但收到以下错误：

AttributeError: 'GroupedData' object has no attribute 'map'

任何建议都会非常感激！

Answer 1

您正在尝试编写UDAF（用户定义聚合函数）而不是UDF（用户定义函数）。 UDAF是处理按密钥分组的数据的函数。具体来说，他们需要定义如何在单个分区中合并组中的多个值，然后如何跨分区合并键的结果。目前在python中没有办法实现UDAF，它们只能在Scala中实现。

但是，你可以在Python中解决它。您可以使用收集集来收集分组值，然后使用常规UDF来执行您想要的操作。唯一需要注意的是collect_set仅适用于原始值，因此您需要将它们编码为字符串。

from pyspark.sql.types import StringType
from pyspark.sql.functions import col, collect_list, concat_ws, udf

def myFunc(data_list):
    for val in data_list:
        b, c = data.split(',')
        # do something

    return <whatever>

myUdf = udf(myFunc, StringType())

df.withColumn('data', concat_ws(',', col('B'), col('C'))) \
  .groupBy('A').agg(collect_list('data').alias('data'))
  .withColumn('data', myUdf('data'))

如果要进行重复数据删除，请使用collect_set。此外，如果您的某些密钥有很多值，这将会很慢，因为密钥的所有值都需要在集群中的某个分区中收集。如果你的最终结果是你通过以某种方式组合每个键的值来构建的值（例如将它们相加），那么使用RDD aggregateByKey方法实现它可能会更快，它允许你为每个键构建一个中间值。在对数据进行混洗之前的分区。

编辑：11/21/2018

由于这个答案是写的，pyspark使用Pandas增加了对UDAF'的支持。使用Panda的UDF和UDAF比使用RDD的直接python函数有一些不错的性能改进。在引擎盖下，它会对列进行矢量化（将多行中的值批处理在一起以优化处理和压缩）。请查看here以获得更好的解释，或者查看下面user6910411的答案。

Answer 2

自Spark 2.3起，您可以使用pandas_udf。 GROUPED_MAP使用Callable[[pandas.DataFrame], pandas.DataFrame]或换句话说是一个函数，该函数从与输入相同形状的Pandas DataFrame映射到输出DataFrame。

例如，如果数据如下所示：

df = spark.createDataFrame(
    [("a", 1, 0), ("a", -1, 42), ("b", 3, -1), ("b", 10, -2)],
    ("key", "value1", "value2")
)

并且您想要计算value1 value2之间成对最小值的平均值，您必须定义输出模式：

from pyspark.sql.types import *

schema = StructType([
    StructField("key", StringType()),
    StructField("avg_min", DoubleType())
])

pandas_udf：

import pandas as pd

from pyspark.sql.functions import pandas_udf
from pyspark.sql.functions import PandasUDFType

@pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
def g(df):
    result = pd.DataFrame(df.groupby(df.key).apply(
        lambda x: x.loc[:, ["value1", "value2"]].min(axis=1).mean()
    ))
    result.reset_index(inplace=True, drop=False)
    return result

并应用它：

df.groupby("key").apply(g).show()

+---+-------+
|key|avg_min|
+---+-------+
|  b|   -1.5|
|  a|   -0.5|
+---+-------+

排除架构定义和装饰器，您可以按原样应用当前的Pandas代码。

自Spark 2.4.0起，还有GROUPED_AGG变体，需要Callable[[pandas.Series, ...], T]，其中T是原始标量：

import numpy as np

@pandas_udf(DoubleType(), functionType=PandasUDFType.GROUPED_AGG)
def f(x, y):
    return np.minimum(x, y).mean()

可与标准group_by / agg构造一起使用：

df.groupBy("key").agg(f("value1", "value2").alias("avg_min")).show()

+---+-------+
|key|avg_min|
+---+-------+
|  b|   -1.5|
|  a|   -0.5|
+---+-------+

请注意，GROUPED_MAP和GROUPPED_AGG pandas_udf的行为与UserDefinedAggregateFunction或Aggregator的行为相同，并且更接近groupByKey或窗口函数与无界框架。首先对数据进行混洗，然后才应用UDF。

为了优化执行，您应该implement Scala UserDefinedAggregateFunction和add Python wrapper。

另见User defined function to be applied to Window in PySpark?

Answer 3

我要超越答案。

因此您可以使用@pandas_udf在pyspark中实现类似pandas.groupby（）。apply的逻辑这是向量化方法，比简单的udf更快。

from pyspark.sql.functions import pandas_udf,PandasUDFType

df3 = spark.createDataFrame(
[("a", 1, 0), ("a", -1, 42), ("b", 3, -1), ("b", 10, -2)],
("key", "value1", "value2")
)

from pyspark.sql.types import *

schema = StructType([
    StructField("key", StringType()),
    StructField("avg_value1", DoubleType()),
    StructField("avg_value2", DoubleType()),
    StructField("sum_avg", DoubleType()),
    StructField("sub_avg", DoubleType())
])

@pandas_udf(schema, functionType=PandasUDFType.GROUPED_MAP)
def g(df):
    gr = df['key'].iloc[0]
    x = df.value1.mean()
    y = df.value2.mean()
    w = df.value1.mean() + df.value2.mean()
    z = df.value1.mean() - df.value2.mean()
    return pd.DataFrame([[gr]+[x]+[y]+[w]+[z]])

df3.groupby("key").apply(g).show()

您将得到以下结果：

+---+----------+----------+-------+-------+
|key|avg_value1|avg_value2|sum_avg|sub_avg|
+---+----------+----------+-------+-------+
|  b|       6.5|      -1.5|    5.0|    8.0|
|  a|       0.0|      21.0|   21.0|  -21.0|
+---+----------+----------+-------+-------+

因此，您可以在分组数据中的其他字段之间进行更多计算，并将它们以列表格式添加到数据框中。

Answer 4

PySpark 3.0.0 版中的另一个扩展新功能： applyInPandas

df = spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], 
                            ("id", "v"))  

def mean_func(key, pdf):
# key is a tuple of one numpy.int64, which is the value
# of 'id' for the current group
return pd.DataFrame([key + (pdf.v.mean(),)])

df.groupby('id').applyInPandas(mean_func, schema="id long, v double").show()

结果：

+---+---+
| id|  v|
+---+---+
|  1|1.5|
|  2|6.0|
+---+---+

欲知更多详情，请参阅：https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.GroupedData.applyInPandas.html

在PySpark中的GroupedData上应用UDF（具有正常运行的python示例）

4 个答案: