30. Higher-order functions on arrays

Higher-order functions transform an array in place — without exploding it into rows and regrouping. They take a small Python lambda that runs on each element. Let's use a simple basket of quantities:

baskets = spark.createDataFrame(
    [(1, [10, 20, 30]), (2, [5, 5]), (3, [100])],
    ["basket_id", "quantities"],
)
baskets.show(truncate=False)

transform — map over each element

transform(array, lambda) returns a new array with the lambda applied to every element:

baskets.select(
    "basket_id",
    F.transform("quantities", lambda x: x * 2).alias("doubled"),
).show(truncate=False)

filter — keep some elements

filter(array, lambda) keeps only elements where the lambda is true — note this is the array F.filter, different from the DataFrame .filter method:

baskets.select(
    "basket_id",
    F.filter("quantities", lambda x: x >= 10).alias("big_only"),
).show(truncate=False)

"Does any element match?" — there's no exists here, but size(filter(...)) > 0 says the same thing and runs everywhere:

baskets.select(
    "basket_id",
    (F.size(F.filter("quantities", lambda x: x >= 50)) > 0).alias("has_big"),
).show()

aggregate — reduce an array to one value

aggregate(array, start, lambda) folds the array down — a sum, product, max, built by hand:

baskets.select(
    "basket_id",
    F.aggregate("quantities", F.lit(0), lambda acc, x: acc + x).alias("total_qty"),
).show()

These beat explode-then-regroup when you want to stay at one row per record — no shuffle, no join back.

Your turn

For each basket, produce an array of the quantities increased by 1 (alias bumped), using transform. Return basket_id and bumped. Assign to result.

baskets = spark.createDataFrame(
    [(1, [10, 20, 30]), (2, [5, 5]), (3, [100])],
    ["basket_id", "quantities"],
)

result = baskets  # <- select basket_id, transform(quantities, x -> x + 1) AS bumped
result.show(truncate=False)

result = baskets.select(
    "basket_id",
    F.transform("quantities", lambda x: x + 1).alias("bumped"),
)