使用 PySpark 在列表列表上应用函数

如何解决使用 PySpark 在列表列表上应用函数

我正在尝试在我的 df 上应用一个函数。我的变量是一个列表列表，对于它们中的每一个，我都会应用我的函数，以获得每个结果的列表。

dff = spark.createDataFrame([
('a',[[1,2,3,4],[1,1,3],[3,5,9,6],[12,4,2]]),('b',[[4,0.2,0.3,0.7],1],[2,7,9]]),('c',9],8],[5,[4,6]]),('d',[[2,2],8,6]])],['num','list_apply'])

因此我写了这个函数，我注册了它，但我找不到任何结果，只出现了错误。

def calc(data):
  a,b,c,d = data
  dist = a+b+c*d
  if dist < 10:
    return True
  else:
    return False
  
calc_udf = f.udf(lambda x: calc(x),ArrayType(BooleanType()))
spark.udf.register("calc_udf",calc,T.ArrayType(T.BooleanType()))

dff = dff.withColumn("dist",f.expr("TRANSFORM(list_apply,x -> calc_udf(x))"))

我想获得的是：

+-----+------------+-------------------------------------+
| num | list_apply |                 dist                |
+-----+------------+-------------------------------------+
| a   |    ...     | [False,True,False,False]         |
| b   |    ...     | [True,False]                 |
| c   |    ...     | [False,False] |
| d   |    ...     | [True,False]                |

解决方法

您的方法的问题在于，在 transform 函数中，您使用 udf 将列表作为参数，这是行不通的，因为 udf 应该采用列（或列名）作为论据。可以修改 calc_udf 取一个列表类型列如下，然后直接在 list_apply 列上使用：

# instead of calc(x) we iterate through lst so we don't need transform
calc_udf = f.udf(lambda lst: [calc(x) for x in lst],ArrayType(BooleanType()))
dff = dff.withColumn("dist",calc_udf("list_apply"))

dff.show()
+---+--------------------+--------------------+
|num|          list_apply|                dist|
+---+--------------------+--------------------+
|  a|[[1.0,2.0,3.0,...|[false,true,fal...|
|  b|[[4.0,0.2,0.3,...| [true,false]|
|  c|[[1.0,1.0,false,fa...|
|  d|[[2.0,...|[true,false]|
+---+--------------------+--------------------+