使用 pyspark pandas 的自定义聚合

Question

我有一个 pyspark pandas 数据框。我想使用一些预定义函数执行自定义聚合，并为了简单起见使用 numpy.nanmean 我收到以下错误“aggs 必须是从列名到聚合函数（字符串或字符串列表）的字典映射”

我在databricks中尝试了这样的最小示例

import pandas as pd
import pyspark.pandas as ps
import numpy as np

# Create a sample PySpark Pandas DataFrame
data = {'Category': ['A', 'B', 'A', 'B', 'A', 'B'],
        'Value1': [10, 20, 30, 40, 50, 60],
        'Value2': [100, np.nan, 300, 400, 500, 600]}  
sdf = spark.createDataFrame(pd.DataFrame(data))

# Convert PySpark DataFrame to PySpark Pandas DataFrame
pdf = ps.DataFrame(sdf)

# Group by 'Category' and apply nanmean_udf to 'Value2'
result = pdf.groupby('Category').agg(
    Sum_Value1=('Value1', 'sum'),
    Mean_val2=('Value2', np.nanmean)
)

print(result)

感谢您的帮助。

Answer 1

您可以使用字典作为

agg

的参数来获取结果

result = pdf.groupby('Category').agg({
    'Value1': 'sum',
    'Value2': np.nanmean
})