Spark Add Constant Column to DataFrame

Spark SQL functions lit()and typedLit() are used to add a new constant column to DataFrame by assigning a literal or constant value. Both of these functions are available in Spark by importing org.apache.spark.sql.functions and they return Column type.

lit() Function - Syntax:


// lit() Function - Syntax:
lit(literal : scala.Any) : org.apache.spark.sql.Column

Creates a [[Column]] of literal value. The passed in object is returned directly if it is already a [[Column]]. If the object is a Scala Symbol, it is converted into a [[Column]] also. Otherwise, a new [[Column]] is created to represent the literal value

Let’s take a look at some Spark code.


import spark.sqlContext.implicits._
import org.apache.spark.sql.functions._
  
val data = Seq(("111",50000),("222",60000),("333",40000))
val df = data.toDF("EmpId","Salary")

1. Spark Add Constant Column using lit() function

Let's see a scala example of how to create a new column with constant value using lit() Spark SQL function. On the below snippet, we are creating a new column by adding a literal '1' to Spark DataFrame.


// Spark Add Constant Column using lit() function
val df2 = df.select(col("EmpId"),col("Salary"),lit("1").as("lit_value1"))
df2.show()


// Output:
+-----+------+----------+
|EmpId|Salary|lit_value1|
+-----+------+----------+
|  111| 50000|         1|
|  222| 60000|         1|
|  333| 40000|         1|
+-----+------+----------+

Adding the same constant literal to all records in DataFrame may not be real-time use full so let's see another example.

2. lit() function with withColumn

The following scala code example shows how to use lit() Spark sql function, using withColumn to derive a new column based on some conditions.


// lit() function with withColumn
val df3 = df2.withColumn("lit_value2",
    when(col("Salary") >=40000 && col("Salary") <= 50000, 
       lit("100").cast(IntegerType))
      .otherwise(lit("200").cast(IntegerType))
  )
  df3.show()

Below is the output for the above code snippet.


// Output:
+-----+------+----------+----------+
|EmpId|Salary|lit_value1|lit_value2|
+-----+------+----------+----------+
|  111| 50000|         1|       100|
|  222| 60000|         1|       200|
|  333| 40000|         1|       100|
+-----+------+----------+----------+

3. typedLit() Function - Syntax


def typedLit[T : TypeTag](literal: T): org.apache.spark.sql.Column

Difference between lit and typedLit is that this function can handle parameterized scala types e.g.: List, Seq and Map.

Example 3 : typedLit() with Seq and Map

Following example shows on how to create a new column with collection using typedLit() sql function. On below snippet, we are creating a new column by adding a collection literal Seq(1, 2, 3), Map("a" -> 1, "b" -> 2) and structure ("a", 2, 1.0) to spark DataFrame.


val df4 = df3.withColumn("typedLit_seq",typedLit(Seq(1, 2, 3)))
    .withColumn("typedLit_map",typedLit(Map("a" -> 1, "b" -> 2)))
    .withColumn("typedLit_struct",typedLit(("a", 2, 1.0)))
df4.printSchema()
df4.show()

df4.printSchema() outputs below schema


// Output:
root
 |-- EmpId: string (nullable = true)
 |-- Salary: integer (nullable = false)
 |-- lit_value1: string (nullable = false)
 |-- lit_value2: integer (nullable = true)
 |-- typedLit_seq: array (nullable = false)
 |    |-- element: integer (containsNull = false)
 |-- typedLit_map: map (nullable = false)
 |    |-- key: string
 |    |-- value: integer (valueContainsNull = false)
 |-- typedLit_struct: struct (nullable = false)
 |    |-- _1: string (nullable = true)
 |    |-- _2: integer (nullable = false)
 |    |-- _3: double (nullable = false)

df4.show() outputs below data frame.


// Output:
+-----+------+----------+----------+------------+----------------+--------------+
|EmpId|Salary|lit_value1|lit_value2|typedLit_seq|    typedLit_map|typedLit_struct|
+-----+------+----------+----------+------------+----------------+--------------+
|  111| 50000|         1|       200|   [1, 2, 3]|[a -> 1, b -> 2]|   [a, 2, 1.0]|
|  222| 60000|         1|       200|   [1, 2, 3]|[a -> 1, b -> 2]|   [a, 2, 1.0]|
|  333| 40000|         1|       200|   [1, 2, 3]|[a -> 1, b -> 2]|   [a, 2, 1.0]|
+-----+------+----------+----------+------------+----------------+--------------+

4. Complete Example:


package com.aistudylabs.spark.dataframe.functions
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.types.IntegerType
object litTypeLit extends App {
  val spark = SparkSession.builder()
    .appName("aistudylabs.com")
    .master("local")
    .getOrCreate()
  import spark.sqlContext.implicits._
  import org.apache.spark.sql.functions._
  val data = Seq(("111",50000),("222",60000),("333",40000))
  val df = data.toDF("EmpId","Salary")
  val df2 = df.select(col("EmpId"),col("Salary"),lit("1").as("lit_value1"))
  df2.show()
  val df3 = df2.withColumn("lit_value2",
    when(col("Salary") >=40000 && col("Salary")  1, "b" -> 2)))
    .withColumn("typedLit_struct",typedLit(("a", 2, 1.0)))
  df4.printSchema()
  df4.show()
}

Conclusion:

You have learned multiple ways to add a constant literal value to DataFrame using Spark SQL lit() function and have learned the difference between lit and typedLit functions.

When possible try to use predefined Spark SQL functions as they are a little bit more compile-time safety and perform better when compared to user-defined functions. If your application is critical on performance try to avoid using custom UDF functions as these are not guarantee on performance.

Happy Learning !!