обратный эффект функции взрыва

Question

В scala с spark-2.4 я хотел бы отфильтровать значение внутри массивов в столбце.

От

+---+------------+
| id|      letter|
+---+------------+
|  1|[x, xxx, xx]|
|  2|[yy, y, yyy]|
+---+------------+

До

+---+-------+
| id| letter|
+---+-------+
|  1|[x, xx]|
|  2|[yy, y]|
+---+-------+

Я думал об использовании explode + filter

val res = Seq(("1", Array("x", "xxx", "xx")), ("2", Array("yy", "y", "yyy"))).toDF("id", "letter")
res.withColumn("tmp", explode(col("letter"))).filter(length(col("tmp")) < 3).drop(col("letter")).show()

И я получаю

+---+---+
| id|tmp|
+---+---+
|  1|  x|
|  1| xx|
|  2| yy|
|  2|  y|
+---+---+

Как мне вернуть zip / groupBy по id?

Или, может быть, есть лучшее, более оптимизированное решение?

Answer 1

Вы можете отфильтровать массив без explode() в Spark 2.4:

res.withColumn("letter", expr("filter(letter, x -> length(x) < 3)")).show()

Вывод:

+---+-------+
| id| letter|
+---+-------+
|  1|[x, xx]|
|  2|[yy, y]|
+---+-------+

Answer 2

В Spark 2.4+, функции более высокого порядка - это путь (filter), в качестве альтернативы используйте collect_list:

res.withColumn("tmp",explode(col("letter")))
  .filter(length(col("tmp")) < 3)
  .drop(col("letter"))
  // aggregate back
  .groupBy($"id")
  .agg(collect_list($"tmp").as("letter"))
  .show()

:

+---+-------+
| id| letter|
+---+-------+
|  1|[x, xx]|
|  2|[yy, y]|
+---+-------+

Как этовводит случайное перемешивание, для этого лучше использовать UDF:

def filter_arr(maxLength:Int)= udf((arr:Seq[String]) => arr.filter(str => str.size<=maxLength))

res
  .select($"id",filter_arr(maxLength = 2)($"letter").as("letter"))
  .show()

дает:

+---+-------+
| id| letter|
+---+-------+
|  1|[x, xx]|
|  2|[yy, y]|
+---+-------+