PySpark在追加超列名称的同时展平数据框

时间:2018-03-13 18:31:43

标签: python apache-spark pyspark spark-dataframe

假设我有一个PySpark数据帧df

>>> df.printSchema()
root
 |-- a: struct
      |-- alpha: integer
      |-- beta: string
      |-- gamma: boolean
 |-- b: string
 |-- c: struct
      |-- delta: string
      |-- epsilon: struct
           |-- omega: string
           |-- psi: boolean

我知道我可以压扁数据框:

select_col_list = [col.replace("a", "a.*").replace("c", "c.*") for col in df.columns]
flat_df = df.select(*select_col_list)

这会产生如下架构:

root
 |-- alpha: integer
 |-- beta: string
 |-- gamma: boolean
 |-- b: string
 |-- delta: string
 |-- epsilon: struct
      |-- omega: string
      |-- psi: boolean

但是我想在我扁平化时将超级列的名称附加到子列,所以我希望得到的模式如下:

root
 |-- a_alpha: integer
 |-- a_beta: string
 |-- a_gamma: boolean
 |-- b: string
 |-- c_delta: string
 |-- c_epsilon: struct
      |-- omega: string
      |-- psi: boolean

我该怎么做?

1 个答案:

答案 0 :(得分:1)

我不认为这是一种直截了当的方式,但这是我提出的一个黑客解决方案。

  1. 定义要展开的列的列表,并使用pyspark.sql.functions.monotonically_increasing_id()创建临时id列。
  2. 遍历数据框中的所有列,并为每个列创建临时数据框。
    • 如果列位于cols_to_expand中:请使用.*展开列。然后使用alias()使用相应的前缀重命名结果(临时)数据帧中的所有字段(id除外)。
    • 如果列不在cols_to_expand中:选择该列并id并将其存储在临时数据框中。
  3. temp_df存储在列表中。
  4. 使用id加入列表中的所有数据框,然后删除id列。
  5. 代码:

    df = df.withColumn('id', f.monotonically_increasing_id())
    cols_to_expand = ['a', 'c']
    flat_dfs = []
    for col in df.columns:
        if col in cols_to_expand:
            temp_df = df.select('id', col+".*")
            temp_df = temp_df.select(
                [
                    f.col(c).alias(col+"_"+c if c != 'id' else c) for c in temp_df.columns
                ]
            )
        else:
            temp_df = df.select('id', col)
    
        flat_dfs.append(temp_df)
    
    flat_df = reduce(lambda x, y: x.join(y, on='id'), flat_dfs)
    
    flat_df = flat_df.drop('id')
    flat_df.printSchema()
    

    结果架构:

    flat_df.printSchema()
    #root
    # |-- a_alpha: integer (nullable = true)
    # |-- a_beta: string (nullable = true)
    # |-- a_gamma: boolean (nullable = true)
    # |-- b: string (nullable = true)
    # |-- c_delta: string (nullable = true)
    # |-- c_epsilon: struct (nullable = true)
    # |    |-- omega: string (nullable = true)
    # |    |-- psi: boolean (nullable = true)