Pyspark 将所有数据帧值增加 1

Question

我试图将数据框中的所有值增加 1，除了 ID 列之外。

例子：

结果：

这是我到目前为止所拥有的，但是当我有很多列要做时（例如 50），它会变得有点长。

df_add = df.select(
  'Id',
  (df['col_a'] + 1).alias('col_a'),
  ..
  ..
)

有没有更Pythonic的方法来达到相同的结果？

Answer 1

编辑（基于@Daniel评论）：您可以直接使用lit函数

from pyspark.sql.functions import col, lit

for column in plus_one_cols:
    df = df.withColumn(column, col(column) + lit(1))

之前的回答：

在列中添加“1”是一种列式操作，更适合pandas_udf

from pyspark.sql.functions import col, pandas_udf, PandasUDFType

@pandas_udf('double', PandasUDFType.SCALAR)
def plus_one(v):
    return v + 1

plus_one_cols = [x for x in df.columns if x != "Id"]

for column in plus_one_cols:
    df = df.withColumn(column, plus_one(col(column)))

这将比逐行操作快得多。您还可以参考PySpark 的 Pandas UDF 简介 - Databricks