Pandas API on Spark (pyspark.pandas): когда и как

# Обычный pandas - выполняется на одной машине
import pandas as pd
df = pd.read_csv("data.csv")
df["revenue"].sum()

# pyspark.pandas - тот же синтаксис, выполняется на кластере
import pyspark.pandas as ps
df = ps.read_csv("s3://bucket/data.csv")
df["revenue"].sum()

# В pandas: мгновенная мутация, O(1) по памяти
import pandas as pd
df = pd.read_csv("data.csv")
df["tax"] = df["price"] * 0.2     # мутация на месте
df.iloc[5]                         # мгновенный доступ по индексу

# В pyspark.pandas: тот же синтаксис, но под капотом - новый DataFrame
import pyspark.pandas as ps
df = ps.read_csv("s3://bucket/data.csv")
df["tax"] = df["price"] * 0.2     # создаётся НОВЫЙ DataFrame, перестраивается план
df.iloc[5]                         # вызывает полный Shuffle для построения индекса!

spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
# Проверить текущее значение
print(spark.conf.get("spark.sql.execution.arrow.pyspark.enabled"))

import pyspark.pandas as ps

df = ps.read_parquet("s3://bucket/events/")  # Action - Spark читает метаданные

# Это всё трансформации - план строится, вычислений нет
filtered  = df[df["amount"] > 100]
enriched  = filtered.copy()
enriched["category"]   = enriched["type"].str.upper()
enriched["amount_tax"] = enriched["amount"] * 1.2

# Action - только здесь Spark реально обрабатывает данные
result = enriched.groupby("category")["amount_tax"].sum()
print(result)   # ещё один Action (вывод требует данных)

# Плохо: enriched вычисляется дважды
total = enriched["amount"].sum()   # Action 1: полный scan
count = enriched["amount"].count() # Action 2: ещё один полный scan

# Хорошо: пересчитать один раз через нативный Spark
spark_df  = enriched.to_spark().cache()
ps_cached = spark_df.pandas_api()
total     = ps_cached["amount"].sum()   # Action 1: из кеша
count     = ps_cached["amount"].count() # Action 2: из кеша

import pyspark.pandas as ps

# Рекомендуется явно установить в начале сессии:
ps.set_option("compute.default_index_type", "distributed")

# Теперь создание DF не будет триггерить дорогой Shuffle
df = ps.read_parquet("s3://bucket/events/")

# Если индекс не нужен - явно отбросить при операциях
result = df.groupby("user_id")["amount"].sum().reset_index(drop=True)

import pyspark.pandas as ps
df = ps.read_parquet("s3://bucket/events/")  # 50 GB данных

# ОПАСНО: весь датасет едет в память Driver
records = df.to_dict()           # OutOfMemoryError на 50 GB
values  = df.values              # то же - возвращает numpy array
rows    = list(df.iterrows())    # то же - итерирует построчно

# ОПАСНО: iloc с отрицательным индексом требует знания общей длины
last_row = df.iloc[-1]           # Spark вынужден собрать всё для нумерации

# БЕЗОПАСНО: остаётся распределённым
filtered   = df[df["amount"] > 1000]                          # Spark Filter
aggregated = df.groupby("user_id")["amount"].sum()            # Spark GroupBy
spark_df   = df.to_spark()                                    # без перемещения данных

# Потенциально опасно - поведение зависит от функции
df.apply(lambda row: custom_python_func(row), axis=1)

# Безопаснее и явно быстрее - нативный UDF
from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType

@udf(returnType=DoubleType())
def custom_spark_udf(value):
    return custom_python_func(value)

df.to_spark().withColumn("result", custom_spark_udf("amount"))

import pyspark.pandas as ps
df = ps.read_parquet("s3://bucket/events/")

# Каждое присваивание создаёт новый DataFrame и перестраивает план
df["tax"]      = df["amount"] * 0.2        # DataFrame #1, план глубиной 1
df["net"]      = df["amount"] - df["tax"]  # DataFrame #2, план глубиной 2
df["category"] = df["type"].str.upper()    # DataFrame #3, план глубиной 3

# После многих мутаций план становится очень глубоким
# При explain() увидите длинную цепочку Project→Filter→Project→...
print(df.to_spark().explain())

# ПЛОХО: 100 итераций = 100 вложенных планов = потенциальный StackOverflowError
for feature_name in feature_list:
    df[feature_name] = compute_feature(df, feature_name)

# ХОРОШО: накопить все вычисления через нативный Spark за один раз
from pyspark.sql import functions as F

spark_df = df.to_spark()
for feature_name, feature_expr in feature_expressions.items():
    spark_df = spark_df.withColumn(feature_name, feature_expr)
# Catalyst оптимизирует всё в один физический план

import pyspark.pandas as ps
ps.set_option("compute.default_index_type", "distributed")

df = ps.read_parquet("s3://datalake/events/year=2024/")

# Знакомые методы работают на терабайтах
print(df.describe())                 # статистика по всем колонкам
print(df["revenue"].value_counts())  # топ значений
print(df.isnull().sum())             # количество пропусков по колонкам

import pyspark.pandas as ps
from pyspark.ml.feature import VectorAssembler

df = ps.read_parquet("s3://datalake/features/")
df["log_amount"] = ps.np.log1p(df["amount"])
df["amount_z"]   = (df["amount"] - df["amount"].mean()) / df["amount"].std()
df["day_of_week"] = df["date"].dt.dayofweek

# Конвертация в нативный Spark для обучения модели
spark_df  = df.to_spark()
assembler = VectorAssembler(
    inputCols=["log_amount", "amount_z", "day_of_week"],
    outputCol="features"
)
ml_df = assembler.transform(spark_df)

import pyspark.pandas as ps

# Чтение
df = ps.read_parquet("s3://bucket/events/")
df = ps.read_csv("data.csv")
df = ps.from_pandas(pd_df)       # из обычного pandas DF (данные идут на кластер)

# DataFrame операции (pandas-синтаксис)
df_filtered = df[df["amount"] > 100]
df["category"] = df["type"].str.upper()
df["rank"] = df.groupby("user_id")["amount"].rank(method="dense")

# Конвертация
pd_df    = df.to_pandas()        # собрать на Driver - осторожно с объёмом!
spark_df = df.to_spark()         # вернуться в нативный Spark DataFrame (без перемещения)
ps_df    = spark_df.pandas_api() # Spark DF в pyspark.pandas DF (без перемещения)

# Безопасный паттерн: работать в ps, при необходимости переключаться в Spark
ps_df    = ps.read_parquet("s3://bucket/events/")
spark_df = ps_df.to_spark()             # бесплатно: просто меняем обёртку
cached   = spark_df.cache()             # кешируем нативными средствами Spark
ps_cached = cached.pandas_api()         # возвращаемся в ps для дальнейших операций

total_per_user = ps_cached.groupby("user_id")["amount"].sum()

# Индекс по умолчанию - распределённый (нет гарантии порядка)
# pandas: целочисленный монотонный индекс
# pyspark.pandas: по умолчанию sequence (глобальный sort) или distributed (без порядка)

# Порядок строк не гарантирован
df.sort_values("date")  # OK, но не меняет физический порядок партиций

# Некоторые методы требуют сбора на Driver
df.to_dict()   # Опасно! Весь DF уходит в память Driver
df.values      # То же

# Явно включить режим совместимости (замедляет - разрешает join разных DF)
ps.set_option("compute.ops_on_diff_frames", True)

# В pandas это работает: автовыравнивание по индексу
pd_df1 = pd.DataFrame({"a": [1, 2, 3]})
pd_df2 = pd.DataFrame({"b": [4, 5, 6]})
result = pd_df1["a"] + pd_df2["b"]   # [5, 7, 9] - выравнивание по индексу

# В pyspark.pandas по умолчанию это ошибка - нет общего индекса
ps_df1 = ps.DataFrame({"a": [1, 2, 3]})
ps_df2 = ps.DataFrame({"b": [4, 5, 6]})
# result = ps_df1["a"] + ps_df2["b"]   # ValueError!

# Включить (медленно - требует join по индексу = Shuffle)
ps.set_option("compute.ops_on_diff_frames", True)
result = ps_df1["a"] + ps_df2["b"]   # работает, но дорого

# Лучше: явный join на нативном Spark
spark_result = ps_df1.to_spark().join(ps_df2.to_spark(), on="key")

# Нативный PySpark
from pyspark.sql.functions import col
spark_df.filter(col("amount") > 100)
spark_df.where("amount > 100")

# pyspark.pandas
ps_df[ps_df["amount"] > 100]
ps_df.query("amount > 100")   # тоже работает, транслируется в WHERE

# Нативный PySpark
from pyspark.sql.functions import upper
spark_df.withColumn("tax", col("amount") * 0.2)
spark_df.withColumn("category", upper(col("type")))

# pyspark.pandas (создаёт новый DataFrame под капотом)
ps_df["tax"]      = ps_df["amount"] * 0.2
ps_df["category"] = ps_df["type"].str.upper()

# Нативный PySpark
from pyspark.sql.functions import sum as spark_sum, avg, count, max as spark_max

spark_df.groupBy("user_id").agg(
    spark_sum("amount").alias("total"),
    avg("amount").alias("avg_amount"),
    count("*").alias("cnt"),
    spark_max("amount").alias("max_amount")
)

# pyspark.pandas - лаконичнее для простых случаев
ps_df.groupby("user_id")["amount"].agg(["sum", "mean", "count", "max"])
# Или словарь:
ps_df.groupby("user_id").agg({"amount": ["sum", "mean"], "events": "count"})

# Нативный PySpark
orders.join(users, orders["user_id"] == users["id"], "left")

# pyspark.pandas
orders_ps.merge(users_ps, left_on="user_id", right_on="id", how="left")

# pandas-style groupby apply - работает, но не оптимально
def normalize_group(group):
    group["normalized"] = (group["amount"] - group["amount"].mean()) / group["amount"].std()
    return group

# Транслируется в pandas_udf под капотом (если можно)
result = ps_df.groupby("category").apply(normalize_group)

# Вместо groupby().apply() с Python-функцией:
ps_df.groupby("category").apply(normalize_group)  # медленно

# Используйте Window functions через нативный Spark:
from pyspark.sql import functions as F
from pyspark.sql.window import Window

spark_df = ps_df.to_spark()
w = Window.partitionBy("category")
spark_df.withColumn(
    "normalized",
    (F.col("amount") - F.avg("amount").over(w)) / F.stddev("amount").over(w)
)
# Catalyst оптимизирует, выполняется в JVM без Python overhead

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import DoubleType
import pandas as pd

# Pandas UDF: функция над батчем значений одной колонки
@pandas_udf(DoubleType())
def risk_score(amounts: pd.Series) -> pd.Series:
    return amounts.apply(lambda x: x * 0.1 if x > 1000 else x * 0.05)

# Применяется в нативном PySpark
spark_df.withColumn("risk", risk_score("amount"))

import pyspark.pandas as ps

ps_df = ps.read_parquet("s3://bucket/events/")
# Работает со всем DataFrame, не только с одной колонкой
result = ps_df.groupby("user_id").agg({
    "amount": ["sum", "mean", "std"],
    "date":   ["min", "max"]
})

from pyspark.sql import functions as F

# Spark: NULL и NaN - разные вещи!
spark_df = spark.createDataFrame([
    (1, None),           # NULL
    (2, float("nan")),   # NaN - не NULL!
    (3, 42.0)
], ["id", "value"])

spark_df.filter(F.col("value").isNull()).show()
# Результат: только строка с None (id=1)
# NaN (id=2) НЕ считается NULL!

spark_df.filter(F.isnan("value")).show()
# Результат: только строка с NaN (id=2)

# Правильная обработка: поймать и то и другое
spark_df.filter(
    F.col("value").isNull() | F.isnan("value")
).show()

import pyspark.pandas as ps
import numpy as np

ps_df = ps.DataFrame({
    "value": [1.0, None, np.nan, 4.0]
})

# pyspark.pandas: pandas-семантика
print(ps_df["value"].isnull().sum())   # 2 (None и nan оба считаются)

# Нормализовать NaN в NULL перед to_spark():
ps_df["value"] = ps_df["value"].fillna(None)  # все nan → None → NULL в Spark

# Нормализовать NaN в нативном Spark:
spark_df = spark_df.withColumn(
    "value",
    F.when(F.isnan("value"), None).otherwise(F.col("value"))
)

import pyspark.pandas as ps

df = ps.read_parquet("s3://bucket/orders/")

# Накопительная сумма по пользователю (работает)
df["cumulative_amount"] = df.groupby("user_id")["amount"].cumsum()

# Ранжирование (работает)
df["rank"] = df.groupby("user_id")["amount"].rank(method="dense", ascending=False)

# Процент от суммы группы (работает)
group_total = df.groupby("category")["amount"].transform("sum")
df["pct_of_category"] = df["amount"] / group_total

from pyspark.sql.window import Window
from pyspark.sql import functions as F

# Нативный Spark: lag/lead с гибким окном
spark_df = df.to_spark()
w = Window.partitionBy("user_id").orderBy("date")
spark_df = spark_df.withColumn("prev_amount", F.lag("amount", 1).over(w))
spark_df = spark_df.withColumn("next_amount", F.lead("amount", 1).over(w))
spark_df = spark_df.withColumn(
    "rolling_sum_3",
    F.sum("amount").over(w.rowsBetween(-2, 0))  # скользящее окно в 3 строки
)

# Вернуться в pyspark.pandas при желании
result_ps = spark_df.pandas_api()

import pyspark.pandas as ps

# Намеренно медленный тип - для демонстрации
ps.set_option("compute.default_index_type", "sequence")

df = ps.read_parquet("s3://bucket/events/")
result = df.groupby("user_id")["amount"].sum()

# Получить план выполнения
result.to_spark().explain(mode="extended")

== Physical Plan ==
*(4) Sort [__index_level_0__ ASC NULLS FIRST], true, 0
+- Exchange rangepartitioning(__index_level_0__ ASC NULLS FIRST, 200)
   +- *(3) HashAggregate(keys=[user_id], functions=[sum(amount)])
      +- Exchange hashpartitioning(user_id, 200)
         +- *(2) HashAggregate(keys=[user_id], functions=[partial_sum(amount)])
            +- *(1) ColumnarToRow
               +- Scan parquet s3://bucket/events/

== Physical Plan ==
*(3) HashAggregate(keys=[user_id], functions=[sum(amount)])
+- Exchange hashpartitioning(user_id, 200)
   +- *(2) HashAggregate(keys=[user_id], functions=[partial_sum(amount)])
      +- *(1) ColumnarToRow
         +- Scan parquet s3://bucket/events/

# ПЛОХО: df_enriched сканируется дважды
total = df_enriched["amount"].sum()
count = df_enriched["amount"].count()

# ХОРОШО:
df_cached = df_enriched.to_spark().cache().pandas_api()
total = df_cached["amount"].sum()
count = df_cached["amount"].count()
df_cached.to_spark().unpersist()  # освободить кеш

# ПЛОХО: 50 GB → Driver RAM → OOM
pd_result = large_ps_df.to_pandas()

# ХОРОШО: агрегировать сначала, потом конвертировать
summary  = large_ps_df.groupby("category")["amount"].sum()
pd_summary = summary.to_pandas()   # маленький агрегат - безопасно

# ПЛОХО: каждое создание DF с sequence index = Shuffle
ps.set_option("compute.default_index_type", "sequence")

# ХОРОШО:
ps.set_option("compute.default_index_type", "distributed")

# ПЛОХО: нет векторизации, медленная сериализация строк
df["score"] = df["text"].apply(lambda x: ml_model.predict(x))

# ХОРОШО: pandas_udf для векторизованной обработки батчами
@pandas_udf(DoubleType())
def batch_predict(texts: pd.Series) -> pd.Series:
    return pd.Series(ml_model.predict(texts.tolist()))

spark_df.withColumn("score", batch_predict("text"))

# ПЛОХО: глубокий план, потенциальный StackOverflow при 100+ итерациях
for col in many_columns:
    ps_df[col] = compute(ps_df, col)

# ХОРОШО: выйти в нативный Spark
spark_df = ps_df.to_spark()
for col_name, expr in column_expressions.items():
    spark_df = spark_df.withColumn(col_name, expr)
result_ps = spark_df.pandas_api()

# ПЛОХО: включено глобально, все операции замедляются
ps.set_option("compute.ops_on_diff_frames", True)

# ХОРОШО: только там где нужно, или через нативный join
ps.set_option("compute.ops_on_diff_frames", False)  # по умолчанию
joined = ps_df1.merge(ps_df2, on="key")  # явный join

import pyspark.pandas as ps

# Чтение из Iceberg таблицы через Spark catalog
spark_df = spark.read.table("iceberg.events.orders")
ps_df    = spark_df.pandas_api()

# Аналитика в pandas-стиле
monthly_summary = (
    ps_df
    .assign(month=ps_df["created_at"].dt.to_period("M"))
    .groupby(["month", "category"])["revenue"]
    .agg(["sum", "count", "mean"])
    .reset_index()
)

# Запись обратно в Iceberg
monthly_summary.to_spark().writeTo("iceberg.analytics.monthly_summary") \
    .tableProperty("write.format.default", "parquet") \
    .createOrReplace()

# Iceberg отклонит запись если типы не совпадают
ps_df["amount"] = ps_df["amount"].astype(str)   # конвертировали в string
# ps_df.to_spark().write.saveAsTable("iceberg.events.orders")
# AnalysisException: Cannot write incompatible data to table

# Запись с контролем партиционирования - через нативный Spark
ps_df.to_spark() \
    .write \
    .partitionBy("year", "month") \
    .mode("overwrite") \
    .parquet("s3://bucket/output/")

import pyspark.pandas as ps

# Стартовая конфигурация для миграции
ps.set_option("compute.default_index_type", "distributed")
ps.set_option("compute.ops_on_diff_frames", False)
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

# Было (pandas):
# import pandas as pd
# df = pd.read_csv("data.csv")

# Стало (pyspark.pandas):
df = ps.read_csv("s3://bucket/data.csv")  # или ps.read_parquet(...)

# Весь остальной pandas-код остаётся почти без изменений
df_clean    = df.dropna(subset=["amount", "user_id"])
df_filtered = df_clean[df_clean["amount"] > 0]
result      = df_filtered.groupby("category")["amount"].agg(["sum", "count", "mean"])

# Узкое место - переписать на нативный Spark
complex_result = df_filtered.to_spark() \
    .join(reference_df, "product_id", "left") \
    .groupBy("category", "subcategory") \
    .agg(F.sum("amount").alias("total"))

import pyspark.pandas as ps

# Проблема 1: sequence индекс по умолчанию - лишние Shuffle
df = ps.read_parquet("s3://datalake/user_events/")

# Проблема 2: in-place мутации в цикле создают глубокий план
feature_cols = ["page_views", "session_time", "clicks", "purchases", "cart_adds"]
for col in feature_cols:
    df[col + "_log"] = ps.np.log1p(df[col])  # каждая итерация = новый DataFrame

# Проблема 3: z-score через apply = Python overhead на каждую строку
for col in feature_cols:
    mean = df[col].mean()
    std  = df[col].std()
    df[col + "_z"] = df[col].apply(lambda x: (x - mean) / std if std > 0 else 0)

# Проблема 4: to_pandas() на большом DF перед записью
pd_result = df.to_pandas()       # OOM риск
pd_result.to_parquet("output.parquet")

import pyspark.pandas as ps
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# Исправление 1: distributed индекс
ps.set_option("compute.default_index_type", "distributed")
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

df = ps.read_parquet("s3://datalake/user_events/")

feature_cols = ["page_views", "session_time", "clicks", "purchases", "cart_adds"]
spark_df = df.to_spark()

# Исправление 2: log-трансформации за один выход в нативный Spark
for col in feature_cols:
    spark_df = spark_df.withColumn(f"{col}_log", F.log1p(F.col(col)))

# Исправление 3: z-score через Window (JVM, без Python overhead)
for col in feature_cols:
    w = Window.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
    mean_val = F.avg(col).over(w)
    std_val  = F.stddev(col).over(w)
    spark_df = spark_df.withColumn(
        f"{col}_z",
        F.when(std_val > 0, (F.col(col) - mean_val) / std_val).otherwise(F.lit(0.0))
    )

# Исправление 4: записать через нативный Spark без to_pandas()
spark_df.write \
    .mode("overwrite") \
    .parquet("s3://datalake/user_features/")

import pyspark.pandas as ps

# Показать все опции
ps.get_option("display.max_rows")
ps.set_option("display.max_rows", 100)

# Разрешить операции на DF из разных источников (замедляет - добавляет join по индексу)
ps.set_option("compute.ops_on_diff_frames", True)

# Дефолтный тип индекса (ключевая настройка производительности)
ps.set_option("compute.default_index_type", "distributed")          # рекомендуется
# ps.set_option("compute.default_index_type", "distributed-sequence")  # монотонный, без глобальной сортировки
# ps.set_option("compute.default_index_type", "sequence")            # медленно, только для отладки

# Количество строк для repr
ps.set_option("display.max_rows", 1000)
ps.set_option("display.max_columns", 50)

# Arrow для передачи данных (критично для производительности)
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "10000")  # размер батча

# Количество партиций после Shuffle
spark.conf.set("spark.sql.shuffle.partitions", "200")   # дефолт для production

# Для тестов - уменьшить чтобы не гонять 196 пустых партиций
spark.conf.set("spark.sql.shuffle.partitions", "4")