如何将Python代码转为纯PySpark实现？已有Spark DataFrame的等效操作

1. 生成类似make_blobs的聚类数据 #

用Spark内置随机函数生成带聚类中心的浮点数据，模拟sklearnmake_blobs的效果：

from pyspark.sql import SparkSession
from pyspark.sql.functions import randn, lit, explode, array_repeat

spark = SparkSession.builder.appName("BlobGenerator").getOrCreate()

# 定义3个聚类中心（可根据需求调整）
centers = [(2.0, 3.0), (-1.0, -2.0), (5.0, -1.0)]
samples_per_center = 333  # 每个中心生成约333条样本，总样本量接近1000

# 生成带噪声的聚类样本
df_centers = spark.createDataFrame(centers, ["center_col1", "center_col2"])
df_blob = df_centers.withColumn("repeat", lit(samples_per_center)) \
    .withColumn("dummy", explode(array_repeat(lit(1), col("repeat")))) \
    .withColumn("col1", col("center_col1") + randn() * 0.8)  # 0.8为噪声标准差，控制聚类紧凑度
    .withColumn("col2", col("center_col2") + randn() * 0.8) \
    .drop("center_col1", "center_col2", "repeat", "dummy")

2. 添加ID并格式化数值（替代Python循环） #

用Spark内置函数完成ID生成和数值格式化，无需Python循环：

from pyspark.sql.functions import format_number, row_number
from pyspark.sql.window import Window

# 生成连续自增ID（从1开始，和原循环逻辑一致）
# 若无特定排序需求，用orderBy(lit(1))保证顺序与生成顺序一致
window = Window.orderBy(lit(1))
df_processed = df_blob.withColumn("id", row_number().over(window)) \
    .withColumn("col1", format_number(col("col1"), 2).cast("float")) \
    .withColumn("col2", format_number(col("col2"), 2).cast("float"))

df_processed.show(5)

如果已有包含两列浮点值的Spark DataFrame（例：名为existing_df，列名为col1、col2），直接用以下代码实现原Python循环的逻辑：

from pyspark.sql.functions import format_number, row_number
from pyspark.sql.window import Window

# 定义窗口，若需严格匹配原循环顺序，需用原始数据的排序字段替代lit(1)
window = Window.orderBy(lit(1))

# 添加ID+格式化数值
df_final = existing_df.withColumn("id", row_number().over(window)) \
    .withColumn("col1", format_number(col("col1"), 2).cast("float")) \
    .withColumn("col2", format_number(col("col2"), 2).cast("float"))

df_final.show(5)

关键注意事项 #

• 排序一致性：原Python循环按数组顺序生成ID，Spark中若未指定明确排序字段，row_number()的结果顺序可能不确定。若需严格匹配原顺序，需确保DataFrame包含可标识原始顺序的字段，并将其用于窗口的orderBy。
• 数值格式化：format_number会自动四舍五入到指定小数位，与Pythonround()行为一致；若需截断而非四舍五入，可改用floor(col("col1")*100)/100或ceil(col("col1")*100)/100实现。

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