Pandas用户定义函数（UDF）-是否可以返回布尔值？

Question

我正在尝试将函数编写为Pandas UDF，它将检查字符串数组的任何元素是否以特定值开头。我正在寻找的结果将是这样的：

df.filter(list_contains(val, df.stringArray_column)).show()

函数list_contains将在df.stringArray的任何元素以val开头的每一行上返回True。

仅举一个例子：

df = spark.read.csv(path)
display(df.filter(list_contains('50', df.stringArray_column)))

上面的代码将显示df的每一行，其中stringArray列的元素以50开头。

我已经在python中编写了一个函数，这非常慢

    def list_contains(val):
    # Perfom what ListContains generated
  def list_contains_udf(column_list):
    for element in column_list:
      if element.startswith(val):
        return True
    return False
  return udf(list_contains_udf, BooleanType())

谢谢您的帮助。

Answer 1

如果我们放松一种情况，没有UDF可能会这样做。目标单词必须是数组列中任何单词的子字符串，而不必是前缀。让我们尝试以下方法。

from pyspark.sql import SparkSession
import pyspark.sql.functions as f

spark = SparkSession.builder.appName('substr_finder').getOrCreate()

# sample data creation
my_df = spark.createDataFrame(
        [('scooby', ['cartoon', 'kids']),
         ('joker', ['dark', 'cartoon']),
         ('tom', ['fun', 'kids']),
         ('car', ['kids', 'cartoon'])
         ]
        , schema=('character', 'tags'))

数据帧my_df看起来如下：

+---------+--------------------+
|character|                tags|
+---------+--------------------+
|   scooby|     [cartoon, kids]|
|    joker|     [dark, cartoon]|
|      tom|[cartoon, fun, kids]|
|      car|     [kids, cartoon]|
+---------+--------------------+

在我们遍历每行逻辑之后，这里仅返回第四行，因为car是cartoon的子字符串（和前缀）。

这里是本机spark操作来实现的。

my_df2 = my_df.withColumn('con_tags', f.concat_ws('_', my_df.tags)) # to concatenate items in the array column

数据帧my_df2看起来像：

+---------+--------------------+----------------+
|character|                tags|        con_tags|
+---------+--------------------+----------------+
|   scooby|     [cartoon, kids]|    cartoon_kids|
|    joker|     [dark, cartoon]|    dark_cartoon|
|      tom|[cartoon, fun, kids]|cartoon_fun_kids|
|      car|     [kids, cartoon]|    kids_cartoon|
+---------+--------------------+----------------+

让我们在my_df2上应用正则表达式匹配器

my_df3 = my_df2.withColumn('matcher', f.expr(r"regexp_extract(con_tags, character, 0)"))

my_df3类似于以下内容：

+---------+--------------------+----------------+-------+
|character|                tags|        con_tags|matcher|
+---------+--------------------+----------------+-------+
|   scooby|     [cartoon, kids]|    cartoon_kids|       |
|    joker|     [dark, cartoon]|    dark_cartoon|       |
|      tom|[cartoon, fun, kids]|cartoon_fun_kids|       |
|      car|     [kids, cartoon]|    kids_cartoon|    car|
+---------+--------------------+----------------+-------+

现在，我们需要过滤列matcher为非空的行。

my_df4 = my_df3.filter(my_df3.matcher != "").drop('con_tags', 'matcher')

这里是最终数据帧：

+---------+---------------+
|character|           tags|
+---------+---------------+
|      car|[kids, cartoon]|
+---------+---------------+

注：如前所述，这些操作也将为row (art, [kids, cartoon])提供肯定的结果，因为art是的子字符串。如果我们可以确保列表列中没有满足此条件的条目，则可以使用以下一系列步骤。也许还有其他使用正则表达式匹配器的方法，但是目前还没有想到。