CODEWITHSUNDEEP
pythonapache-sparkpysparkuser-defined-functions
spark-health-learn
spark-health-learn

Reputation: 49

How to create a UDF that creates a new column AND modifies an existing column

I have a dataframe like this:

id | color
---| -----
1  | red-dark
2  | green-light
3  | red-light
4  | blue-sky
5  | green-dark

I would like to create a UDF such that my dataframe becomes:

id | color | shade
---| ----- | -----
1  | red   |  dark
2  | green |  light
3  | red   |  light
4  | blue  |  sky
5  | green |  dark

I've written a UDF for this:

def my_function(data_str):
    return ",".join(data_str.split("-"))

my_function_udf = udf(my_function, StringType())

#apply the UDF

df = df.withColumn("shade", my_function_udf(df['color']))

However, this doesn't transform the dataframe as I intend it to be. instead it turns it into:

id | color      | shade
---| ---------- | -----
1  | red-dark   |  red,dark
2  | green-dark |  green,light
3  | red-light  |  red,light
4  | blue-sky   |  blue,sky
5  | green-dark |  green,dark

How can I transform the dataframe as I want it in pyspark?

Tried based on suggested question

schema = ArrayType(StructType([
    StructField("color", StringType(), False),
    StructField("shade", StringType(), False)
]))

color_shade_udf = udf(
    lambda s: [tuple(s.split("-"))],
    schema
)

df = df.withColumn("colorshade", color_shade_udf(df['color']))

#Gives the following

id | color      | colorshade
---| ---------- | -----
1  | red-dark   |  [{"color":"red","shade":"dark"}]
2  | green-dark |  [{"color":"green","shade":"dark"}]
3  | red-light  |  [{"color":"red","shade":"light"}]
4  | blue-sky   |  [{"color":"blue","shade":"sky"}]
5  | green-dark |  [{"color":"green","shade":"dark"}]

I feel like I am getting closer

Upvotes: 1

Views: 962

Answers (1)

mtoto
mtoto

Reputation: 24178

You can use the built-in function split():

from pyspark.sql.functions import split, col

df.withColumn("arr", split(df.color, "\\-")) \
  .select("id", 
          col("arr")[0].alias("color"),
          col("arr")[1].alias("shade")) \
  .drop("arr") \
  .show()
+---+-----+-----+
| id|color|shade|
+---+-----+-----+
|  1|  red| dark|
|  2|green|light|
|  3|  red|light|
|  4| blue|  sky|
|  5|green| dark|
+---+-----+-----+

Upvotes: 2

Related Questions