Column Expressions

Column Expressions

Column Expressions and Common DataFrame Operations

A Column is an expression, not data: F.col("total") >= 50 builds a tree that Catalyst compiles. Combine conditions with &, | and ~ (with parentheses), and look in pyspark.sql.functions before writing a Python function of your own (When to Reach for a UDF at All).

Column expressions on BookNest orders

orders = spark.read.parquet("data/orders.parquet")
shaped = (orders
          .withColumns({
              "month": F.date_format("order_ts", "yyyy-MM"),
              "n_lines": F.size("items"),
              "copies": F.aggregate("items", F.lit(0), lambda acc, i: acc + i.qty),
              "size": F.when(F.col("total") >= 50, "large")
                       .when(F.col("total") >= 20, "medium").otherwise("small")})
          .withColumnRenamed("channel", "app")
          .where(F.col("status").isin("delivered", "shipped") & F.col("coupon").isNull())
          .select("order_id", "month", "app", "n_lines", "copies", "total", "size"))
shaped.show(3)

Output

+--------+-------+-------+-------+------+-----+------+
|order_id|  month|    app|n_lines|copies|total|  size|
+--------+-------+-------+-------+------+-----+------+
|       1|2025-01|android|      1|     1|14.99| small|
|       2|2025-01|android|      2|     3|46.18|medium|
|       3|2025-01|android|      1|     2|32.40|medium|
+--------+-------+-------+-------+------+-----+------+
only showing top 3 rows

withColumns adds several columns in one projection; long chains of withColumn calls inflate the plan and slow analysis. F.aggregate is a higher-order function: its lambda is translated into a JVM expression over the items array and never runs in Python. when/otherwise is SQL's CASE; F.expr(...) accepts SQL text.

MENU