Присвоение значения элементам в группе на основе других элементов в группе

Question

Я пытаюсь присвоить значение каждой строке в группе в зависимости от того, находится ли входная строка в начале другой строки в группе.Например, «hello» находится в начале «hello world», поэтому обоим из них присваивается значение «ВМЕСТЕ».«До свидания» не присваивается значение «ВМЕСТЕ», поскольку оно не находится в начале другой строки в группе.

input = [('1', 'hello'),('1', 'goodbye'),('1', 'hola'),('1', 'hi'),('1', 'hello world'),('1', 'hello world, goodbye'),('1', 'hello world, goodbye earth'),('2', 'bonjour'),('2', 'gracias'),('2','bonjour madame')]
spark.createDataFrame(input,['group','input']).show(10,truncate=False)

Вот входные данные:

Ниже приведен желаемый результат:

Я использую pyspark, но если кто-то знает, как это реализоватьв Python я могу перевести его в pyspark.

Answer 1

Сначала мы используем оконную функцию для разделения DF по столбцу ' group '. Мы снова группируем каждую из этих групп по первому слову каждой строки столбца ' input '. Мы делаем это с помощью функции substring_index . После создания оконной функции мы добавляем столбец в DF, подсчитывая размер каждой группы, созданной оконной функцией, - проверяя, имеет ли группа одну строку или более.

from pyspark.sql import Window
from pyspark.sql.functions import *

w = Window.partitionBy(df.group, substring_index(df.input, ' ', 1))
df.withColumn('output', when(count('*').over(w) > 1, 'TOGETHER').otherwise('STAND ALONE'))