在 PySpark 数据框中重命名已透视和聚合的列

Question

以下是一个dataframe示例：

from pyspark.sql.functions import avg, first

rdd = sc.parallelize(
    [
        (0, "A", 223,"201603", "PORT"), 
        (0, "A", 22,"201602", "PORT"), 
        (0, "A", 422,"201601", "DOCK"), 
        (1,"B", 3213,"201602", "DOCK"), 
        (1,"B", 3213,"201601", "PORT"), 
        (2,"C", 2321,"201601", "DOCK")
    ]
)
df_data = sqlContext.createDataFrame(rdd, ["id","type", "cost", "date", "ship"])

df_data.show()

我对它进行了数据透视表操作，

df_data.groupby(df_data.id, df_data.type).pivot("date").agg(avg("cost"), first("ship")).show()

+---+----+----------------+--------------------+----------------+--------------------+----------------+--------------------+
| id|type|201601_avg(cost)|201601_first(ship)()|201602_avg(cost)|201602_first(ship)()|201603_avg(cost)|201603_first(ship)()|
+---+----+----------------+--------------------+----------------+--------------------+----------------+--------------------+
|  2|   C|          2321.0|                DOCK|            null|                null|            null|                null|
|  0|   A|           422.0|                DOCK|            22.0|                PORT|           223.0|                PORT|
|  1|   B|          3213.0|                PORT|          3213.0|                DOCK|            null|                null|
+---+----+----------------+--------------------+----------------+--------------------+----------------+--------------------+

但是我得到的这些列名非常复杂。在聚合操作中应用别名(alias)通常有效，但由于此处使用了透视(pivot)，因此列名变得更加糟糕:

+---+----+--------------------------------------------------------------+------------------------------------------------------------------+--------------------------------------------------------------+------------------------------------------------------------------+--------------------------------------------------------------+------------------------------------------------------------------+
| id|type|201601_(avg(cost),mode=Complete,isDistinct=false) AS cost#1619|201601_(first(ship)(),mode=Complete,isDistinct=false) AS ship#1620|201602_(avg(cost),mode=Complete,isDistinct=false) AS cost#1619|201602_(first(ship)(),mode=Complete,isDistinct=false) AS ship#1620|201603_(avg(cost),mode=Complete,isDistinct=false) AS cost#1619|201603_(first(ship)(),mode=Complete,isDistinct=false) AS ship#1620|
+---+----+--------------------------------------------------------------+------------------------------------------------------------------+--------------------------------------------------------------+------------------------------------------------------------------+--------------------------------------------------------------+------------------------------------------------------------------+
|  2|   C|                                                        2321.0|                                                              DOCK|                                                          null|                                                              null|                                                          null|                                                              null|
|  0|   A|                                                         422.0|                                                              DOCK|                                                          22.0|                                                              PORT|                                                         223.0|                                                              PORT|
|  1|   B|                                                        3213.0|                                                              PORT|                                                        3213.0|                                                              DOCK|                                                          null|                                                              null|
+---+----+--------------------------------------------------------------+------------------------------------------------------------------+--------------------------------------------------------------+------------------------------------------------------------------+--------------------------------------------------------------+------------------------------------------------------------------+ 

有没有一种方法可以在透视和聚合时即时更改列名？

- Ivan

Answer 1

您可以直接为聚合设置别名：

pivoted = df_data \
    .groupby(df_data.id, df_data.type) \
    .pivot("date") \
    .agg(
       avg('cost').alias('cost'),
       first("ship").alias('ship')
    )

pivoted.printSchema()
##root
##|-- id: long (nullable = true)
##|-- type: string (nullable = true)
##|-- 201601_cost: double (nullable = true)
##|-- 201601_ship: string (nullable = true)
##|-- 201602_cost: double (nullable = true)
##|-- 201602_ship: string (nullable = true)
##|-- 201603_cost: double (nullable = true)
##|-- 201603_ship: string (nullable = true)

- jastingo

---

我认为这应该是被接受的答案。 - Mattias Moser

Answer 2

一个简单的正则表达式就可以解决问题：

import re

def clean_names(df):
    p = re.compile("^(\w+?)_([a-z]+)\((\w+)\)(?:\(\))?")
    return df.toDF(*[p.sub(r"\1_\3", c) for c in df.columns])

pivoted = df_data.groupby(...).pivot(...).agg(...)

clean_names(pivoted).printSchema()
## root
##  |-- id: long (nullable = true)
##  |-- type: string (nullable = true)
##  |-- 201601_cost: double (nullable = true)
##  |-- 201601_ship: string (nullable = true)
##  |-- 201602_cost: double (nullable = true)
##  |-- 201602_ship: string (nullable = true)
##  |-- 201603_cost: double (nullable = true)
##  |-- 201603_ship: string (nullable = true)

如果你想保留函数名称，你可以将替换模式改为例如\1_\2_\3。

- zero323

---

当你有许多列，尤其是当你事先不知道列名时，例如由于某些ETL操作时，你必须使用@zero323的方法。 - Kai

Answer 3

一个简单的方法是在聚合函数之后使用别名。 我从你创建的df_data spark数据框开始。

df_data.groupby(df_data.id, df_data.type).pivot("date").agg(avg("cost").alias("avg_cost"), first("ship").alias("first_ship")).show()
+---+----+---------------+-----------------+---------------+-----------------+---------------+-----------------+
| id|type|201601_avg_cost|201601_first_ship|201602_avg_cost|201602_first_ship|201603_avg_cost|201603_first_ship|
+---+----+---------------+-----------------+---------------+-----------------+---------------+-----------------+
|  1|   B|         3213.0|             PORT|         3213.0|             DOCK|           null|             null|
|  2|   C|         2321.0|             DOCK|           null|             null|           null|             null|
|  0|   A|          422.0|             DOCK|           22.0|             PORT|          223.0|             PORT|
+---+----+---------------+-----------------+---------------+-----------------+---------------+-----------------+

列名的形式将是“原始列名_别名列名”。对于您的情况，原始列名将是201601，别名列名将是avg_cost，而列名是201601_avg_cost（由下划线“_”链接）。

- Frank

---

只有当你有多个聚合时，这才有效。在我的情况下，在旋转之前，我选择了该列并进行了转换。即 df_data.select(($"date" + "_avg_cost").as("date"), $"cost")).pivot("date").agg(avg("cost"))。 - jgaw

如何在使用alias的agg时将下划线更改为空格？ - rakeeee

Answer 4

我写了一个简单快速的函数来完成这个任务。享受吧！ :)

# This function efficiently rename pivot tables' urgly names
def rename_pivot_cols(rename_df, remove_agg):
    """change spark pivot table's default ugly column names at ease.
        Option 1: remove_agg = True: `2_sum(sum_amt)` --> `sum_amt_2`.
        Option 2: remove_agg = False: `2_sum(sum_amt)` --> `sum_sum_amt_2`
    """
    for column in rename_df.columns:
        if remove_agg == True:
            start_index = column.find('(')
            end_index = column.find(')')
            if (start_index > 0 and end_index > 0):
                rename_df = rename_df.withColumnRenamed(column, column[start_index+1:end_index]+'_'+column[:1])
        else:
            new_column = column.replace('(','_').replace(')','')
            rename_df = rename_df.withColumnRenamed(column, new_column[2:]+'_'+new_column[:1])   
    return rename_df

- Victor Z

Answer 5

来自 zero323 的修改版本，适用于 Spark 2.4

import re

def clean_names(df):
    p = re.compile("^(\w+?)_([a-z]+)\((\w+)(,\s\w+)\)(:\s\w+)?")
    return df.toDF(*[p.sub(r"\1_\3", c) for c in df.columns])

当前列名类似于 0_first(is_flashsale, false): int

- Mithril

Answer 6

在進行樞紐分析之前，您可以更改欄位的名稱：

.withColumn("ship", F.concat(F.lit("ship_"), "ship"))

- DANIEL SILVA RAMOS

---

如果您能提供完整的示例，包括您的建议和源代码格式化，那就太理想了。 - Cwt