Индексатор строк, CountVectorizer Pyspark в одной строке

Question

Привет, я столкнулся с проблемой, когда у меня есть строки с двумя столбцами массива слов.

column1, column2
["a", "b" ,"b", "c"], ["a","b", "x", "y"]

В основном я хочу подсчитать появление каждого слова между столбцами, чтобы получить два массива:

[1, 2, 1, 0, 0], 
[1, 1, 0, 1, 1]

Таким образом, "a" появляется один раз в каждом массиве, "b" появляется дважды в столбце 1 и один раз в столбце 2, "c" появляется только в столбце 1, "x" и "y" только в столбце 2. Так далее и тому подобное.

Я попытался взглянуть на функцию CountVectorizer из библиотеки ml, но не уверен, работает ли она построчно, массивы могут быть очень большими в каждом столбце? И значения 0 (где одно слово появляется в одном столбце, но не в другом), похоже, не переносятся.

Любая помощь приветствуется.

Answer 1

Для Spark 2.4+ это можно сделать с помощью DataFrame API и встроенных функций массива.

Сначала получите все слова для каждой строки, используя array_union функция. Затем используйте transform функция для преобразования массива слов, где для каждого элемента вычисляется количество вхождений в каждом столбце, используя size а также array_remove функции:

df = spark.createDataFrame([(["a", "b", "b", "c"], ["a", "b", "x", "y"])], ["column1", "column2"])

df.withColumn("words", array_union("column1", "column2")) \
  .withColumn("occ_column1",
              expr("transform(words, x -> size(column1) - size(array_remove(column1, x)))")) \
  .withColumn("occ_column2",
              expr("transform(words, x -> size(column2) - size(array_remove(column2, x)))")) \
  .drop("words") \
  .show(truncate=False)

Выход:

+------------+------------+---------------+---------------+
|column1     |column2     |occ_column1    |occ_column2    |
+------------+------------+---------------+---------------+
|[a, b, b, c]|[a, b, x, y]|[1, 2, 1, 0, 0]|[1, 1, 0, 1, 1]|
+------------+------------+---------------+---------------+