Apache spark dealing with case statements

by

These are few ways to write If-Else / When-Then-Else / When-Otherwise expression in pyspark.

Sample dataframe

df = spark.createDataFrame([(1,1),(2,2),(3,3)],['id','value'])

df.show()

#+---+-----+
#| id|value|
#+---+-----+
#|  1|    1|
#|  2|    2|
#|  3|    3|
#+---+-----+

#Desired Output:
#+---+-----+----------+
#| id|value|value_desc|
#+---+-----+----------+
#|  1|    1|       one|
#|  2|    2|       two|
#|  3|    3|     other|
#+---+-----+----------+

Option#1: withColumn() using when-otherwise

from pyspark.sql.functions import when

df.withColumn("value_desc",when(df.value == 1, 'one').when(df.value == 2, 'two').otherwise('other')).show()

Option#2: select() using when-otherwise

from pyspark.sql.functions import when

df.select("*",when(df.value == 1, 'one').when(df.value == 2, 'two').otherwise('other').alias('value_desc')).show()

Option3: selectExpr() using SQL equivalent CASE expression

df.selectExpr("*","CASE WHEN value == 1 THEN  'one' WHEN value == 2 THEN  'two' ELSE 'other' END AS value_desc").show()

SQL like expression can also be written in withColumn() and select() using pyspark.sql.functions.expr function. Here are examples.

Option4: select() using expr function

from pyspark.sql.functions import expr 

df.select("*",expr("CASE WHEN value == 1 THEN  'one' WHEN value == 2 THEN  'two' ELSE 'other' END AS value_desc")).show()

Option5: withColumn() using expr function

from pyspark.sql.functions import expr 

df.withColumn("value_desc",expr("CASE WHEN value == 1 THEN  'one' WHEN value == 2 THEN  'two' ELSE 'other' END AS value_desc")).show()

Output:

#+---+-----+----------+
#| id|value|value_desc|
#+---+-----+----------+
#|  1|    1|       one|
#|  2|    2|       two|
#|  3|    3|     other|
#+---+-----+----------+

More Related Contents:

  1. TypeError: Column is not iterable – How to iterate over ArrayType()?
  2. Finding duplicates from large data set using Apache Spark
  3. Using a column value as a parameter to a spark DataFrame function
  4. How do I split an RDD into two or more RDDs?
  5. How to improve performance for slow Spark jobs using DataFrame and JDBC connection?
  6. Python Spark Cumulative Sum by Group Using DataFrame
  7. How to save/insert each DStream into a permanent table
  8. Spark parquet partitioning : Large number of files
  9. Apache Spark: What is the equivalent implementation of RDD.groupByKey() using RDD.aggregateByKey()?
  10. Pyspark: Pass multiple columns in UDF
  11. Spark ALS predictAll returns empty
  12. Why is the fold action necessary in Spark?
  13. Find maximum row per group in Spark DataFrame
  14. How to find median and quantiles using Spark
  15. How to access element of a VectorUDT column in a Spark DataFrame?
  16. Spark RDD to DataFrame python
  17. Link Spark with iPython Notebook
  18. How DAG works under the covers in RDD?
  19. Spark load data and add filename as dataframe column
  20. Spark groupByKey alternative
  21. How to extract an element from a array in pyspark
  22. PySpark: how to resample frequencies
  23. Explode array data into rows in spark [duplicate]
  24. PySpark DataFrames – way to enumerate without converting to Pandas?
  25. How to find count of Null and Nan values for each column in a PySpark dataframe efficiently?
  26. pyspark: count distinct over a window
  27. How spark read a large file (petabyte) when file can not be fit in spark’s main memory
  28. Serialize a custom transformer using python to be used within a Pyspark ML pipeline
  29. Not able to cat dbfs file in databricks community edition cluster. FileNotFoundError: [Errno 2] No such file or directory:
  30. pyspark: rolling average using timeseries data

Leave a Comment