syntax = "proto3";

service SparkSamplingService {
    rpc SamplingJob(SamplingRequest) returns (SamplingResponse) {}
    rpc RelationSamplingJob(RelationSamplingRequest) returns (RelationSamplingResponse) {}
    rpc CancelJob(CancelRequest) returns (CancelResponse) {}
}

enum SamplingMethodType{
    UNKNOWN_METHOD = 0;
    RANDOM_SAMPLING_METHOD = 1;
    STRATIFIED_SAMPLING_METHOD = 2;
    SIMPLE_SAMPLING_METHOD = 3;
    CLUSTER_SAMPLING_METHOD = 4;
}

enum FileFormatType{
    UNKNOWN_FORMAT = 0;
    FILE_FORMAT_CSV = 1;
}

/**
 * Response message for SMS tasks
 */
message SamplingRequest {
    SamplingMethodType sampling_method = 1;
    FileFormatType file_format = 2;
    SamplingConf sampling_conf = 3;
    FileFormatConf format_conf = 4;
    string input_path = 5;
    string output_path = 6;
    string job_id = 7;
}

message CancelRequest{
    string job_id = 1;
}


/**
 * Response message for canceled tasks
 * Need to determine if cancellation is successful based on whether the task reports an error or not
 */
message CancelResponse {
    // Status code, return 0 for success
    int32 code = 1;
    // message
    string message = 2;
}

/**
 * Response message for SMS tasks
 * When a task is cancelled it returns code:5000
 * Other error codes: error.py
 */
message SamplingResponse {
    // Status code, return 0 for success
    int32 code = 1;
    string message = 2;
    ResponseData data = 3;
    message ResponseData{
        // Original SMS request
        SamplingRequest parent_request = 1;
        // Sample output file path
        string sampled_path = 2;
        repeated HookMessage hook_msg = 3;
    }
}

message dict {
    string name = 1;
    string value = 2;
}

/**
 * config sampling for SMS tasks
 */
message SamplingConf {
    // Sampling rate
    // 'Float str' for random sampling
    // 'Json str' for stratified sampling like {str:float}
    // eg. '{"label_A":0.3, "label_B":0.7}' means to draw label_A and label_B in the ratio of 3:7
    // if 'Float str' is passed for stratified sampling, it is considered to be the same for each stratum
    // same as {"label_A":fraction, "label_B":fraction}
    string fraction = 1;
    // sampling with replacement(Bootstrap Sample), default is False
    bool with_replacement = 2;
    uint32 seed = 3;
    // Key for stratified sampling, necessary parameters for stratified sampling
    string stratified_key = 4;
    // Maximum number of samples drawn
    // SIMPLE_SAMPLING_METHOD requires not to be 0
    uint32 count = 5;
    // Columns to be retained, empty or null means to retain all
    // Sampled dataframe only has columns with sampling_col
    // In relationSampling, 'output_col' is a better choice if you don't know which columns need to be used in the association
    repeated string sampling_col = 6;

    // For cluster sampling, column names to distinguish between groups
    string group_by = 7;
    // The number of clusters to be sampled for whole cluster sampling
    // If the number of clusters to be extracted is not clear, use fraction as the scale input
    // When used with fraction, group_by is sampled in scale, limiting the maximum number of groups to group_num
    uint32 group_num = 8;

    // When sampling in stratification, passing in True will guarantee at least one sample per stratum
    // Will ignore count if it's True
    // set env FORCE_STRATIFILED_ENSURE_COL 'true' or 'True' will force set it's True, otherwise default False
    bool ensure_col = 9;
}

message FileFormatConf {
    // csv have a header, default is False
    bool with_header = 1;
    // File separator, default is '\001', using a comma as the default separator requires passing ','
    string sep = 2;
}

message RelationSamplingRequest{
    message SamplingStage{
        // input_path or input_name determines the dataframe involved in the association relationship
        // One of input_name or input_path must be passed in
        // name has higher priority than path, if no name, using path to construct dataframe
        string input_path = 1;
        string input_name = 2;
        string output_path = 3;
        // Columns to be retained when output, only work when export dataframe, empty or null means to retain all
        repeated string output_col = 4;
        string name = 5;
        SamplingMethodType sampling_method = 6;
        SamplingConf sampling_conf = 7;
        FileFormatType file_format = 8;
        FileFormatConf format_conf = 9;
    }

    message Relation{
        // similar with input_path and input_name, determines the dataframe
        // One of relation_path or relation_name must be passed in
        // name has higher priority than path, if no name, using path to construct dataframe
        string relation_path = 1;
        string relation_name = 2;

        // One of relation_path or relation_name or both must be passed in
        string relation_col = 3;
        string input_col = 4;
        // when using path to construct dataframe, allow read as custom file format
        FileFormatType file_format = 5;
        FileFormatConf format_conf = 6;
    }

    message RelationStage{
        // input_path or input_name determines the dataframe involved in the association relationship
        // One of input_name or input_path must be passed in
        // name has higher priority than path, if no name, using path to construct dataframe
        string input_path = 1;
        string input_name = 2;
        // If output_path exists, this dataframe will be exported
        string output_path = 3;
        // Columns to be retained when output, only work when export dataframe, empty or null means to retain all
        repeated string output_col = 4;
        // Name of the dataframe, using in response and other stage
        string name = 5;
        // input dataframe semi-join relation dataframe by input_col == relation_col
        repeated Relation relations = 6;
        FileFormatType file_format = 7;
        FileFormatConf format_conf = 8;
    }
    repeated SamplingStage sampling_stages = 1;
    repeated RelationStage relation_stages = 2;
    string job_id = 3;

    // Global default configuration, used when the stage is not specifically configured
    // If a relation or sample using name not path, FileFormat* only applied when output
    SamplingMethodType default_sampling_method = 4;
    SamplingConf default_sampling_conf = 5;
    FileFormatType default_file_format = 6;
    FileFormatConf default_format_conf = 7;

    // Only build dataframe, not export dataframe
    bool dry_run = 8;
}


message RelationSamplingResponse{
    message SampleResult{
        string name = 1;
        string sampled_path = 2;
        repeated HookMessage hook_msg = 3;
    }

    // Status code, return 0 for success
    int32 code = 1;
    string message = 2;
    ResponseData data = 3;
    message ResponseData{
        // Original request
        RelationSamplingRequest parent_request = 1;
        // Sample output file path
        repeated SampleResult results = 2;
    }
}


message HookMessage{
    string hook_name = 1;
    string period = 2;
    repeated pair meta = 3;
}

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