/*
 * 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
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * 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 ballista.protobuf;

option java_multiple_files = true;
option java_package = "org.ballistacompute.protobuf";
option java_outer_classname = "BallistaProto";

import "datafusion.proto";

///////////////////////////////////////////////////////////////////////////////////////////////////
// Ballista Logical Plan
///////////////////////////////////////////////////////////////////////////////////////////////////

// LogicalPlan is a nested type
message LogicalPlanNode {
  oneof LogicalPlanType {
    ListingTableScanNode listing_scan = 1;
    ProjectionNode projection = 3;
    SelectionNode selection = 4;
    LimitNode limit = 5;
    AggregateNode aggregate = 6;
    JoinNode join = 7;
    SortNode sort = 8;
    RepartitionNode repartition = 9;
    EmptyRelationNode empty_relation = 10;
    CreateExternalTableNode create_external_table = 11;
    ExplainNode explain = 12;
    WindowNode window = 13;
    AnalyzeNode analyze = 14;
    CrossJoinNode cross_join = 15;
    ValuesNode values = 16;
    LogicalExtensionNode extension = 17;
    CreateCatalogSchemaNode create_catalog_schema = 18;
    UnionNode union = 19;
    CreateCatalogNode create_catalog = 20;
    SubqueryAliasNode subquery_alias = 21;
    CreateViewNode create_view = 22;
  }
}

message LogicalExtensionNode {
  bytes node = 1;
  repeated LogicalPlanNode inputs = 2;
}

message ProjectionColumns {
  repeated string columns = 1;
}

message Statistics {
  int64 num_rows = 1;
  int64 total_byte_size = 2;
  repeated ColumnStats column_stats = 3;
  bool is_exact = 4;
}

message FileRange {
  int64 start = 1;
  int64 end = 2;
}

message PartitionedFile {
  string path = 1;
  uint64 size = 2;
  uint64 last_modified_ns = 3;
  repeated datafusion.ScalarValue partition_values = 4;
  FileRange range = 5;
}

message CsvFormat {
  bool has_header = 1;
  string delimiter = 2;
}

message ParquetFormat {
  bool enable_pruning = 1;
}

message AvroFormat {}

message ListingTableScanNode {
  string table_name = 1;
  string path = 2;
  string file_extension = 3;
  ProjectionColumns projection = 4;
  datafusion.Schema schema = 5;
  repeated datafusion.LogicalExprNode filters = 6;
  repeated string table_partition_cols = 7;
  bool collect_stat = 8;
  uint32 target_partitions = 9;
  oneof FileFormatType {
    CsvFormat csv = 10;
    ParquetFormat parquet = 11;
    AvroFormat avro = 12;
  }
}

message ProjectionNode {
  LogicalPlanNode input = 1;
  repeated datafusion.LogicalExprNode expr = 2;
  oneof optional_alias {
    string alias = 3;
  }
}

message SelectionNode {
  LogicalPlanNode input = 1;
  datafusion.LogicalExprNode expr = 2;
}

message SortNode {
  LogicalPlanNode input = 1;
  repeated datafusion.LogicalExprNode expr = 2;
}

message RepartitionNode {
  LogicalPlanNode input = 1;
  oneof partition_method {
    uint64 round_robin = 2;
    HashRepartition hash = 3;
  }
}

message HashRepartition {
  repeated datafusion.LogicalExprNode hash_expr = 1;
  uint64 partition_count = 2;
}

message EmptyRelationNode {
  bool produce_one_row = 1;
}

message CreateExternalTableNode {
  string name = 1;
  string location = 2;
  FileType file_type = 3;
  bool has_header = 4;
  datafusion.DfSchema schema = 5;
  repeated string table_partition_cols = 6;
  bool if_not_exists = 7;
  string delimiter = 8;
}

message CreateCatalogSchemaNode {
  string schema_name = 1;
  bool if_not_exists = 2;
  datafusion.DfSchema schema = 3;
}

message CreateCatalogNode {
  string catalog_name = 1;
  bool if_not_exists = 2;
  datafusion.DfSchema schema = 3;
}

message CreateViewNode {
  string name = 1;
  LogicalPlanNode input = 2;
  bool or_replace = 3;
}

// a node containing data for defining values list. unlike in SQL where it's two dimensional, here
// the list is flattened, and with the field n_cols it can be parsed and partitioned into rows
message ValuesNode {
  uint64 n_cols = 1;
  repeated datafusion.LogicalExprNode values_list = 2;
}

enum FileType {
  NdJson = 0;
  Parquet = 1;
  CSV = 2;
  Avro = 3;
}

message AnalyzeNode {
  LogicalPlanNode input = 1;
  bool verbose = 2;
}

message ExplainNode {
  LogicalPlanNode input = 1;
  bool verbose = 2;
}

message AggregateNode {
  LogicalPlanNode input = 1;
  repeated datafusion.LogicalExprNode group_expr = 2;
  repeated datafusion.LogicalExprNode aggr_expr = 3;
}

message WindowNode {
  LogicalPlanNode input = 1;
  repeated datafusion.LogicalExprNode window_expr = 2;
}

enum JoinType {
  INNER = 0;
  LEFT = 1;
  RIGHT = 2;
  FULL = 3;
  SEMI = 4;
  ANTI = 5;
}

enum JoinConstraint {
  ON = 0;
  USING = 1;
}

message JoinNode {
  LogicalPlanNode left = 1;
  LogicalPlanNode right = 2;
  JoinType join_type = 3;
  JoinConstraint join_constraint = 4;
  repeated datafusion.Column left_join_column = 5;
  repeated datafusion.Column right_join_column = 6;
  bool null_equals_null = 7;
}

message UnionNode {
  repeated LogicalPlanNode inputs = 1;
}

message CrossJoinNode {
  LogicalPlanNode left = 1;
  LogicalPlanNode right = 2;
}

message LimitNode {
  LogicalPlanNode input = 1;
  uint32 limit = 2;
}

message SelectionExecNode {
  datafusion.LogicalExprNode expr = 1;
}

message SubqueryAliasNode {
  LogicalPlanNode input = 1;
  string alias = 2;
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Ballista Physical Plan
///////////////////////////////////////////////////////////////////////////////////////////////////

// PhysicalPlanNode is a nested type
message PhysicalPlanNode {
  oneof PhysicalPlanType {
    ParquetScanExecNode parquet_scan = 1;
    CsvScanExecNode csv_scan = 2;
    EmptyExecNode empty = 3;
    ProjectionExecNode projection = 4;
    GlobalLimitExecNode global_limit = 6;
    LocalLimitExecNode local_limit = 7;
    AggregateExecNode aggregate = 8;
    HashJoinExecNode hash_join = 9;
    ShuffleReaderExecNode shuffle_reader = 10;
    SortExecNode sort = 11;
    CoalesceBatchesExecNode coalesce_batches = 12;
    FilterExecNode filter = 13;
    CoalescePartitionsExecNode merge = 14;
    UnresolvedShuffleExecNode unresolved = 15;
    RepartitionExecNode repartition = 16;
    WindowAggExecNode window = 17;
    ShuffleWriterExecNode shuffle_writer = 18;
    CrossJoinExecNode cross_join = 19;
    AvroScanExecNode avro_scan = 20;
    PhysicalExtensionNode extension = 21;
    UnionExecNode union = 22;
    ExplainExecNode explain = 23;
  }
}

message PhysicalExtensionNode {
  bytes node = 1;
  repeated PhysicalPlanNode inputs = 2;
}

// physical expressions
message PhysicalExprNode {
  oneof ExprType {
    // column references
    PhysicalColumn column = 1;

    datafusion.ScalarValue literal = 2;

    // binary expressions
    PhysicalBinaryExprNode binary_expr = 3;

    // aggregate expressions
    PhysicalAggregateExprNode aggregate_expr = 4;

    // null checks
    PhysicalIsNull is_null_expr = 5;
    PhysicalIsNotNull is_not_null_expr = 6;
    PhysicalNot not_expr = 7;

    PhysicalCaseNode case_ = 8;
    PhysicalCastNode cast = 9;
    PhysicalSortExprNode sort = 10;
    PhysicalNegativeNode negative = 11;
    PhysicalInListNode in_list = 12;
    PhysicalScalarFunctionNode scalar_function = 13;
    PhysicalTryCastNode try_cast = 14;

    // window expressions
    PhysicalWindowExprNode window_expr = 15;

    PhysicalScalarUdfNode scalar_udf = 16;
  }
}

message PhysicalScalarUdfNode {
  string name = 1;
  repeated PhysicalExprNode args = 2;
  datafusion.ArrowType return_type = 4;
}

message PhysicalAggregateExprNode {
  datafusion.AggregateFunction aggr_function = 1;
  repeated PhysicalExprNode expr = 2;
}

message PhysicalWindowExprNode {
  oneof window_function {
    datafusion.AggregateFunction aggr_function = 1;
    datafusion.BuiltInWindowFunction built_in_function = 2;
    // udaf = 3
  }
  PhysicalExprNode expr = 4;
}

message PhysicalIsNull {
  PhysicalExprNode expr = 1;
}

message PhysicalIsNotNull {
  PhysicalExprNode expr = 1;
}

message PhysicalNot {
  PhysicalExprNode expr = 1;
}

message PhysicalAliasNode {
  PhysicalExprNode expr = 1;
  string alias = 2;
}

message PhysicalBinaryExprNode {
  PhysicalExprNode l = 1;
  PhysicalExprNode r = 2;
  string op = 3;
}

message PhysicalSortExprNode {
  PhysicalExprNode expr = 1;
  bool asc = 2;
  bool nulls_first = 3;
}

message PhysicalWhenThen {
  PhysicalExprNode when_expr = 1;
  PhysicalExprNode then_expr = 2;
}

message PhysicalInListNode {
  PhysicalExprNode expr = 1;
  repeated PhysicalExprNode list = 2;
  bool negated = 3;
}

message PhysicalCaseNode {
  PhysicalExprNode expr = 1;
  repeated PhysicalWhenThen when_then_expr = 2;
  PhysicalExprNode else_expr = 3;
}

message PhysicalScalarFunctionNode {
  string name = 1;
  datafusion.ScalarFunction fun = 2;
  repeated PhysicalExprNode args = 3;
  datafusion.ArrowType return_type = 4;
}

message PhysicalTryCastNode {
  PhysicalExprNode expr = 1;
  datafusion.ArrowType arrow_type = 2;
}

message PhysicalCastNode {
  PhysicalExprNode expr = 1;
  datafusion.ArrowType arrow_type = 2;
}

message PhysicalNegativeNode {
  PhysicalExprNode expr = 1;
}

message UnresolvedShuffleExecNode {
  uint32 stage_id = 1;
  datafusion.Schema schema = 2;
  uint32 input_partition_count = 3;
  uint32 output_partition_count = 4;
}

message FilterExecNode {
  PhysicalPlanNode input = 1;
  PhysicalExprNode expr = 2;
}

message FileGroup {
  repeated PartitionedFile files = 1;
}

message ScanLimit {
  // wrap into a message to make it optional
  uint32 limit = 1;
}

message FileScanExecConf {
  repeated FileGroup file_groups = 1;
  datafusion.Schema schema = 2;
  repeated uint32 projection = 4;
  ScanLimit limit = 5;
  Statistics statistics = 6;
  repeated string table_partition_cols = 7;
}

message ParquetScanExecNode {
  FileScanExecConf base_conf = 1;
  datafusion.LogicalExprNode pruning_predicate = 2;
}

message CsvScanExecNode {
  FileScanExecConf base_conf = 1;
  bool has_header = 2;
  string delimiter = 3;
}

message AvroScanExecNode {
  FileScanExecConf base_conf = 1;
}

enum PartitionMode {
  COLLECT_LEFT = 0;
  PARTITIONED = 1;
}

message HashJoinExecNode {
  PhysicalPlanNode left = 1;
  PhysicalPlanNode right = 2;
  repeated JoinOn on = 3;
  JoinType join_type = 4;
  PartitionMode partition_mode = 6;
  bool null_equals_null = 7;
}

message UnionExecNode {
  repeated PhysicalPlanNode inputs = 1;
}

message ExplainExecNode {
  datafusion.Schema schema = 1;
  repeated datafusion.StringifiedPlan stringified_plans = 2;
  bool verbose = 3;
}

message CrossJoinExecNode {
  PhysicalPlanNode left = 1;
  PhysicalPlanNode right = 2;
}

message PhysicalColumn {
  string name = 1;
  uint32 index = 2;
}

message JoinOn {
  PhysicalColumn left = 1;
  PhysicalColumn right = 2;
}

message EmptyExecNode {
  bool produce_one_row = 1;
  datafusion.Schema schema = 2;
}

message ProjectionExecNode {
  PhysicalPlanNode input = 1;
  repeated PhysicalExprNode expr = 2;
  repeated string expr_name = 3;
}

enum AggregateMode {
  PARTIAL = 0;
  FINAL = 1;
  FINAL_PARTITIONED = 2;
}

message WindowAggExecNode {
  PhysicalPlanNode input = 1;
  repeated PhysicalExprNode window_expr = 2;
  repeated string window_expr_name = 3;
  datafusion.Schema input_schema = 4;
}

message AggregateExecNode {
  repeated PhysicalExprNode group_expr = 1;
  repeated PhysicalExprNode aggr_expr = 2;
  AggregateMode mode = 3;
  PhysicalPlanNode input = 4;
  repeated string group_expr_name = 5;
  repeated string aggr_expr_name = 6;
  // we need the input schema to the partial aggregate to pass to the final aggregate
  datafusion.Schema input_schema = 7;
}

message ShuffleWriterExecNode {
  //TODO it seems redundant to provide job and stage id here since we also have them
  // in the TaskDefinition that wraps this plan
  string job_id = 1;
  uint32 stage_id = 2;
  PhysicalPlanNode input = 3;
  PhysicalHashRepartition output_partitioning = 4;
}

message ShuffleReaderExecNode {
  repeated ShuffleReaderPartition partition = 1;
  datafusion.Schema schema = 2;
}

message ShuffleReaderPartition {
  // each partition of a shuffle read can read data from multiple locations
  repeated PartitionLocation location = 1;
}

message GlobalLimitExecNode {
  PhysicalPlanNode input = 1;
  uint32 limit = 2;
}

message LocalLimitExecNode {
  PhysicalPlanNode input = 1;
  uint32 limit = 2;
}

message SortExecNode {
  PhysicalPlanNode input = 1;
  repeated PhysicalExprNode expr = 2;
}

message CoalesceBatchesExecNode {
  PhysicalPlanNode input = 1;
  uint32 target_batch_size = 2;
}

message CoalescePartitionsExecNode {
  PhysicalPlanNode input = 1;
}

message PhysicalHashRepartition {
  repeated PhysicalExprNode hash_expr = 1;
  uint64 partition_count = 2;
}

message RepartitionExecNode{
  PhysicalPlanNode input = 1;
  oneof partition_method {
    uint64 round_robin = 2;
    PhysicalHashRepartition hash = 3;
    uint64 unknown = 4;
  }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
// Ballista Scheduling
///////////////////////////////////////////////////////////////////////////////////////////////////

message KeyValuePair {
  string key = 1;
  string value = 2;
}

message Action {

  oneof ActionType {
    // Fetch a partition from an executor
    FetchPartition fetch_partition = 3;
  }

  // configuration settings
  repeated KeyValuePair settings = 100;
}

message ExecutePartition {
  string job_id = 1;
  uint32 stage_id = 2;
  repeated uint32 partition_id = 3;
  PhysicalPlanNode plan = 4;
  // The task could need to read partitions from other executors
  repeated PartitionLocation partition_location = 5;
  // Output partition for shuffle writer
  PhysicalHashRepartition output_partitioning = 6;
}

message FetchPartition {
  string job_id = 1;
  uint32 stage_id = 2;
  uint32 partition_id = 3;
  string path = 4;
}

// Mapping from partition id to executor id
message PartitionLocation {
  PartitionId partition_id = 1;
  ExecutorMetadata executor_meta = 2;
  PartitionStats partition_stats = 3;
  string path = 4;
}

// Unique identifier for a materialized partition of data
message PartitionId {
  string job_id = 1;
  uint32 stage_id = 2;
  uint32 partition_id = 4;
}

message PartitionStats {
  int64 num_rows = 1;
  int64 num_batches = 2;
  int64 num_bytes = 3;
  repeated ColumnStats column_stats = 4;
}

message ColumnStats {
  datafusion.ScalarValue min_value = 1;
  datafusion.ScalarValue max_value = 2;
  uint32 null_count = 3;
  uint32 distinct_count = 4;
}

// Used by scheduler
message ExecutorMetadata {
  string id = 1;
  string host = 2;
  uint32 port = 3;
  uint32 grpc_port = 4;
  ExecutorSpecification specification = 5;
}

// Used by grpc
message ExecutorRegistration {
  string id = 1;
  // "optional" keyword is stable in protoc 3.15 but prost is still on 3.14 (see https://github.com/tokio-rs/prost/issues/430 and https://github.com/tokio-rs/prost/pull/455)
  // this syntax is ugly but is binary compatible with the "optional" keyword (see https://stackoverflow.com/questions/42622015/how-to-define-an-optional-field-in-protobuf-3)
  oneof optional_host {
    string host = 2;
  }
  uint32 port = 3;
  uint32 grpc_port = 4;
  ExecutorSpecification specification = 5;
}

message ExecutorHeartbeat {
  string executor_id = 1;
  // Unix epoch-based timestamp in seconds
  uint64 timestamp = 2;
  ExecutorState state = 3;
}

message ExecutorState {
  repeated ExecutorMetric metrics = 1;
}

message ExecutorMetric {
  // TODO add more metrics
  oneof metric {
    uint64 available_memory = 1;
  }
}

message ExecutorSpecification {
  repeated ExecutorResource resources = 1;
}

message ExecutorResource {
  // TODO add more resources
  oneof resource {
    uint32 task_slots = 1;
  }
}

message ExecutorData {
  string executor_id = 1;
  repeated ExecutorResourcePair resources = 2;
}

message ExecutorResourcePair {
  ExecutorResource total = 1;
  ExecutorResource available = 2;
}

message RunningTask {
  string executor_id = 1;
}

message FailedTask {
  string error = 1;
}

message CompletedTask {
  string executor_id = 1;
  // TODO tasks are currently always shuffle writes but this will not always be the case
  // so we might want to think about some refactoring of the task definitions
  repeated ShuffleWritePartition partitions = 2;
}

message ShuffleWritePartition {
  uint64 partition_id = 1;
  string path = 2;
  uint64 num_batches = 3;
  uint64 num_rows = 4;
  uint64 num_bytes = 5;
}

message TaskStatus {
  PartitionId task_id = 1;
  oneof status {
    RunningTask running = 2;
    FailedTask failed = 3;
    CompletedTask completed = 4;
  }
}

message PollWorkParams {
  ExecutorRegistration metadata = 1;
  bool can_accept_task = 2;
  // All tasks must be reported until they reach the failed or completed state
  repeated TaskStatus task_status = 3;
}

message TaskDefinition {
  PartitionId task_id = 1;
  bytes plan = 2;
  // Output partition for shuffle writer
  PhysicalHashRepartition output_partitioning = 3;
  string session_id = 4;
  repeated KeyValuePair props = 5;
}

message JobSessionConfig {
  string session_id = 1;
  repeated KeyValuePair configs = 2;
}

message PollWorkResult {
  TaskDefinition task = 1;
}

message RegisterExecutorParams {
  ExecutorRegistration metadata = 1;
}

message RegisterExecutorResult {
  bool success = 1;
}

message HeartBeatParams {
  string executor_id = 1;
  ExecutorState state = 2;
}

message HeartBeatResult {
  // TODO it's from Spark for BlockManager
  bool reregister = 1;
}

message StopExecutorParams {
}

message StopExecutorResult {
}

message UpdateTaskStatusParams {
  string executor_id = 1;
  // All tasks must be reported until they reach the failed or completed state
  repeated TaskStatus task_status = 2;
}

message UpdateTaskStatusResult {
  bool success = 1;
}

message ExecuteQueryParams {
  oneof query {
    bytes logical_plan = 1;
    string sql = 2;
  }
  oneof optional_session_id {
    string session_id = 3;
  }
  repeated KeyValuePair settings = 4;
}

message ExecuteSqlParams {
  string sql = 1;
}

message ExecuteQueryResult {
  string job_id = 1;
  string session_id = 2;
}

message GetJobStatusParams {
  string job_id = 1;
}

message CompletedJob {
  repeated PartitionLocation partition_location = 1;
}

message QueuedJob {}

// TODO: add progress report
message RunningJob {}

message FailedJob {
  string error = 1;
}

message JobStatus {
  oneof status {
    QueuedJob queued = 1;
    RunningJob running = 2;
    FailedJob failed = 3;
    CompletedJob completed = 4;
  }
}

message GetJobStatusResult {
  JobStatus status = 1;
}

message GetFileMetadataParams {
  string path = 1;
  FileType file_type = 2;
}

message GetFileMetadataResult {
  datafusion.Schema schema = 1;
}

message FilePartitionMetadata {
  repeated string filename = 1;
}

message LaunchTaskParams {
  // Allow to launch a task set to an executor at once
  repeated TaskDefinition task = 1;
}

message LaunchTaskResult {
  bool success = 1;
  // TODO when part of the task set are scheduled successfully
}

service SchedulerGrpc {
  // Executors must poll the scheduler for heartbeat and to receive tasks
  rpc PollWork (PollWorkParams) returns (PollWorkResult) {}

  rpc RegisterExecutor(RegisterExecutorParams) returns (RegisterExecutorResult) {}

  // Push-based task scheduler will only leverage this interface
  // rather than the PollWork interface to report executor states
  rpc HeartBeatFromExecutor (HeartBeatParams) returns (HeartBeatResult) {}

  rpc UpdateTaskStatus (UpdateTaskStatusParams) returns (UpdateTaskStatusResult) {}

  rpc GetFileMetadata (GetFileMetadataParams) returns (GetFileMetadataResult) {}

  rpc ExecuteQuery (ExecuteQueryParams) returns (ExecuteQueryResult) {}

  rpc GetJobStatus (GetJobStatusParams) returns (GetJobStatusResult) {}
}

service ExecutorGrpc {
  rpc LaunchTask (LaunchTaskParams) returns (LaunchTaskResult) {}

  rpc StopExecutor (StopExecutorParams) returns (StopExecutorResult) {}
}
