从 groupby 和已知序列向量创建 Spark 数据帧

Question

我有以下数据框

| Col1 | Col2 | Value  |
|------|------|--------|
| Key1 | Key2 | 45.0   |
| Key2 | Key2 | 34.0   |
| Key2 | Key3 | 121.0  |

我有值序列[1,2,3]。我想生成以下数据框

| Col1 | Col2 | Value  |
|------|------|--------|
| Key1 | Key2 |  1     |
| Key1 | Key2 |  2     |
| Key1 | Key2 |  3     |
| Key2 | Key2 |  1     |
| Key2 | Key2 |  2     |
| Key2 | Key2 |  3     |
| Key2 | Key3 |  1     |
| Key2 | Key3 |  2     |
| Key2 | Key3 |  3     |

(Col1,Col2) 对是从 groupBy 生成的，在我的新数据框中，我希望 GroupBy 返回的每一对重复三次，一次针对序列的每个值。然后，我将使用 (Col1,Col2,Value) 的组合将该数据框与另一个数据框左连接。

我想到的一个解决方案可能是采用三个不同数据帧的并集，每个数据帧的行包含不同的 (Col1,Col2) 和给定的序列值，即

df.groupBy(['Col1','Col2']).agg(F.lit(1).alias('Value')).union(df.groupBy(['Col1','Col2']).agg(F.lit(2).alias('Value')))

我想找到一个更优雅的解决方案，因为这个看起来与我在 SQL 中所做的非常相似。

Answer 1

只需创建一个文字数组列并将其展开：

from pyspark.sql import functions as F

df = spark.createDataFrame([
    ("Key1", "Key2", 45.0), ("Key2", "Key2", 34.0),
    ("Key2", "Key3", 121.0)], ("Col1", "Col2", "Value"))

values = [1, 2, 3]
df1 = df.withColumn("Value", F.explode(F.array(*[F.lit(x) for x in values])))

df1.show()
#+----+----+-----+
#|Col1|Col2|Value|
#+----+----+-----+
#|Key1|Key2|    1|
#|Key1|Key2|    2|
#|Key1|Key2|    3|
#|Key2|Key2|    1|
#|Key2|Key2|    2|
#|Key2|Key2|    3|
#|Key2|Key3|    1|
#|Key2|Key3|    2|
#|Key2|Key3|    3|
#+----+----+-----+

从 groupby 和已知序列向量创建 Spark 数据帧

Create Spark dataframe from groupby and known sequence vector

dataframe

apache-spark

pyspark

apache-spark-sql