时间戳

首页 > 解决方案 > 如何炸开多列、不同类型、不同长度？

python - 如何炸开多列、不同类型、不同长度？

问题描述

我有一个包含不同时间周期（1/6、3/6、6/6 等）列的 DF，并且想“分解”所有列以创建一个新的 DF，其中每一行都是 1/ 6个循环。

from pyspark import Row 
from pyspark.sql import SparkSession 
from pyspark.sql.functions import explode, arrays_zip, col

spark = SparkSession.builder \
    .appName('DataFrame') \
    .master('local[*]') \
    .getOrCreate()

df = spark.createDataFrame([Row(a=1, b=[1, 2, 3, 4, 5, 6], c=[11, 22, 33], d=['foo'])])

|  a|                 b|           c|    d|
+---+------------------+------------+-----+
|  1|[1, 2, 3, 4, 5, 6]|[11, 22, 33]|[foo]|
+---+------------------+------------+-----+

我正在做爆炸：

df2 = (df.withColumn("tmp", arrays_zip("b", "c", "d"))
       .withColumn("tmp", explode("tmp"))
       .select("a", col("tmp.b"), col("tmp.c"), "d"))

但输出不是我想要的：

|  a|  b|   c|    d|
+---+---+----+-----+
|  1|  1|  11|[foo]|
|  1|  2|  22|[foo]|
|  1|  3|  33|[foo]|
|  1|  4|null|[foo]|
|  1|  5|null|[foo]|
|  1|  6|null|[foo]|
+---+---+----+-----+

我希望它看起来像这样：

|  a|  b|  c|  d|
+---+---+---+---+
|  1|  1| 11|foo|
|   |  2|   |   |
|   |  3| 22|   |
|   |  4|   |   |
|   |  5| 33|   |
|   |  6|   |   |
+---+---+---+---+

我是 Spark 的新手，从一开始我就有复杂的话题！:)

2019-07-15 更新：也许有人有不使用 UDF 的解决方案？-> 由@jxc 回答

2019-07-17 更新：也许有人有一个解决方案，如何以更复杂的顺序更改 null <-> 值序列？就像我们希望在列第一个值中c那样或更复杂的情况，然后是：Null, 11, Null, 22, Null, 33dNullfooNull, Null, Null

|  a|  b|  c|  d|
+---+---+---+---+
|  1|  1|   |   |
|   |  2| 11|foo|
|   |  3|   |   |
|   |  4| 22|   |
|   |  5|   |   |
|   |  6| 33|   |
+---+---+---+---+

标签： pythonpyspark

解决方案

这是不使用 udf 的一种方法：

2019/07/17 更新：调整 SQL stmt 并添加 N=6 作为 SQL 的参数。

2019/07/16 更新：删除了临时列t，替换array(0,1,2,3,4,5)为转换函数中的常量。在这种情况下，我们可以直接对数组元素的值而不是它们的索引进行操作。

更新：我删除了使用字符串函数并将数组元素中的数据类型全部转换为字符串的原始方法，效率较低。Spark 2.4+的Spark SQL高阶函数应该比原来的方法好。

设置

from pyspark.sql import functions as F, Row

df = spark.createDataFrame([ Row(a=1, b=[1, 2, 3, 4, 5, 6], c=['11', '22', '33'], d=['foo'], e=[111,222]) ])

>>> df.show()
+---+------------------+------------+-----+----------+
|  a|                 b|           c|    d|         e|
+---+------------------+------------+-----+----------+
|  1|[1, 2, 3, 4, 5, 6]|[11, 22, 33]|[foo]|[111, 222]|
+---+------------------+------------+-----+----------+

# columns you want to do array-explode
cols = df.columns

# number of array elements to set
N = 6

使用 SQL 高阶函数：transform

使用 Spark SQL 高阶函数：transform()，执行以下操作：

  1. 创建以下 Spark SQL 代码，其中{0}将替换为 column_name，{1}将替换为N：

    stmt = '''
   CASE
      WHEN '{0}' in ('d') THEN
        transform(sequence(0,{1}-1), x -> IF(x == 1, `{0}`[0], NULL))
      WHEN size(`{0}`) <= {1}/2 AND size(`{0}`) > 1 THEN
        transform(sequence(0,{1}-1), x -> IF(((x+1)*size(`{0}`))%{1} == 0, `{0}`[int((x-1)*size(`{0}`)/{1})], NULL))
      ELSE `{0}`
    END AS `{0}`
'''

    注意：数组转换仅在数组包含多个（除非在单独的WHEN子句中指定）和<= N/2元素（在本例中为1 < size <= 3）时定义。其他大小的数组将保持原样。

  2. 使用selectExpr()对所有必需的列运行上述 SQL

    df1 = df.withColumn('a', F.array('a')) \
        .selectExpr(*[ stmt.format(c,N) for c in cols ])

>>> df1.show()
+---+------------------+----------------+-----------+---------------+
|  a|                 b|               c|          d|              e|
+---+------------------+----------------+-----------+---------------+
|[1]|[1, 2, 3, 4, 5, 6]|[, 11,, 22,, 33]|[, foo,,,,]|[,, 111,,, 222]|
+---+------------------+----------------+-----------+---------------+

  3. 运行arrays_zip并爆炸：

    df_new = df1.withColumn('vals', F.explode(F.arrays_zip(*cols))) \
            .select('vals.*') \
            .fillna('', subset=cols)

>>> df_new.show()
+----+---+---+---+----+
|   a|  b|  c|  d|   e|
+----+---+---+---+----+
|   1|  1|   |   |null|
|null|  2| 11|foo|null|
|null|  3|   |   | 111|
|null|  4| 22|   |null|
|null|  5|   |   |null|
|null|  6| 33|   | 222|
+----+---+---+---+----+

    注意：fillna('', subset=cols)仅更改包含字符串的列

在一个方法链中：

df_new = df.withColumn('a', F.array('a')) \
           .selectExpr(*[ stmt.format(c,N) for c in cols ]) \
           .withColumn('vals', F.explode(F.arrays_zip(*cols))) \
           .select('vals.*') \
           .fillna('', subset=cols)

用变换函数解释：

转换功能（下面列出，反映旧版本的要求）

transform(sequence(0,5), x -> IF((x*size({0}))%6 == 0, {0}[int(x*size({0})/6)], NULL))

如帖子中所述，{0}将替换为列名。这里我们c以包含 3 个元素的 column- 为例：

  • 在转换函数中，sequence(0,5)创建一个array(0,1,2,3,4,5)包含 6 个元素的常量数组，其余的设置 lambda 函数，其中一个参数x具有元素的值。
  • IF(condition, true_value, false_value) : 是一个标准的 SQL 函数

  • 我们应用的条件是：(x*size(c))%6 == 0where size(c)=3，如果这个条件为真，则返回c[int(x*size(c)/6)]，否则返回NULL。所以x从 0 到 5，我们将有：

    ((0*3)%6)==0) true   -->  c[int(0*3/6)] = c[0]
((1*3)%6)==0) false  -->  NULL
((2*3)%6)==0) true   -->  c[int(2*3/6)] = c[1]
((3*3)%6)==0) false  -->  NULL
((4*3)%6)==0) true   -->  c[int(4*3/6)] = c[2]
((5*3)%6)==0) false  -->  NULL

类似于 column-e，它包含一个 2 元素数组。

推荐阅读