在Pyspark中有条件地将列添加到数据框

时间:2017-01-20 02:15:37

标签: python apache-spark dataframe pyspark multiple-columns

我在PySpark中有一个数据框。我想有条件地在数据框中添加一列。

说如果数据框没有列,则添加一个null值的列。 如果列存在则不执行任何操作并返回与新数据帧相同的数据帧

如何在PySpark中传递条件语句

1 个答案:

答案 0 :(得分:1)

这并不难,但你需要的不仅仅是一个列名来做正确的事。必需的进口

from pyspark.sql import types as t
from pyspark.sql.functions import lit
from pyspark.sql import DataFrame

示例数据:

df = sc.parallelize([("a", 1, [1, 2, 3])]).toDF(["x", "y", "z"])

辅助函数(用于遗留Python版本的条带类型注释):

def add_if_not_present(df: DataFrame, name: str, dtype: t.DataType) -> DataFrame:
    return (df if name in df.columns 
        else df.withColumn(name, lit(None).cast(dtype)))

使用示例:

add_if_not_present(df, "foo", t.IntegerType())
DataFrame[x: string, y: bigint, z: array<bigint>, foo: int]
add_if_not_present(df, "x", t.IntegerType())
DataFrame[x: string, y: bigint, z: array<bigint>]
add_if_not_present(df, "foobar", 
  t.StructType([
      t.StructField("foo", t.IntegerType()), 
      t.StructField("bar", t.IntegerType())]))
DataFrame[x: string, y: bigint, z: array<bigint>, foobar: struct<foo:int,bar:int>]
相关问题