// Copyright 2020 The Merlin Authors
//
// Licensed 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 merlin.batch.spec;

option go_package = "github.com/caraml-dev/merlin-pyspark-app/pkg/spec";
option java_multiple_files = true;
option java_outer_classname = "PredictionJobProto";
option java_package = "com.gojek.merlin.batch.spec";

message PredictionJob {
    string version = 1;
    string kind = 2;
    string name = 3;
    oneof source {
        BigQuerySource bigquery_source = 11;
        GcsSource gcs_source = 12;
        MaxComputeSource maxcompute_source = 13;
    }
    Model model = 21;
    oneof sink {
        BigQuerySink bigquery_sink = 31;
        GcsSink gcs_sink = 32;
        MaxComputeSink maxcompute_sink = 33;
    }
}

message BigQuerySource {
    string table = 1;
    repeated string features = 2;
    map<string, string> options = 3;
}

message GcsSource {
    FileFormat format = 1;
    string uri = 2;
    repeated string features = 3;
    map<string, string> options = 4;
}

message MaxComputeSource {
    string table = 1;
    string endpoint = 2;
    repeated string features = 3;
    map<string, string> options = 4;
}

message Model {
    ModelType type = 1;
    string uri = 2;
    ModelResult result = 3;
    map<string, string> options = 4;

    message ModelResult {
        ResultType type = 1;
        // only if type is array
        ResultType item_type = 2;
    }
}

message BigQuerySink {
    string table = 1;
    string staging_bucket = 2;
    string result_column = 3;
    SaveMode save_mode = 4;
    map<string, string> options = 5;
}

message GcsSink {
    FileFormat format = 1;
    string uri = 2;
    string result_column = 3;
    SaveMode save_mode = 4;
    map<string, string> options = 5;
}

message MaxComputeSink {
    string table = 1;
    string endpoint = 2;
    string result_column = 3;
    SaveMode save_mode = 4;
    map<string, string> options = 5;
}

enum ResultType {
    DOUBLE = 0;
    FLOAT = 1;
    INTEGER = 2;
    LONG = 3;
    STRING = 4;

    ARRAY = 10;
}

enum ModelType {
    INVALID_MODEL_TYPE = 0;
    XGBOOST = 1;
    TENSORFLOW = 2;
    SKLEARN = 3;
    PYTORCH = 4;
    ONNX = 5;
    PYFUNC = 6;
    PYFUNC_V2 = 7;
}

enum FileFormat {
    INVALID_FILE_FORMAT = 0;
    CSV = 1;
    PARQUET = 2;
    AVRO = 3;
    JSON = 4;
}

enum SaveMode {
    ERRORIFEXISTS = 0;
    OVERWRITE = 1;
    APPEND = 2;
    IGNORE = 3;
    ERROR = 4;
}