Azure环境下DataFrame空值赋值随机数问题求助

场景1：使用Pandas DataFrame #

首先导入依赖库：

import pandas as pd
import numpy as np

同一行的rn1和rn2使用同一个随机数

# 筛选rn列为空的行，生成对应数量的随机数
random_values = np.random.randint(0, 1001, size=df['rn'].isna().sum())
# 批量赋值给rn1和rn2列
df.loc[df['rn'].isna(), ['rn1', 'rn2']] = random_values[:, np.newaxis]

注：np.random.randint(0,1001)生成0到1000的整数（左闭右开区间，所以上限设为1001）；[:, np.newaxis]将一维数组转为二维，确保同一随机数能同时赋值给同一行的两列。

rn1和rn2各自独立生成随机数

# 分别给rn1、rn2赋值独立随机数
df.loc[df['rn'].isna(), 'rn1'] = np.random.randint(0, 1001, size=df['rn'].isna().sum())
df.loc[df['rn'].isna(), 'rn2'] = np.random.randint(0, 1001, size=df['rn'].isna().sum())

场景2：使用Azure Databricks PySpark DataFrame #

如果是在Azure Databricks中使用PySpark，可按以下方式处理：
首先导入依赖函数：

from pyspark.sql import functions as F

rn1和rn2各自独立生成随机数

df = df.withColumn(
    "rn1",
    F.when(F.col("rn").isNull(), F.randint(0, 1000)).otherwise(F.col("rn1"))
).withColumn(
    "rn2",
    F.when(F.col("rn").isNull(), F.randint(0, 1000)).otherwise(F.col("rn2"))
)

同一行的rn1和rn2使用同一个随机数

先生成临时随机数列，再同步赋值给rn1和rn2：

df = df.withColumn(
    "temp_random",
    F.when(F.col("rn").isNull(), F.randint(0, 1000)).otherwise(F.lit(None))
).withColumn(
    "rn1",
    F.when(F.col("rn").isNull(), F.col("temp_random")).otherwise(F.col("rn1"))
).withColumn(
    "rn2",
    F.when(F.col("rn").isNull(), F.col("temp_random")).otherwise(F.col("rn2"))
).drop("temp_random")

注：F.randint(0,1000)在Spark 3.0及以上版本支持，生成0到1000的整数。

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