// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.



syntax = "proto3";

package spark.spark_operator;

import "expr.proto";
import "partitioning.proto";
import "types.proto";

option java_package = "org.apache.comet.serde";

// The basic message representing a Spark operator.
message Operator {
  // The child operators of this
  repeated Operator children = 1;

  // Spark plan ID
  uint32 plan_id = 2;

  oneof op_struct {
    Scan scan = 100;
    Projection projection = 101;
    Filter filter = 102;
    Sort sort = 103;
    HashAggregate hash_agg = 104;
    Limit limit = 105;
    ShuffleWriter shuffle_writer = 106;
    Expand expand = 107;
    SortMergeJoin sort_merge_join = 108;
    HashJoin hash_join = 109;
    Window window = 110;
    NativeScan native_scan = 111;
    IcebergScan iceberg_scan = 112;
    ParquetWriter parquet_writer = 113;
    Explode explode = 114;
    CsvScan csv_scan = 115;
  }
}

message SparkPartitionedFile {
  string file_path = 1;
  int64 start = 2;
  int64 length = 3;
  int64 file_size = 4;
  repeated spark.spark_expression.Expr partition_values = 5;
}

// This name and the one above are not great, but they correspond to the (unfortunate) Spark names.
// I prepended "Spark" since I think there's a name collision on the native side, but we can revisit.
message SparkFilePartition {
  repeated SparkPartitionedFile partitioned_file = 1;
}

message SparkStructField {
  string name = 1;
  spark.spark_expression.DataType data_type = 2;
  bool nullable = 3;
}

message Scan {
  repeated spark.spark_expression.DataType fields = 1;
  // The source of the scan (e.g. file scan, broadcast exchange, shuffle, etc). This
  // is purely for informational purposes when viewing native query plans in
  // debug mode.
  string source = 2;
  // Whether native code can assume ownership of batches that it receives
  bool arrow_ffi_safe = 3;
}

// Common data shared by all partitions in split mode (sent once at planning)
message NativeScanCommon {
  repeated SparkStructField required_schema = 1;
  repeated SparkStructField data_schema = 2;
  repeated SparkStructField partition_schema = 3;
  repeated spark.spark_expression.Expr data_filters = 4;
  repeated int64 projection_vector = 5;
  string session_timezone = 6;
  repeated spark.spark_expression.Expr default_values = 7;
  repeated int64 default_values_indexes = 8;
  bool case_sensitive = 9;
  map<string, string> object_store_options = 10;
  bool encryption_enabled = 11;
  string source = 12;
  repeated spark.spark_expression.DataType fields = 13;
}

message NativeScan {
  // Common data shared across partitions (schemas, filters, projections, config)
  NativeScanCommon common = 1;

  // Single partition's file list (injected at execution time)
  SparkFilePartition file_partition = 2;
}

message CsvScan {
  repeated SparkStructField data_schema = 1;
  repeated SparkStructField partition_schema = 2;
  repeated int32 projection_vector = 3;
  repeated SparkFilePartition file_partitions = 4;
  map<string, string> object_store_options = 5;
  CsvOptions csv_options = 6;
}

message CsvOptions {
  bool has_header = 1;
  string delimiter = 2;
  string quote = 3;
  string escape = 4;
  optional string comment = 5;
  string terminator = 7;
  bool truncated_rows = 8;
}

// Partition value for Iceberg partition data
message PartitionValue {
  int32 field_id = 1;
  oneof value {
    int32 int_val = 2;
    int64 long_val = 3;
    int64 date_val = 4;           // days since epoch
    int64 timestamp_val = 5;       // microseconds since epoch
    int64 timestamp_tz_val = 6;    // microseconds with timezone
    string string_val = 7;
    double double_val = 8;
    float float_val = 9;
    bytes decimal_val = 10;        // unscaled BigInteger bytes
    bool bool_val = 11;
    bytes uuid_val = 12;
    bytes fixed_val = 13;
    bytes binary_val = 14;
  }
  bool is_null = 15;
}

// Collection of partition values for a single partition
message PartitionData {
  repeated PartitionValue values = 1;
}

// Common data shared by all partitions in split mode (sent once, captured in closure)
message IcebergScanCommon {
  // Catalog-specific configuration for FileIO (credentials, S3/GCS config, etc.)
  map<string, string> catalog_properties = 1;

  // Table metadata file path for FileIO initialization
  string metadata_location = 2;

  // Schema to read
  repeated SparkStructField required_schema = 3;

  // Deduplication pools (must contain all entries for cross-partition deduplication)
  repeated string schema_pool = 4;
  repeated string partition_type_pool = 5;
  repeated string partition_spec_pool = 6;
  repeated string name_mapping_pool = 7;
  repeated ProjectFieldIdList project_field_ids_pool = 8;
  repeated PartitionData partition_data_pool = 9;
  repeated DeleteFileList delete_files_pool = 10;
  repeated spark.spark_expression.Expr residual_pool = 11;

  // Number of data files to read concurrently within a single task
  uint32 data_file_concurrency_limit = 12;
}

message IcebergScan {
  // Common data shared across partitions (pools, metadata, catalog props)
  IcebergScanCommon common = 1;

  // Single partition's file scan tasks
  repeated IcebergFileScanTask file_scan_tasks = 2;
}

// Helper message for deduplicating field ID lists
message ProjectFieldIdList {
  repeated int32 field_ids = 1;
}

// Helper message for deduplicating delete file lists
message DeleteFileList {
  repeated IcebergDeleteFile delete_files = 1;
}

// Iceberg FileScanTask containing data file, delete files, and residual filter
message IcebergFileScanTask {
  // Data file path (e.g., s3://bucket/warehouse/db/table/data/00000-0-abc.parquet)
  string data_file_path = 1;

  // Byte range to read (for split files)
  uint64 start = 2;
  uint64 length = 3;

  // Record count if reading entire file
  optional uint64 record_count = 4;

  // Total file size from the manifest entry, used to skip stat/HEAD calls
  uint64 file_size_in_bytes = 5;

  // Indices into IcebergScan deduplication pools
  uint32 schema_idx = 15;
  optional uint32 partition_type_idx = 16;
  optional uint32 partition_spec_idx = 17;
  optional uint32 name_mapping_idx = 18;
  uint32 project_field_ids_idx = 19;
  optional uint32 partition_data_idx = 20;
  optional uint32 delete_files_idx = 21;
  optional uint32 residual_idx = 22;
}

// Iceberg delete file for MOR tables (positional or equality deletes)
// Positional: (file_path, row_position) pairs to skip
// Equality: Column values to filter out (specified by equality_ids)
message IcebergDeleteFile {
  // Delete file path
  string file_path = 1;

  // POSITION_DELETES or EQUALITY_DELETES
  string content_type = 2;

  // Partition spec ID
  int32 partition_spec_id = 3;

  // Equality field IDs (empty for positional deletes)
  repeated int32 equality_ids = 4;

  // // Total file size from the manifest entry, used to skip stat/HEAD calls
  uint64 file_size_in_bytes = 5;
}

message Projection {
  repeated spark.spark_expression.Expr project_list = 1;
}

message Filter {
  spark.spark_expression.Expr predicate = 1;
}

message Sort {
  repeated spark.spark_expression.Expr sort_orders = 1;
  optional int32 fetch = 3;
  optional int32 skip = 4;
}

message HashAggregate {
  repeated spark.spark_expression.Expr grouping_exprs = 1;
  repeated spark.spark_expression.AggExpr agg_exprs = 2;
  repeated spark.spark_expression.Expr result_exprs = 3;
  AggregateMode mode = 5;
}

message Limit {
  int32 limit = 1;
  int32 offset = 2;
}

enum CompressionCodec {
  None = 0;
  Zstd = 1;
  Lz4 = 2;
  Snappy = 3;
}

message ShuffleWriter {
  spark.spark_partitioning.Partitioning partitioning = 1;
  string output_data_file = 3;
  string output_index_file = 4;
  CompressionCodec codec = 5;
  int32 compression_level = 6;
  bool tracing_enabled = 7;
  // Size of the write buffer in bytes used when writing shuffle data to disk.
  // Larger values may improve write performance but use more memory.
  int32 write_buffer_size = 8;
}

message ParquetWriter {
  string output_path = 1;
  CompressionCodec compression = 2;
  repeated string column_names = 4;
  // Working directory for temporary files (used by FileCommitProtocol)
  // If not set, files are written directly to output_path
  optional string work_dir = 5;
  // Job ID for tracking this write operation
  optional string job_id = 6;
  // Task attempt ID for this specific task
  optional int32 task_attempt_id = 7;
  // Options for configuring object stores such as AWS S3, GCS, etc. The key-value pairs are taken
  // from Hadoop configuration for compatibility with Hadoop FileSystem implementations of object
  // stores.
  // The configuration values have hadoop. or spark.hadoop. prefix trimmed. For instance, the
  // configuration value "spark.hadoop.fs.s3a.access.key" will be stored as "fs.s3a.access.key" in
  // the map.
  map<string, string> object_store_options = 8;
}

enum AggregateMode {
  Partial = 0;
  Final = 1;
}

message Expand {
  repeated spark.spark_expression.Expr project_list = 1;
  int32 num_expr_per_project = 3;
}

message Explode {
  // The array expression to explode into multiple rows
  spark.spark_expression.Expr child = 1;
  // Whether this is explode_outer (produces null row for empty/null arrays)
  bool outer = 2;
  // Expressions for other columns to project alongside the exploded values
  repeated spark.spark_expression.Expr project_list = 3;
}

message HashJoin {
  repeated spark.spark_expression.Expr left_join_keys = 1;
  repeated spark.spark_expression.Expr right_join_keys = 2;
  JoinType join_type = 3;
  optional spark.spark_expression.Expr condition = 4;
  BuildSide build_side = 5;
}

message SortMergeJoin {
  repeated spark.spark_expression.Expr left_join_keys = 1;
  repeated spark.spark_expression.Expr right_join_keys = 2;
  JoinType join_type = 3;
  repeated spark.spark_expression.Expr sort_options = 4;
  optional spark.spark_expression.Expr condition = 5;
}

enum JoinType {
  Inner = 0;
  LeftOuter = 1;
  RightOuter = 2;
  FullOuter = 3;
  LeftSemi = 4;
  LeftAnti = 5;
}

enum BuildSide {
  BuildLeft = 0;
  BuildRight = 1;
}

message WindowExpr {
  spark.spark_expression.Expr built_in_window_function = 1;
  spark.spark_expression.AggExpr agg_func = 2;
  WindowSpecDefinition spec = 3;
}

enum WindowFrameType {
  Rows = 0;
  Range = 1;
}

message WindowFrame {
  WindowFrameType frame_type = 1;
  LowerWindowFrameBound lower_bound = 2;
  UpperWindowFrameBound upper_bound = 3;
}

message LowerWindowFrameBound {
  oneof lower_frame_bound_struct {
    UnboundedPreceding unboundedPreceding = 1;
    Preceding preceding = 2;
    CurrentRow currentRow = 3;
  }
}

message UpperWindowFrameBound {
  oneof upper_frame_bound_struct {
    UnboundedFollowing unboundedFollowing = 1;
    Following following = 2;
    CurrentRow currentRow = 3;
  }
}

message Preceding {
  int64 offset = 1;
}

message Following {
  int64 offset = 1;
}

message UnboundedPreceding {}
message UnboundedFollowing {}
message CurrentRow {}

message WindowSpecDefinition {
  repeated spark.spark_expression.Expr partitionSpec = 1;
  repeated spark.spark_expression.Expr orderSpec = 2;
  WindowFrame frameSpecification = 3;
}

message Window {
  repeated WindowExpr window_expr = 1;
  repeated spark.spark_expression.Expr order_by_list = 2;
  repeated spark.spark_expression.Expr partition_by_list = 3;
  Operator child = 4;
}
