Pyspark concat column based on other columns values

Question

I am working with pyspark below is my dataframe

col_list   | col_1| col_2 |...|col_n
col_1,col2 | x    | y     |...|n
col_2,col_n| a    | b     |...|z

I want concat columns which are present in my col_list column values

Expected DF:

col_list   | col_1| col_2 |...|col_n | concat_cols
col_1,col2 | x    | y     |...|n     | [col_1:x,col_2:y]
col_2,col_n| a    | b     |...|z     | [col_2:b,col_n:z]

Trying with below approaches:

Creating udf to pass col_list and create expression to create array col like array(df['col_1'],df['col_2']) , but this will create column with string , not sure how to execute expression
Using create map passing col_list in loop and create column, need to check on this.

mck · Accepted Answer · 2021-04-19 14:04:11Z

2

You can use some map methods:

import pyspark.sql.functions as F

df2 = df.withColumn(
    'col_map', 
    F.map_from_arrays(
        F.array([F.lit(c) for c in df.columns[1:]]), 
        F.array(df.columns[1:])
    )
).withColumn(
    'concat_cols', 
    F.expr("""
        map_from_arrays(
            split(col_list, ','), 
            transform(split(col_list, ','), x -> col_map[x])
        )
    """)
).drop('col_map')

df2.show(truncate=False)
+-----------+-----+-----+-----+------------------------+
|col_list   |col_1|col_2|col_n|concat_cols             |
+-----------+-----+-----+-----+------------------------+
|col_1,col_2|x    |y    |n    |[col_1 -> x, col_2 -> y]|
|col_2,col_n|a    |b    |z    |[col_2 -> b, col_n -> z]|
+-----------+-----+-----+-----+------------------------+

If you prefer a string representation rather than a map type column, you can do

import pyspark.sql.functions as F

df2 = df.withColumn(
    'col_map', 
    F.map_from_arrays(
        F.array([F.lit(c) for c in df.columns[1:]]), 
        F.array(df.columns[1:])
    )
).withColumn(
    'concat_cols', 
    F.expr("""
        concat_ws(',', 
            transform(split(col_list, ','), x -> concat_ws(':', x, col_map[x]))
        )
    """)
).drop('col_map')

df2.show(truncate=False)
+-----------+-----+-----+-----+---------------+
|col_list   |col_1|col_2|col_n|concat_cols    |
+-----------+-----+-----+-----+---------------+
|col_1,col_2|x    |y    |n    |col_1:x,col_2:y|
|col_2,col_n|a    |b    |z    |col_2:b,col_n:z|
+-----------+-----+-----+-----+---------------+

answered Apr 19, 2021 at 14:04

mck

42.7k13 gold badges44 silver badges62 bronze badges

Sign up to request clarification or add additional context in comments.

6 Comments

Prabhanj Over a year ago

i am getting [col_1 -> x, col_2 -> ] in concat_cols

mck Over a year ago

did you have spaces in your column names? that may cause issues

Prabhanj Over a year ago

there are chances, happening to few columns

mck Over a year ago

Then you need to make sure to remove spaces. See stackoverflow.com/questions/57318519/…

Prabhanj Over a year ago

getting spaces for some records in col_lis causing issue, fixed by chaining col_map[x] to col_map[trim(x)]

|

Collectives™ on Stack Overflow

Pyspark concat column based on other columns values

1 Answer 1

6 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

6 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related