Выполнять сложные агрегаты по двум спискам столбцов?

Question

Я новичок в программировании, и я очищаю и упрощаю свой код для выполнения групповой работы и агрегации на фрейме данных pyspark. Я пытаюсь сделать вещи проще для понимания и работаю над рефакторингом некоторых из моего кода. Когда я пробую следующий код, я получаю сообщение об ошибке:

TypeError: Invalid argument, not a string or column: 

Вот мой код:

groupBy = ['ColA']

convert_to_list = ['Col1', 'Col2', 'Col3',]
convert_to_set = ['Col4', 'Col5', 'Col6',]

fun_list = [F.collect_list]
funs_set = [F.collect_set]

exprs = F.concat(
    [f(F.col(c)) for f in fun_list for c in convert_to_list], 
    [f(F.col(c)) for f in funs_set for c in convert_to_set]
)

df = df.groupby(*groupBy).agg(*exprs)

Очень ценю вашу помощь. Я не уверен, как передать правильные столбцы в функцию agg.

Пример ввода и ожидаемый результат

Answer 1

Ваша проблема - вызов concat - вы передаете 2 списка, когда он ожидает, что аргументы будут столбцами. Похоже, вы хотели объединить эти списки вместе, что можно сделать с помощью оператора сложения в python.

Попробуйте:

exprs = (
    [f(F.col(c)).alias(c) for f in fun_list for c in convert_to_list] + 
    [f(F.col(c)).alias(c) for f in funs_set for c in convert_to_set]
)

Я также добавил в вызов alias, чтобы имена столбцов сохранялись после агрегации.