python - 如何映射列以在 spark sql 数据框中创建新列？

Question

在 python 和 pandas 中，我可以像这样创建一个新列：

使用 pandas 数据框中的两列来创建一个字典。

 dict1 = dict(zip(data["id"], data["duration"]))

然后我可以应用这个字典在第二个数据框中创建一个新列。

df['id_duration'] = df['id'].map(lambda x: dict1[x] if x in dict1.keys() else -1)

如果我有一个数据框（有两列：和）和一个数据框（有一个列） id_duration，如何在 spark sql 数据框中创建一个新列？dataiddurationdfid

标签： pythonpandasdataframeapache-spark

Answer 1

使用字典会很可惜，因为您需要将整个数据帧收集data到驱动程序上，这对性能非常不利，并可能导致 OOM 错误。

您可以简单地在两个数据帧之间执行左外连接，并na.fill使用-1.

data = spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ['id', 'duration'])
df = spark.createDataFrame([(1, 2), (3, 4)], ['id', 'x'])

df\
    .join(data.withColumnRenamed("duration", "id_duration"), ['id'], 'left')\
    .na.fill(-1).show()

+---+---+-----------+
| id|  x|id_duration|
+---+---+-----------+
|  5|  6|         -1|
|  1|  2|         10|
|  3|  4|         30|
+---+---+-----------+