Dataframe 在PySpark中将列值组合添加到数据帧的更好方法

我有一个包含3列的数据集，

id

，

day

，

value

。对于

id

和

day

的所有组合，我需要在

value

中添加带零的行

# Simplified version of my data frame
data = [("1", "2020-04-01", 5), 
        ("2", "2020-04-01", 5), 
        ("3", "2020-04-02", 4)]
df = spark.createDataFrame(data,['id','day', 'value'])

我想到的是：

# Create all combinations of id and day
ids= df.select('id').distinct()
days = df.select('day').distinct()
full = ids.crossJoin(days)

# Add combinations back to df filling value with zeros
df_full = df.join(full, ['id', 'day'], 'rightouter')\
    .na.fill(value=0,subset=['value'])

输出我需要的内容：

>>> df_full.orderBy(['id','day']).show()
+---+----------+-----+
| id|       day|value|
+---+----------+-----+
|  1|2020-04-01|    5|
|  1|2020-04-02|    0|
|  2|2020-04-01|    5|
|  2|2020-04-02|    0|
|  3|2020-04-01|    0|
|  3|2020-04-02|    4|
+---+----------+-----+

问题是这两种操作在计算上都非常昂贵。当我用我的全部数据运行它时，它给了我一个比通常需要几个小时才能运行的任务大一个数量级的任务

---

有没有更有效的方法？还是我遗漏了什么？

这就是我要实现的方式。需要指出的是，两个数据帧必须具有相同的模式，否则

stack

函数将引发错误

导入pyspark.sql.f函数
#我的数据框的简化版本
数据=[（“1”，“2020-04-01”，5），
("2", "2020-04-01", 5), 
("3", "2020-04-02", 4)]
df=spark.createDataFrame（数据，['id'，'day'，'value']）
#创建包含所有不同日期的数据帧
df_days=df.select（f.col（'day'）。别名（'r_day'））。distinct（）
#自联接以查找所有组合
df_final=df.join（df_days，on=df['day'！=df_days['r_day']））
# +---+----------+-----+----------+
#| id | day | value | r|u day|
# +---+----------+-----+----------+
# |  1|2020-04-01|    5|2020-04-02|
# |  2|2020-04-01|    5|2020-04-02|
# |  3|2020-04-02|    4|2020-04-01|
# +---+----------+-----+----------+
#取消Pivot数据帧
df_final=df_final。选择（'id'，f.expr（'stack（2，day，value，r_day，cast（0作为bigint））作为（day，value）'））
df_final.orderBy（'id'，'day'）.show（）

输出：

+---+----------+-----+
| id|       day|value|
+---+----------+-----+
|  1|2020-04-01|    5|
|  1|2020-04-02|    0|
|  2|2020-04-01|    5|
|  2|2020-04-02|    0|
|  3|2020-04-01|    0|
|  3|2020-04-02|    4|
+---+----------+-----+

---

像这样的。你可以，我把第一排分开，因为它更清楚发生了什么。不过，您可以将其添加到“主循环”中

数据=[
（“1”，日期（2020年4月1日），5日），
（“2”，日期（2020年4月2日），5日），
（“3”，日期（2020年4月3日），5），
（“1”，日期（2020年4月3日）和5日），
]
df=spark.createDataFrame（数据，[“id”、“日期”、“值”]）
row_dates=df.select（“date”）.distinct（）.collect（）
日期=[item.asDict（）[“date”]用于第_行日期中的项目]
def map_行（日期：List[date]）->可调用[[Iterator[row]]，Iterator[row]]：
dates.sort（）
def内部（分区）：
最后一行=无
对于分区中的行：
#为分区中的第一行填写缺少的日期
如果最后一行为“无”：
对于日期中的日期：
如果日期<行日期：
收益率行（行id，天，0）
其他：
#将当前行设置为最后一行，生成当前行并中断循环
最后一行=行
产量行
打破
其他：
#如果当前行与最后一行具有相同的id
如果last_row.id==row.id:
#上次和当前之间的收益日期
对于日期中的日期：
如果日期>最后一行日期和日期<行日期：
收益率行（行id，天，0）
#将电流设置为最后一个电流，并产生电流
最后一行=行
产量行
其他：
#如果当前行是新id
对于日期中的日期：
#为最后一行id运行可能的剩余日期
如果日期>最后一行日期：
收益率行（最后一行id，天，0）
对于日期中的日期：
#在row.date之前填写缺少的日期
如果日期<行日期：
收益率行（行id，天，0）
其他：
#苏韦特先生
最后一行=行
产量行
打破
返回内部
rdd=(
df.重新划分（1，“id”）
.sortWithinPartitions（“id”、“日期”）
.rdd.mapPartitions（地图行（日期））
)
new_df=spark.createDataFrame（rdd）
新测向显示（10，错误）

这似乎并没有更快，至少在简化数据的情况下是如此。我将运行一个基准比较来确定。但我很想了解一下为什么这应该更快的背后的直觉（请原谅我的无知）。非常感谢。也许我正在避免一个更像

交叉连接

的步骤会更快，但当然，只测试真实数据。@LeonardoViotti我做了一个更新来简化我的代码。请再检查一次，让我知道它是否比您的解决方案快。我做了一些基准测试，不幸的是，您的答案花费了将近7倍的时间。2000年对2000年。我的问题代码是300。这是我的本地设置，仅使用模拟3x3数据帧：/@LeonardoViotti可能

stack

操作并不快，因为我认为您对数据一无所知。例如，日期或id是连续的吗？还是一切都是随机的？您的真实数据集有多大，即大约有多少行和唯一日期/id:s？Hi@Molotch，数据约为5x10^9行，id为500万，时间为61天。但我不确定在这种情况下连续意味着什么（我尝试过谷歌搜索，但没有多大成功），也就是说，如果第一个日期是2021-01-01，最后一个日期是2021-01-10，那么您知道这两个日期之间的所有日期都在数据集中。因为您知道这些日期，而且没有那么多，所以我会尝试按id重新划分数据集，以便每个id的所有行都位于同一分区中。然后按id和日期对分区进行排序，使其顺序正确。最后，在传递给mapPartitions调用的函数中应用已知日期。该函数只会为每个缺失的日期添加一行。嗨@Molotch，这实际上很有意义！实际上，我还没有尝试实现它，但是我不确定在

mapPartitions上使用什么函数
data = [
    ("1", date(2020, 4, 1), 5),
    ("2", date(2020, 4, 2), 5),
    ("3", date(2020, 4, 3), 5),
    ("1", date(2020, 4, 3), 5),
]


df = spark.createDataFrame(data, ["id", "date", "value"])

row_dates = df.select("date").distinct().collect()

dates = [item.asDict()["date"] for item in row_dates]


def map_row(dates: List[date]) -> Callable[[Iterator[Row]], Iterator[Row]]:
    dates.sort()

    def inner(partition):
        last_row = None

        for row in partition:
            # fill in missing dates for first row in partition
            if last_row is None:
                for day in dates:
                    if day < row.date:
                        yield Row(row.id, day, 0)
                    else:
                        # set current row as last row, yield current row and break out of the loop
                        last_row = row
                        yield row
                        break
            else:
                # if current row has same id as last row
                if last_row.id == row.id:
                    # yield dates between last and current
                    for day in dates:
                        if day > last_row.date and day < row.date:
                            yield Row(row.id, day, 0)
                    
                    # set current as last and yield current
                    last_row = row
                    yield row

                else:
                    # if current row is new id
                    for day in dates:
                        # run potential remaining dates for last_row.id
                        if day > last_row.date:
                            yield Row(last_row.id, day, 0)

                    for day in dates:
                        # fill in missing dates before row.date
                        if day < row.date:
                            yield Row(row.id, day, 0)                    
                        else:
                            # unt so weiter
                            last_row = row
                            yield row
                            break

    return inner


rdd = (
    df.repartition(1, "id")
    .sortWithinPartitions("id", "date")
    .rdd.mapPartitions(map_row(dates))
)
new_df = spark.createDataFrame(rdd)
new_df.show(10, False)