test: verify writes after RocksDB upgrade

Author: zhourrr

scripts/rocksdb_patch_compat.py

@@ -21,6 +21,10 @@
 
     # Validate with a separately installed source build.
     python scripts/rocksdb_patch_compat.py validate var/compatibility/v051
+
+    # Copy, query, extend, and reopen the old fixture with the source build.
+    python scripts/rocksdb_patch_compat.py upgrade \
+        var/compatibility/v051 var/compatibility/v051-upgraded --overwrite
 """
 
 from __future__ import annotations
@@ -47,6 +51,8 @@
 )
 
 DOC_COUNT = 256
+UPGRADE_DOC_COUNT = 64
+UPGRADED_DOC_COUNT = DOC_COUNT + UPGRADE_DOC_COUNT
 VECTOR_DIMENSION = 8
 GENERATOR_VERSION = "0.5.1"
 
@@ -57,20 +63,31 @@ def zvec_version() -> str:
 
 def embedding(ordinal: int) -> list[float]:
     """Return a deterministic, non-degenerate vector for an ordinal."""
+    if ordinal >= DOC_COUNT:
+        upgrade_ordinal = ordinal - DOC_COUNT
+        return [
+            -1.0 - ((upgrade_ordinal * (dimension + 3)) % 67) / 67.0
+            for dimension in range(VECTOR_DIMENSION)
+        ]
     return [
         ((ordinal * (dimension + 3)) % 257) / 257.0
         for dimension in range(VECTOR_DIMENSION)
     ]
 
 
-def documents() -> list[Doc]:
+def documents(
+    start: int = 0,
+    stop: int = DOC_COUNT,
+    extra_terms: tuple[str, ...] = (),
+) -> list[Doc]:
     docs = []
-    for ordinal in range(DOC_COUNT):
+    for ordinal in range(start, stop):
         terms = ["compatibility", "fixture", f"document{ordinal}"]
         if ordinal % 2 == 0:
             terms.append("evenmarker")
         if ordinal % 5 == 0:
             terms.append("fivemarker")
+        terms.extend(extra_terms)
         docs.append(
             Doc(
                 id=f"doc-{ordinal:03d}",
@@ -132,6 +149,22 @@ def sidecar_path(collection_path: Path) -> Path:
     return collection_path.with_name(f"{collection_path.name}.fixture.json")
 
 
+def sst_files(collection_path: Path) -> list[str]:
+    return sorted(
+        str(path.relative_to(collection_path))
+        for path in collection_path.rglob("*.sst")
+    )
+
+
+def read_metadata(collection_path: Path) -> dict:
+    metadata = json.loads(sidecar_path(collection_path).read_text(encoding="utf-8"))
+    if metadata["generator_version"] != GENERATOR_VERSION:
+        raise AssertionError(
+            f"expected a {GENERATOR_VERSION} fixture, got {metadata}"
+        )
+    return metadata
+
+
 def generate(collection_path: Path, overwrite: bool) -> None:
     if zvec_version() != GENERATOR_VERSION:
         raise RuntimeError(
@@ -171,11 +204,8 @@ def generate(collection_path: Path, overwrite: bool) -> None:
     del collection
     gc.collect()
 
-    sst_files = sorted(
-        str(path.relative_to(collection_path))
-        for path in collection_path.rglob("*.sst")
-    )
-    if not sst_files:
+    persisted_sst_files = sst_files(collection_path)
+    if not persisted_sst_files:
         raise AssertionError("fixture contains no persisted RocksDB SST files")
 
     sidecar.write_text(
@@ -184,7 +214,7 @@ def generate(collection_path: Path, overwrite: bool) -> None:
                 "generator_version": zvec_version(),
                 "doc_count": DOC_COUNT,
                 "vector_dimension": VECTOR_DIMENSION,
-                "sst_files": sst_files,
+                "sst_files": persisted_sst_files,
             },
             indent=2,
             sort_keys=True,
@@ -194,25 +224,15 @@ def generate(collection_path: Path, overwrite: bool) -> None:
     )
     print(
         f"generated {collection_path} with zvec {zvec_version()}: "
-        f"{DOC_COUNT} docs, {len(sst_files)} SST files"
+        f"{DOC_COUNT} docs, {len(persisted_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:
+def assert_collection_contents(collection, expected_doc_count: int) -> None:
+    if collection.stats.doc_count != expected_doc_count:
         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}"
+            f"expected {expected_doc_count} documents, "
+            f"got {collection.stats.doc_count}"
         )
 
     fetched = collection.fetch(
@@ -223,9 +243,14 @@ def validate(collection_path: Path) -> None:
     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_start = expected_doc_count - 16
+    scalar_hits = collection.query(
+        filter=f"ordinal>={scalar_start}", topk=expected_doc_count
+    )
     scalar_ids = {doc.id for doc in scalar_hits}
-    expected_scalar_ids = {f"doc-{ordinal:03d}" for ordinal in range(240, 256)}
+    expected_scalar_ids = {
+        f"doc-{ordinal:03d}" for ordinal in range(scalar_start, expected_doc_count)
+    }
     if scalar_ids != expected_scalar_ids:
         raise AssertionError(
             f"inverted-index mismatch: expected {expected_scalar_ids}, "
@@ -236,10 +261,12 @@ def validate(collection_path: Path) -> None:
         queries=Query(
             field_name="content", fts=Fts(match_string="evenmarker")
         ),
-        topk=DOC_COUNT,
+        topk=expected_doc_count,
     )
     expected_fts_ids = {
-        f"doc-{ordinal:03d}" for ordinal in range(DOC_COUNT) if ordinal % 2 == 0
+        f"doc-{ordinal:03d}"
+        for ordinal in range(expected_doc_count)
+        if ordinal % 2 == 0
     }
     fts_ids = {doc.id for doc in fts_hits}
     if fts_ids != expected_fts_ids:
@@ -256,18 +283,129 @@ def validate(collection_path: Path) -> None:
     if len(vector_hits) != 1 or vector_hits[0].id != "doc-127":
         raise AssertionError(f"vector-index mismatch: {vector_hits}")
 
+    if expected_doc_count > DOC_COUNT:
+        upgrade_hits = collection.query(
+            queries=Query(
+                field_name="content", fts=Fts(match_string="upgrademark")
+            ),
+            topk=expected_doc_count,
+        )
+        expected_upgrade_ids = {
+            f"doc-{ordinal:03d}"
+            for ordinal in range(DOC_COUNT, expected_doc_count)
+        }
+        if {doc.id for doc in upgrade_hits} != expected_upgrade_ids:
+            raise AssertionError(f"upgraded FTS mismatch: {upgrade_hits}")
+
+        last_ordinal = expected_doc_count - 1
+        new_vector_hits = collection.query(
+            queries=Query(
+                field_name="embedding", vector=embedding(last_ordinal)
+            ),
+            topk=1,
+            include_vector=True,
+        )
+        if len(new_vector_hits) != 1 or new_vector_hits[0].id != (
+            f"doc-{last_ordinal:03d}"
+        ):
+            raise AssertionError(f"upgraded vector-index mismatch: {new_vector_hits}")
+
+
+def validate(collection_path: Path) -> None:
+    metadata = read_metadata(collection_path)
+    expected_doc_count = metadata["doc_count"]
+    collection = zvec.open(
+        path=str(collection_path),
+        option=CollectionOption(read_only=True, enable_mmap=True),
+    )
+    assert_collection_contents(collection, expected_doc_count)
+
     del collection
     gc.collect()
     print(
         f"validated zvec {metadata['generator_version']} fixture with "
-        f"zvec {zvec_version()}: fetch, inverted, FTS, and vector indexes pass"
+        f"zvec {zvec_version()}: {expected_doc_count} docs; fetch, inverted, "
+        "FTS, and vector indexes pass"
+    )
+
+
+def upgrade(source_path: Path, destination_path: Path, overwrite: bool) -> None:
+    if zvec_version() == GENERATOR_VERSION:
+        raise RuntimeError(
+            "upgrade requires the new source build, not the 0.5.1 generator"
+        )
+    if source_path == destination_path:
+        raise ValueError("upgrade destination must differ from the source fixture")
+
+    metadata = read_metadata(source_path)
+    if metadata["doc_count"] != DOC_COUNT:
+        raise AssertionError("upgrade source must be the untouched 0.5.1 fixture")
+
+    destination_sidecar = sidecar_path(destination_path)
+    if destination_path.exists():
+        if not overwrite:
+            raise FileExistsError(
+                f"{destination_path} already exists; pass --overwrite to replace it"
+            )
+        shutil.rmtree(destination_path)
+    if destination_sidecar.exists() and overwrite:
+        destination_sidecar.unlink()
+
+    destination_path.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copytree(source_path, destination_path)
+    shutil.copy2(sidecar_path(source_path), destination_sidecar)
+
+    collection = zvec.open(
+        path=str(destination_path),
+        option=CollectionOption(read_only=False, enable_mmap=True),
+    )
+
+    # Query the release-created indexes before changing anything.
+    assert_collection_contents(collection, DOC_COUNT)
+
+    require_ok(
+        collection.insert(
+            documents(
+                DOC_COUNT,
+                UPGRADED_DOC_COUNT,
+                extra_terms=("upgrademark",),
+            )
+        )
+    )
+    collection.flush()
+    collection.optimize()
+    collection.flush()
+
+    # Query both old and newly inserted data through every index type.
+    assert_collection_contents(collection, UPGRADED_DOC_COUNT)
+    del collection
+    gc.collect()
+
+    metadata.update(
+        {
+            "doc_count": UPGRADED_DOC_COUNT,
+            "sst_files": sst_files(destination_path),
+            "upgraded_with_version": zvec_version(),
+        }
+    )
+    destination_sidecar.write_text(
+        json.dumps(metadata, indent=2, sort_keys=True) + "\n",
+        encoding="utf-8",
+    )
+
+    # Reopen read-only to prove the mixed-version collection persisted cleanly.
+    validate(destination_path)
+    print(
+        f"upgraded a {GENERATOR_VERSION} fixture with zvec {zvec_version()}: "
+        f"inserted {UPGRADE_DOC_COUNT} docs and persisted {UPGRADED_DOC_COUNT}"
     )
 
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("mode", choices=("generate", "validate"))
+    parser.add_argument("mode", choices=("generate", "validate", "upgrade"))
     parser.add_argument("collection_path", type=Path)
+    parser.add_argument("destination_path", nargs="?", type=Path)
     parser.add_argument(
         "--overwrite",
         action="store_true",
@@ -279,9 +417,21 @@ def parse_args() -> argparse.Namespace:
 def main() -> None:
     args = parse_args()
     if args.mode == "generate":
+        if args.destination_path is not None:
+            raise ValueError("generate accepts one collection path")
         generate(args.collection_path.resolve(), args.overwrite)
-    else:
+    elif args.mode == "validate":
+        if args.destination_path is not None:
+            raise ValueError("validate accepts one collection path")
         validate(args.collection_path.resolve())
+    else:
+        if args.destination_path is None:
+            raise ValueError("upgrade requires source and destination paths")
+        upgrade(
+            args.collection_path.resolve(),
+            args.destination_path.resolve(),
+            args.overwrite,
+        )
 
 
 if __name__ == "__main__":