包含 StringType 元素的 ArrayType 列上的 UDF 函数答案

【问题标题】：UDF function on ArrayType column that contains StringType elements包含 StringType 元素的 ArrayType 列上的 UDF 函数
【发布时间】：2021-05-12 17:06:28
【问题描述】：

我需要一个 udf 函数来输入数据框的数组列，并对其中的两个字符串元素执行相等性检查。我的数据框有这样的架构。

ID	date	options
1	2021-01-06	['red', 'green']
2	2021-01-07	['Blue', 'Blue']
3	2021-01-08	['Blue', 'Yellow']
4	2021-01-09	nan

我已经试过了：

def equality_check(options: list):
  try:
   if options[0] == options[1]:
     return 1
   else:
     return 0
  except:
     return -1

equality_udf = f.udf(equality_check, t.IntegerType())

但它抛出了索引错误。我相信选项列是字符串数组。期望是这样的：

ID	date	options	equality_check
1	2021-01-06	['red', 'green']	0
2	2021-01-07	['Blue', 'Blue']	1
3	2021-01-08	['Blue', 'Yellow']	0
4	2021-01-09	nan	-1

【问题讨论】：

标签： apache-spark pyspark apache-spark-sql user-defined-functions pyspark-dataframes

【解决方案1】：

您可以检查options 列表是否已定义或其长度小于2，而不是使用try/except。这是一个工作示例：

from pyspark.sql import functions as F
from pyspark.sql.types import IntegerType

data = [
    (1, "2021-01-06", ['red', 'green']),
    (2, "2021-01-07", ['Blue', 'Blue']),
    (3, "2021-01-08", ['Blue', 'Yellow']),
    (4, "2021-01-09", None),
]
df = spark.createDataFrame(data, ["ID", "date", "options"])

def equality_check(options: list):
    if not options or len(options) < 2:
        return -1

    return int(options[0] == options[1])

equality_udf = F.udf(equality_check, IntegerType())

df1 = df.withColumn("equality_check", equality_udf(F.col("options")))
df1.show()

#+---+----------+--------------+--------------+
#| ID|      date|       options|equality_check|
#+---+----------+--------------+--------------+
#|  1|2021-01-06|  [red, green]|             0|
#|  2|2021-01-07|  [Blue, Blue]|             1|
#|  3|2021-01-08|[Blue, Yellow]|             0|
#|  4|2021-01-09|          null|            -1|
#+---+----------+--------------+--------------+

但是，我建议您不要使用 UDF，因为您可以只使用内置函数来做同样的事情：

df1 = df.withColumn(
    "equality_check",
    F.when(F.size(F.col("options")) < 2, -1)
        .when(F.col("options")[0] == F.col("options")[1], 1)
        .otherwise(0)
)

【讨论】：