Introduction

In Apache Spark using Java, the withColumn() method of a DataFrame is utilized to add a new column to the DataFrame, change the value of an existing column, convert the datatype of a column, or derive a new column from an existing column.

Spark withColumn() Syntax and Usage

The withColumn() method in Spark's DataFrame functions as a transformation operation used to manipulate the column values of all or selected rows in a DataFrame.

Calling the withColumn() method returns a new DataFrame after processing actions like adding, updating, or deriving columns.

Below is the syntax for the Java withColumn() function:

DataFrame df = df.withColumn(String colName, Column col);

• colName:String - The name of the new column you wish to create or update.
• col:Column - The column expression used for derivation or update.

Since withColumn() is a transformation function, execution only occurs when an action is invoked.

Using the Spark withColumn() method adds a projection internally, which may cause large plans and performance issues if used extensively, such as within loops. In such cases, it is better to use the select method to work with multiple columns efficiently.

First, let's create a DataFrame to work with using Java Spark:

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.*;

import java.util.Arrays;
import java.util.List;

SparkSession spark = SparkSession.builder()
        .appName("Spark withColumn Example")
        .master("local")
        .getOrCreate();

List<Row> data = Arrays.asList(
        RowFactory.create(RowFactory.create("James", "", "Smith"), "36636", "M", "3000"),
        RowFactory.create(RowFactory.create("Michael", "Rose", ""), "40288", "M", "4000"),
        RowFactory.create(RowFactory.create("Robert", "", "Williams"), "42114", "M", "4000"),
        RowFactory.create(RowFactory.create("Maria", "Anne", "Jones"), "39192", "F", "4000"),
        RowFactory.create(RowFactory.create("Jen", "Mary", "Brown"), "", "F", "-1")
);

StructType schema = new StructType(new StructField[]{
        new StructField("name", new StructType(new StructField[]{
                new StructField("firstname", DataTypes.StringType, false, Metadata.empty()),
                new StructField("middlename", DataTypes.StringType, false, Metadata.empty()),
                new StructField("lastname", DataTypes.StringType, false, Metadata.empty())
        }), false, Metadata.empty()),
        new StructField("dob", DataTypes.StringType, false, Metadata.empty()),
        new StructField("gender", DataTypes.StringType, false, Metadata.empty()),
        new StructField("salary", DataTypes.StringType, false, Metadata.empty())
});

Dataset<Row> df = spark.createDataFrame(data, schema);

1. Add a New Column to DataFrame

To add a new column, use the withColumn() method, specifying the column name and expression. Ensure the column doesn’t pre-exist; otherwise, it will update the column.

import org.apache.spark.sql.functions;

df = df.withColumn("Country", functions.lit("USA"));

// Chaining to operate on multiple columns
df = df.withColumn("Country", functions.lit("USA"))
       .withColumn("AnotherColumn", functions.lit("AnotherValue"));

While this approach is acceptable for a small number of columns, chaining withColumn() is not recommended for multiple columns; use select() instead for better performance.

2. Change Value of an Existing Column

The withColumn() method can update existing column values by using the column’s current name and desired expression.

// Change Value of an Existing Column
df = df.withColumn("salary", df.col("salary").multiply(100));

This example multiplies the value of “salary” by 100 and updates the “salary” column.

3. Derive a New Column From an Existing Column

To derive a new column, specify the desired column name and apply an operation on an existing column with an expression.

// Derive New Column From an Existing Column
df = df.withColumn("CopiedColumn", df.col("salary").multiply(-1));

This creates a column “CopiedColumn” by multiplying “salary” by -1.

4. Change Column Data Type

Using the cast() function with withColumn(), you can change a column’s datatype. This snippet changes the “salary” from String to Integer.

// Change Column Data Type
df = df.withColumn("salary", df.col("salary").cast("integer"));

5. Split Column into Multiple Columns

Though not utilizing withColumn(), splitting a DataFrame column into multiple columns can be achieved using the map() transformation in Java Spark.

// Split Column into Multiple Columns
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.Encoders;
import scala.Tuple6;

Dataset<Row> dfFromData = spark.createDataFrame(Arrays.asList(
        RowFactory.create("Robert, Smith", "1 Main st, Newark, NJ, 92537"),
        RowFactory.create("Maria, Garcia", "3456 Walnut st, Newark, NJ, 94732")),
        new StructType(new StructField[]{
                new StructField("name", DataTypes.StringType, false, Metadata.empty()),
                new StructField("address", DataTypes.StringType, false, Metadata.empty())
        }));

Dataset<Tuple6<String, String, String, String, String, String>> newDF = dfFromData.map(
        (MapFunction<Row, Tuple6<String, String, String, String, String, String>>) row -> {
            String[] nameSplit = row.getString(0).split(", ");
            String[] addSplit = row.getString(1).split(", ");
            return new Tuple6<>(nameSplit[0], nameSplit[1], addSplit[0], addSplit[1], addSplit[2], addSplit[3]);
        }, Encoders.tuple(Encoders.STRING(), Encoders.STRING(), Encoders.STRING(), Encoders.STRING(), Encoders.STRING(), Encoders.STRING()));

Dataset<Row> finalDF = newDF.toDF("First Name", "Last Name", "Address Line1", "City", "State", "ZipCode");
finalDF.printSchema();
finalDF.show(false);

This splits the “name” column into “First Name” and “Last Name” and the “address” column into “Address Line1,” “City,” “State,” and “ZipCode.”

The output of the final DataFrame will be:

root
 |-- First Name: string (nullable = true)
 |-- Last Name: string (nullable = true)
 |-- Address Line1: string (nullable = true)
 |-- City: string (nullable = true)
 |-- State: string (nullable = true)
 |-- ZipCode: string (nullable = true)

+----------+---------+--------------+-------+-----+-------+
|First Name|Last Name|Address Line1 |City   |State|ZipCode|
+----------+---------+--------------+-------+-----+-------+
|Robert    | Smith   |1 Main st     | Newark| NJ  | 92537 |
|Maria     | Garcia  |3456 Walnut st| Newark| NJ  | 94732 |
+----------+---------+--------------+-------+-----+-------+

The complete code can be downloaded from GitHub.