how to make a new column by pairing elements of the other column?

user_16831723 0
Viewed 217

I have a big data dataframe and I want to make pairs from elements of the other column.

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

output: the pair of the above rows:

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

Note that tuple will work instead of list. like ('summer','book').

I know in pandas I can do this:

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

but not sure in pyspark.

pyspark
1 Answers

You can use a UDF to do the same as you would do in python. Then cast the output to an array of array of strings.

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)
user_48262950
Related