cases.py

import json
import logging
from enum import Enum, auto

from vectordb_bench import config
from vectordb_bench.backend.clients.api import MetricType
from vectordb_bench.backend.filter import Filter, FilterOp, IntFilter, LabelFilter, NewIntFilter, NonFilter, non_filter
from vectordb_bench.backend.payload import PayloadProfile
from vectordb_bench.base import BaseModel
from vectordb_bench.frontend.components.custom.getCustomConfig import CustomDatasetConfig

from .dataset import (
    CustomDataset,
    Dataset,
    DatasetManager,
    DatasetWithSizeType,
    FtsDatasetManager,
    FtsDatasetWithSizeType,
)

log = logging.getLogger(__name__)

FTS_FILTER_ID_FIELD = "filter_id"
FTS_FILTER_RATES = (0.5, 0.75, 0.9, 0.95, 0.99)


def _format_filter_rate(filter_rate: float) -> str:
    return f"{filter_rate * 100:g}%"


def _is_supported_fts_filter_rate(filter_rate: float) -> bool:
    return any(abs(filter_rate - supported_rate) < 1e-9 for supported_rate in FTS_FILTER_RATES)


class CaseType(Enum):
    """
    Example:
        >>> case_cls = CaseType.CapacityDim128.case_cls
        >>> assert c is not None
        >>> CaseType.CapacityDim128.case_name
        "Capacity Test (128 Dim Repeated)"
    """

    CapacityDim128 = 1
    CapacityDim960 = 2

    Performance768D100M = 3
    Performance768D10M = 4
    Performance768D1M = 5

    Performance768D10M1P = 6
    Performance768D1M1P = 7
    Performance768D10M99P = 8
    Performance768D1M99P = 9

    Performance1536D500K = 10
    Performance1536D5M = 11

    Performance1536D500K1P = 12
    Performance1536D5M1P = 13
    Performance1536D500K99P = 14
    Performance1536D5M99P = 15

    Performance1024D1M = 17
    Performance1024D10M = 20

    Performance1536D50K = 50

    Custom = 100
    PerformanceCustomDataset = 101

    StreamingPerformanceCase = 200
    StreamingCustomDataset = 201

    LabelFilterPerformanceCase = 300

    NewIntFilterPerformanceCase = 400
    CloudPayloadSearchCase = 500
    FTSBm25Performance = 503
    CloudInsertCase = 600
    CloudColdLatencyCase = 700
    CloudMultiTenantSearchCase = 800

    def case_cls(self, custom_configs: dict | None = None) -> type["Case"]:
        if custom_configs is None:
            return type2case.get(self)()
        return type2case.get(self)(**custom_configs)

    def case_name(self, custom_configs: dict | None = None) -> str:
        c = self.case_cls(custom_configs)
        if c is not None:
            return c.name
        raise ValueError("Case unsupported")

    def case_description(self, custom_configs: dict | None = None) -> str:
        c = self.case_cls(custom_configs)
        if c is not None:
            return c.description
        raise ValueError("Case unsupported")


class CaseLabel(Enum):
    Load = auto()
    Performance = auto()
    Streaming = auto()
    CloudInsert = auto()
    CloudColdLatency = auto()
    FullTextSearchPerformance = auto()


class Case(BaseModel):
    """Undefined case

    Fields:
        case_id(CaseType): default 9 case type plus one custom cases.
        label(CaseLabel): performance or load.
        dataset(DataSet): dataset for this case runner.
        filter_rate(float | None): one of 99% | 1% | None
        filters(dict | None): filters for search
    """

    case_id: CaseType
    label: CaseLabel
    name: str
    description: str
    dataset: DatasetManager

    load_timeout: float | int | None = None
    optimize_timeout: float | int | None = None

    filter_rate: float | None = None
    payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY

    @property
    def filters(self) -> Filter:
        return non_filter

    def estimated_payload_bytes_per_query(self, k: int | None) -> int:
        if k is None:
            k = config.K_DEFAULT
        return self.payload_profile.estimated_bytes_per_query(k=k, dim=self.dataset.data.dim)

    @property
    def is_multitenant(self) -> bool:
        return False

    @property
    def with_scalar_labels(self) -> bool:
        return self.filters.type == FilterOp.StrEqual or self.payload_profile == PayloadProfile.SCALAR_LABEL

    def check_scalar_labels(self) -> None:
        if self.with_scalar_labels and not self.dataset.data.with_scalar_labels:
            msg = f"Case init failed: no scalar_labels data in current dataset ({self.dataset.data.full_name})"
            raise ValueError(msg)

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.check_scalar_labels()


class CapacityCase(Case):
    label: CaseLabel = CaseLabel.Load
    filter_rate: float | None = None
    load_timeout: float | int = config.CAPACITY_TIMEOUT_IN_SECONDS
    optimize_timeout: float | int | None = None


class PerformanceCase(Case):
    label: CaseLabel = CaseLabel.Performance
    filter_rate: float | None = None
    load_timeout: float | int = config.LOAD_TIMEOUT_DEFAULT
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_DEFAULT
    int_value: float | None = None


class CapacityDim960(CapacityCase):
    case_id: CaseType = CaseType.CapacityDim960
    dataset: DatasetManager = Dataset.GIST.manager(100_000)
    name: str = "Capacity Test (960 Dim Repeated)"
    description: str = """This case tests the vector database's loading capacity by repeatedly inserting large-dimension
     vectors (GIST 100K vectors, <b>960 dimensions</b>) until it is fully loaded. Number of inserted vectors will be
     reported."""


class CapacityDim128(CapacityCase):
    case_id: CaseType = CaseType.CapacityDim128
    dataset: DatasetManager = Dataset.SIFT.manager(500_000)
    name: str = "Capacity Test (128 Dim Repeated)"
    description: str = """This case tests the vector database's loading capacity by repeatedly inserting small-dimension
     vectors (SIFT 100K vectors, <b>128 dimensions</b>) until it is fully loaded. Number of inserted vectors will be
     reported."""


class Performance768D10M(PerformanceCase):
    case_id: CaseType = CaseType.Performance768D10M
    dataset: DatasetManager = Dataset.COHERE.manager(10_000_000)
    name: str = "Search Performance Test (10M Dataset, 768 Dim)"
    description: str = """This case tests the search performance of a vector database with a large dataset
    (<b>Cohere 10M vectors</b>, 768 dimensions) at varying parallel levels.
    Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_768D_10M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_768D_10M


class IntFilterPerformanceCase(PerformanceCase):
    @property
    def filters(self) -> Filter:
        int_field = self.dataset.data.train_id_field
        int_value = int(self.dataset.data.size * self.filter_rate)
        return IntFilter(filter_rate=self.filter_rate, int_field=int_field, int_value=int_value)


class Performance768D1M(PerformanceCase):
    case_id: CaseType = CaseType.Performance768D1M
    dataset: DatasetManager = Dataset.COHERE.manager(1_000_000)
    name: str = "Search Performance Test (1M Dataset, 768 Dim)"
    description: str = """This case tests the search performance of a vector database with a medium dataset
    (<b>Cohere 1M vectors</b>, 768 dimensions) at varying parallel levels.
    Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_768D_1M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_768D_1M


class Performance768D10M1P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance768D10M1P
    filter_rate: float | int | None = 0.01
    dataset: DatasetManager = Dataset.COHERE.manager(10_000_000)
    name: str = "Filtering Search Performance Test (10M Dataset, 768 Dim, Filter 1%)"
    description: str = """This case tests the search performance of a vector database with a large dataset
    (<b>Cohere 10M vectors</b>, 768 dimensions) under a low filtering rate (<b>1% vectors</b>), at varying parallel
    levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_768D_10M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_768D_10M


class Performance768D1M1P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance768D1M1P
    filter_rate: float | int | None = 0.01
    dataset: DatasetManager = Dataset.COHERE.manager(1_000_000)
    name: str = "Filtering Search Performance Test (1M Dataset, 768 Dim, Filter 1%)"
    description: str = """This case tests the search performance of a vector database with a medium dataset
    (<b>Cohere 1M vectors</b>, 768 dimensions) under a low filtering rate (<b>1% vectors</b>),
    at varying parallel levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_768D_1M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_768D_1M


class Performance768D10M99P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance768D10M99P
    filter_rate: float | int | None = 0.99
    dataset: DatasetManager = Dataset.COHERE.manager(10_000_000)
    name: str = "Filtering Search Performance Test (10M Dataset, 768 Dim, Filter 99%)"
    description: str = """This case tests the search performance of a vector database with a large dataset
    (<b>Cohere 10M vectors</b>, 768 dimensions) under a high filtering rate (<b>99% vectors</b>),
    at varying parallel levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_768D_10M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_768D_10M


class Performance768D1M99P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance768D1M99P
    filter_rate: float | int | None = 0.99
    dataset: DatasetManager = Dataset.COHERE.manager(1_000_000)
    name: str = "Filtering Search Performance Test (1M Dataset, 768 Dim, Filter 99%)"
    description: str = """This case tests the search performance of a vector database with a medium dataset
    (<b>Cohere 1M vectors</b>, 768 dimensions) under a high filtering rate (<b>99% vectors</b>),
    at varying parallel levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_768D_1M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_768D_1M


class Performance768D100M(PerformanceCase):
    case_id: CaseType = CaseType.Performance768D100M
    filter_rate: float | int | None = None
    dataset: DatasetManager = Dataset.LAION.manager(100_000_000)
    name: str = "Search Performance Test (100M Dataset, 768 Dim)"
    description: str = """This case tests the search performance of a vector database with a large 100M dataset
    (<b>LAION 100M vectors</b>, 768 dimensions), at varying parallel levels. Results will show index building time,
    recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_768D_100M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_768D_100M


class Performance1536D500K(PerformanceCase):
    case_id: CaseType = CaseType.Performance1536D500K
    filter_rate: float | int | None = None
    dataset: DatasetManager = Dataset.OPENAI.manager(500_000)
    name: str = "Search Performance Test (500K Dataset, 1536 Dim)"
    description: str = """This case tests the search performance of a vector database with a medium 500K dataset
    (<b>OpenAI 500K vectors</b>, 1536 dimensions), at varying parallel levels. Results will show index building time,
    recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1536D_500K
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_500K


class Performance1536D5M(PerformanceCase):
    case_id: CaseType = CaseType.Performance1536D5M
    filter_rate: float | int | None = None
    dataset: DatasetManager = Dataset.OPENAI.manager(5_000_000)
    name: str = "Search Performance Test (5M Dataset, 1536 Dim)"
    description: str = """This case tests the search performance of a vector database with a medium 5M dataset
    (<b>OpenAI 5M vectors</b>, 1536 dimensions), at varying parallel levels. Results will show index building time,
    recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1536D_5M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_5M


class Performance1536D500K1P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance1536D500K1P
    filter_rate: float | int | None = 0.01
    dataset: DatasetManager = Dataset.OPENAI.manager(500_000)
    name: str = "Filtering Search Performance Test (500K Dataset, 1536 Dim, Filter 1%)"
    description: str = """This case tests the search performance of a vector database with a large dataset
    (<b>OpenAI 500K vectors</b>, 1536 dimensions) under a low filtering rate (<b>1% vectors</b>),
    at varying parallel levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1536D_500K
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_500K


class Performance1536D5M1P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance1536D5M1P
    filter_rate: float | int | None = 0.01
    dataset: DatasetManager = Dataset.OPENAI.manager(5_000_000)
    name: str = "Filtering Search Performance Test (5M Dataset, 1536 Dim, Filter 1%)"
    description: str = """This case tests the search performance of a vector database with a large dataset
    (<b>OpenAI 5M vectors</b>, 1536 dimensions) under a low filtering rate (<b>1% vectors</b>),
    at varying parallel levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1536D_5M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_5M


class Performance1536D500K99P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance1536D500K99P
    filter_rate: float | int | None = 0.99
    dataset: DatasetManager = Dataset.OPENAI.manager(500_000)
    name: str = "Filtering Search Performance Test (500K Dataset, 1536 Dim, Filter 99%)"
    description: str = """This case tests the search performance of a vector database with a medium dataset
    (<b>OpenAI 500K vectors</b>, 1536 dimensions) under a high filtering rate (<b>99% vectors</b>),
    at varying parallel levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1536D_500K
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_500K


class Performance1536D5M99P(IntFilterPerformanceCase):
    case_id: CaseType = CaseType.Performance1536D5M99P
    filter_rate: float | int | None = 0.99
    dataset: DatasetManager = Dataset.OPENAI.manager(5_000_000)
    name: str = "Filtering Search Performance Test (5M Dataset, 1536 Dim, Filter 99%)"
    description: str = """This case tests the search performance of a vector database with a medium dataset
    (<b>OpenAI 5M vectors</b>, 1536 dimensions) under a high filtering rate (<b>99% vectors</b>),
    at varying parallel levels. Results will show index building time, recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1536D_5M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1536D_5M


class Performance1024D1M(PerformanceCase):
    case_id: CaseType = CaseType.Performance1024D1M
    filter_rate: float | int | None = None
    dataset: DatasetManager = Dataset.BIOASQ.manager(1_000_000)
    name: str = "Search Performance Test (1M Dataset, 1024 Dim)"
    description: str = """This case tests the search performance of a vector database with a medium 1M dataset
    (<b>Bioasq 1M vectors</b>, 1024 dimensions), at varying parallel levels. Results will show index building time,
    recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1024D_1M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1024D_1M


class Performance1024D10M(PerformanceCase):
    case_id: CaseType = CaseType.Performance1024D10M
    filter_rate: float | int | None = None
    dataset: DatasetManager = Dataset.BIOASQ.manager(10_000_000)
    name: str = "Search Performance Test (10M Dataset, 1024 Dim)"
    description: str = """This case tests the search performance of a vector database with a large 10M dataset
    (<b>Bioasq 10M vectors</b>, 1024 dimensions), at varying parallel levels. Results will show index building time,
    recall, and maximum QPS."""
    load_timeout: float | int = config.LOAD_TIMEOUT_1024D_10M
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_1024D_10M


class Performance1536D50K(PerformanceCase):
    case_id: CaseType = CaseType.Performance1536D50K
    filter_rate: float | int | None = None
    dataset: DatasetManager = Dataset.OPENAI.manager(50_000)
    name: str = "Search Performance Test (50K Dataset, 1536 Dim)"
    description: str = """This case tests the search performance of a vector database with a medium 50K dataset
    (<b>OpenAI 50K vectors</b>, 1536 dimensions), at varying parallel levels. Results will show index building time,
    recall, and maximum QPS."""
    load_timeout: float | int = 3600
    optimize_timeout: float | int | None = config.OPTIMIZE_TIMEOUT_DEFAULT


def metric_type_map(s: str) -> MetricType:
    if s.lower() == "cosine":
        return MetricType.COSINE
    if s.lower() == "l2" or s.lower() == "euclidean":
        return MetricType.L2
    if s.lower() == "ip":
        return MetricType.IP
    err_msg = f"Not support metric_type: {s}"
    log.error(err_msg)
    raise RuntimeError(err_msg)


class PerformanceCustomDataset(PerformanceCase):
    case_id: CaseType = CaseType.PerformanceCustomDataset
    name: str = "Performance With Custom Dataset"
    description: str = ""
    gt_file: str
    dataset: DatasetManager
    label_percentage: float | None = None
    use_filter: bool

    def __init__(
        self,
        name: str,
        description: str,
        load_timeout: float,
        optimize_timeout: float,
        dataset_config: dict,
        label_percentage: float | None = None,
        use_filter: bool = False,
        **kwargs,
    ):
        dataset_config = CustomDatasetConfig(**dataset_config)
        dataset = CustomDataset(
            name=dataset_config.name,
            size=dataset_config.size,
            dim=dataset_config.dim,
            metric_type=metric_type_map(dataset_config.metric_type),
            use_shuffled=dataset_config.use_shuffled,
            with_gt=dataset_config.with_gt,
            dir=dataset_config.dir,
            file_num=dataset_config.file_count,
            train_file=dataset_config.train_name,
            test_file=f"{dataset_config.test_name}.parquet",
            train_id_field=dataset_config.train_id_name,
            train_vector_field=dataset_config.train_col_name,
            test_vector_field=dataset_config.test_col_name,
            gt_neighbors_field=dataset_config.gt_col_name,
            scalar_labels_file=f"{dataset_config.scalar_labels_name}.parquet",
        )
        filter_rate = (1.0 - label_percentage) if (use_filter and label_percentage is not None) else None
        super().__init__(
            name=name,
            description=description,
            load_timeout=load_timeout,
            optimize_timeout=optimize_timeout,
            gt_file=f"{dataset_config.gt_name}.parquet",
            dataset=DatasetManager(data=dataset),
            use_filter=use_filter,
            label_percentage=label_percentage,
            filter_rate=filter_rate,
        )

    @property
    def filters(self) -> Filter:
        if self.use_filter is True:
            return LabelFilter(label_percentage=self.label_percentage)
        return NonFilter(gt_file_name=self.gt_file)


class StreamingPerformanceCase(Case):
    case_id: CaseType = CaseType.StreamingPerformanceCase
    label: CaseLabel = CaseLabel.Streaming
    dataset_with_size_type: DatasetWithSizeType
    insert_rate: int
    search_stages: list[float]
    concurrencies: list[int]
    optimize_after_write: bool = True
    read_dur_after_write: int = 30

    def __init__(
        self,
        dataset_with_size_type: DatasetWithSizeType | str = DatasetWithSizeType.CohereSmall.value,
        insert_rate: int = 500,
        search_stages: list[float] | str = (0.5, 0.8),
        concurrencies: list[int] | str = (5, 10),
        **kwargs,
    ):
        num_per_batch = config.NUM_PER_BATCH
        if insert_rate % config.NUM_PER_BATCH != 0:
            _insert_rate = max(
                num_per_batch,
                insert_rate // num_per_batch * num_per_batch,
            )
            log.warning(
                f"[streaming_case init] insert_rate(={insert_rate}) should be "
                f"divisible by NUM_PER_BATCH={num_per_batch}), reset to {_insert_rate}",
            )
            insert_rate = _insert_rate
        if not isinstance(dataset_with_size_type, DatasetWithSizeType):
            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
        dataset = dataset_with_size_type.get_manager()
        name = f"Streaming-Perf - {dataset_with_size_type.value}, {insert_rate} rows/s"
        description = (
            "This case tests the search performance of vector database while maintaining "
            f"a fixed insertion speed. (dataset: {dataset_with_size_type.value})"
        )

        if isinstance(search_stages, str):
            search_stages = json.loads(search_stages)
        if isinstance(concurrencies, str):
            concurrencies = json.loads(concurrencies)

        super().__init__(
            name=name,
            description=description,
            dataset=dataset,
            dataset_with_size_type=dataset_with_size_type,
            insert_rate=insert_rate,
            search_stages=search_stages,
            concurrencies=concurrencies,
            **kwargs,
        )


class StreamingCustomDataset(Case):
    case_id: CaseType = CaseType.StreamingCustomDataset
    label: CaseLabel = CaseLabel.Streaming
    name: str = "Streaming Performance With Custom Dataset"
    description: str = ""
    dataset: DatasetManager
    insert_rate: int
    search_stages: list[float]
    concurrencies: list[int]
    optimize_after_write: bool = True
    read_dur_after_write: int = 30

    def __init__(
        self,
        description: str,
        dataset_config: dict,
        insert_rate: int = 500,
        search_stages: list[float] | str = (0.5, 0.8),
        concurrencies: list[int] | str = (5, 10),
        optimize_after_write: bool = True,
        read_dur_after_write: int = 30,
        **kwargs,
    ):
        num_per_batch = config.NUM_PER_BATCH
        if insert_rate % config.NUM_PER_BATCH != 0:
            _insert_rate = max(
                num_per_batch,
                insert_rate // num_per_batch * num_per_batch,
            )
            log.warning(
                f"[streaming_case init] insert_rate(={insert_rate}) should be "
                f"divisible by NUM_PER_BATCH={num_per_batch}), reset to {_insert_rate}",
            )
            insert_rate = _insert_rate

        dataset_config = CustomDatasetConfig(**dataset_config)
        dataset = CustomDataset(
            name=dataset_config.name,
            size=dataset_config.size,
            dim=dataset_config.dim,
            metric_type=metric_type_map(dataset_config.metric_type),
            use_shuffled=dataset_config.use_shuffled,
            with_gt=dataset_config.with_gt,
            dir=dataset_config.dir,
            file_num=dataset_config.file_count,
            train_file=dataset_config.train_name,
            test_file=f"{dataset_config.test_name}.parquet",
            train_id_field=dataset_config.train_id_name,
            train_vector_field=dataset_config.train_col_name,
            test_vector_field=dataset_config.test_col_name,
            gt_neighbors_field=dataset_config.gt_col_name,
            scalar_labels_file=f"{dataset_config.scalar_labels_name}.parquet",
        )
        name = f"Streaming-Perf - Custom - {dataset_config.name}, {insert_rate} rows/s"
        description = (
            description
            if description
            else f"This case tests the search performance of vector database while maintaining "
            f"a fixed insertion speed. (dataset: Custom - {dataset_config.name})"
        )

        if isinstance(search_stages, str):
            search_stages = json.loads(search_stages)
        if isinstance(concurrencies, str):
            concurrencies = json.loads(concurrencies)

        super().__init__(
            name=name,
            description=description,
            dataset=DatasetManager(data=dataset),
            insert_rate=insert_rate,
            search_stages=search_stages,
            concurrencies=concurrencies,
            optimize_after_write=optimize_after_write,
            read_dur_after_write=read_dur_after_write,
            **kwargs,
        )


class NewIntFilterPerformanceCase(PerformanceCase):
    case_id: CaseType = CaseType.NewIntFilterPerformanceCase
    dataset_with_size_type: DatasetWithSizeType
    filter_rate: float

    def __init__(
        self,
        dataset_with_size_type: DatasetWithSizeType | str,
        filter_rate: float,
        int_value: float | None = 0,
        **kwargs,
    ):
        if not isinstance(dataset_with_size_type, DatasetWithSizeType):
            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
        name = f"Int-Filter-{filter_rate*100:.1f}% - {dataset_with_size_type.value}"
        description = f"Int-Filter-{filter_rate*100:.1f}% Performance Test ({dataset_with_size_type.value})"
        dataset = dataset_with_size_type.get_manager()
        load_timeout = dataset_with_size_type.get_load_timeout()
        optimize_timeout = dataset_with_size_type.get_optimize_timeout()
        filters = IntFilter(filter_rate=filter_rate, int_value=int_value)
        filter_rate = filters.filter_rate
        super().__init__(
            name=name,
            description=description,
            dataset=dataset,
            load_timeout=load_timeout,
            optimize_timeout=optimize_timeout,
            filter_rate=filter_rate,
            int_value=int_value,
            dataset_with_size_type=dataset_with_size_type,
            **kwargs,
        )

    @property
    def filters(self) -> Filter:
        int_field = self.dataset.data.train_id_field
        int_value = int(self.dataset.data.size * self.filter_rate)
        return NewIntFilter(filter_rate=self.filter_rate, int_field=int_field, int_value=int_value)


class CloudPayloadSearchCase(PerformanceCase):
    case_id: CaseType = CaseType.CloudPayloadSearchCase
    dataset_with_size_type: DatasetWithSizeType | None = None
    payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY
    filter_rate: float | None = None
    label_percentage: float | None = None

    def __init__(
        self,
        dataset_with_size_type: DatasetWithSizeType | str | None = None,
        payload_profile: PayloadProfile | str = PayloadProfile.IDS_ONLY,
        filter_rate: float | None = None,
        label_percentage: float | None = None,
        **kwargs,
    ):
        if filter_rate is not None and label_percentage is not None:
            msg = "CloudPayloadSearchCase supports only one filter type per run"
            raise ValueError(msg)
        if dataset_with_size_type is not None and not isinstance(dataset_with_size_type, DatasetWithSizeType):
            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
        if not isinstance(payload_profile, PayloadProfile):
            payload_profile = PayloadProfile(payload_profile)

        if dataset_with_size_type is None:
            dataset = Dataset.LAION.manager(100_000_000)
            load_timeout = config.LOAD_TIMEOUT_768D_100M
            optimize_timeout = config.OPTIMIZE_TIMEOUT_768D_100M
            dataset_name = "LAION 100M (768dim)"
        else:
            dataset = dataset_with_size_type.get_manager()
            load_timeout = dataset_with_size_type.get_load_timeout()
            optimize_timeout = dataset_with_size_type.get_optimize_timeout()
            dataset_name = dataset_with_size_type.value

        name = f"Cloud Payload Search - {payload_profile.value} - {dataset_name}"
        description = (
            "Cloud leaderboard search envelope case with explicit response payload profile. "
            f"Payload profile: {payload_profile.value}; dataset: {dataset_name}."
        )
        super().__init__(
            name=name,
            description=description,
            dataset=dataset,
            load_timeout=load_timeout,
            optimize_timeout=optimize_timeout,
            dataset_with_size_type=dataset_with_size_type,
            payload_profile=payload_profile,
            filter_rate=filter_rate,
            label_percentage=label_percentage,
            **kwargs,
        )

    @property
    def filters(self) -> Filter:
        if self.label_percentage is not None:
            return LabelFilter(label_percentage=self.label_percentage)
        if self.filter_rate is None:
            return non_filter
        int_field = self.dataset.data.train_id_field
        int_value = int(self.dataset.data.size * self.filter_rate)
        return NewIntFilter(filter_rate=self.filter_rate, int_field=int_field, int_value=int_value)


class CloudColdLatencyCase(Case):
    case_id: CaseType = CaseType.CloudColdLatencyCase
    label: CaseLabel = CaseLabel.CloudColdLatency
    dataset_with_size_type: DatasetWithSizeType | None = None
    payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY
    filter_rate: float | None = None
    label_percentage: float | None = None
    query_count: int = 1000

    def __init__(
        self,
        dataset_with_size_type: DatasetWithSizeType | str | None = None,
        payload_profile: PayloadProfile | str = PayloadProfile.IDS_ONLY,
        filter_rate: float | None = None,
        label_percentage: float | None = None,
        query_count: int = 1000,
        **kwargs,
    ):
        if filter_rate is not None and label_percentage is not None:
            msg = "CloudColdLatencyCase supports only one filter type per run"
            raise ValueError(msg)
        if query_count <= 0:
            msg = "query_count must be positive"
            raise ValueError(msg)
        if dataset_with_size_type is not None and not isinstance(dataset_with_size_type, DatasetWithSizeType):
            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
        if not isinstance(payload_profile, PayloadProfile):
            payload_profile = PayloadProfile(payload_profile)

        if dataset_with_size_type is None:
            dataset = Dataset.LAION.manager(100_000_000)
            load_timeout = config.LOAD_TIMEOUT_768D_100M
            optimize_timeout = config.OPTIMIZE_TIMEOUT_768D_100M
            dataset_name = "LAION 100M (768dim)"
        else:
            dataset = dataset_with_size_type.get_manager()
            load_timeout = dataset_with_size_type.get_load_timeout()
            optimize_timeout = dataset_with_size_type.get_optimize_timeout()
            dataset_name = dataset_with_size_type.value

        name = f"Cloud Cold Latency - {payload_profile.value} - {dataset_name}"
        description = (
            "Cloud leaderboard cold/warm serial latency case with explicit response payload profile. "
            f"Payload profile: {payload_profile.value}; dataset: {dataset_name}; query count: {query_count}."
        )
        super().__init__(
            name=name,
            description=description,
            dataset=dataset,
            load_timeout=load_timeout,
            optimize_timeout=optimize_timeout,
            dataset_with_size_type=dataset_with_size_type,
            payload_profile=payload_profile,
            filter_rate=filter_rate,
            label_percentage=label_percentage,
            query_count=query_count,
            **kwargs,
        )

    @property
    def filters(self) -> Filter:
        if self.label_percentage is not None:
            return LabelFilter(label_percentage=self.label_percentage)
        if self.filter_rate is None:
            return non_filter
        int_field = self.dataset.data.train_id_field
        int_value = int(self.dataset.data.size * self.filter_rate)
        return NewIntFilter(filter_rate=self.filter_rate, int_field=int_field, int_value=int_value)


class CloudInsertCase(Case):
    case_id: CaseType = CaseType.CloudInsertCase
    label: CaseLabel = CaseLabel.CloudInsert
    batch_size: int
    duration: float | None = None
    readiness_timeout: float | None = config.CLOUD_INSERT_READINESS_TIMEOUT
    readiness_poll_interval: float = config.CLOUD_INSERT_READINESS_POLL_INTERVAL
    dataset_with_size_type: DatasetWithSizeType | None = None

    def __init__(
        self,
        batch_size: int,
        duration: float | None = None,
        readiness_timeout: float | None = config.CLOUD_INSERT_READINESS_TIMEOUT,
        readiness_poll_interval: float = config.CLOUD_INSERT_READINESS_POLL_INTERVAL,
        dataset_with_size_type: DatasetWithSizeType | str | None = None,
        **kwargs,
    ):
        if dataset_with_size_type is not None and not isinstance(dataset_with_size_type, DatasetWithSizeType):
            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
        dataset = (
            Dataset.LAION.manager(100_000_000)
            if dataset_with_size_type is None
            else dataset_with_size_type.get_manager()
        )
        super().__init__(
            name=f"Cloud Insert - batch {batch_size}",
            description="Cloud leaderboard insert-only case with readiness polling.",
            dataset=dataset,
            batch_size=batch_size,
            duration=duration,
            readiness_timeout=readiness_timeout,
            readiness_poll_interval=readiness_poll_interval,
            dataset_with_size_type=dataset_with_size_type,
            **kwargs,
        )


class CloudMultiTenantSearchCase(PerformanceCase):
    case_id: CaseType = CaseType.CloudMultiTenantSearchCase
    dataset_with_size_type: DatasetWithSizeType = DatasetWithSizeType.CohereLarge
    tenant_count: int = 1000
    tenant_prefix: str = "tenant_"
    tenant_id_width: int = 4
    tenant_distribution: str = "uniform_by_id_mod"
    measure_recall: bool = False
    payload_profile: PayloadProfile = PayloadProfile.IDS_ONLY
    filter_rate: float | None = None
    label_percentage: float | None = None

    def __init__(
        self,
        dataset_with_size_type: DatasetWithSizeType | str = DatasetWithSizeType.CohereLarge,
        tenant_count: int = 1000,
        tenant_prefix: str = "tenant_",
        tenant_id_width: int = 4,
        payload_profile: PayloadProfile | str = PayloadProfile.IDS_ONLY,
        filter_rate: float | None = None,
        label_percentage: float | None = None,
        **kwargs,
    ):
        if filter_rate is not None and label_percentage is not None:
            msg = "CloudMultiTenantSearchCase supports only one filter type per run"
            raise ValueError(msg)
        if not isinstance(dataset_with_size_type, DatasetWithSizeType):
            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
        if not isinstance(payload_profile, PayloadProfile):
            payload_profile = PayloadProfile(payload_profile)
        if tenant_count <= 0:
            msg = "tenant_count must be greater than 0"
            raise ValueError(msg)
        if tenant_id_width <= 0:
            msg = "tenant_id_width must be greater than 0"
            raise ValueError(msg)

        dataset = dataset_with_size_type.get_manager()
        super().__init__(
            name=f"Cloud Multi-Tenant Search - {dataset_with_size_type.value}, {tenant_count} tenants",
            description=(
                "Multi-tenant QPS/latency benchmark with deterministic tenant routing "
                f"({dataset_with_size_type.value}, {tenant_count} tenants)."
            ),
            dataset=dataset,
            load_timeout=dataset_with_size_type.get_load_timeout(),
            optimize_timeout=dataset_with_size_type.get_optimize_timeout(),
            dataset_with_size_type=dataset_with_size_type,
            tenant_count=tenant_count,
            tenant_prefix=tenant_prefix,
            tenant_id_width=tenant_id_width,
            payload_profile=payload_profile,
            filter_rate=filter_rate,
            label_percentage=label_percentage,
            **kwargs,
        )

    @property
    def is_multitenant(self) -> bool:
        return True

    def tenant_for_id(self, row_id: int) -> str:
        tenant_id = int(row_id) % self.tenant_count
        return f"{self.tenant_prefix}{tenant_id:0{self.tenant_id_width}d}"

    def tenant_labels_for_ids(self, row_ids: list[int]) -> list[str]:
        return [self.tenant_for_id(row_id) for row_id in row_ids]

    def tenant_labels(self) -> list[str]:
        return [f"{self.tenant_prefix}{tenant_id:0{self.tenant_id_width}d}" for tenant_id in range(self.tenant_count)]

    @property
    def filters(self) -> Filter:
        if self.label_percentage is not None:
            return LabelFilter(label_percentage=self.label_percentage)
        if self.filter_rate is None:
            return non_filter
        int_field = self.dataset.data.train_id_field
        int_value = int(self.dataset.data.size * self.filter_rate)
        return NewIntFilter(filter_rate=self.filter_rate, int_field=int_field, int_value=int_value)


class LabelFilterPerformanceCase(PerformanceCase):
    case_id: CaseType = CaseType.LabelFilterPerformanceCase
    dataset_with_size_type: DatasetWithSizeType
    label_percentage: float

    def __init__(
        self,
        dataset_with_size_type: DatasetWithSizeType | str,
        label_percentage: float,
        **kwargs,
    ):
        if not isinstance(dataset_with_size_type, DatasetWithSizeType):
            dataset_with_size_type = DatasetWithSizeType(dataset_with_size_type)
        name = f"Label-Filter-{label_percentage*100:.1f}% - {dataset_with_size_type.value}"
        description = f"Label-Filter-{label_percentage*100:.1f}% Performance Test ({dataset_with_size_type.value})"
        dataset = dataset_with_size_type.get_manager()
        load_timeout = dataset_with_size_type.get_load_timeout()
        optimize_timeout = dataset_with_size_type.get_optimize_timeout()
        filters = LabelFilter(label_percentage=label_percentage)
        filter_rate = filters.filter_rate
        super().__init__(
            name=name,
            description=description,
            dataset=dataset,
            load_timeout=load_timeout,
            optimize_timeout=optimize_timeout,
            filter_rate=filter_rate,
            dataset_with_size_type=dataset_with_size_type,
            label_percentage=label_percentage,
            **kwargs,
        )

    @property
    def filters(self) -> Filter:
        return LabelFilter(label_percentage=self.label_percentage)


class FtsPerformanceCase(Case):
    """Base class for full-text search BM25 performance cases."""

    label: CaseLabel = CaseLabel.FullTextSearchPerformance
    dataset: FtsDatasetManager

    filter_rate: float | None = None

    @property
    def filters(self) -> Filter:
        if self.filter_rate is None:
            return non_filter
        int_value = int(self.dataset.data.size * self.filter_rate)
        return NewIntFilter(filter_rate=self.filter_rate, int_field=FTS_FILTER_ID_FIELD, int_value=int_value)

    def estimated_payload_bytes_per_query(self, k: int | None) -> int:
        if k is None:
            k = config.K_DEFAULT
        return self.payload_profile.estimated_bytes_per_query(k=k, dim=0)


class FTSBm25Performance(FtsPerformanceCase):
    case_id: CaseType = CaseType.FTSBm25Performance
    dataset_with_size_type: FtsDatasetWithSizeType = FtsDatasetWithSizeType.MSMarcoSmall

    def __init__(
        self,
        dataset_with_size_type: FtsDatasetWithSizeType | str = FtsDatasetWithSizeType.MSMarcoSmall,
        filter_rate: float | None = None,
        **kwargs,
    ):
        if not isinstance(dataset_with_size_type, FtsDatasetWithSizeType):
            dataset_with_size_type = FtsDatasetWithSizeType(dataset_with_size_type)
        if filter_rate is not None:
            if not dataset_with_size_type.is_advanced:
                msg = "FTS filter cases are only supported for MS MARCO Large and HotpotQA Large"
                raise ValueError(msg)
            if not _is_supported_fts_filter_rate(filter_rate):
                supported_rates = ", ".join(_format_filter_rate(rate) for rate in FTS_FILTER_RATES)
                msg = f"FTS filter_rate must be one of: {supported_rates}"
                raise ValueError(msg)
        dataset = dataset_with_size_type.get_manager()
        filter_suffix = f", Filter {_format_filter_rate(filter_rate)}" if filter_rate is not None else ""
        name = f"FTS BM25 Performance - {dataset_with_size_type.value}{filter_suffix}"
        description = (
            f"This case tests native BM25 full-text search performance on {dataset_with_size_type.value}. "
            "It measures index building time, recall, serial latency, and search QPS."
        )
        if filter_rate is not None:
            description += (
                f" The FTS filter case searches only documents with {FTS_FILTER_ID_FIELD} >= "
                f"int(dataset_size * {filter_rate})."
            )
        super().__init__(
            name=name,
            description=description,
            dataset=dataset,
            dataset_with_size_type=dataset_with_size_type,
            filter_rate=filter_rate,
            load_timeout=dataset_with_size_type.get_load_timeout(),
            optimize_timeout=dataset_with_size_type.get_optimize_timeout(),
            **kwargs,
        )


type2case = {
    CaseType.CapacityDim960: CapacityDim960,
    CaseType.CapacityDim128: CapacityDim128,
    CaseType.Performance768D100M: Performance768D100M,
    CaseType.Performance768D10M: Performance768D10M,
    CaseType.Performance768D1M: Performance768D1M,
    CaseType.Performance768D10M1P: Performance768D10M1P,
    CaseType.Performance768D1M1P: Performance768D1M1P,
    CaseType.Performance768D10M99P: Performance768D10M99P,
    CaseType.Performance768D1M99P: Performance768D1M99P,
    CaseType.Performance1536D500K: Performance1536D500K,
    CaseType.Performance1536D5M: Performance1536D5M,
    CaseType.Performance1536D500K1P: Performance1536D500K1P,
    CaseType.Performance1536D5M1P: Performance1536D5M1P,
    CaseType.Performance1536D500K99P: Performance1536D500K99P,
    CaseType.Performance1536D5M99P: Performance1536D5M99P,
    CaseType.Performance1024D1M: Performance1024D1M,
    CaseType.Performance1024D10M: Performance1024D10M,
    CaseType.Performance1536D50K: Performance1536D50K,
    CaseType.PerformanceCustomDataset: PerformanceCustomDataset,
    CaseType.StreamingPerformanceCase: StreamingPerformanceCase,
    CaseType.StreamingCustomDataset: StreamingCustomDataset,
    CaseType.NewIntFilterPerformanceCase: NewIntFilterPerformanceCase,
    CaseType.LabelFilterPerformanceCase: LabelFilterPerformanceCase,
    CaseType.CloudPayloadSearchCase: CloudPayloadSearchCase,
    CaseType.FTSBm25Performance: FTSBm25Performance,
    CaseType.CloudInsertCase: CloudInsertCase,
    CaseType.CloudColdLatencyCase: CloudColdLatencyCase,
    CaseType.CloudMultiTenantSearchCase: CloudMultiTenantSearchCase,
}