/*
 * 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.connect;

import "spark/connect/relations.proto";
import "spark/connect/expressions.proto";
import "spark/connect/ml_common.proto";

option java_multiple_files = true;
option java_package = "org.apache.spark.connect.proto";
option go_package = "internal/generated";

// Command for ML
message MlCommand {
  oneof command {
    Fit fit = 1;
    Fetch fetch = 2;
    Delete delete = 3;
    Write write = 4;
    Read read = 5;
    Evaluate evaluate = 6;
    CleanCache clean_cache = 7;
    GetCacheInfo get_cache_info = 8;
    CreateSummary create_summary = 9;
    GetModelSize get_model_size = 10;
  }

  // Command for estimator.fit(dataset)
  message Fit {
    // (Required) Estimator information (its type should be OPERATOR_TYPE_ESTIMATOR)
    MlOperator estimator = 1;
    // (Optional) parameters of the Estimator
    optional MlParams params = 2;
    // (Required) the training dataset
    Relation dataset = 3;
  }

  // Command to delete the cached objects which could be a model
  // or summary evaluated by a model
  message Delete {
    repeated ObjectRef obj_refs = 1;
    // if set `evict_only` to true, only evict the cached model from memory,
    // but keep the offloaded model in Spark driver local disk.
    optional bool evict_only = 2;
  }

  // Force to clean up all the ML cached objects
  message CleanCache { }

  // Get the information of all the ML cached objects
  message GetCacheInfo { }

  // Command to write ML operator
  message Write {
    // It could be an estimator/evaluator or the cached model
    oneof type {
      // Estimator or evaluator
      MlOperator operator = 1;
      // The cached model
      ObjectRef obj_ref = 2;
    }
    // (Optional) The parameters of operator which could be estimator/evaluator or a cached model
    optional MlParams params = 3;
    // (Required) Save the ML instance to the path
    string path = 4;
    // (Optional) Overwrites if the output path already exists.
    optional bool should_overwrite = 5;
    // (Optional) The options of the writer
    map<string, string> options = 6;
  }

  // Command to load ML operator.
  message Read {
    // (Required) ML operator information
    MlOperator operator = 1;
    // (Required) Load the ML instance from the input path
    string path = 2;
  }

  // Command for evaluator.evaluate(dataset)
  message Evaluate {
    // (Required) Evaluator information (its type should be OPERATOR_TYPE_EVALUATOR)
    MlOperator evaluator = 1;
    // (Optional) parameters of the Evaluator
    optional MlParams params = 2;
    // (Required) the evaluating dataset
    Relation dataset = 3;
  }

  // This is for re-creating the model summary when the model summary is lost
  // (model summary is lost when the model is offloaded and then loaded back)
  message CreateSummary {
    ObjectRef model_ref = 1;
    Relation dataset = 2;
  }

  // This is for query the model estimated in-memory size
  message GetModelSize {
    ObjectRef model_ref = 1;
  }
}

// The result of MlCommand
message MlCommandResult {
  oneof result_type {
    // The result of the attribute
    Expression.Literal param = 1;
    // Evaluate a Dataset in a model and return the cached ID of summary
    string summary = 2;
    // Operator information
    MlOperatorInfo operator_info = 3;
  }

  // Represents an operator info
  message MlOperatorInfo {
    oneof type {
      // The cached object which could be a model or summary evaluated by a model
      ObjectRef obj_ref = 1;
      // Operator name
      string name = 2;
    }
    // (Optional) the 'uid' of a ML object
    // Note it is different from the 'id' of a cached object.
    optional string uid = 3;
    // (Optional) parameters
    optional MlParams params = 4;
    // (Optional) warning message generated during the ML command execution
    optional string warning_message = 5;
  }
}
