时间戳

首页 > 解决方案 > 如何缩放火花数据框中的数据子集

python - 如何缩放火花数据框中的数据子集

问题描述

我有一个包含几列的火花数据框，重要的是一列带有日期时间戳，另一列带有值。如果该值对应于某个时间之前的日期时间，我想将值列中的所有值缩放一个常数。

例如

id | datetime   | value | other_col |
1  | 2020-10-1  | 0.1   | 200       |
1  | 2020-10-20 | 0.01  | 100       |
1  | 2020-10-14 | 0.2   | 200       |
1  | 2020-10-25 | 1     | 200       |

应该转变为

id | datetime   | value | other_col |
1  | 2020-10-1  | 1     | 200       |
1  | 2020-10-20 | 0.01  | 100       |
1  | 2020-10-14 | 2     | 200       |
1  | 2020-10-25 | 1     | 200       |

如果日期时间 < 2020-10-20。只有值列中的值应该改变，数据框的其余部分应该保持不变。

我尝试使用 map 函数来做到这一点，但我不熟悉 spark 数据帧，目前我只得到一个数字输出而不是整个数据帧。

from pyspark.sql.types import Row


def scale_data(x):
    if x.datetime<= "2020-10-20 08:00:00.00":
        return x.value * 10
    return x


df2 = df.rdd.map(lambda x: Row(scale_data(x))).toDF()

我希望返回的整个数据框仅包含满足条件的数据子集已被转换。

标签： pythondataframeapache-sparkpyspark

解决方案

你可以使用 when 函数来解决这个问题。您不必编写可能会减慢代码速度的 UDF。

from pyspark.sql.functions import *

df2 = df.withColumn("scaled", when(df["datetime"] <= lit("2020-10-20 08:00:00.00"), df["value"]*lit(10))
.otherwise(df["value"])))

推荐阅读