python - 如何迭代 pyspark.sql.Column？

Question

我有一个 pyspark DataFrame，我想获取一个特定的列并迭代它的值。例如:

userId    itemId
1         2
2         2
3         7
4         10

我通过 df.userId 获取 userId 列，并且我想为该列中的每个 userId 应用一个方法。我怎样才能做到这一点？

Answer 1

您的问题对于您要应用的功能类型不是很具体，因此我创建了一个示例，根据 itemId 的值添加项目描述。

首先让我们导入相关库并创建数据:

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

df = spark.createDataFrame([(1,2),(2,2),(3,7),(4,10)], ['userId', 'itemId'])

其次，创建函数并将其转换为 PySpark 可以使用的 UDF 函数:

def item_description(itemId):
    items = {2  : "iPhone 8",
             7  : "Apple iMac",
             10 : "iPad"}
    return items[itemId]

item_description_udf = udf(item_description,StringType())

最后，为 ItemDescription 添加一个新列，并用 item_description_udf 函数返回的值填充它:

df = df.withColumn("ItemDescription",item_description_udf(df.itemId))    
df.show()

这给出了以下输出:

+------+------+---------------+
|userId|itemId|ItemDescription|
+------+------+---------------+
|     1|     2|       iPhone 8|
|     2|     2|       iPhone 8|
|     3|     7|     Apple iMac|
|     4|    10|           iPad|
+------+------+---------------+