从UDF返回StructType的ArrayType时出错（并在多个UDF中使用单个函数）

Question

（编辑）更改了字段名称（从foo、bar……改为name和city），因为旧的名称令人困惑

我需要在多个UDF中使用单个函数，并根据输入返回不同的结构。

我的实现的这个简化版本基本上完成了我正在寻找的内容：

from pyspark.sql.types import IntegerType, StructType, StringType
from pyspark.sql.functions import when, col

df = spark.createDataFrame([1, 2, 3], IntegerType()).toDF('id')

struct_one = StructType().add('name', StringType(), True)
struct_not_one = StructType().add('city', StringType(), True)

def select(id):
  if id == 1:
    return {'name': 'Alice'}
  else:
    return {'city': 'Seattle'}

one_udf = udf(select, struct_one)
not_one_udf = udf(select, struct_not_one)

df = df.withColumn('one', when((col('id') == 1), one_udf(col('id'))))\
       .withColumn('not_one', when((col('id') != 1), not_one_udf(col('id'))))

display(df)   

（编辑）输出：

id  one               not_one
1   {"name":"Alice"}  null
2   null              {"city":"Seattle"}
3   null              {"city":"Seattle"}

但是，返回StructType的ArrayType的代码不幸失败：

from pyspark.sql.types import IntegerType, StructType, StringType, ArrayType
from pyspark.sql.functions import when, col

df = spark.createDataFrame([1, 2, 3], IntegerType()).toDF('id')

struct_one = StructType().add('name', StringType(), True)
struct_not_one = ArrayType(StructType().add('city', StringType(), True))

def select(id):
  if id == 1:
    return {'name': 'Alice'}
  else:
    return [{'city': 'Seattle'}, {'city': 'Milan'}]

one_udf = udf(select, struct_one)
not_one_udf = udf(select, struct_not_one)

df = df.withColumn('one', when((col('id') == 1), one_udf(col('id'))))\
       .withColumn('not_one', when((col('id') != 1), not_one_udf(col('id'))))

display(df)      

错误消息是：

ValueError：结构类型的意外元组“名称”

（编辑）所需输出为：

id  one                 not_one
1   {"name":"Alice"}    null
2   null                [{"city":"Seattle"},{"city":"Milan"}]
3   null                [{"city":"Seattle"},{"city":"Milan"}]

但是，其他类型（StringType、ENTgerType、...）的返回和ArrayType可以使用。

在多个UDF中不使用单个函数时，也可以返回StructType数组：

from pyspark.sql.types import IntegerType, StructType, StringType, ArrayType
from pyspark.sql.functions import when, col

df = spark.createDataFrame([1, 2, 3], IntegerType()).toDF('id')

struct_not_one = ArrayType(StructType().add('city', StringType(), True))

def select(id):
    return [{'city': 'Seattle'}, {'city': 'Milan'}]

not_one_udf = udf(select, struct_not_one)

df = df.withColumn('not_one', when((col('id') != 1), not_one_udf(col('id'))))

display(df)   

（编辑）输出：

id  not_one
1   null
2   [{"city":"Seattle"},{"city":"Milan"}]
3   [{"city":"Seattle"},{"city":"Milan"}]

为什么返回结构类型的数组类型并使用多个UDF和一个函数不起作用？

谢谢

Answer 1

“Spark SQL（包括SQL以及DataFrame和Dataset API）不保证子表达式的求值顺序。。。因此，依赖布尔表达式的副作用或求值顺序以及WHERE和HAVING子句的顺序是危险的，因为这些表达式和子句可以在查询优化和规划期间重新排序。具体来说，如果UDF依赖SQL中的短路语义进行空检查，则不能保证在调用UDF之前会进行空检查。"

请参阅评估顺序和空值检查

为了保持自定义项的通用性，您可以将“when filter”推入自定义项中：

from pyspark.sql.types import IntegerType, StructType, StringType, ArrayType
from pyspark.sql.functions import when, col, lit

df = spark.createDataFrame([1, 2, 3], IntegerType()).toDF('id')

struct_one = StructType().add('name', StringType(), True)
struct_not_one = ArrayType(StructType().add('city', StringType(), True))

def select(id, test):

  if eval(test.format(id)) is False:
    return None

  if id == 1:
    return {'name': 'Alice'}
  else:
    return [{'city': 'Seattle'}, {'city': 'Milan'}]

one_udf = udf(select, struct_one)
not_one_udf = udf(select, struct_not_one)

df = df.withColumn('one', one_udf(col('id'), lit('{} == 1')))\
       .withColumn('not_one', not_one_udf(col('id'), lit('{} != 1')))

display(df)