create new pyspark dataframe columns by splitting values on white space

Question

I have a pyspark dataframe like the input data below. I would like to split the values in the productname column on white space. I'd then like to create new columns with the first 3 values. I have example input and output data below. Can someone please suggest how to do this with pyspark?

input data:

+------+-------------------+
|id    |productname        |
+------+-------------------+
|235832|EXTREME BERRY Sweet|             
|419736|BLUE CHASER SAUCE  |             
|124513|LAAVA C2L5         |
+------+-------------------+

output:

+------+-------------------+-------------+-------------+-------------+
|id    |productname        |product1     |product2     |product3     |
+------+-------------------+-------------+-------------+-------------+
|235832|EXTREME BERRY Sweet|EXTREME      |BERRY        |Sweet        |
|419736|BLUE CHASER SAUCE  |BLUE         |CHASER       |SAUCE        |
|124513|LAAVA C2L5         |LAAVA        |C2L5         |             |
+------+-------------------+-------------+-------------+-------------+

can we assume u only need three more columns?(1,2,3) or could there be more depending on productname? — murtihash, Mar 25 '20 at 01:27

notNull · Accepted Answer · 2020-03-25T02:02:39.173

Split the productname column then create new columns using element_at (or) .getItem() on index value.

df.withColumn("tmp",split(col("productname"),"\s+")).\
withColumn("product1",element_at(col("tmp"),1)).\
withColumn("product2",element_at(col("tmp"),2)).\
withColumn("product3",coalesce(element_at(col("tmp"),3),lit(""))).drop("tmp").show()

#or

df.withColumn("tmp",split(col("productname"),"\s+")).\
withColumn("product1",col("tmp").getItem(0)).\
withColumn("product2",col("tmp").getItem(1)).\
withColumn("product3",coalesce(col("tmp").getItem(2),lit(""))).drop("tmp").show()
#+------+-------------------+--------+--------+--------+
#|    id|        productname|product1|product2|product3|
#+------+-------------------+--------+--------+--------+
#|235832|EXTREME BERRY Sweet| EXTREME|   BERRY|   Sweet|
#|     4|  BLUE CHASER SAUCE|    BLUE|  CHASER|   SAUCE|
#|     1|         LAAVA C2L5|   LAAVA|    C2L5|        |
#+------+-------------------+--------+--------+--------+

To do more dynamic way:

df.show()
#+------+-------------------+
#|    id|        productname|
#+------+-------------------+
#|235832|EXTREME BERRY Sweet|
#|     4|  BLUE CHASER SAUCE|
#|     1|         LAAVA C2L5|
#+------+-------------------+
#caluculate array max size and store into variable
arr=int(df.select(size(split(col("productname"),"\s+")).alias("size")).orderBy(desc("size")).collect()[0][0])

#loop through arr variable and add the columns replace null with ""
(df.withColumn('temp', split('productname', '\s+')).select("*",*(coalesce(col('temp').getItem(i),lit("")).alias('product{}'.format(i+1)) for i in range(arr))).drop("temp").show())

#+------+-------------------+--------+--------+--------+
#|    id|        productname|product1|product2|product3|
#+------+-------------------+--------+--------+--------+
#|235832|EXTREME BERRY Sweet| EXTREME|   BERRY|   Sweet|
#|     4|  BLUE CHASER SAUCE|    BLUE|  CHASER|   SAUCE|
#|     1|         LAAVA C2L5|   LAAVA|    C2L5|        |
#+------+-------------------+--------+--------+--------+

thank you, I like the dynamic version. I did have to edit the part using orderBy though to ".orderBy("size",ascending=False)". I don't know if that has to do with the version of pyspark I'm using or not. — user3476463, Mar 26 '20 at 01:16
@user3476463, **`.orderBy(desc("size"))`** is doing the same thing.. like we are ordering by descending! — notNull, Mar 26 '20 at 01:48

murtihash · Answer 2 · 2020-03-25T01:46:22.140

You can use split, element_at, and when/otherwise clause with array_union to put empty strings.

from pyspark.sql import functions as F
from pyspark.sql.functions import when
df.withColumn("array", F.split("productname","\ "))\
  .withColumn("array", F.when(F.size("array")==2, F.array_union(F.col("array"),F.array(F.lit(""))))\
                        .when(F.size("array")==1, F.array_union(F.col("array"),F.array(F.lit(" "),F.lit(""))))\
                        .otherwise(F.col("array")))\
  .withColumn("product1", F.element_at("array",1))\
  .withColumn("product2", F.element_at("array",2))\
  .withColumn("product3", F.element_at("array",3)).drop("array")\
  .show(truncate=False)

+------+-------------------+--------+--------+--------+
|id    |productname        |product1|product2|product3|
+------+-------------------+--------+--------+--------+
|235832|EXTREME BERRY Sweet|EXTREME |BERRY   |Sweet   |
|419736|BLUE CHASER SAUCE  |BLUE    |CHASER  |SAUCE   |
|124513|LAAVA C2L5         |LAAVA   |C2L5    |        |
|123455|LAVA               |LAVA    |        |        |
+------+-------------------+--------+--------+--------+

create new pyspark dataframe columns by splitting values on white space

2 Answers2