Question

我想使用data.groupby.apply（）将函数应用于每个组的Pyspark数据框的每一行。

我使用了分组地图熊猫UDF。但是我不知道如何在函数中添加另一个参数。

我尝试将参数用作全局变量，但函数无法重新识别（我的参数是pyspark数据框）
我还尝试了此问题（针对熊猫数据框）Use Pandas groupby() + apply() with arguments

中提出的解决方案

      @pandas_udf(schema,PandasUDFType.GROUPED_MAP)
        def function(key,data, interval):
            interval_df=interval.filter(interval["var"]==key).toPandas()
            for value in interval_df:
                  #Apply some operations

        return Data.groupBy("msn").apply(calc_diff, ('arg1'))

或

 @pandas_udf(schema,PandasUDFType.GROUPED_MAP)
        def function(key,data, interval):
            interval_df=interval.filter(interval["var"]==key).toPandas()
            for value in interval_df:
                  #Apply some operations

        return Data.groupBy("msn").apply(lambda x: calc_diff(x,'arg1'))

但是我得到了错误：

ValueError：无效函数：函数类型为GROUPED_MAP的pandas_udfs必须采用一个参数（数据）或两个参数（键，数据）。

任何人都可以帮助我解决上述问题。

谢谢

Answer 1

我喜欢@hwrd的想法，但可以将其变成生成器模式，以便像@Feng的示例中那样更易于集成：

def function_generator(key):
    @pandas_udf(schema,PandasUDFType.GROUPED_MAP)
    def function(interval):
        interval_df=interval.filter(interval["var"]==key).toPandas()
        for value in interval_df:
              #Apply some operations
    return function

calc_diff = function_generator('arg1')
output = Data.groupBy("msn").apply(calc_diff)

Answer 2

您可以在函数内部创建pandas udf，以便在创建函数时就知道函数参数。（或者，您可以导入functools并使用部分函数评估来做同样的事情。）这是PySpark文档中的example，经过修改可以传入一些参数：

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 my_function(df, by="id", column="v", value=1.0):
    schema = "{} long, {} double".format(by, column)

    @pandas_udf(schema, PandasUDFType.GROUPED_MAP)
    def subtract_value(pdf):
        # pdf is a pandas.DataFrame
        v = pdf[column]
        g = pdf[by]
        return pdf.assign(v = v - g * value)

    return df.groupby(by).apply(subtract_value)

my_function(df, by="id", column="v", value=2.0).show()

Answer 3

我认为您可以做这样的事情

def myfun(data, key, interval):
    #Apply some operations
    return something

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def myfun_udf(data):
    return myfun(data=data, key=mykey, interval=myinterval)


mykey=1
myinterval=2

Data.groupBy("msn").apply(myfun_udf)

Answer 4

所有答案似乎都很有用，但没有对正在发生的事情进行正式描述。所以，我从所有人，特别是@sifta 中获取了点点滴滴，并试图将其解释为。也许它可以在未来帮助某人。

假设我有一个 PySpark DF，如下所示

# test = pd.DataFrame({
# 'c1': ['a', 'a', 'b', 'b', 'b'],
# 'c2': ['a1', 'a2', 'b1', 'b1', 'b2']})
# test = spark.createDataFrame(test)

+---+---+
| c1| c2|
+---+---+
|  a| a1|
|  a| a2|
|  b| b1|
|  b| b1|
|  b| b2|
+---+---+

我的目标是创建另一个列 c3，它可以是 group count + some fixed value。好吧，这绝对不是最好的例子，但让我们尝试使用 groupby 来解决它。我们需要传递一些不直接支持的参数（固定值）。

所以，根据答案，我们可以想出

schema = t.StructType([
  t.StructField('c1', t.StringType()),
  t.StructField('c2', t.StringType()),
  t.StructField('c3', t.IntegerType()),
])

def fn_wrapper(df, val):

  @f.pandas_udf(schema, f.PandasUDFType.GROUPED_MAP)
  def fn(pdf):
    pdf['c3'] = pdf.shape[0] + val
    return pdf

  return df.groupby('c1', 'c2').apply(fn)

fn_wrapper(test, 7).show()

但这到底是什么意思？

我们有用于映射 fn (return pdf) 返回的 Pandas DF 的架构。

那么，让我们了解一下这个技巧是如何工作的。我们已经定义了一个名为 fn_wrapper 的普通 UDF，它采用 Pyspark DF 和要在核心 Pandas groupby 中使用的参数。我们在 fn_wrapper(test, 7).show() 中调用它。现在，当我们在 fn_wrapper 中时，我们只是在它里面有一个函数体，它只是暂时编译而不执行。

接下来，执行语句return df.groupby('c1', 'c2').apply(fn)。看，我们将函数 fn 定义为 Pandas UDF，它没有任何参数。然而，当 scope of fn_wrapper 并且它定义了 fn is called as apply 时，我们仍然在 val 中。因此，我们可以轻松地在 pdf['c3'] = pdf.shape[0] + val 中引用 val，其中我们将分组数据作为 Pandas DataFrame 呈现。

我想这有助于理解将函数包装为函数内部的函数以及如何使用参数 Pandas UDF。

结合使用带有参数的分组地图熊猫UDF

4 个答案: