[WIP] Add Binary Case Support

tests/test_dataset.py

@@ -1,12 +1,16 @@
-from vectordb_bench.backend.dataset import Dataset
+import struct
+
+from vectordb_bench import config
+from vectordb_bench.backend.clients import MetricType
+from vectordb_bench.backend.dataset import BinaryDatasetManager, Dataset, DatasetWithSizeType, SIFTBinary
 import logging
 import pytest
 from pydantic import ValidationError
 from vectordb_bench.backend.data_source import DatasetSource
 
-
 log = logging.getLogger("vectordb_bench")
 
+
 class TestDataSet:
     def test_iter_dataset(self):
         for ds in Dataset:
@@ -29,6 +33,7 @@ def test_iter_cohere(self):
         cohere_10m.prepare()
 
         import time
+
         before = time.time()
         for i in cohere_10m:
             log.debug(i.head(1))
@@ -40,9 +45,11 @@ def test_iter_cohere(self):
     def test_iter_laion(self):
         laion_100m = Dataset.LAION.manager(100_000_000)
         from vectordb_bench.backend.data_source import DatasetSource
+
         laion_100m.prepare(source=DatasetSource.AliyunOSS)
 
         import time
+
         before = time.time()
         for i in laion_100m:
             log.debug(i.head(1))
@@ -75,3 +82,37 @@ def test_download_small(self):
             local_ds_root=openai_50k.data_dir,
         )
 
+    def test_sift_binary_dataset_type(self):
+        dataset = DatasetWithSizeType.SIFTBinary1M.get_manager()
+
+        assert isinstance(dataset, BinaryDatasetManager)
+        assert dataset.data.name == "SIFTBinary"
+        assert dataset.data.dim == 128
+        assert dataset.data.metric_type == MetricType.HAMMING
+
+    def test_binary_dataset_manager_reads_vec_tool_bin_layout(self, tmp_path, monkeypatch):
+        monkeypatch.setattr(config, "DATASET_LOCAL_DIR", tmp_path)
+        vectors = [
+            bytes(range(16)),
+            bytes([255] * 16),
+            bytes([170] * 16),
+        ]
+        truth = [[2, 1], [0, 2]]
+
+        data = SIFTBinary(size=1_000_000)
+        manager = BinaryDatasetManager(data=data)
+        manager.data_dir.mkdir(parents=True)
+        manager.data_dir.joinpath("base.bin").write_bytes(struct.pack("<II", len(vectors), 128) + b"".join(vectors))
+        manager.data_dir.joinpath("query.bin").write_bytes(struct.pack("<II", 1, 128) + vectors[0])
+        manager.data_dir.joinpath("truth.ibin").write_bytes(
+            struct.pack("<II", len(truth), len(truth[0])) + b"".join(struct.pack("<i", v) for row in truth for v in row)
+        )
+
+        assert manager._read_binary_vectors("query.bin") == [vectors[0].hex()]
+        assert manager._read_truth_ids("truth.ibin") == truth
+
+        batches = list(manager.iter_batches(batch_size=2))
+        assert batches[0]["id"].tolist() == [0, 1]
+        assert batches[0]["emb"].tolist() == [vectors[0].hex(), vectors[1].hex()]
+        assert batches[1]["id"].tolist() == [2]
+        assert batches[1]["emb"].tolist() == [vectors[2].hex()]

tests/test_elasticsearch_binary_config.py

@@ -0,0 +1,11 @@
+from vectordb_bench.backend.clients import MetricType
+from vectordb_bench.backend.clients.elastic_cloud.config import ElasticCloudIndexConfig, ESElementType
+
+
+def test_elasticsearch_hamming_metric_uses_bit_dense_vector():
+    config = ElasticCloudIndexConfig(metric_type=MetricType.HAMMING, element_type=ESElementType.float)
+
+    index_param = config.index_param()
+
+    assert index_param["element_type"] == "bit"
+    assert index_param["similarity"] == "l2_norm"

tests/test_oss_opensearch_fts.py

@@ -0,0 +1,180 @@
+import sys
+import types
+
+import pytest
+
+
+class _FakeOpenSearch:
+    pass
+
+
+sys.modules.setdefault("opensearchpy", types.SimpleNamespace(OpenSearch=_FakeOpenSearch))
+
+from vectordb_bench.backend.clients import DB
+from vectordb_bench.backend.clients.api import IndexType
+from vectordb_bench.backend.clients.oss_opensearch.config import OSSOpenSearchFtsConfig
+from vectordb_bench.backend.clients.oss_opensearch.oss_opensearch import OSSOpenSearch
+from vectordb_bench.backend.payload import PayloadProfile
+
+
+def make_fts_db():
+    db = OSSOpenSearch.__new__(OSSOpenSearch)
+    db.index_name = "idx"
+    db.id_col_name = "doc_id"
+    db.text_col_name = "text"
+    db._is_fts = True
+    db.client = object()
+    return db
+
+
+def test_oss_opensearch_fts_config_defaults():
+    config = OSSOpenSearchFtsConfig()
+
+    assert config.index_param()["properties"]["doc_id"] == {"type": "keyword"}
+    assert config.index_param()["properties"]["text"] == {"type": "text"}
+    assert config.search_param() == {}
+
+
+def test_oss_opensearch_fts_config_supports_bm25_similarity():
+    config = OSSOpenSearchFtsConfig(bm25_k1=1.2, bm25_b=0.75)
+
+    assert config.index_param()["properties"]["text"]["similarity"] == "vdbbench_bm25"
+    assert config.similarity_settings() == {
+        "similarity": {
+            "vdbbench_bm25": {
+                "type": "BM25",
+                "k1": 1.2,
+                "b": 0.75,
+            }
+        }
+    }
+
+
+def test_oss_opensearch_declares_full_text_support():
+    assert OSSOpenSearch.supports_full_text_search() is True
+    assert DB.OSSOpenSearch.case_config_cls(IndexType.FTS) is OSSOpenSearchFtsConfig
+
+
+def test_oss_opensearch_create_index_fts_uses_text_mappings_and_settings():
+    db = OSSOpenSearch.__new__(OSSOpenSearch)
+    db._is_fts = True
+    db.case_config = OSSOpenSearchFtsConfig(
+        number_of_shards=2,
+        number_of_replicas=1,
+        refresh_interval="10s",
+    )
+    db.index_name = "idx"
+    calls = {}
+
+    class Indices:
+        def create(self, **kwargs):
+            calls.update(kwargs)
+
+    class Client:
+        indices = Indices()
+
+    db._create_index(Client())
+
+    assert calls == {
+        "index": "idx",
+        "body": {
+            "settings": {
+                "index": {
+                    "number_of_shards": 2,
+                    "number_of_replicas": 1,
+                    "refresh_interval": "10s",
+                }
+            },
+            "mappings": {
+                "properties": {
+                    "doc_id": {"type": "keyword"},
+                    "text": {"type": "text"},
+                }
+            },
+        },
+    }
+
+
+def test_oss_opensearch_insert_documents_builds_bulk_body():
+    db = make_fts_db()
+    captured = {}
+
+    class Client:
+        def bulk(self, **kwargs):
+            captured.update(kwargs)
+            return {"errors": False}
+
+    db.client = Client()
+
+    assert db.insert_documents(["alpha", "beta"], ["d1", "d2"]) == (2, None)
+    assert captured["body"] == [
+        {"index": {"_index": "idx", "_id": "d1"}},
+        {"doc_id": "d1", "text": "alpha"},
+        {"index": {"_index": "idx", "_id": "d2"}},
+        {"doc_id": "d2", "text": "beta"},
+    ]
+
+
+def test_oss_opensearch_insert_documents_validates_lengths():
+    db = make_fts_db()
+
+    class Client:
+        def bulk(self, **kwargs):
+            raise AssertionError("bulk should not be called")
+
+    db.client = Client()
+
+    with pytest.raises(ValueError, match="Mismatch between texts .* and doc_ids .* lengths"):
+        db.insert_documents(["alpha", "beta"], ["d1"])
+
+
+def test_oss_opensearch_search_documents_builds_match_query():
+    db = make_fts_db()
+    calls = {}
+
+    class Client:
+        def search(self, **kwargs):
+            calls.update(kwargs)
+            return {"hits": {"hits": [{"fields": {"doc_id": ["d1"]}}]}}
+
+    db.client = Client()
+
+    assert db.search_documents("hello world", k=3) == ["d1"]
+    assert calls["index"] == "idx"
+    assert calls["body"] == {"query": {"match": {"text": "hello world"}}}
+    assert calls["size"] == 3
+    assert calls["stored_fields"] == "_none_"
+    assert calls["filter_path"] == ["hits.hits._id", "hits.hits.fields.doc_id"]
+
+
+def test_oss_opensearch_search_documents_requests_text_payload():
+    db = make_fts_db()
+    calls = {}
+
+    class Client:
+        def search(self, **kwargs):
+            calls.update(kwargs)
+            return {"hits": {"hits": [{"_id": "d1", "_source": {"text": "hello"}}]}}
+
+    db.client = Client()
+
+    assert db.search_documents("hello world", k=3, payload_profile=PayloadProfile.TEXT) == ["d1"]
+    assert calls["_source"] == ["text"]
+    assert "stored_fields" not in calls
+    assert calls["filter_path"] == [
+        "hits.hits._id",
+        "hits.hits.fields.doc_id",
+        "hits.hits._source.text",
+    ]
+
+
+def test_oss_opensearch_document_methods_guard_non_fts_mode():
+    db = OSSOpenSearch.__new__(OSSOpenSearch)
+    db._is_fts = False
+    db.client = object()
+
+    with pytest.raises(RuntimeError, match="OSSOpenSearch full-text insert requires OSSOpenSearchFtsConfig"):
+        db.insert_documents(["alpha"], ["d1"])
+
+    with pytest.raises(RuntimeError, match="OSSOpenSearch full-text search requires OSSOpenSearchFtsConfig"):
+        db.search_documents("alpha")

vectordb_bench/backend/clients/__init__.py

@@ -538,6 +538,10 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912, PLR0915
             return AWSOpenSearchIndexConfig
 
         if self == DB.OSSOpenSearch:
+            if index_type == IndexType.FTS:
+                from .oss_opensearch.config import OSSOpenSearchFtsConfig
+
+                return OSSOpenSearchFtsConfig
             from .oss_opensearch.config import OSSOpenSearchIndexConfig
 
             return OSSOpenSearchIndexConfig

vectordb_bench/backend/clients/elastic_cloud/cli.py

@@ -200,8 +200,8 @@ class ElasticCloudHNSWTypedDict(CommonTypedDict, ElasticCloudTypedDict):
         str,
         click.option(
             "--element-type",
-            type=click.Choice(["float", "byte"], case_sensitive=False),
-            help="Element type for vectors (float: 4 bytes, byte: 1 byte)",
+            type=click.Choice(["float", "byte", "bit"], case_sensitive=False),
+            help="Element type for vectors (float: 4 bytes, byte: 1 byte, bit: packed binary)",
             required=False,
             default="float",
             show_default=True,

vectordb_bench/backend/clients/elastic_cloud/config.py

@@ -68,6 +68,7 @@ def to_dict(self) -> dict:
 class ESElementType(StrEnum):
     float = "float"  # 4 byte
     byte = "byte"  # 1 byte, -128 to 127
+    bit = "bit"  # packed binary vector
 
 
 class ElasticCloudIndexConfig(BaseModel, DBCaseConfig):
@@ -87,6 +88,12 @@ class ElasticCloudIndexConfig(BaseModel, DBCaseConfig):
     M: int | None = None
     num_candidates: int | None = None
 
+    def _metric_value(self) -> str | None:
+        metric = self.metric_type
+        if metric is None:
+            return None
+        return metric.value if hasattr(metric, "value") else str(metric)
+
     def __eq__(self, obj: any):
         return (
             self.index == obj.index
@@ -111,17 +118,24 @@ def __hash__(self) -> int:
         )
 
     def parse_metric(self) -> str:
+        if self._metric_value() == MetricType.HAMMING.value:
+            return "l2_norm"
         if self.metric_type == MetricType.L2:
             return "l2_norm"
         if self.metric_type == MetricType.IP:
             return "dot_product"
         return "cosine"
 
+    def parse_element_type(self) -> ESElementType:
+        if self._metric_value() == MetricType.HAMMING.value:
+            return ESElementType.bit
+        return self.element_type
+
     def index_param(self) -> dict:
         return {
             "type": "dense_vector",
             "index": True,
-            "element_type": self.element_type.value,
+            "element_type": self.parse_element_type().value,
             "similarity": self.parse_metric(),
             "index_options": {
                 "type": self.index.value,