Pyspark: применять Pivot только к столбцам Dataframe ⇐ Python

[Anonymous]

У меня есть следующий фрейм данных.

Код: Выделить всё

df = spark.createDataFrame(
[
[1, 'AB', 12, '2022-01-01']
, [1, 'AA', 22, '2022-01-10']
, [1, 'AC', 11, '2022-01-11']
, [2, 'AB', 22, '2022-02-01']
, [2, 'AA', 28, '2022-02-10']
, [2, 'AC', 25, '2022-02-22']
]
, 'code: int, doc_type: string, amount: int, load_date: string'
)
df = df.withColumn('load_date', F.to_date('load_date'))

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

Код: Выделить всё

(
df.groupBy('code')
.pivot('doc_type', ['AB', 'AA', 'AC'])
.agg(F.sum('amount').alias('amnt'), F.first('load_date').alias('ldt'))
.show()
)

+----+-------+----------+-------+----------+-------+----------+
|code|AB_amnt|    AB_ldt|AA_amnt|    AA_ldt|AC_amnt|    AC_ldt|
+----+-------+----------+-------+----------+-------+----------+
|   1|     12|2022-01-01|     22|2022-01-10|     11|2022-01-11|
|   2|     22|2022-02-01|     28|2022-02-10|     25|2022-02-22|
+----+-------+----------+-------+----------+-------+----------+

Это то, что я хочу.

Код: Выделить всё

(
df.groupBy('code')
.agg(
F.sum(F.when(F.col('doc_type') == 'AB', F.col('amount'))).alias('AB_amnt')
, F.sum(F.when(F.col('doc_type') == 'AA', F.col('amount'))).alias('AA_amnt')
, F.sum(F.when(F.col('doc_type') == 'AC', F.col('amount'))).alias('AC_amnt')
, F.first('load_date').alias('load_date')
)
.show()
)

+----+-------+-------+-------+----------+
|code|AB_amnt|AA_amnt|AC_amnt| load_date|
+----+-------+-------+-------+----------+
|   1|     12|     22|     11|2022-01-01|
|   2|     22|     28|     25|2022-02-01|
+----+-------+-------+-------+----------+

Есть ли более простой способ сделать это? У меня есть несколько столбцов, которые нужно поместить в сводную таблицу, а также в не сводную таблицу.
Я использую Databricks 14.3 LTS со Spark 3.5.0

Подробнее здесь: https://stackoverflow.com/questions/794 ... me-columns