python-2.7 - 如何在 pyspark 中创建具有两个输入的 UDF

Question

我是 pyspark 的新手，我正在尝试创建一个简单的 udf，它必须采用两个输入列，检查第二列是否有空格，如果有，则将第一个列拆分为两个值并覆盖原始列。这就是我所做的:

def split(x, y):
if x == "EXDRA" and y == "":
    return ("EXT", "DCHA")
if x == "EXIZQ" and y == "":
    return ("EXT", "IZDA")

udf_split = udf(split, ArrayType())

df = df \
.withColumn("x", udf_split(df['x'], df['y'])[1]) \
.withColumn("y", udf_split(df['x'], df['y'])[0])

但是当我运行此代码时，出现以下错误:

File "<stdin>", line 1, in <module>
TypeError: __init__() takes at least 2 arguments (1 given)

我做错了什么？

谢谢你，阿尔瓦罗

Answer 1

我不确定你想做什么，但根据我的理解，我会这样做:

from pyspark.sql.types import *
from pyspark.sql.functions import udf, col

def split(x, y):
    if x == "EXDRA" and y == "":
        return ("EXT", "DCHA")
    if x == "EXIZQ" and y == "":
        return ("EXT", "IZDA")

schema = StructType([StructField("x1", StringType(), False), StructField("y1", StringType(), False)])
udf_split = udf(split, schema)

df = spark.createDataFrame([("EXDRA", ""), ("EXIZQ", ""), ("", "foo")], ("x", "y"))

df.show()

# +-----+---+
# |    x|  y|
# +-----+---+
# |EXDRA|   |
# |EXIZQ|   |
# |     |foo|
# +-----+---+

df = df \
.withColumn("split", udf_split(df['x'], df['y'])) \
.withColumn("x", col("split.x1")) \
.withColumn("y", col("split.y1"))

df.printSchema()

# root
#  |-- x: string (nullable = true)
#  |-- y: string (nullable = true)
#  |-- split: struct (nullable = true)
#  |    |-- x1: string (nullable = false)
#  |    |-- y1: string (nullable = false)


df.show()

# +----+----+----------+
# |   x|   y|     split|
# +----+----+----------+
# | EXT|DCHA|[EXT,DCHA]|
# | EXT|IZDA|[EXT,IZDA]|
# |null|null|      null|
# +----+----+----------+