syntax = "proto3";

package chroma;

import "chromadb/proto/chroma.proto";

message ScanOperator {
    Collection collection = 1;
    // Reserve for deprecated fields
    reserved 2, 3, 4;
    Segment knn = 5;
    Segment metadata = 6;
    Segment record = 7;
}

message FilterOperator {
    optional UserIds ids = 1;
    optional Where where = 2;
    optional WhereDocument where_document = 3;
}

message KNNOperator {
    repeated Vector embeddings = 1;
    uint32 fetch = 2;
}

message LimitOperator {
    uint32 offset = 1;
    optional uint32 limit = 2;
}

message ProjectionOperator {
    bool document = 1;
    bool embedding = 2;
    bool metadata = 3;
}

message KNNProjectionOperator {
    ProjectionOperator projection = 1;
    bool distance = 2;
}

message CountPlan {
    ScanOperator scan = 1;
}

message CountResult {
    uint32 count = 1;
    uint64 pulled_log_bytes = 2;
}

message GetPlan {
    ScanOperator scan = 1;
    FilterOperator filter = 2;
    LimitOperator limit = 3;
    ProjectionOperator projection = 4;
}

message ProjectionRecord {
    string id = 1;
    optional string document = 2;
    optional Vector embedding = 3;
    optional UpdateMetadata metadata = 4;
}

message GetResult {
    repeated ProjectionRecord records = 1;
    uint64 pulled_log_bytes = 2;
}

message KNNPlan {
    ScanOperator scan = 1;
    FilterOperator filter = 2;
    KNNOperator knn = 3;
    KNNProjectionOperator projection = 4;
}

message KNNProjectionRecord {
    ProjectionRecord record = 1;
    optional float distance = 2;
}

message KNNResult {
    repeated KNNProjectionRecord records = 1;
}

message KNNBatchResult {
    repeated KNNResult results = 1;
    uint64 pulled_log_bytes = 2;
}

message QueryVector {
    oneof vector {
        Vector dense = 1;
        SparseVector sparse = 2;
    }
}

message Rank {
    message Division {
        Rank left = 1;
        Rank right = 2;
    }

    message Knn {
        QueryVector embedding = 1;
        string key = 2;
        uint32 limit = 3;
        optional float default = 4;
        bool ordinal = 5;
    }
    
    message Subtraction {
        Rank left = 1;
        Rank right = 2;
    }

    message RankList {
        repeated Rank ranks = 1;
    }

    oneof rank {
        Rank absolute = 1;
        Division division = 2;
        Rank exponentiation = 3;
        Knn knn = 4;
        Rank logarithm = 5;
        RankList maximum = 6;
        RankList minimum = 7;
        RankList multiplication = 8;
        Subtraction subtraction = 9;
        RankList summation = 10;
        float value = 11;
    }
}

message SelectOperator {
    repeated string fields = 1;
}

message SearchPayload {
    FilterOperator filter = 1;
    Rank rank = 2;
    LimitOperator limit = 3;
    SelectOperator select = 4;
}

message SearchPlan {
    ScanOperator scan = 1;
    repeated SearchPayload payloads = 2;
}

message SearchRecord {
    string id = 1;
    optional string document = 2;
    optional Vector embedding = 3;
    optional UpdateMetadata metadata = 4;
    optional float score = 5;
}

message SearchPayloadResult {
    repeated SearchRecord records = 1;
}

message SearchResult {
    repeated SearchPayloadResult results = 1;
    uint64 pulled_log_bytes = 2;
}

service QueryExecutor {
    rpc Count(CountPlan) returns (CountResult) {}
    rpc Get(GetPlan) returns (GetResult) {}
    rpc KNN(KNNPlan) returns (KNNBatchResult) {}
    rpc Search(SearchPlan) returns (SearchResult) {}
}

