如何将此UDF转换为Pandas UDF

Question

我遇到了这个函数的一些性能问题，它的目标是如果字符串数组的字符串与val参数匹配，则返回True。我想将其转换为Pandas UDF。

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

受@jxc注释的启发，尝试在Databricks的单元格中使用下面的sql。

%sql
SELECT exists(column_list, element -> substr(element, 1, length(val)) == val)

我使用SQL理解的代码是使用substr来获取字符串element的头N (length(val))长度，以及它是否等于val本身。

否则，请参考PySpark文档的pyspark.sql.UDFRegistration(sparkSession)类，将相似的函数注册为UDF，组合使用。