Spark Plugin API: SparkPlugin и QueryStage интерфейс для замены executor

# ──────────────────────────────────────────────────────────
# Регистрация плагина через SparkSession конфигурацию
# ──────────────────────────────────────────────────────────
# spark.plugins - список полных имён Java/Scala-классов через запятую.
# Каждый класс должен реализовывать интерфейс SparkPlugin.
# Spark инстанцирует эти классы дважды:
#   1) В JVM-процессе Driver-а при создании SparkContext
#   2) В JVM-процессе каждого Executor-а при его старте
spark = SparkSession.builder \
    .master("spark://host:7077") \
    .config("spark.plugins", "com.mycompany.NativeEnginePlugin") \
    .config("spark.plugins", "io.glutenproject.GlutenPlugin") \  # Gluten/Velox
    # Или несколько плагинов через запятую:
    # .config("spark.plugins", "com.a.PluginA,com.b.PluginB") \
    .getOrCreate()

# Дополнительные расширения Catalyst (отдельный конфиг):
# SparkSessionExtensions - это другой механизм, работающий
# совместно с Plugin API для внедрения правил оптимизатора
spark = SparkSession.builder \
    .config("spark.sql.extensions",
            "io.glutenproject.GlutenSparkExtensions") \
    .config("spark.plugins", "io.glutenproject.GlutenPlugin") \
    .getOrCreate()

// Интерфейс, который должен реализовывать ваш плагин-класс
trait SparkPlugin {
  // Возвращает объект для Driver-стороны.
  // Вызывается один раз при создании SparkContext на Driver.
  // Может вернуть null если Driver-логика не нужна.
  def driverPlugin(): DriverPlugin

  // Возвращает объект для Executor-стороны.
  // Вызывается при старте каждого Executor.
  // Может вернуть null если Executor-логика не нужна.
  def executorPlugin(): ExecutorPlugin
}

// Интерфейс DriverPlugin (Spark 3.0+)
trait DriverPlugin {

  // Вызывается при создании SparkContext на Driver.
  // sc - SparkContext (доступ ко всему Spark)
  // pluginContext - вспомогательный контекст плагина
  // Возвращает Map[String, String] - конфиг, который будет
  // передан на Executor-ы через RPC при их старте.
  def init(
    sc: SparkContext,
    pluginContext: PluginContext
  ): java.util.Map[String, String]

  // Вызывается когда Driver получает RPC-сообщение от
  // ExecutorPlugin через pluginContext.ask() или send().
  // Позволяет организовать двунаправленный обмен данными.
  def receive(message: AnyRef): AnyRef

  // Вызывается при завершении приложения.
  // Используйте для освобождения Driver-сторонних ресурсов.
  def shutdown(): Unit
}

// LoggingDriverPlugin.java
import org.apache.spark.SparkContext;
import org.apache.spark.api.plugin.DriverPlugin;
import org.apache.spark.api.plugin.PluginContext;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.execution.QueryExecution;
import org.apache.spark.sql.util.QueryExecutionListener;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.HashMap;
import java.util.Map;

public class LoggingDriverPlugin implements DriverPlugin {

    private static final Logger LOG =
        LoggerFactory.getLogger(LoggingDriverPlugin.class);

    @Override
    public Map<String, String> init(SparkContext sc, PluginContext pluginContext) {
        LOG.info("[Plugin] DriverPlugin.init() called on Driver");

        // Получаем SparkSession из SparkContext
        SparkSession spark = SparkSession.active();

        // Регистрируем QueryExecutionListener - он вызывается
        // ПОСЛЕ каждого SQL-запроса (для мониторинга и анализа).
        // Это не то же самое что Catalyst rules - это post-execution hook.
        spark.listenerManager().register(new QueryExecutionListener() {

            @Override
            public void onSuccess(
                    String funcName,
                    QueryExecution qe,
                    long durationNs) {
                // qe.sparkPlan() - физический план ДО оптимизации AQE
                // qe.executedPlan() - финальный план С оптимизацией AQE
                LOG.info("[Plugin] Query '{}' succeeded in {}ms",
                    funcName, durationNs / 1_000_000);
                LOG.debug("[Plugin] Physical plan:\n{}",
                    qe.executedPlan().toString());
            }

            @Override
            public void onFailure(
                    String funcName,
                    QueryExecution qe,
                    Exception exception) {
                LOG.error("[Plugin] Query '{}' FAILED: {}",
                    funcName, exception.getMessage());
            }
        });

        // Возвращаем конфиг, который будет передан на все Executor-ы.
        // Это Map<String, String> с произвольными параметрами.
        // Executor-ы получат его в ExecutorPlugin.init().
        Map<String, String> config = new HashMap<>();
        config.put("plugin.driver.host", sc.master());
        config.put("plugin.version", "1.0.0");
        config.put("plugin.log.level", "DEBUG");
        return config;
    }

    @Override
    public Object receive(Object message) {
        // Здесь обрабатываем сообщения от Executor-ов.
        // Например, ExecutorPlugin может отправить статистику
        // использования нативной памяти на Driver.
        LOG.info("[Plugin] Received message from Executor: {}", message);
        return null; // ответ Executor-у (опционально)
    }

    @Override
    public void shutdown() {
        LOG.info("[Plugin] DriverPlugin.shutdown() - cleaning up");
    }
}

// Интерфейс ExecutorPlugin (Spark 3.0+)
trait ExecutorPlugin {

  // Вызывается при старте Executor до первой задачи.
  // pluginContext - контекст для отправки сообщений на Driver
  // extraConf - Map из DriverPlugin.init() (Driver → Executor transfer)
  def init(
    pluginContext: PluginContext,
    extraConf: java.util.Map[String, String]
  ): Unit

  // Вызывается ПЕРЕД запуском каждой Task на этом Executor.
  // taskContext - информация о текущей задаче
  def onTaskStart(): Unit

  // Вызывается ПОСЛЕ завершения каждой Task (успешно или с ошибкой).
  // succeeded - true если задача завершилась без исключений
  def onTaskSucceeded(): Unit
  def onTaskFailed(failureReason: TaskFailedReason): Unit

  // Вызывается при завершении Executor (shutdown).
  def shutdown(): Unit
}

// NativeExecutorPlugin.java
import org.apache.spark.api.plugin.ExecutorPlugin;
import org.apache.spark.api.plugin.PluginContext;
import org.apache.spark.TaskContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

public class NativeExecutorPlugin implements ExecutorPlugin {

    private static final Logger LOG =
        LoggerFactory.getLogger(NativeExecutorPlugin.class);

    // Счётчик аллоцированной нативной памяти (Off-Heap).
    // AtomicLong потому что несколько задач могут работать параллельно.
    private final AtomicLong nativeMemoryAllocated = new AtomicLong(0);

    // Ссылка на контекст плагина для отправки сообщений на Driver
    private PluginContext pluginContext;

    // Указатель на нативный контекст (в реальности - long-адрес в памяти)
    // В JNI это типичная практика: хранить void* как long в Java
    private long nativeContextPtr = 0L;

    @Override
    public void init(
            PluginContext pluginContext,
            Map<String, String> extraConf) {

        this.pluginContext = pluginContext;

        LOG.info("[Plugin] ExecutorPlugin.init() - initializing native runtime");

        // Читаем конфигурацию, переданную от DriverPlugin.init()
        String logLevel = extraConf.getOrDefault("plugin.log.level", "INFO");
        String driverHost = extraConf.getOrDefault("plugin.driver.host", "unknown");
        LOG.info("[Plugin] Driver host: {}, log level: {}", driverHost, logLevel);

        // В реальном плагине (Gluten, RAPIDS) здесь происходит:
        // 1. Загрузка нативной библиотеки через System.loadLibrary()
        // 2. Инициализация нативного memory allocator
        // 3. Регистрация JNI-коллбэков

        // Пример загрузки нативной библиотеки:
        // System.loadLibrary("native_spark_plugin");  // libNative_spark_plugin.so
        // this.nativeContextPtr = initNativeContext(offHeapSizeBytes);

        // Мы имитируем инициализацию:
        long offHeapSizeMb = Long.parseLong(
            extraConf.getOrDefault("plugin.offheap.mb", "1024")
        );
        LOG.info("[Plugin] Reserving {}MB Off-Heap for native allocator", offHeapSizeMb);
        this.nativeContextPtr = 1L; // имитация указателя

        // Отправляем сообщение на Driver о готовности Executor-а
        // pluginContext.send() - асинхронная отправка
        pluginContext.send("EXECUTOR_READY:" + offHeapSizeMb + "MB");
    }

    @Override
    public void onTaskStart() {
        // Вызывается перед каждой Task.
        // Полезно для сброса per-task счётчиков или
        // инициализации task-level нативных ресурсов.
        LOG.debug("[Plugin] Task starting on executor");
    }

    @Override
    public void onTaskSucceeded() {
        // Вызывается после успешной Task.
        // Можно логировать статистику использования нативной памяти.
        long allocatedMb = nativeMemoryAllocated.get() / (1024 * 1024);
        if (allocatedMb > 0) {
            LOG.info("[Plugin] Task completed, native memory peak: {}MB", allocatedMb);
        }
        // В реальном плагине: освобождение task-level нативных буферов
    }

    @Override
    public void onTaskFailed(TaskFailedReason failureReason) {
        // Критически важно: при падении Task нативные ресурсы
        // могут не освободиться автоматически (нет GC для нативной памяти).
        // Здесь ОБЯЗАТЕЛЬНО нужно освобождать нативные буферы.
        LOG.error("[Plugin] Task FAILED: {} - releasing native resources",
            failureReason);
        // В реальном плагине: освободить нативную память принудительно
        // freeNativeTaskResources(nativeContextPtr);
        nativeMemoryAllocated.set(0);
    }

    @Override
    public void shutdown() {
        LOG.info("[Plugin] ExecutorPlugin.shutdown() - freeing native context");
        // В реальном плагине: освободить нативный контекст полностью
        // destroyNativeContext(nativeContextPtr);
        this.nativeContextPtr = 0L;
    }
}

// Регистрация расширений Catalyst в SparkSessionExtensions
// (обычно делается внутри DriverPlugin.init() или отдельным классом)

spark.withExtensions { extensions =>

    // injectOptimizerRule - добавление правила в Catalyst LOGICAL optimizer.
    // Работает с Logical Plan (до выбора физических операторов).
    // Используется для: predicate pushdown, constant folding,
    // join reordering и других логических оптимизаций.
    extensions.injectOptimizerRule { session =>
        new MyCustomLogicalOptimizationRule(session)
    }

    // injectPlannerStrategy - добавление стратегии Physical Planner.
    // Определяет, как конкретный Logical оператор превращается
    // в Physical оператор. Используется для замены стандартных
    // физических операторов на нативные.
    extensions.injectPlannerStrategy { session =>
        new NativePlannerStrategy(session)
    }

    // injectQueryStagePrepRule - правила, применяемые к КАЖДОМУ
    // QueryStage ПЕРЕД его материализацией (в режиме AQE).
    // Это самая поздняя точка, где можно изменить план до выполнения.
    extensions.injectQueryStagePrepRule { session =>
        new PreStageNativeOptimizationRule(session)
    }

    // injectColumnar - специальные правила для замены ROW-based
    // операторов на COLUMNAR операторы (работающие с Arrow ColumnarBatch).
    // Это главный механизм, через который Gluten/RAPIDS заменяют операторы.
    extensions.injectColumnar { session =>
        NativeColumnarOverrideRules(session)
    }

    // injectPostHocResolutionRule - правила после анализа,
    // но до оптимизации. Используется редко.
    extensions.injectPostHocResolutionRule { session =>
        new MyPostHocRule(session)
    }
}

// Упрощённый пример ColumnarRule для замены операторов
// (реальная реализация в Gluten значительно сложнее)
class NativeColumnarOverrideRules(session: SparkSession)
    extends ColumnarRule {

  // Вызывается с физическим планом.
  // Возвращает трансформированный план.
  override def preColumnarTransitions(plan: SparkPlan): SparkPlan = {
    // transformDown обходит дерево план сверху вниз
    // и применяет partial function к каждому узлу
    plan.transformDown {

      // Если встречаем стандартный HashAggregateExec (JVM) —
      // заменяем на нативный NativeHashAggregateExec
      case agg: HashAggregateExec
          if canNativelyExecute(agg) =>
        NativeHashAggregateExec(
          agg.requiredChildDistributionExpressions,
          agg.groupingExpressions,
          agg.aggregateExpressions,
          agg.aggregateAttributes,
          agg.initialInputBufferOffset,
          agg.resultExpressions,
          agg.child
        )

      // SortMergeJoinExec → NativeSortMergeJoinExec
      case join: SortMergeJoinExec
          if canNativelyExecute(join) =>
        NativeSortMergeJoinExec(join.leftKeys, join.rightKeys,
          join.joinType, join.condition, join.left, join.right)

      // FilterExec → NativeFilterExec
      case filter: FilterExec
          if canNativelyExecute(filter) =>
        NativeFilterExec(filter.condition, filter.child)

      // Все остальные операторы - оставляем JVM-реализацию.
      // Это принципиальная особенность plug-in подхода:
      // незамена → автоматический fallback на JVM, не ошибка.
    }
  }

  // Проверяет, может ли нативный движок выполнить оператор.
  // Если нет - оставляем JVM-версию (graceful degradation).
  private def canNativelyExecute(plan: SparkPlan): Boolean = {
    // Проверяем типы данных всех колонок
    val unsupportedTypes = plan.output.exists { attr =>
      attr.dataType match {
        case _: MapType => true       // MapType не поддерживается нативно
        case _: UserDefinedType[_] => true  // кастомные типы - нет
        case _ => false
      }
    }
    !unsupportedTypes
  }
}

Пример физического плана с AQE:

SortMergeJoinExec
├── Exchange (ShuffleExchangeExec)          ← граница QueryStage 2
│   └── FilterExec + ParquetScan           ← QueryStage 1 (orders)
└── Exchange (BroadcastExchangeExec)        ← граница QueryStage 3
    └── FilterExec + ParquetScan           ← QueryStage 0 (customers)

AQE выполняет:
1. QueryStage 0 и 1 параллельно
2. После завершения: видит что customers < 10MB → меняет SortMergeJoin на BroadcastHashJoin
3. Выполняет финальный join с broadcast

// QueryStagePrepRule - применяется к каждому QueryStage перед выполнением
class NativeQueryStagePrepRule(session: SparkSession)
    extends Rule[SparkPlan] {

  override def apply(plan: SparkPlan): SparkPlan = {
    // plan - физический план ОДНОГО QueryStage
    // (уже с AQE-оптимизациями)

    // Логируем план для диагностики
    LOG.info("QueryStage plan before native transformation:\n{}",
      plan.treeString)

    // Применяем нативные трансформации к этому stage
    val nativePlan = applyNativeTransformations(plan)

    LOG.info("QueryStage plan after native transformation:\n{}",
      nativePlan.treeString)

    nativePlan
  }

  private def applyNativeTransformations(plan: SparkPlan): SparkPlan = {
    plan.transformDown {
      // Замена операторов - аналогично ColumnarRule
      case agg: HashAggregateExec if isNativeSupported(agg) =>
        NativeHashAggregateExec.from(agg)

      // ... другие замены
    }
  }
}

// NativeHashAggregateExec - упрощённый учебный пример
// Реальный код в Gluten/RAPIDS значительно сложнее
public class NativeHashAggregateExec extends SparkPlan
        implements ColumnarSupport {

    private final HashAggregateExec originalPlan;

    // Флаг: используем ли columnar (Arrow) или row-based путь
    @Override
    public boolean supportsColumnar() {
        return true;  // говорим Spark: мы работаем с ColumnarBatch
    }

    @Override
    public RDD<ColumnarBatch> executeColumnar() {
        // Получаем входные данные в columnar формате от child оператора
        RDD<ColumnarBatch> childRDD = child().executeColumnar();

        return childRDD.mapPartitions(iter -> {
            // Собираем все батчи партиции в нативной памяти
            List<Long> arrowPtrs = new ArrayList<>();
            while (iter.hasNext()) {
                ColumnarBatch batch = iter.next();
                // extractArrowBufPtr() возвращает long-указатель
                // на Arrow буфер в Off-Heap памяти
                arrowPtrs.add(ArrowUtils.extractArrowBufPtr(batch));
            }

            // Один JNI-вызов для всей партиции
            // nativeAggregate() - метод, объявленный как native
            long resultPtr = nativeAggregate(
                arrowPtrs.stream().mapToLong(Long::longValue).toArray(),
                buildAggregateConfig()
            );

            // Оборачиваем результат в ColumnarBatch без копирования
            ColumnarBatch resultBatch = ArrowUtils.wrapArrowPtr(resultPtr);
            return Collections.singletonList(resultBatch).iterator();
        });
    }

    // JNI-объявление: этот метод реализован в libNativePlugin.so
    // При вызове JVM передаёт управление в нативный код
    private native long nativeAggregate(
        long[] arrowInputPtrs,
        byte[] aggregateConfig  // protobuf-сериализованный конфиг
    );
}

Архитектура Gluten:

PySpark / SparkSQL
        ↓
Spark Catalyst (Logical Plan → Physical Plan)
        ↓
GlutenPlugin (ColumnarRule):
  SortMergeJoinExec     → VeloxSortMergeJoinExec (C++ JNI)
  HashAggregateExec     → VeloxHashAggregateExec (C++ JNI)
  BroadcastHashJoinExec → VeloxBroadcastHashJoin (C++ JNI)
  ParquetScan           → VeloxParquetScan (C++ JNI)
  [UDFs, MLlib, ...]    → JVM fallback (без изменений)
        ↓
Velox Execution Engine (C++)
  - Velox expression evaluator (vectorized)
  - Velox hash table (cache-friendly)
  - Velox Parquet reader (columnar, no JVM deserialization)
  - AVX-512/AVX2 SIMD через Velox operators

spark = SparkSession.builder \
    .config("spark.plugins", "io.glutenproject.GlutenPlugin") \
    .config("spark.sql.extensions",
            "io.glutenproject.GlutenSparkExtensions") \
    # Выбор нативного backend: velox (CPU) или другие
    .config("spark.gluten.sql.columnar.backend.lib", "velox") \
    # Off-Heap размер для Velox - Velox не использует JVM heap
    .config("spark.memory.offHeap.enabled", "true") \
    .config("spark.memory.offHeap.size", "8g") \
    # Включаем конкретные операторы (можно контролировать фолбек)
    .config("spark.gluten.sql.columnar.hashagg", "true") \
    .config("spark.gluten.sql.columnar.sort", "true") \
    .config("spark.gluten.sql.columnar.shuffleManager",
            "org.apache.spark.shuffle.sort.ColumnarShuffleManager") \
    .getOrCreate()

spark = SparkSession.builder \
    # RAPIDS Plugin для Spark
    .config("spark.plugins", "com.nvidia.spark.SQLPlugin") \
    # GPU resource discovery
    .config("spark.worker.resource.gpu.discoveryScript",
            "/opt/sparkRapidsPlugin/getGpusResources.sh") \
    .config("spark.executor.resource.gpu.amount", "1") \
    .config("spark.task.resource.gpu.amount", "0.5") \
    # Размер GPU-батча
    .config("spark.rapids.sql.batchSizeBytes", "512m") \
    # Fallback для неподдерживаемых операций
    .config("spark.rapids.sql.incompatibleOps.enabled", "false") \
    .getOrCreate()

spark = SparkSession.builder \
    .config("spark.plugins", "org.apache.spark.CometPlugin") \
    .config("spark.sql.extensions",
            "org.apache.comet.CometSparkSessionExtensions") \
    # Включаем нативное выполнение
    .config("spark.comet.enabled", "true") \
    # Нативный shuffle (Arrow IPC между воркерами)
    .config("spark.comet.exec.shuffle.enabled", "true") \
    # Размер батча для нативных операторов
    .config("spark.comet.batchSize", "8192") \
    .getOrCreate()

# Python-сторона: как ColumnarBatch виден в PySpark
# (обычно прозрачно, но можно увидеть в explain())

# Запрос, который вызывает переходы Row ↔ Columnar:
df.filter("amount > 1000") \
  .groupBy("region") \
  .agg({"amount": "sum"}) \
  .explain(extended=True)

# В плане при включённом нативном плагине (например Gluten):
# == Physical Plan ==
# VeloxColumnarToRowExec            ← конвертация из Columnar в Row
# +- VeloxHashAggregateExec         ← нативная агрегация (Velox, columnar)
#    +- VeloxColumnarShuffleExchange  ← нативный shuffle (Arrow IPC)
#    +- VeloxHashAggregateExec       ← partial aggregate (нативная)
#       +- VeloxFilterExec           ← нативный filter (SIMD)
#          +- VeloxParquetScan       ← нативное чтение Parquet
#
# vs без плагина:
# == Physical Plan ==
# *(2) HashAggregate(...)
# +- Exchange hashpartitioning(region, 200)
#    +- *(1) HashAggregate(...)
#       +- *(1) Filter (amount > 1000)
#          +- *(1) ColumnarToRow      ← конвертация Arrow → Row
#             +- FileScan parquet

# ──────────────────────────────────────────────────────────
# Запуск с Comet (наиболее простой для локального тестирования)
# ──────────────────────────────────────────────────────────
# Установка: pip install pyspark apache-comet
# (или скачать JAR: https://github.com/apache/datafusion-comet/releases)

import os
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# Путь к JAR-файлу Comet
comet_jar = "/path/to/comet-spark-spark3.5_2.12-0.4.0.jar"

spark = SparkSession.builder \
    .master("local[4]") \
    .config("spark.driver.extraClassPath", comet_jar) \
    .config("spark.executor.extraClassPath", comet_jar) \
    .config("spark.plugins", "org.apache.spark.CometPlugin") \
    .config("spark.sql.extensions",
            "org.apache.comet.CometSparkSessionExtensions") \
    .config("spark.comet.enabled", "true") \
    .config("spark.comet.exec.shuffle.enabled", "true") \
    # Off-Heap для нативных буферов
    .config("spark.memory.offHeap.enabled", "true") \
    .config("spark.memory.offHeap.size", "2g") \
    # Уровень логирования плагина
    .config("spark.comet.logLevel", "INFO") \
    .getOrCreate()

spark.sparkContext.setLogLevel("WARN")
print(f"Spark version: {spark.version}")
print(f"Comet enabled: {spark.conf.get('spark.comet.enabled')}")

# ──────────────────────────────────────────────────────────
# Тестовый запрос: тяжёлая аналитика
# ──────────────────────────────────────────────────────────
# Создаём DataFrame для тестирования
import random
from pyspark.sql.types import *

# Генерируем данные: транзакции с нескольких регионов
n_rows = 1_000_000
data = [(
    i,
    random.choice(["North", "South", "East", "West", "Central"]),
    random.choice(["Electronics", "Clothing", "Food", "Sports"]),
    random.uniform(10.0, 10000.0),
    random.randint(1, 100)
) for i in range(n_rows)]

schema = StructType([
    StructField("id", LongType(), False),
    StructField("region", StringType(), True),
    StructField("category", StringType(), True),
    StructField("amount", DoubleType(), True),
    StructField("quantity", IntegerType(), True),
])

df = spark.createDataFrame(data, schema).cache()
df.count()  # форсируем кэширование

# ──────────────────────────────────────────────────────────
# Запрос с JOIN + GROUP BY + оконными функциями
# ──────────────────────────────────────────────────────────
from pyspark.sql.window import Window

# Таблица регионов (для JOIN)
regions = spark.createDataFrame([
    ("North", "Moscow"),
    ("South", "Krasnodar"),
    ("East", "Vladivostok"),
    ("West", "Saint Petersburg"),
    ("Central", "Voronezh")
], ["region", "city"])

# Сложный аналитический запрос
query = df \
    .join(regions, on="region", how="left") \
    .filter(F.col("amount") > 100) \
    .groupBy("region", "city", "category") \
    .agg(
        F.sum("amount").alias("total_revenue"),
        F.count("*").alias("order_count"),
        F.avg("amount").alias("avg_order"),
        F.sum(F.col("amount") * F.col("quantity")).alias("gross_revenue")
    ) \
    .withColumn(
        "revenue_rank",
        F.rank().over(
            Window.partitionBy("region").orderBy(F.desc("total_revenue"))
        )
    )

# ──────────────────────────────────────────────────────────
# Анализ плана: ищем нативные операторы
# ──────────────────────────────────────────────────────────
print("=" * 60)
print("PHYSICAL PLAN WITH COMET PLUGIN:")
print("=" * 60)
query.explain(extended=True)

# Что искать в плане с Comet:
# - CometScanExec или CometParquetScan - нативное чтение
# - CometFilterExec - нативная фильтрация (SIMD)
# - CometHashAggregateExec - нативная агрегация (DataFusion)
# - CometSortMergeJoinExec или CometBroadcastHashJoinExec
# - CometColumnarToRowExec - переход из нативного в JVM row-based
#   (присутствует только на финальном выходе или у неподдерживаемых операторов)

# Что НЕ должно быть много в плане с плагином:
# - ColumnarToRow / RowToColumnar в СЕРЕДИНЕ плана
#   (это означает неподдерживаемый оператор создал разрыв)

# ──────────────────────────────────────────────────────────
# Бенчмарк: Spark vs Spark+Comet
# ──────────────────────────────────────────────────────────
import time

def run_and_time(label, query_fn, n_runs=3):
    """Запускает запрос n_runs раз, возвращает медиану."""
    times = []
    for i in range(n_runs):
        start = time.time()
        result = query_fn()
        result.count()  # форсируем полное выполнение
        elapsed = time.time() - start
        times.append(elapsed)
        print(f"  {label} run {i+1}: {elapsed:.2f}s")
    return sorted(times)[n_runs // 2]

# Запускаем с плагином (Comet включён в конфиге)
print("\n--- С Comet Plugin ---")
comet_time = run_and_time("Comet", lambda: query)

# Отключаем Comet на лету (не перезапуская SparkSession)
# и запускаем без нативного движка
spark.conf.set("spark.comet.enabled", "false")
print("\n--- Без плагина (JVM Spark) ---")
jvm_time = run_and_time("JVM Spark", lambda: query)

# Включаем обратно
spark.conf.set("spark.comet.enabled", "true")

print(f"\nSpeedUp: {jvm_time / comet_time:.1f}×")
print(f"JVM Spark: {jvm_time:.2f}s")
print(f"Comet:     {comet_time:.2f}s")

# ──────────────────────────────────────────────────────────
# Диагностика fallback: почему оператор не пошёл в нативный движок
# ──────────────────────────────────────────────────────────

# Включаем детальное логирование Comet
spark.conf.set("spark.comet.logLevel", "DEBUG")

# Запрос с UDF - UDF не может быть нативным (Python UDF остаётся в JVM)
from pyspark.sql.types import DoubleType

@F.udf(returnType=DoubleType())
def custom_score(amount, quantity):
    """Кастомная Python-логика - остаётся в JVM."""
    return amount * quantity * 1.1

query_with_udf = df \
    .withColumn("score", custom_score("amount", "quantity")) \
    .groupBy("region") \
    .agg(F.sum("score").alias("total_score"))

print("Plan с UDF (часть в JVM, часть в нативном):")
query_with_udf.explain(extended=True)

# В логах Comet при DEBUG уровне:
# [DEBUG] CometSparkSessionExtensions: Operator BatchEvalPython not supported
#         by Comet - falling back to JVM execution for this subtree
# [DEBUG] CometSparkSessionExtensions: Inserted RowToColumnarExec before
#         CometHashAggregateExec to handle JVM→Native transition

# В плане это выглядит как:
# CometColumnarToRowExec
#   CometHashAggregateExec  ← нативная агрегация после UDF
#     RowToColumnarExec      ← конвертация JVM → Arrow перед нативным
#       BatchEvalPython [custom_score] ← UDF остаётся в JVM
#         CometFilterExec    ← нативный filter перед UDF
#           CometScanExec    ← нативное чтение

# Это показывает "разрыв" плана из-за UDF:
# [нативный] → [JVM UDF] → [конвертация] → [нативный]
# Каждый разрыв = ColumnarToRow + RowToColumnar = дополнительный overhead

spark.conf.set("spark.comet.logLevel", "WARN")

# ──────────────────────────────────────────────────────────
# Мониторинг Off-Heap памяти (нативной памяти плагина)
# ──────────────────────────────────────────────────────────

# Spark UI показывает Off-Heap использование в разделе Executor:
# http://localhost:4040/executors/

# Через Spark Metrics API:
def get_executor_memory_stats(spark):
    """Возвращает статистику памяти всех executors."""
    sc = spark.sparkContext
    status_tracker = sc.statusTracker()

    # executorInfos() доступен через JVM через Py4J
    executor_info = sc._jvm.org.apache.spark.SparkContext \
        .getOrCreate() \
        .statusStore() \
        .executorList(True)

    # В реальном коде: парсим ExecutorSummary из REST API
    # GET http://localhost:4040/api/v1/applications/{appId}/executors
    import urllib.request
    import json

    app_id = sc.applicationId
    url = f"http://localhost:4040/api/v1/applications/{app_id}/executors"
    try:
        with urllib.request.urlopen(url) as response:
            executors = json.loads(response.read())
            for ex in executors:
                print(f"Executor {ex['id']}:")
                print(f"  JVM Heap Used:    {ex['memoryUsed'] / 1024 / 1024:.0f} MB")
                print(f"  Off-Heap Used:    "
                      f"{ex.get('peakMemoryMetrics', {}).get('OffHeapMemoryUsed', 0) / 1024 / 1024:.0f} MB")
                print(f"  JVM Heap Max:     "
                      f"{ex['maxMemory'] / 1024 / 1024:.0f} MB")
    except Exception as e:
        print(f"Spark UI not accessible: {e}")

get_executor_memory_stats(spark)

// MySparkPlugin.scala
package com.example.plugin

import org.apache.spark.SparkContext
import org.apache.spark.api.plugin._
import org.slf4j.LoggerFactory
import java.util

// ── Главный класс плагина - точка регистрации ─────────────
// Должен быть в spark.plugins конфиге
class MySparkPlugin extends SparkPlugin {

  override def driverPlugin(): DriverPlugin = new MyDriverPlugin()
  override def executorPlugin(): ExecutorPlugin = new MyExecutorPlugin()
}

// ── Driver-сторона ─────────────────────────────────────────
class MyDriverPlugin extends DriverPlugin {
  private val log = LoggerFactory.getLogger(classOf[MyDriverPlugin])

  override def init(
      sc: SparkContext,
      pluginContext: PluginContext
  ): util.Map[String, String] = {

    log.info("=== [Driver Plugin] INIT ===")
    log.info("Spark master: {}", sc.master)
    log.info("App name: {}", sc.appName)
    log.info("Executor memory: {} bytes", sc.executorMemory)

    // Инициализируем Driver-сторонние ресурсы
    // (в реальном плагине: нативные библиотеки, connection pools)

    // Передаём конфигурацию на все Executor-ы
    val config = new util.HashMap[String, String]()
    config.put("plugin.driver.initialized.at",
      System.currentTimeMillis().toString)
    config.put("plugin.driver.app.name", sc.appName)
    config.put("plugin.native.enabled", "true")
    config
  }

  override def receive(message: AnyRef): AnyRef = {
    log.info("=== [Driver Plugin] Received from Executor: {} ===", message)
    // Возвращаем ответ Executor-у (опционально)
    "ACK:" + message
  }

  override def shutdown(): Unit = {
    log.info("=== [Driver Plugin] SHUTDOWN ===")
    // Освобождаем Driver-сторонние ресурсы
  }
}

// ── Executor-сторона ───────────────────────────────────────
class MyExecutorPlugin extends ExecutorPlugin {
  private val log = LoggerFactory.getLogger(classOf[MyExecutorPlugin])

  private var pluginContext: PluginContext = _
  private var executorId: String = "unknown"
  private var taskCount: Long = 0

  override def init(
      pluginContext: PluginContext,
      extraConf: util.Map[String, String]
  ): Unit = {

    this.pluginContext = pluginContext

    // Получаем ID этого Executor (если доступен через контекст)
    log.info("=== [Executor Plugin] INIT ===")

    // Читаем конфиг, переданный DriverPlugin
    val initTime = extraConf.get("plugin.driver.initialized.at")
    val appName = extraConf.get("plugin.driver.app.name")
    val nativeEnabled = extraConf.get("plugin.native.enabled")

    log.info("Config from Driver: appName={}, nativeEnabled={}, driverInitTime={}",
      appName, nativeEnabled, initTime)

    // Инициализируем Executor-сторонние ресурсы
    // В реальном плагине: System.loadLibrary("native_plugin_lib")
    log.info("Native runtime initialized (simulated)")

    // Отправляем сообщение на Driver
    pluginContext.send("EXECUTOR_READY:executor-" + Thread.currentThread().getId)
  }

  override def onTaskStart(): Unit = {
    taskCount += 1
    log.debug("[Executor Plugin] Task #{} starting", taskCount)
  }

  override def onTaskSucceeded(): Unit = {
    log.debug("[Executor Plugin] Task #{} succeeded", taskCount)
  }

  override def onTaskFailed(failureReason: TaskFailedReason): Unit = {
    log.warn("[Executor Plugin] Task #{} FAILED: {}", taskCount, failureReason)
    // КРИТИЧЕСКИ ВАЖНО: освободить нативные ресурсы при ошибке
    // В реальном плагине: freeNativeTaskResources()
  }

  override def shutdown(): Unit = {
    log.info("[Executor Plugin] SHUTDOWN after {} tasks", taskCount)
    // Освобождаем нативный контекст
    // В реальном плагине: destroyNativeContext()
  }
}

# Использование учебного плагина из Python
spark = SparkSession.builder \
    .master("local[2]") \
    .config("spark.driver.extraClassPath", "/path/to/my-plugin.jar") \
    .config("spark.executor.extraClassPath", "/path/to/my-plugin.jar") \
    .config("spark.plugins", "com.example.plugin.MySparkPlugin") \
    .config("spark.driver.extraJavaOptions",
            "-Dlog4j.logger.com.example.plugin=DEBUG") \
    .getOrCreate()

# Запускаем запрос - плагин логирует каждый шаг
df = spark.range(100000).groupBy((F.col("id") % 10).alias("group")).count()
df.show()

# В логах увидим:
# [Driver Plugin] INIT
# [Executor Plugin] INIT (для каждого Executor)
# [Executor Plugin] Received: EXECUTOR_READY:executor-123
# [Executor Plugin] Task #1 starting
# [Executor Plugin] Task #1 succeeded
# ...
# [Driver Plugin] SHUTDOWN
# [Executor Plugin] SHUTDOWN after N tasks

SELECT
    region,
    category,
    SUM(amount) as total,
    COUNT(*) as orders,
    AVG(amount) as avg_amount,
    RANK() OVER (PARTITION BY region ORDER BY SUM(amount) DESC) as rank
FROM transactions
WHERE amount > 0
GROUP BY region, category

[INFO]  GlutenExecutorPlugin: Initialized Velox runtime, offHeap=4g
[INFO]  GlutenExecutorPlugin: Task#142 starting (Stage 3, queryStage=5)
[INFO]  VeloxHashAggregateExec: Processing partition 0/8, inputRows=14_500_000
[INFO]  VeloxHashJoinExec: Building hash table, rightSide=2_100_000 rows
[WARN]  VeloxMemoryPool: Memory pressure: 92% of offHeap used
[INFO]  VeloxHashJoinExec: Building hash table, rightSide=2_100_000 rows (RETRY)
[ERROR] VeloxMemoryPool: OOM - failed to allocate 512MB, available=180MB
[ERROR] GlutenExecutorPlugin: Native OOM in Task#142, Stage 3 —
        executor will be terminated
[ERROR] KubernetesExternalClusterManager: Executor pod killed (OOMKilled)