如何在Databricks的PySpark DataFrame中添加随机整数列rating

需求说明 #

现有如下结构的PySpark DataFrame：

需要新增一列rating，值为4到9（包含两端）的随机整数，最终得到包含该列的DataFrame。

实现脚本 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import randint

# 初始化SparkSession
spark = SparkSession.builder.appName("AddRandomRating").getOrCreate()

# 创建示例DataFrame
data = [("A", "xyz"), ("B", "opq"), ("C", "mns")]
df = spark.createDataFrame(data, ["Customer", "Name"])

# 添加随机rating列：生成4到9（包含）的整数
df_with_rating = df.withColumn("rating", randint(4, 10))

# 展示结果
df_with_rating.show()

关键说明 #

• randint(low, high)函数会生成[low, high)区间的随机整数，因此要包含9的话，high参数设为10即可。
• 如果不想用randint，也可以通过rand()函数转换实现：

  from pyspark.sql.functions import rand
  df_with_rating = df.withColumn("rating", (rand() * 6).cast("integer") + 4)
  

  原理是：rand()生成0到1的浮点数，乘以6后得到0到6的浮点数，转为整数后是0到5，加4就得到4到9的整数。

内容的提问来源于stack exchange，提问作者Emma