Introduction

Spark DataFrame provides a drop() method to drop a column/field from a DataFrame/Dataset. The drop() method is also used to remove multiple columns at a time from a Spark DataFrame/Dataset.

First, let’s create a DataFrame.

import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import java.util.Arrays;
import java.util.List;

SparkSession spark = SparkSession.builder()
    .appName("Spark Drop Column")
    .master("local")
    .getOrCreate();

List<Row> structureData = Arrays.asList(
    RowFactory.create("James", "", "Smith", "36636", "NewYork", 3100),
    RowFactory.create("Michael", "Rose", "", "40288", "California", 4300),
    RowFactory.create("Robert", "", "Williams", "42114", "Florida", 1400),
    RowFactory.create("Maria", "Anne", "Jones", "39192", "Florida", 5500),
    RowFactory.create("Jen", "Mary", "Brown", "34561", "NewYork", 3000)
);

StructType structureSchema = new StructType(new StructField[]{
    DataTypes.createStructField("firstname", DataTypes.StringType, true),
    DataTypes.createStructField("middlename", DataTypes.StringType, true),
    DataTypes.createStructField("lastname", DataTypes.StringType, true),
    DataTypes.createStructField("id", DataTypes.StringType, true),
    DataTypes.createStructField("location", DataTypes.StringType, true),
    DataTypes.createStructField("salary", DataTypes.IntegerType, true)
});

Dataset<Row> df = spark.createDataFrame(structureData, structureSchema);

df.printSchema();

This yields the output:

root
 |-- firstname: string (nullable = true)
 |-- middlename: string (nullable = true)
 |-- lastname: string (nullable = true)
 |-- id: string (nullable = true)
 |-- location: string (nullable = true)
 |-- salary: integer (nullable = true)

1. drop() syntax

Spark drop() has 3 different signatures. I’ve explained using all these signatures with examples in the below sections.

// Spark DataFrame drop() syntax
1) drop(String colName)
2) drop(String... colNames)
3) drop(Column col)

2. Drop one column

The first and third signature takes the column name as String type and Column type, respectively. When you use the third signature make sure you import org.apache.spark.sql.functions.col

// Drop one column from DataFrame
Dataset<Row> df2 = df.drop("firstname"); // First signature
df2.printSchema();

df.drop(df.col("firstname")).printSchema();

import static org.apache.spark.sql.functions.col;
df.drop(col("firstname")).printSchema(); // Third signature

The above 3 examples remove the column “firstname” from the DataFrame. You can use either one of these, depending on your needs.

root
 |-- middlename: string (nullable = true)
 |-- lastname: string (nullable = true)
 |-- id: string (nullable = true)
 |-- location: string (nullable = true)
 |-- salary: integer (nullable = true)

3. Drop Multiple Columns from DataFrame

This uses the second signature of the drop(), which removes more than one column from a DataFrame.

// Referencing more than one column
df.drop("firstname", "middlename", "lastname")
  .printSchema();

// Using array of columns
String[] cols = {"firstname", "middlename", "lastname"};
df.drop(cols)
  .printSchema();

The above two examples remove more than one column at a time from DataFrame. These both yield the same output.

root
 |-- id: string (nullable = true)
 |-- location: string (nullable = true)
 |-- salary: integer (nullable = true)

The complete example is also available at Spark Examples Github project for references.