如何融化Spark DataFrame？

PySpark 或至少 Scala 中的 Apache Spark 中是否有等效的 Pandas Melt 函数？

到目前为止，我一直在 Python 中运行示例数据集，现在我想对整个数据集使用 Spark。

火花 >= 3.4

在 Spark 3.4 或更高版本中，您可以使用内置melt method

(sdf
    .melt(
        ids=['A'], values=['B', 'C'], 
        variableColumnName="variable", 
        valueColumnName="value")
    .show())

+---+--------+-----+
|  A|variable|value|
+---+--------+-----+
|  a|       B|    1|
|  a|       C|    2|
|  b|       B|    3|
|  b|       C|    4|
|  c|       B|    5|
|  c|       C|    6|
+---+--------+-----+

此方法适用于所有 API，因此可以在 Scala 中使用

sdf.melt(Array($"A"), Array($"B", $"C"), "variable", "value")

or SQL

SELECT * FROM sdf UNPIVOT (val FOR col in (col_1, col_2))

火花3.2（仅限 Python，需要 Pandas 和 pyarrow）

(sdf
    .to_koalas()
    .melt(id_vars=['A'], value_vars=['B', 'C'])
    .to_spark()
    .show())

+---+--------+-----+
|  A|variable|value|
+---+--------+-----+
|  a|       B|    1|
|  a|       C|    2|
|  b|       B|    3|
|  b|       C|    4|
|  c|       B|    5|
|  c|       C|    6|
+---+--------+-----+

火花

没有内置函数（如果您使用 SQL 并启用了 Hive 支持，则可以使用stack功能，但它没有在 Spark 中公开，也没有本机实现），但自己推出是很简单的。所需进口：

from pyspark.sql.functions import array, col, explode, lit, struct
from pyspark.sql import DataFrame
from typing import Iterable 

实施示例：

def melt(
        df: DataFrame, 
        id_vars: Iterable[str], value_vars: Iterable[str], 
        var_name: str="variable", value_name: str="value") -> DataFrame:
    """Convert :class:`DataFrame` from wide to long format."""

    # Create array<struct<variable: str, value: ...>>
    _vars_and_vals = array(*(
        struct(lit(c).alias(var_name), col(c).alias(value_name)) 
        for c in value_vars))

    # Add to the DataFrame and explode
    _tmp = df.withColumn("_vars_and_vals", explode(_vars_and_vals))

    cols = id_vars + [
            col("_vars_and_vals")[x].alias(x) for x in [var_name, value_name]]
    return _tmp.select(*cols)

和一些测试（基于Pandas 文档测试):

import pandas as pd

pdf = pd.DataFrame({'A': {0: 'a', 1: 'b', 2: 'c'},
                   'B': {0: 1, 1: 3, 2: 5},
                   'C': {0: 2, 1: 4, 2: 6}})

pd.melt(pdf, id_vars=['A'], value_vars=['B', 'C'])

   A variable  value
0  a        B      1
1  b        B      3
2  c        B      5
3  a        C      2
4  b        C      4
5  c        C      6

sdf = spark.createDataFrame(pdf)
melt(sdf, id_vars=['A'], value_vars=['B', 'C']).show()

+---+--------+-----+
|  A|variable|value|
+---+--------+-----+
|  a|       B|    1|
|  a|       C|    2|
|  b|       B|    3|
|  b|       C|    4|
|  c|       B|    5|
|  c|       C|    6|
+---+--------+-----+

注意：要与旧版 Python 版本一起使用，请删除类型注释。

Related:

• R SparkR - 相当于熔化函数
• 聚集在星光熠熠