如何在Python函数内更新pyspark数据框

Question

我有一个 Python 函数，它接收 pyspark 数据帧并检查它是否具有脚本中使用的其他函数所需的所有列。特别是，如果缺少列

'weight'

，我想通过为其分配新列来更新用户传递的数据帧。

例如：

from pyspark.sql import functions as F

def verify_cols(df):
    if 'weight' not in df.columns:
        df = df.withColumn('weight', F.lit(1))  # Can I update `df` inside this function?

如您所见，我希望该功能更新

df

。我怎样才能实现这个目标？如果可能的话，我想避免使用

return

声明。

这篇文章非常相似，但使用了pandas的

inplace

参数。

Answer 1

为了避免 return 语句，您可以声明一个类并将 df 保留为成员字段。

from pyspark.sql import functions as F
from pyspark.sql.DataFrame import DataFrame
class Validator:
    def __init__(self, df: DataFrame):
        self.df = df

    def verify_cols(self):
        if 'weight' not in self.df.columns:
            self.df = self.df.withColumn('weight', F.lit(1))

调用verify_cols方法后，字段df将会更新。