test: add RocksDB patch compatibility fixture

Author: zhourrr

scripts/rocksdb_patch_compat.py

@@ -0,0 +1,288 @@
+#!/usr/bin/env python3
+# Copyright 2025-present the zvec project
+#
+# 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.
+"""Generate and validate a persisted collection across zvec versions.
+
+Example:
+    # Generate with the published release.
+    python -m pip install zvec==0.5.1
+    python scripts/rocksdb_patch_compat.py generate var/compatibility/v051
+
+    # Validate with a separately installed source build.
+    python scripts/rocksdb_patch_compat.py validate var/compatibility/v051
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import shutil
+from pathlib import Path
+
+import zvec
+from zvec import (
+    CollectionOption,
+    DataType,
+    Doc,
+    FieldSchema,
+    Fts,
+    FtsIndexParam,
+    HnswIndexParam,
+    InvertIndexParam,
+    MetricType,
+    Query,
+    VectorSchema,
+)
+
+DOC_COUNT = 256
+VECTOR_DIMENSION = 8
+GENERATOR_VERSION = "0.5.1"
+
+
+def zvec_version() -> str:
+    return getattr(zvec, "__version__", "unknown")
+
+
+def embedding(ordinal: int) -> list[float]:
+    """Return a deterministic, non-degenerate vector for an ordinal."""
+    return [
+        ((ordinal * (dimension + 3)) % 257) / 257.0
+        for dimension in range(VECTOR_DIMENSION)
+    ]
+
+
+def documents() -> list[Doc]:
+    docs = []
+    for ordinal in range(DOC_COUNT):
+        terms = ["compatibility", "fixture", f"document{ordinal}"]
+        if ordinal % 2 == 0:
+            terms.append("evenmarker")
+        if ordinal % 5 == 0:
+            terms.append("fivemarker")
+        docs.append(
+            Doc(
+                id=f"doc-{ordinal:03d}",
+                fields={
+                    "ordinal": ordinal,
+                    "category": f"category-{ordinal % 4}",
+                    "content": " ".join(terms),
+                },
+                vectors={"embedding": embedding(ordinal)},
+            )
+        )
+    return docs
+
+
+def collection_schema():
+    return zvec.CollectionSchema(
+        name="rocksdb_patch_compat_v051",
+        fields=[
+            FieldSchema(
+                "ordinal",
+                DataType.INT64,
+                nullable=False,
+                index_param=InvertIndexParam(enable_range_optimization=True),
+            ),
+            FieldSchema(
+                "category",
+                DataType.STRING,
+                nullable=False,
+                index_param=InvertIndexParam(),
+            ),
+            FieldSchema(
+                "content",
+                DataType.STRING,
+                nullable=False,
+                index_param=FtsIndexParam(
+                    tokenizer_name="standard",
+                    filters=["lowercase"],
+                ),
+            ),
+        ],
+        vectors=[
+            VectorSchema(
+                "embedding",
+                DataType.VECTOR_FP32,
+                dimension=VECTOR_DIMENSION,
+                index_param=HnswIndexParam(metric_type=MetricType.L2),
+            )
+        ],
+    )
+
+
+def require_ok(statuses) -> None:
+    for status in statuses if isinstance(statuses, list) else [statuses]:
+        if not status.ok():
+            raise RuntimeError(f"zvec operation failed: {status}")
+
+
+def sidecar_path(collection_path: Path) -> Path:
+    return collection_path.with_name(f"{collection_path.name}.fixture.json")
+
+
+def generate(collection_path: Path, overwrite: bool) -> None:
+    if zvec_version() != GENERATOR_VERSION:
+        raise RuntimeError(
+            f"fixture generation requires zvec=={GENERATOR_VERSION}, "
+            f"loaded {zvec_version()}"
+        )
+
+    if collection_path.exists():
+        if not overwrite:
+            raise FileExistsError(
+                f"{collection_path} already exists; pass --overwrite to replace it"
+            )
+        shutil.rmtree(collection_path)
+
+    sidecar = sidecar_path(collection_path)
+    if sidecar.exists() and overwrite:
+        sidecar.unlink()
+
+    collection_path.parent.mkdir(parents=True, exist_ok=True)
+    collection = zvec.create_and_open(
+        path=str(collection_path),
+        schema=collection_schema(),
+        option=CollectionOption(read_only=False, enable_mmap=True),
+    )
+    require_ok(collection.insert(documents()))
+    collection.flush()
+    collection.optimize()
+    collection.flush()
+
+    if collection.stats.doc_count != DOC_COUNT:
+        raise AssertionError(
+            f"expected {DOC_COUNT} documents, got {collection.stats.doc_count}"
+        )
+
+    # Collection has no explicit close API in 0.5.1. Releasing the last Python
+    # reference closes the native collection before inspecting its files.
+    del collection
+    gc.collect()
+
+    sst_files = sorted(
+        str(path.relative_to(collection_path))
+        for path in collection_path.rglob("*.sst")
+    )
+    if not sst_files:
+        raise AssertionError("fixture contains no persisted RocksDB SST files")
+
+    sidecar.write_text(
+        json.dumps(
+            {
+                "generator_version": zvec_version(),
+                "doc_count": DOC_COUNT,
+                "vector_dimension": VECTOR_DIMENSION,
+                "sst_files": sst_files,
+            },
+            indent=2,
+            sort_keys=True,
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    print(
+        f"generated {collection_path} with zvec {zvec_version()}: "
+        f"{DOC_COUNT} docs, {len(sst_files)} SST files"
+    )
+
+
+def validate(collection_path: Path) -> None:
+    sidecar = sidecar_path(collection_path)
+    metadata = json.loads(sidecar.read_text(encoding="utf-8"))
+    if metadata["generator_version"] != GENERATOR_VERSION:
+        raise AssertionError(
+            f"expected a {GENERATOR_VERSION} fixture, got {metadata}"
+        )
+
+    collection = zvec.open(
+        path=str(collection_path),
+        option=CollectionOption(read_only=True, enable_mmap=True),
+    )
+    if collection.stats.doc_count != DOC_COUNT:
+        raise AssertionError(
+            f"expected {DOC_COUNT} documents, got {collection.stats.doc_count}"
+        )
+
+    fetched = collection.fetch(
+        ["doc-000", "doc-127", "doc-255"], include_vector=True
+    )
+    if set(fetched) != {"doc-000", "doc-127", "doc-255"}:
+        raise AssertionError(f"unexpected fetch result: {set(fetched)}")
+    if fetched["doc-127"].fields["ordinal"] != 127:
+        raise AssertionError("scalar field changed while reopening the fixture")
+
+    scalar_hits = collection.query(filter="ordinal>=240", topk=DOC_COUNT)
+    scalar_ids = {doc.id for doc in scalar_hits}
+    expected_scalar_ids = {f"doc-{ordinal:03d}" for ordinal in range(240, 256)}
+    if scalar_ids != expected_scalar_ids:
+        raise AssertionError(
+            f"inverted-index mismatch: expected {expected_scalar_ids}, "
+            f"got {scalar_ids}"
+        )
+
+    fts_hits = collection.query(
+        queries=Query(
+            field_name="content", fts=Fts(match_string="evenmarker")
+        ),
+        topk=DOC_COUNT,
+    )
+    expected_fts_ids = {
+        f"doc-{ordinal:03d}" for ordinal in range(DOC_COUNT) if ordinal % 2 == 0
+    }
+    fts_ids = {doc.id for doc in fts_hits}
+    if fts_ids != expected_fts_ids:
+        raise AssertionError(
+            f"FTS mismatch: expected {len(expected_fts_ids)} hits, "
+            f"got {len(fts_ids)}"
+        )
+
+    vector_hits = collection.query(
+        queries=Query(field_name="embedding", vector=embedding(127)),
+        topk=1,
+        include_vector=True,
+    )
+    if len(vector_hits) != 1 or vector_hits[0].id != "doc-127":
+        raise AssertionError(f"vector-index mismatch: {vector_hits}")
+
+    del collection
+    gc.collect()
+    print(
+        f"validated zvec {metadata['generator_version']} fixture with "
+        f"zvec {zvec_version()}: fetch, inverted, FTS, and vector indexes pass"
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("mode", choices=("generate", "validate"))
+    parser.add_argument("collection_path", type=Path)
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="replace an existing collection when generating",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.mode == "generate":
+        generate(args.collection_path.resolve(), args.overwrite)
+    else:
+        validate(args.collection_path.resolve())
+
+
+if __name__ == "__main__":
+    main()