将 Spark 数据帧列转换为行

Question

我有一个像这样的 Spark 数据框：

    +-----+-----+-------+--------+
    |  A  |B    |2017Q1 | 2017Q2 |
    +-----+----------------------+
    |  1  |  101|  0.6  |  0.8   |
    |  2  |  102|  0.7  |  0.9   |
    |  3  |  103|  0.9  |  0.4   |
    |  ...|  ...|  ...  |  ...   |

这里的年和季度列可以是动态的。 这意味着我可能还会获得 2017Q3、2017Q4 等的列。 我想将2017Q1和2017Q2的列值转换为行，如下所示：

+-----+-----+-------+--------+----+
|  A  |B    |Year   | Quarter|Val |
+-----+----------------------+----+
|  1  |  101| 2017  |  1     |0.6 |
|  1  |  101| 2017  |  2     |0.8 |
|  2  |  102| 2017  |  1     |0.7 |
|  ...|  ...|  ...  |  ...   |.   |

有人可以帮我解决这个问题吗？ 我正在使用 Spark 2.4.4

Answer 1

在 Spark SQL 中，您可以创建一个包含两个季度的值的数组。 由于您需要记住哪个值对应于哪个季度，因此您可以创建一个结构体来将季度的索引与其值绑定。 要使其动态化，您可以使用季度列表。 为了使它更通用，我们可以从数据框的列名中提取季度，如下所示。

val quarters = df.columns
    .filter( _.matches("[0-9]{4}Q[1-4]") ) // all the columns matching the regex
    .sorted

df.withColumn("value", explode(array(
       quarters.indices.map(i =>
            struct(lit(i+1) as "val", col(quarters(i)) as "quarter")
       ) : _*
  )))
  .withColumn("Quarter", $"value.quarter")
  .withColumn("Val", $"value.val")
  .drop( quarters :+ "value" : _*)

编辑因为在评论中被问到，这里是一个（痛苦的）Java 版本：

// a function to generate struct(lit(i+1) as "val", col(quarters(i)) as "quarter")
public static Column generateStruct(int i, List<String> quarters) {
    Column[] columns = {
        lit(i+1).alias("val"),
        col(quarters.get(i)).alias("quarter")
    };
    Seq<Column> columnSeq = JavaConverters
        .asScalaBufferConverter(Arrays.asList(columns))
        .asScala().toSeq();
    return struct(columnSeq);
}

List<String> quarters = Arrays.stream(df.columns())
    .filter(x -> x.matches("[0-9]{4}Q[1-4]"))
    .sorted().collect(Collectors.toList());

List<Column> quarterColumns = new ArrayList<>();
for(int i = 0; i < quarters.size(); i++)
    quarterColumns.add(generateStruct(i, quarters));


df
    .withColumn("value", explode(array(
        JavaConverters.asScalaBufferConverter(quarterColumns).asScala().toSeq()
    )))
    .withColumn("Quarter", col("value.quarter"))
    .withColumn("Val", col("value.val"))
    .drop(JavaConverters.asScalaBufferConverter(quarters).asScala().toSeq())
    .drop("value");