DataFrame API: select, filter, withColumn, alias и Column expressions

Как устроен DataFrame API изнутри: декларативный подход, объект Column как expression tree, трансформации select/filter/withColumn, оптимизации Catalyst и anti-patterns производственного кода.

core

Почему DataFrame API - не просто «pandas для кластера»¶

Когда Data Engineer впервые видит PySpark DataFrame API, первый инстинкт - сравнить его с pandas. Цепочки методов, знакомые названия select, filter, groupBy. Но под похожим синтаксисом скрывается принципиально иная модель выполнения.

Разница в одном слове: декларативность.

В pandas операция df["col"] + 1 немедленно проходит по всему массиву в памяти и возвращает новый Series. В PySpark df["col"] + 1 создаёт объект-выражение - дерево из узлов логического плана. Никакие данные при этом не перемещаются. Spark не знает даже, сколько строк в DataFrame.

Это не тонкость реализации - это главное архитектурное решение, которое делает возможными все оптимизации: predicate pushdown, column pruning, join reordering, code generation.

Код разработчика попадает на вход многоступенчатого конвейера. На каждом шаге Catalyst что-то улучшает - применяет правила оптимизации, проверяет схему, выбирает физические стратегии. К моменту реального выполнения исходный код уже не узнать: Spark мог переставить фильтры, убрать лишние колонки и выбрать другой тип соединения. Задача разработчика - написать логически корректный код. Задача Catalyst - сделать его быстрым.

Immutability: DataFrame никогда не меняется¶

Любая трансформация создаёт новый объект DataFrame. Старый остаётся нетронутым.

df0 = spark.read.parquet("/data/orders")

df1 = df0.filter(col("status") == "completed")   # df0 не изменён
df2 = df1.withColumn("tax", col("amount") * 0.2) # df1 не изменён
df3 = df2.select("id", "amount", "tax")           # df2 не изменён

Каждый df - это ссылка на новый логический план. В памяти нет копий данных - только дерево операций. Физические данные вообще не трогаются до вызова Action (collect, show, count, write).

Это свойство принципиально важно для надёжности pipeline: промежуточный df можно передавать в функции, хранить в переменных, использовать несколько раз - он никогда не изменится под рукой.

Создание DataFrame¶

Чтение из источника¶

Типичный источник данных в реальных проектах - файлы в хранилище (Parquet, JSON, CSV) или таблицы в каталоге:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, lit, when, to_date, regexp_replace

spark = SparkSession.builder.appName("df-api-lesson").getOrCreate()

# Parquet - рекомендуемый формат для аналитических нагрузок
orders = spark.read.parquet("/data/orders/")

# JSON с явной схемой (inferSchema дорого на больших файлах)
from pyspark.sql.types import StructType, StructField, StringType, LongType, DoubleType

schema = StructType([
    StructField("order_id",   LongType(),   nullable=False),
    StructField("customer_id", LongType(),  nullable=True),
    StructField("amount",     DoubleType(), nullable=True),
    StructField("status",     StringType(), nullable=True),
    StructField("created_at", StringType(), nullable=True),
])

raw = spark.read.schema(schema).json("/data/raw/orders/2024-01/")

Создание из Python-данных (для тестов и отладки)¶

from pyspark.sql import Row

# Способ 1: createDataFrame из списка кортежей + columns
df = spark.createDataFrame(
    [(1, "Alice", 1500.0), (2, "Bob", 750.0), (3, "Carol", None)],
    schema=["customer_id", "name", "amount"]
)

# Способ 2: через Row - явное именование полей
data = [
    Row(order_id=101, status="completed", amount=320.0),
    Row(order_id=102, status="pending",   amount=150.0),
    Row(order_id=103, status="cancelled", amount=None),
]
orders = spark.createDataFrame(data)

Второй способ удобен в тестах: Row делает поля самодокументирующими и позволяет создать DataFrame с той же структурой, что в prod-коде.

Объект Column: дерево выражений¶

Прежде чем разбирать трансформации, важно понять, что такое Column в Spark.

col("amount") - это не значение и не ссылка на массив данных. Это узел дерева выражений (Expression), который Catalyst умеет анализировать и оптимизировать.

from pyspark.sql.functions import col

# Каждая из этих строк создаёт объект Column (Expression node)
c1 = col("amount")              # UnresolvedAttribute("amount")
c2 = col("amount") * 1.2       # Multiply(UnresolvedAttribute, Literal(1.2))
c3 = col("amount") > 100       # GreaterThan(UnresolvedAttribute, Literal(100))
c4 = c3 & (col("status") == "completed")  # And(GreaterThan, EqualTo)

Catalyst понимает это дерево и умеет его преобразовывать. Например, он может вычислить константное выражение заранее (Literal(1.2) * Literal(100) → Literal(120)) или перенести предикат ближе к источнику. Именно потому Python UDF непрозрачен для оптимизатора - он видит чёрный ящик вместо понятного дерева выражений.

Четыре способа сослаться на колонку¶

В PySpark одну и ту же колонку можно задать четырьмя способами:

# 1. Строка - работает в большинстве методов, но не даёт Column-объект
df.select("amount", "status")

# 2. df["col"] - Column привязана к конкретному DataFrame (безопасно при self-join)
df.select(df["amount"], df["status"])

# 3. col("col") - Column без привязки к DataFrame (чаще всего используется)
from pyspark.sql.functions import col
df.select(col("amount"), col("status"))

# 4. expr("sql expression") - SQL-выражение как строка
from pyspark.sql.functions import expr
df.select(expr("amount * 1.2 as amount_with_vat"))

Когда что использовать:

Способ	Когда применять
строка `"col"`	простые случаи: `select`, `drop`, `orderBy` по имени
`df["col"]`	self-join, когда важно различать DataFrame-источник
`col("col")`	построение сложных выражений, передача в функции
`expr("...")`	перенос SQL-логики без переписывания на Column API

Select: проекция колонок¶

select - это реализация оператора проекции из реляционной алгебры. Он не фильтрует строки, а выбирает и трансформирует колонки. Catalyst превращает select в Projection node логического плана, который при чтении Parquet позволяет пропустить целые колонки на диске.

Базовый select¶

# Выбрать подмножество колонок по именам
slim = orders.select("order_id", "amount", "status")

# То же самое через Column-объекты
slim = orders.select(col("order_id"), col("amount"), col("status"))

# Вычислить новую колонку прямо в select
with_vat = orders.select(
    col("order_id"),
    col("amount"),
    (col("amount") * 0.2).alias("vat"),
    (col("amount") * 1.2).alias("amount_with_vat"),
)

Вычисление alias("vat") происходит без промежуточного withColumn - одним select. Это важный приём: если нужно добавить много колонок, лучше сделать это в одном select, чем последовательными withColumn.

selectExpr: SQL-синтаксис в DataFrame pipeline¶

selectExpr принимает строки с SQL-выражениями. Удобен при миграции SQL-запросов или когда выражение сложно записать через Column API:

cleaned = raw.selectExpr(
    "order_id",
    "customer_id",
    "CAST(amount AS DOUBLE)         AS amount",
    "UPPER(TRIM(status))            AS status",
    "TO_DATE(created_at, 'yyyy-MM-dd') AS created_date",
    "amount * 0.2                   AS vat",
)

selectExpr эквивалентен select(expr(...)) для каждой строки. Catalyst разбирает SQL-строки и строит то же дерево выражений - никакой потери производительности по сравнению с Column API нет.

Alias: именование колонок¶

alias - это не просто косметика. Имена колонок влияют на корректность self-join и читаемость логического плана в explain().

# Без alias: после операции колонка называется "(amount * 0.2)"
df.select(col("amount") * 0.2)         # колонка → "(amount * 0.2)"

# С alias: понятное имя, с которым можно работать дальше
df.select((col("amount") * 0.2).alias("vat"))

Зачем alias нужен при self-join¶

# Self-join: найти повторяющихся клиентов в разных заказах
left  = orders.alias("l")
right = orders.alias("r")

# Без alias df["column"] - Spark не поймёт, из какого DataFrame
result = left.join(right, left["customer_id"] == right["customer_id"]) \
    .filter(left["order_id"] != right["order_id"]) \
    .select(left["order_id"].alias("order_a"), right["order_id"].alias("order_b"))

withColumnRenamed: переименование существующих колонок¶

# Единичное переименование
df = df.withColumnRenamed("customer_id", "cust_id")

# Переименование нескольких - через select + alias эффективнее
df = df.select(
    col("order_id"),
    col("customer_id").alias("cust_id"),
    col("amount").alias("order_amount"),
    col("status").alias("order_status"),
)

Второй вариант через select создаёт один узел плана вместо цепочки переименований.

Filter и Where: фильтрация строк¶

filter и where - это алиасы одного метода. Можно использовать любой по вкусу. В Spark оба создают одинаковый узел Filter в логическом плане.

# filter и where идентичны - используйте то, что читается лучше
completed = orders.filter(col("status") == "completed")
completed = orders.where(col("status") == "completed")

# SQL-строка тоже работает
completed = orders.filter("status = 'completed'")

Все предикатные выражения¶

# Равенство и неравенство
col("status") == "completed"
col("status") != "cancelled"

# Числовые сравнения
col("amount") > 100
col("amount") >= 100
col("amount").between(100, 500)    # включительно с обоих краёв

# Список допустимых значений
col("status").isin("completed", "shipped")
~col("status").isin("cancelled", "refunded")   # NOT IN

# Строковые предикаты
col("email").like("%@gmail.com")       # SQL LIKE, % - wildcard
col("name").startswith("Ив")
col("name").endswith("ов")
col("description").contains("акция")
col("phone").rlike(r"^\+7\d{10}$")    # регулярное выражение

# NULL-проверки (через специальные методы, не через ==)
col("amount").isNull()
col("amount").isNotNull()

Boolean logic: скобки обязательны¶

При комбинировании условий операторами & (AND), | (OR), ~ (NOT) скобки вокруг каждого условия обязательны. Это одна из самых частых ошибок:

# ОШИБКА: Python вычислит col("amount") > 100 | col("status") == "completed"
# как col("amount") > (100 | col("status")) == "completed" - неверно
orders.filter(col("amount") > 100 | col("status") == "completed")

# ПРАВИЛЬНО: каждое условие в скобках
orders.filter(
    (col("amount") > 100) | (col("status") == "completed")
)

# Сложные условия - выносить в переменные для читаемости
is_high_value = col("amount") > 500
is_completed  = col("status") == "completed"
is_not_test   = ~col("email").like("%@test.%")

result = orders.filter(is_completed & is_high_value & is_not_test)

NULL в фильтрации: трёхзначная логика SQL¶

В SQL NULL означает «неизвестно». Из этого следуют нелогичные на первый взгляд правила:

# amount содержит: [100, 200, None, 50]
# Ожидание: filter удалит None строку, остальное сравнится с 100
# Реальность: None != 100 возвращает NULL, а не True - строка теряется

df.filter(col("amount") != 100)   # строка с amount=None потеряна

# ПРАВИЛЬНО: явная проверка на NULL
df.filter(col("amount").isNotNull() & (col("amount") != 100))

# Для сравнения двух колонок, где обе могут быть NULL:
# eqNullSafe (<=>) возвращает True, если оба NULL или оба равны
df.filter(col("col_a").eqNullSafe(col("col_b")))

Правило: при фильтрации данных с потенциальными NULL всегда добавляйте явную проверку isNotNull() или используйте eqNullSafe. Иначе строки с NULL тихо исчезнут из результата.

WithColumn: создание и изменение колонок¶

withColumn(name, expression) добавляет новую колонку или перезаписывает существующую:

# Добавить новую колонку
df = orders.withColumn("vat",        col("amount") * 0.2)
df = df.withColumn("amount_with_vat", col("amount") + col("vat"))

# Перезаписать существующую (привести тип)
df = df.withColumn("amount", col("amount").cast("double"))

# Нормализовать строку
df = df.withColumn("status", F.upper(F.trim(col("status"))))

Проблема итеративного withColumn¶

Частая ошибка в production-коде - добавлять много колонок через цикл:

# АНТИПАТТЕРН: цепочка withColumn в цикле
features = ["a", "b", "c", ..., "z"]   # 50+ колонок
for feat in features:
    df = df.withColumn(f"{feat}_normalized", col(feat) / col(f"{feat}_max"))

Каждый withColumn добавляет новый узел в логический план. При 50+ итерациях план становится деревом глубиной 50, что вызывает StackOverflowError в Catalyst при компиляции. Кроме того, каждый новый узел делает explain() нечитаемым.

Решение: один select со всеми выражениями:

# ПРАВИЛЬНО: все колонки в одном select
from pyspark.sql.functions import col

base_cols = [col(c) for c in df.columns if c not in features]
new_cols  = [
    (col(f) / col(f"{f}_max")).alias(f"{f}_normalized")
    for f in features
]

df = df.select(*base_cols, *new_cols)

Один select - один Projection в плане, независимо от числа выражений.

Column expressions: пространство встроенных функций¶

pyspark.sql.functions - это более 300 встроенных функций, реализованных в JVM и выполняющихся внутри Tungsten. В отличие от Python UDF, они прозрачны для Catalyst и не требуют сериализации данных в Python.

import pyspark.sql.functions as F

# --- Константы ---
df.withColumn("country", F.lit("RU"))            # литеральное значение
df.withColumn("pi",      F.lit(3.14159))

# --- Типизация ---
df.withColumn("amount",  col("amount").cast("double"))
df.withColumn("date",    col("created_at").cast("date"))
df.withColumn("ts",      col("ts_str").cast("timestamp"))

# --- Строковые функции ---
df.withColumn("email",   F.lower(F.trim(col("email"))))
df.withColumn("domain",  F.regexp_replace(col("email"), r".+@", ""))
df.withColumn("parts",   F.split(col("full_name"), " "))    # ArrayType
df.withColumn("fname",   F.split(col("full_name"), " ")[0]) # первый элемент

# Конкатенация строк
df.withColumn("full",    F.concat(col("first"), F.lit(" "), col("last")))
df.withColumn("full",    F.concat_ws(" ", col("first"), col("last")))  # то же, separator-first

# --- Числовые функции ---
df.withColumn("rounded", F.round(col("amount"), 2))
df.withColumn("absval",  F.abs(col("delta")))
df.withColumn("sqrtv",   F.sqrt(col("value")))

# --- Дата и время ---
df.withColumn("date",    F.to_date(col("ts"), "yyyy-MM-dd"))
df.withColumn("year",    F.year(col("date")))
df.withColumn("month",   F.month(col("date")))
df.withColumn("dow",     F.dayofweek(col("date")))          # 1=Sun, 7=Sat
df.withColumn("age_d",   F.datediff(F.current_date(), col("birth_date")))

# --- Null-обработка ---
df.withColumn("amount",  F.coalesce(col("amount"), F.lit(0.0)))  # первое non-null
df.withColumn("status",  F.nvl(col("status"), F.lit("unknown"))) # Spark 3.5+

When / otherwise: CASE WHEN в DataFrame API¶

when(condition, value).otherwise(default) - это Spark-эквивалент SQL CASE WHEN:

# Простая категоризация
df = orders.withColumn(
    "size_category",
    F.when(col("amount") < 100, "small")
     .when(col("amount") < 500, "medium")
     .when(col("amount") < 2000, "large")
     .otherwise("enterprise")
)

# Условная трансформация
df = orders.withColumn(
    "adjusted_amount",
    F.when(col("currency") == "USD", col("amount") * 90)
     .when(col("currency") == "EUR", col("amount") * 98)
     .when(col("currency") == "RUB", col("amount"))
     .otherwise(None)   # неизвестная валюта → NULL
)

# Флаги (boolean-колонки)
df = orders.withColumn(
    "is_vip",
    (col("total_spent") > 10000) & (col("orders_count") > 5)
)

Цепочки .when().when() компилируются в один CASE WHEN ... WHEN ... END на уровне физического плана - не в if-elif Python.

Chaining: читаемый ETL pipeline¶

Все трансформации возвращают новый DataFrame, что позволяет строить цепочки (method chaining):

from pyspark.sql.functions import col, when, to_date, upper, trim, coalesce, lit

silver = (
    raw
    # 1. Привести типы
    .withColumn("amount",     col("amount").cast("double"))
    .withColumn("created_at", to_date(col("created_at"), "yyyy-MM-dd"))

    # 2. Нормализовать строки
    .withColumn("status", upper(trim(col("status"))))
    .withColumn("email",  trim(col("email")))

    # 3. Заполнить пустые значения по умолчанию
    .withColumn("region", coalesce(col("region"), lit("UNKNOWN")))

    # 4. Вычислить производные поля
    .withColumn("vat",            col("amount") * 0.2)
    .withColumn("amount_with_vat", col("amount") + col("vat"))
    .withColumn(
        "size_category",
        when(col("amount") < 100, "small")
        .when(col("amount") < 1000, "medium")
        .otherwise("large")
    )

    # 5. Фильтрация битых записей
    .filter(col("amount").isNotNull())
    .filter(col("amount") > 0)
    .filter(col("order_id").isNotNull())

    # 6. Финальная проекция (только нужные колонки)
    .select(
        "order_id", "customer_id", "created_at",
        "status", "amount", "vat", "amount_with_vat",
        "size_category", "region",
    )
)

Скобки вокруг всего выражения позволяют переносить строки без \. Это стандартный стиль production PySpark-кода.

Под капотом: оптимизации Catalyst¶

Column Pruning (обрезка колонок)¶

Когда вы читаете Parquet и применяете select, Catalyst распространяет информацию о нужных колонках вниз до источника данных. Parquet-ридер откроет только те column chunks, которые реально нужны - остальные пропустит на уровне диска.

Следствие: никогда не делайте select("*") перед фильтрацией и агрегацией - Spark теряет возможность сказать Parquet-ридеру, что 47 из 50 колонок не нужны.

Predicate Pushdown (перенос предикатов)¶

Фильтры из filter() / where() Catalyst переносит как можно ближе к источнику данных. Для Parquet это означает использование row group statistics - блоки с заведомо неподходящими значениями не читаются вообще.

Важное следствие: порядок filter в коде не критичен для производительности - Catalyst сам переставит предикаты. Но пишите фильтры рядом с источником данных - это улучшает читаемость.

Explain: смотрим на план¶

# Три уровня детализации
df.explain()                    # только Physical Plan
df.explain("simple")            # краткий Logical + Physical
df.explain("extended")          # Unresolved → Analyzed → Optimized → Physical

# Пример: что происходит после filter + select на Parquet
result = (
    spark.read.parquet("/data/orders")
    .filter(col("status") == "completed")
    .select("order_id", "amount")
)
result.explain("extended")

В выводе смотрите на:

PushedFilters в FileScan parquet - что Spark передал ридеру
ReadSchema - сколько колонок реально читается
Exchange - shuffle (дорогая операция, старайтесь минимизировать)
BroadcastHashJoin vs SortMergeJoin - стратегия соединения

Практика: Bronze → Silver pipeline¶

Пример реального pipeline очистки данных из Data Lake:

from pyspark.sql.functions import (
    col, to_date, upper, trim, coalesce, lit,
    when, regexp_replace, length, current_timestamp
)

# --- Чтение Bronze данных ---
raw = spark.read.option("multiline", True).json("/data/bronze/orders/2024-01/")

# --- Silver трансформация ---
silver = (
    raw
    # Приведение типов
    .withColumn("order_id",    col("order_id").cast("long"))
    .withColumn("customer_id", col("customer_id").cast("long"))
    .withColumn("amount",      col("amount").cast("double"))
    .withColumn("created_at",  to_date(col("created_at"), "yyyy-MM-dd"))

    # Нормализация строк
    .withColumn("status",      upper(trim(col("status"))))
    .withColumn("email",       trim(col("email")))
    .withColumn("phone",       regexp_replace(col("phone"), r"[^\d+]", ""))

    # Производные поля
    .withColumn("vat",         col("amount") * 0.2)
    .withColumn(
        "size_category",
        when(col("amount") < 100, "SMALL")
        .when(col("amount") < 1000, "MEDIUM")
        .when(col("amount") < 5000, "LARGE")
        .otherwise("ENTERPRISE")
    )
    .withColumn(
        "is_valid_email",
        col("email").rlike(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$")
    )
    .withColumn("ingested_at", current_timestamp())

    # Фильтрация невалидных записей
    .filter(col("order_id").isNotNull())
    .filter(col("amount").isNotNull() & (col("amount") > 0))
    .filter(col("created_at").isNotNull())

    # Финальная проекция
    .select(
        "order_id", "customer_id", "created_at",
        "status", "amount", "vat", "size_category",
        "email", "is_valid_email", "phone",
        "ingested_at",
    )
)

# Запись в Silver
silver.write \
    .mode("overwrite") \
    .partitionBy("created_at") \
    .parquet("/data/silver/orders/2024-01/")

# Проверка плана перед записью
silver.explain("extended")

Anti-patterns в DataFrame API¶

Best practices production ETL¶

Читаемость и структура:

# Разбивайте длинные pipeline на именованные шаги
def cast_types(df: DataFrame) -> DataFrame:
    return df \
        .withColumn("amount",     col("amount").cast("double")) \
        .withColumn("created_at", to_date(col("created_at"), "yyyy-MM-dd"))

def normalize_strings(df: DataFrame) -> DataFrame:
    return df \
        .withColumn("status", upper(trim(col("status")))) \
        .withColumn("email",  trim(col("email")))

def add_derived_fields(df: DataFrame) -> DataFrame:
    return df \
        .withColumn("vat", col("amount") * 0.2) \
        .withColumn("size_category", when(...).otherwise(...))

def filter_invalid(df: DataFrame) -> DataFrame:
    return df \
        .filter(col("order_id").isNotNull()) \
        .filter(col("amount") > 0)

# Pipeline читается как список шагов
silver = (
    raw
    .transform(cast_types)
    .transform(normalize_strings)
    .transform(add_derived_fields)
    .transform(filter_invalid)
)

Метод .transform(func) передаёт текущий DataFrame в функцию - это стандартный способ декомпозиции ETL-кода.

Правила для production:

Правило	Почему
Явный `select` в конце каждого шага	Ограничивает количество колонок, которые несёт следующий шаг
Фильтр как можно раньше	Меньше данных на всех последующих шагах
`F.*` вместо Python UDF	Tungsten codegen, Catalyst-оптимизация
Именованные функции-трансформеры	Тестируемость через `pytest` + `chispa`
`alias` после каждого вычисления	Читаемый план, нет безымянных колонок
Проверка `explain()` перед деплоем	Убедиться в column pruning и predicate pushdown

Задания для практики¶

Задание 1: Типы и очистка

Дан DataFrame с CSV-данными:

order_id,customer_id,amount,status,created_at,email
1,101,"1 500,00",Completed ,2024-01-15,user@gmail.com
2,102,,PENDING,2024-01-16,bad-email
3,103,-50.0,completed,NULL,

Напишите pipeline, который:

Приводит amount к double (убрать пробелы и запятую как разделитель тысяч)
Нормализует status к верхнему регистру без пробелов
Приводит created_at к date
Отфильтровывает строки с amount <= 0 или amount IS NULL
Добавляет флаг is_valid_email через rlike

Решение

from pyspark.sql.functions import (
    col, upper, trim, to_date, regexp_replace, when, lit
)

raw = spark.createDataFrame([
    ("1", "101", "1 500,00", "Completed ", "2024-01-15", "user@gmail.com"),
    ("2", "102", None,       "PENDING",    "2024-01-16", "bad-email"),
    ("3", "103", "-50.0",    "completed",  "NULL",       None),
], ["order_id", "customer_id", "amount", "status", "created_at", "email"])

silver = (
    raw
    # 1. amount: убрать пробелы, заменить запятую-разделитель на точку, кастовать
    .withColumn(
        "amount",
        regexp_replace(
            regexp_replace(col("amount"), r"\s", ""),
            r",", "."
        ).cast("double")
    )
    # 2. status: trim + upper
    .withColumn("status", upper(trim(col("status"))))
    # 3. created_at: строку "NULL" → null, затем to_date
    .withColumn(
        "created_at",
        when(col("created_at") == "NULL", lit(None))
        .otherwise(to_date(col("created_at"), "yyyy-MM-dd"))
    )
    # 4. Отфильтровать amount IS NULL или amount <= 0
    .filter(col("amount").isNotNull() & (col("amount") > 0))
    # 5. Флаг валидного email через rlike
    .withColumn(
        "is_valid_email",
        col("email").rlike(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$")
    )
)
silver.show(truncate=False)
# +--------+-----------+--------+----------+-----------+---+--------+--------------------+
# |order_id|customer_id|  amount|    status|created_at |...|   email|      is_valid_email|
# +--------+-----------+--------+----------+-----------+---+--------+--------------------+
# |       1|        101|  1500.0| COMPLETED|2024-01-15 |...|user@...|                true|
# +--------+-----------+--------+----------+-----------+---+--------+--------------------+
# Строки 2 и 3 отфильтрованы: amount IS NULL и amount=-50 <= 0

Ключевые моменты: regexp_replace вызывается цепочкой - сначала убираем пробелы, затем меняем , на .. Строка "NULL" в CSV - это не SQL NULL, поэтому нужен явный when. Два regexp_replace объединяются Catalyst в один проход по данным.

Задание 2: Сложная бизнес-логика

Для датасета заказов напишите выражение priority_score:

Если amount > 1000 И status = 'COMPLETED' → 3
Если amount > 500 ИЛИ status = 'COMPLETED' → 2
Если status = 'PENDING' → 1
Иначе → 0

Реализуйте через when().otherwise() без Python UDF.

Решение

from pyspark.sql.functions import col, when

orders = spark.createDataFrame([
    (1, 1200.0, "COMPLETED"),
    (2,  600.0, "COMPLETED"),
    (3,  300.0, "PENDING"),
    (4,   50.0, "CANCELLED"),
], ["order_id", "amount", "status"])

result = orders.withColumn(
    "priority_score",
    when(
        (col("amount") > 1000) & (col("status") == "COMPLETED"),
        3
    ).when(
        (col("amount") > 500) | (col("status") == "COMPLETED"),
        2
    ).when(
        col("status") == "PENDING",
        1
    ).otherwise(0)
)

result.show()
# +--------+------+---------+--------------+
# |order_id|amount|   status|priority_score|
# +--------+------+---------+--------------+
# |       1|1200.0|COMPLETED|             3|  # amount>1000 AND completed
# |       2| 600.0|COMPLETED|             2|  # amount>500 OR completed
# |       3| 300.0|  PENDING|             1|  # pending
# |       4|  50.0|CANCELLED|             0|  # ничего не совпало
# +--------+------+---------+--------------+

Условия в when проверяются сверху вниз - первое совпавшее побеждает. Поэтому порядок важен: самое строгое условие (AND) ставим первым, более широкие (OR) - ниже. Все условия компилируются в один CASE WHEN в физическом плане - без Python и без ветвления на уровне JVM-вызовов.

Задание 3: Explain план

Напишите запрос к Parquet-файлу, который:

Читает 10 колонок из файла с 20 колонками
Применяет фильтр
Вычисляет производную колонку

Вызовите explain("extended") и найдите в выводе:

PushedFilters - какие предикаты переданы ридеру
ReadSchema - сколько колонок реально читается

Убедитесь, что Catalyst применил column pruning (в ReadSchema должны быть только выбранные колонки).

Решение

import os, tempfile
from pyspark.sql.functions import col

# Создаём тестовый Parquet с 20 колонками
base_cols = ["id", "amount", "name", "status", "category",
             "region", "currency", "channel", "segment", "score"]
extra_cols = [f"extra_{i}" for i in range(10)]  # ещё 10 колонок

all_cols = base_cols + extra_cols
data = [
    tuple([i, float(i * 100), f"user{i}", "completed" if i % 2 == 0 else "pending",
           f"cat{i%5}", "RU", "RUB", "web", "B2C", float(i) * 1.5]
           + [f"val{j}" for j in range(10)])
    for i in range(1, 11)
]

tmp = tempfile.mkdtemp()
parquet_path = os.path.join(tmp, "orders20")
spark.createDataFrame(data, all_cols).write.parquet(parquet_path)

# Запрос: читаем 5 из 20 колонок, фильтруем, вычисляем производную
result = (
    spark.read.parquet(parquet_path)
    .filter(col("status") == "completed")        # predicate pushdown
    .select("id", "amount", "name", "status", "region")  # column pruning
    .withColumn("doubled_amount", col("amount") * 2)
)

result.explain("extended")

В выводе explain("extended") ищем ключевые строки:

== Physical Plan ==
*(1) Project [id#0L, amount#1, name#2, status#3, region#5,
              (amount#1 * 2.0) AS doubled_amount#21]
+- *(1) Filter (isnotnull(status#3) AND (status#3 = completed))
   +- *(1) FileScan parquet [id#0L,amount#1,name#2,status#3,region#5]
          PushedFilters: [IsNotNull(status), EqualTo(status,completed)]
          ReadSchema: struct<id:bigint,amount:double,name:string,
                             status:string,region:string>

PushedFilters - предикат status = completed передан Parquet-ридеру: блоки строк с неподходящим status не читаются с диска
ReadSchema - только 5 колонок из 20; остальные 15 (extra_0..extra_9, currency, ...) пропущены Parquet-ридером на уровне column chunks
Всё выражение в одном *(1) - WholeStageCodeGen объединил Filter, Project и вычисление doubled_amount в единый байткод-проход