Примените функцию к списку списков с помощью PySpark

#pyspark

Вопрос:

Я пытаюсь применить функцию к своему df. Моя переменная представляет собой список списков, для каждого из которых я бы применил свою функцию, чтобы получить список результатов для каждого из них.

 dff = spark.createDataFrame([
('a', [[1,2,3,4], [1,1,2,3], [3,5,9,6], [12,4,2,2]]), 
('b', [[4,0.2,0.3,0.7],[1,1,1,1],[2,7,5,9]]), 
('c', [[1,1,2,9],[2,5,4,8],[5,7,4,8],[1,2,3,4],[4,4,4,6]]),
('d' ,[[2,2,2,2],[4,8,5,9],[1,5,9,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, True, False]                 |
| c   |    ...     | [False, False, False, False, False] |
| d   |    ...     | [True, False, False]                |


 

Ответ №1:

Проблема с вашим подходом заключается в том, что в 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, true, false]|
|  c|[[1.0, 1.0, 2.0, ...|[false, false, fa...|
|  d|[[2.0, 2.0, 2.0, ...|[true, false, false]|
 --- -------------------- -------------------- 
 

Комментарии:

1. Большое спасибо за вашу помощь, я думал, что transform функция будет принимать каждый элемент основного списка в каждой строке.

2. Знаете ли вы, как я мог бы изменить свою функцию, чтобы взять этот список списка, чтобы в итоге получить одно значение, т. Е. Взять max() из полного списка результатов ? Здесь он берет элемент за элементом, применяет функцию и возвращает список результатов. Можно ли предоставить ему полный список [1,2,3,4], [1,1,2,3]... , получить максимум a b c*d и вернуть одно True или False ? Как эта transform() функция может быть применена ?