Pyspark 按包含列名的列表分組

Question

所以我有一個帶有多個分組變量的 pyspark 數據框我想寫一個 function 用戶可以在其中傳遞數據框，以及分組列的列表=[a, b, Z4A8A08F09D37B748553 列... （要計算基於組的聚合值的列）。 動態大小的列表將由用戶決定它可以包含多少元素我在 pyspark 中找不到語法，我可以在其中使用列列表作為分組變量以及聚合變量。

請幫忙看代碼

Answer 1

在groupBy()中使用列表非常容易，因為它很容易接受列表輸入（列列表）。 至於agg() ，您可以使用列表推導來創建聚合。

假設您只想在所有聚合列（例如max() ）中使用 1 個聚合 function ，您的 function 將接受用於分組的列列表和要聚合的列列表。 請參閱下面的示例 function - 注意.agg()中的理解。

def groupByFunc(sdf, groupByVars=[], aggVars=[]):
    assert len(groupByVars) > 0, "at least 1 grouping variable in the list required"
    assert len(aggVars) > 0, "at least 1 aggregating variable in the list required"

    grouped_sdf = sdf. \
        groupBy(groupByVars). \
        agg(*[func.sum(k).alias(k+'_sum') for k in aggVars])

    return grouped_sdf

# used as following
agg_sdf = groupByFunc(sdf, groupByVars=['col1', 'col2'], aggVars=['col_a', 'col_b', 'col_c')
# the aforementioned return a dataframe using `sdf` 
# grouped on col1 and col2 with the sum of col_a, col_b, col_c individually
# agg_sdf will have 5 fields - col1, col2, col_a_sum, col_b_sum, col_c_sum

如果不同列的聚合不同，可以更改 function 以接受包含列名和聚合 function 作為鍵值對的字典。

def groupByFunc(sdf, groupByVars=[], aggVarFuncDict={}):
    assert len(groupByVars) > 0, "at least 1 grouping variable in the list required"
    assert len(aggVarFuncDict) > 0, "at least 1 aggregating variable and function in the dictionary required"

    grouped_sdf = sdf. \
        groupBy(groupByVars). \
        agg(*[aggVarFuncDict[k](k).alias(k+'_'+str(aggVarFuncDict[k]).split(' ')[1].lower()) for k in aggVarFuncDict.keys()])

    return grouped_sdf

# used as following
agg_sdf = groupByFunc(sdf, groupByVars=['col1', 'col2'], aggVarFuncDict={'col3': func.countDistinct, 'col5': func.sum})
# the aforementioned return a dataframe using `sdf` 
# grouped on col1 and col2 with the distinct count of col3 and sum of col5 individually
# agg_sdf will have 4 fields - col1, col2, col3_countdistinct, col5_sum