|
| 1 | +# Copyright 2025-present the zvec project |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import pickle |
| 18 | + |
| 19 | +import numpy as np |
| 20 | +import pytest |
| 21 | + |
| 22 | +import zvec |
| 23 | +from zvec import ( |
| 24 | + CollectionOption, |
| 25 | + CollectionSchema, |
| 26 | + Doc, |
| 27 | + FieldSchema, |
| 28 | + FlatIndexParam, |
| 29 | + HnswIndexParam, |
| 30 | + HnswQueryParam, |
| 31 | + InvertIndexParam, |
| 32 | + Query, |
| 33 | + VamanaIndexParam, |
| 34 | + VamanaQueryParam, |
| 35 | + VectorSchema, |
| 36 | +) |
| 37 | +from zvec.typing import DataType, MetricType, QuantizeType |
| 38 | + |
| 39 | +DIMENSION = 16 |
| 40 | +NUM_DOCS = 64 |
| 41 | +TOPK = 5 |
| 42 | + |
| 43 | + |
| 44 | +def _index_param(index_type: str, quantize_type: QuantizeType): |
| 45 | + if index_type == "flat": |
| 46 | + return FlatIndexParam( |
| 47 | + metric_type=MetricType.L2, |
| 48 | + quantize_type=quantize_type, |
| 49 | + ) |
| 50 | + if index_type == "hnsw": |
| 51 | + return HnswIndexParam( |
| 52 | + metric_type=MetricType.L2, |
| 53 | + m=16, |
| 54 | + ef_construction=64, |
| 55 | + quantize_type=quantize_type, |
| 56 | + ) |
| 57 | + if index_type == "vamana": |
| 58 | + return VamanaIndexParam( |
| 59 | + metric_type=MetricType.L2, |
| 60 | + max_degree=16, |
| 61 | + search_list_size=32, |
| 62 | + quantize_type=quantize_type, |
| 63 | + ) |
| 64 | + raise AssertionError(f"unknown index type: {index_type}") |
| 65 | + |
| 66 | + |
| 67 | +def _query_param(index_type: str): |
| 68 | + if index_type == "hnsw": |
| 69 | + return HnswQueryParam(ef=64) |
| 70 | + if index_type == "vamana": |
| 71 | + return VamanaQueryParam(ef_search=64) |
| 72 | + return None |
| 73 | + |
| 74 | + |
| 75 | +def _schema( |
| 76 | + name: str, |
| 77 | + index_type: str, |
| 78 | + quantize_type: QuantizeType, |
| 79 | +) -> CollectionSchema: |
| 80 | + return CollectionSchema( |
| 81 | + name=name, |
| 82 | + fields=[ |
| 83 | + FieldSchema( |
| 84 | + "sequence", |
| 85 | + DataType.INT64, |
| 86 | + nullable=False, |
| 87 | + index_param=InvertIndexParam(), |
| 88 | + ) |
| 89 | + ], |
| 90 | + vectors=[ |
| 91 | + VectorSchema( |
| 92 | + "embedding", |
| 93 | + DataType.VECTOR_FP32, |
| 94 | + DIMENSION, |
| 95 | + index_param=_index_param(index_type, quantize_type), |
| 96 | + ) |
| 97 | + ], |
| 98 | + ) |
| 99 | + |
| 100 | + |
| 101 | +def _documents() -> list[Doc]: |
| 102 | + rng = np.random.default_rng(20260720) |
| 103 | + vectors = rng.normal(size=(NUM_DOCS, DIMENSION)).astype(np.float32) |
| 104 | + return [ |
| 105 | + Doc( |
| 106 | + id=str(i), |
| 107 | + fields={"sequence": i}, |
| 108 | + vectors={"embedding": vectors[i].tolist()}, |
| 109 | + ) |
| 110 | + for i in range(NUM_DOCS) |
| 111 | + ] |
| 112 | + |
| 113 | + |
| 114 | +def _query(collection, index_type: str, vector: list[float]): |
| 115 | + query = Query( |
| 116 | + field_name="embedding", |
| 117 | + vector=vector, |
| 118 | + param=_query_param(index_type), |
| 119 | + ) |
| 120 | + results = collection.query(query, topk=TOPK) |
| 121 | + assert results |
| 122 | + return results |
| 123 | + |
| 124 | + |
| 125 | +@pytest.mark.parametrize( |
| 126 | + "quantize_type, value", |
| 127 | + [ |
| 128 | + (QuantizeType.UNIFORM_UINT7, 5), |
| 129 | + (QuantizeType.UNIFORM_UINT8, 6), |
| 130 | + ], |
| 131 | +) |
| 132 | +def test_uniform_quantize_type_surface(quantize_type, value): |
| 133 | + assert quantize_type.value == value |
| 134 | + assert quantize_type.name in {"UNIFORM_UINT7", "UNIFORM_UINT8"} |
| 135 | + assert getattr(zvec.QuantizeType, quantize_type.name) is quantize_type |
| 136 | + |
| 137 | + |
| 138 | +@pytest.mark.parametrize("index_type", ["flat", "hnsw", "vamana"]) |
| 139 | +@pytest.mark.parametrize( |
| 140 | + "quantize_type", |
| 141 | + [QuantizeType.UNIFORM_UINT7, QuantizeType.UNIFORM_UINT8], |
| 142 | +) |
| 143 | +def test_uniform_index_param_pickle_roundtrip(index_type, quantize_type): |
| 144 | + original = _index_param(index_type, quantize_type) |
| 145 | + restored = pickle.loads(pickle.dumps(original)) |
| 146 | + assert restored.quantize_type == quantize_type |
| 147 | + assert restored.to_dict() == original.to_dict() |
| 148 | + |
| 149 | + |
| 150 | +@pytest.mark.parametrize("index_type", ["flat", "hnsw", "vamana"]) |
| 151 | +@pytest.mark.parametrize( |
| 152 | + "quantize_type", |
| 153 | + [QuantizeType.UNIFORM_UINT7, QuantizeType.UNIFORM_UINT8], |
| 154 | +) |
| 155 | +def test_uniform_quantizer_optimize_query_and_reopen( |
| 156 | + tmp_path, index_type, quantize_type |
| 157 | +): |
| 158 | + path = tmp_path / f"{index_type}_{quantize_type.name.lower()}" |
| 159 | + option = CollectionOption(read_only=False, enable_mmap=True) |
| 160 | + collection = zvec.create_and_open( |
| 161 | + path=str(path), |
| 162 | + schema=_schema(path.name, index_type, quantize_type), |
| 163 | + option=option, |
| 164 | + ) |
| 165 | + documents = _documents() |
| 166 | + try: |
| 167 | + for result in collection.insert(documents): |
| 168 | + assert result.ok(), result.message() |
| 169 | + |
| 170 | + probe = documents[7].vector("embedding") |
| 171 | + assert _query(collection, index_type, probe)[0].id == "7" |
| 172 | + |
| 173 | + collection.optimize() |
| 174 | + assert _query(collection, index_type, probe)[0].id == "7" |
| 175 | + |
| 176 | + # Uniform quantization is trained globally during optimize(). New |
| 177 | + # writes remain in the raw FP32 index until the next optimize, so |
| 178 | + # queries must fall back to the complete raw index instead of missing |
| 179 | + # the unquantized tail segment. |
| 180 | + appended = Doc( |
| 181 | + id="appended", |
| 182 | + fields={"sequence": NUM_DOCS}, |
| 183 | + vectors={"embedding": [100.0] * DIMENSION}, |
| 184 | + ) |
| 185 | + for result in collection.insert([appended]): |
| 186 | + assert result.ok(), result.message() |
| 187 | + appended_probe = appended.vector("embedding") |
| 188 | + assert _query(collection, index_type, appended_probe)[0].id == "appended" |
| 189 | + |
| 190 | + collection.optimize() |
| 191 | + assert _query(collection, index_type, appended_probe)[0].id == "appended" |
| 192 | + |
| 193 | + post_optimize = Doc( |
| 194 | + id="post-optimize", |
| 195 | + fields={"sequence": NUM_DOCS + 1}, |
| 196 | + vectors={"embedding": [200.0] * DIMENSION}, |
| 197 | + ) |
| 198 | + for result in collection.insert([post_optimize]): |
| 199 | + assert result.ok(), result.message() |
| 200 | + post_optimize_probe = post_optimize.vector("embedding") |
| 201 | + assert ( |
| 202 | + _query(collection, index_type, post_optimize_probe)[0].id == "post-optimize" |
| 203 | + ) |
| 204 | + finally: |
| 205 | + del collection |
| 206 | + |
| 207 | + reopened = zvec.open(path=str(path), option=option) |
| 208 | + try: |
| 209 | + assert reopened.schema.vectors[0].index_param.quantize_type == quantize_type |
| 210 | + assert _query(reopened, index_type, probe)[0].id == "7" |
| 211 | + assert _query(reopened, index_type, appended_probe)[0].id == "appended" |
| 212 | + assert ( |
| 213 | + _query(reopened, index_type, post_optimize_probe)[0].id == "post-optimize" |
| 214 | + ) |
| 215 | + finally: |
| 216 | + reopened.destroy() |
| 217 | + |
| 218 | + |
| 219 | +@pytest.mark.parametrize( |
| 220 | + "quantize_type", |
| 221 | + [QuantizeType.UNIFORM_UINT7, QuantizeType.UNIFORM_UINT8], |
| 222 | +) |
| 223 | +@pytest.mark.parametrize("metric_type", [MetricType.IP, MetricType.COSINE]) |
| 224 | +def test_uniform_quantizer_rejects_non_l2(tmp_path, quantize_type, metric_type): |
| 225 | + schema = CollectionSchema( |
| 226 | + name="invalid_uniform_metric", |
| 227 | + vectors=[ |
| 228 | + VectorSchema( |
| 229 | + "embedding", |
| 230 | + DataType.VECTOR_FP32, |
| 231 | + DIMENSION, |
| 232 | + index_param=HnswIndexParam( |
| 233 | + metric_type=metric_type, |
| 234 | + quantize_type=quantize_type, |
| 235 | + ), |
| 236 | + ) |
| 237 | + ], |
| 238 | + ) |
| 239 | + with pytest.raises(Exception, match=quantize_type.name): |
| 240 | + zvec.create_and_open(path=str(tmp_path / "invalid"), schema=schema) |
0 commit comments