как создать новый столбец, объединив элементы другого столбца?

#pyspark

Вопрос:

У меня есть фрейм данных больших данных, и я хочу создавать пары из элементов другого столбца.

 col
['summer','book','hot']
['g','o','p']
 

вывод:
пара приведенных выше строк:

 new_col
['summer','book'],['summer','hot'],['hot','book']
['g','o'],['g','p'],['p','o']
 

Обратите внимание, что кортеж будет работать вместо списка. Нравится ('summer','book') .

Я знаю, что в панд я могу это сделать:

 df['col'].apply(lambda x: list(itertools.combinations(x, 2)))
 

но не уверен, что в пыспарке.

Ответ №1:

Вы можете использовать UDF для выполнения того же, что и в python. Затем приведите выходные данные к массиву массив строк.

 import itertools
from pyspark.sql import functions as F

combinations_udf = F.udf(
    lambda x: list(itertools.combinations(x, 2)), "array<array<string>>" 
)

  
df = spark.createDataFrame([(['hot','summer', 'book'],), 
                            (['g', 'o', 'p'], ),
                           ], ['col1'])

df1 = df.withColumn("new_col", combinations_udf(F.col("col1")))

display(df1)