Pyspark conditional fill a column

user_14353779 0
Viewed 33

I want to make a column id1 in ss_df such that if length of id is 13, then take substring from 6th digit to end of digits; else when length of id is 9 take the substring from the 3rd digit to the end. How do I achive this using pyspark? ss_df :

id id1
9999945678909 45678909
789065437 9065437
67898095 67898095

What I tried :

from pyspark.sql.functions import length

if (length(col("id")) == 13):
  ss_df = ss_df.withColumn("id1", substring("id",6,15))
if (length(col("id")) == 9):
  ss_df = ss_df.withColumn("id1", substring("id",3,15))
pyspark
1 Answers

pyspark functions provide utility when in the following manner to create and populate column conditionally

>>> from pyspark.sql import functions as F
>>> df = spark.createDataFrame([("abcde", 1), ("abc", 2), ("a", 3)], ["value", "r"])
>>> df = df.withColumn('id', F.when(F.length(df.value) == 3, F.substring(df.value, 1, 10)).otherwise(F.when(F.length(df.value) == 5, F.substring(df.value, 2, 10)).otherwise(None)))
>>> df.show()
+-----+---+----+                                                                
|value|  r|  id|
+-----+---+----+
|abcde|  1|bcde|
|  abc|  2| abc|
|    a|  3|null|

Note: In Function F.substring(Col, index, len), index is not 0 based but 1 based index.

user_74156340
Related