Core DataFrame Operations

This chapter covers the DataFrame steps you use every day: choose columns and rows, add columns, sort and remove repeats.

Each step returns a new DataFrame, and nothing runs until an action such as show() (chapter 5).

The examples use orders.csv from chapter 3, read with inferSchema=True, and the customers:

Two NULL values matter in this chapter. Dev has no country, and order 105 has no amount.

select keeps the columns you name. filter keeps the rows where a condition is true. col("status") points to a column by its name:

from pyspark.sql.functions import col

(orders
 .select("order_id", "status", "amount")
 .filter(col("status") == "paid")
 .show())

selectExpr takes SQL text instead of column objects. It is handy for a quick new column:

orders.selectExpr("order_id", "amount * 100 AS cents").filter("order_id <= 102").show()

A condition can be true, false or unknown. A comparison with NULL is unknown. filter keeps a row only when the condition is true.

So a NULL row can vanish without an error. Look for the customers who do not live in India (IN):

customers.filter(col("country") != "IN").show()

You might expect the opposite filter to return Dev. It does not:

customers.filter(~(col("country") != "IN")).show()

Five customers went in, but 2 + 2 came out. Dev is in neither result. Ask for the NULL rows on purpose with isNull():

customers.filter((col("country") != "IN") | col("country").isNull()).show()

withColumn adds a column, or replaces one with the same name. withColumnRenamed changes a name, and drop removes a column.

lit("shop") makes a column that holds the same value in every row. when(...).otherwise(...) picks a value by condition, like if and else:

from pyspark.sql.functions import lit, when

(orders
 .withColumn("source", lit("shop"))
 .withColumn("size", when(col("amount") >= 50, "large")
                     .when(col("amount") >= 20, "medium")
                     .otherwise("small"))
 .withColumnRenamed("customer_id", "cust")
 .drop("order_date")
 .show())

Without otherwise, a row that matches no when gets NULL, not its old value. For example, when(col("amount") >= 50, "large") alone sizes only orders 103 and 106. The other four get NULL.

cast: change a type

cast changes the type of a column. Here a made-up table holds amounts as text, and one is broken:

raw_amounts = spark.createDataFrame(
    [(104, "15.50"), (105, ""), (106, "n/a")], "order_id int, amount string")

raw_amounts.withColumn("amount", col("amount").cast("double")).show()

try_cast turns a bad value into NULL instead of stopping:

raw_amounts.withColumn("amount", col("amount").try_cast("double")).show()

distinct removes repeated rows. orderBy sorts. limit(n) keeps the first n rows:

orders.select("status").distinct().orderBy("status").limit(2).show()

distinct and orderBy are wide steps: they move rows between tasks in a shuffle (chapter 5). On big data, filter and select first, so they move less data.

Where do NULL values go in a sort? Here are the amounts, largest first:

orders.orderBy(col("amount").desc()).select("order_id", "amount").show()

dropDuplicates(["customer_id"]) keeps one row for each customer. But it does not promise which row:

orders.dropDuplicates(["customer_id"]).select("order_id", "customer_id").show()

To keep the newest row for each key every time, number the rows with a window function. Chapter 12, Window Functions, shows how.

Put the steps in brackets, one step on each line. Filter first, so later steps see less data, and sort last:

(orders
 .filter(col("status") == "paid")
 .withColumn("amount_cents", (col("amount") * 100).cast("int"))
 .select("order_id", "customer_id", "amount_cents")
 .orderBy(col("amount_cents").desc())
 .show())

To add many columns, do not call withColumn in a long loop. Each call makes the plan bigger, and a very big plan can take minutes to build before any task runs.

Use one select(*columns) with all the new columns, or withColumns({...}), which takes many columns at once (Spark 3.3 and later).

Each step has a cost and one trap to remember:

Check yourself: filter country != "IN" gives 2 customers and its opposite gives 2. Where is the fifth? (Answer: Dev, whose country is NULL.)