[TamimEhsan]

Summary

• Part of Proposal: expose core quantization through a C API and a Python module #183

This adds rabitqlib.quant, a Python module that exposes the quantization building blocks behind the indexes: rotation, quantization, correction factors, packed codes and distance estimation. With it, users can store RaBitQ codes in their own database or file format and score them with their own candidate selection, without building an index.

The module is provisional in the sense of PEP 411: any minor release may change it, with a release note. The status is stated in the module docstring and on the docs page. Importing the module emits no warning. Byte formats are versioned from the first release, and bytes of an unknown version raise quant.FormatError.

Everything is additive. Existing Python and C++ APIs, defaults and saved-index formats are unchanged, and import rabitqlib does not import the new module.

It builds on the portable seeded rotations and the frozen fast-mode scale table from #199.

Example

Encode 10,000 vectors in 10 clusters, then find the top 10 for a query with a 1-bit first pass and 4-bit refinement:

import numpy as np
from rabitqlib import quant

rng = np.random.default_rng(0)
data = rng.standard_normal((10_000, 100)).astype(np.float32)
labels = rng.integers(0, 10, len(data))   # cluster assignments from any k-means

# Rotate everything into the working space (100 -> padded 128 dimensions).
rot = quant.Rotator(100, seed=42)
x = rot.rotate(data)

# Encode each cluster against its centroid: a 1-bit prefix plus 3 extra bits.
lists = []
for j in range(10):
    ids = np.flatnonzero(labels == j)
    c = x[ids].mean(axis=0)
    residuals = x[ids] - c
    codes = quant.quantize(residuals, nbits=4)
    prefix, ex = quant.split(codes, nbits=4, prefix_bits=1)
    lists.append(dict(
        ids=ids,
        centroid=c,
        bitsets=quant.pack_codes(prefix, nbits=1),
        f_prefix=quant.code_factors(residuals, prefix, nbits=1, metric="l2", centroid=c),
        ex=quant.pack_codes(ex, nbits=3),
        f_full=quant.code_factors(residuals, codes, nbits=4, metric="l2", centroid=c),
    ))
# Codes and factors are plain NumPy arrays: store them anywhere, with
# quant.PACKED_CODES_VERSION and quant.FACTORS_VERSION next to them.

# Search: a cheap 1-bit estimate for every vector, then refine the candidates.
q = rot.rotate(rng.standard_normal((1, 100)).astype(np.float32))[0]
query = quant.Query(q, metric="l2")
first = []
for lst in lists:
    query.set_centroid(lst["centroid"])
    first.append(query.estimate_bitset(lst["bitsets"], lst["f_prefix"]))
upper = np.concatenate([e["distance"] + e["error"] for e in first])
threshold = np.partition(upper, 9)[9]      # 10th-smallest upper bound

ids, dists = [], []
for lst, est in zip(lists, first):
    cand = np.flatnonzero(est["distance"] - est["error"] <= threshold)
    query.set_centroid(lst["centroid"])
    dists.append(query.refine_packed(lst["ex"], lst["f_full"], est[cand],
                                     indices=cand, ex_bits=3))
    ids.append(lst["ids"][cand])
ids, dists = np.concatenate(ids), np.concatenate(dists)
top10 = ids[np.argsort(dists)[:10]]

One Query serves every cluster: set_centroid switches lists without preparing the query again. A vector is refined only if the lower bound of its interval, distance - error, is at most the 10th-smallest upper bound, so the true top 10 survive the first pass within the interval's confidence. On this data, the example refines 1,484 of the 10,000 vectors, and 9 of the exact top 10 make its result. The data is unclustered noise, so the estimates are close together and many vectors survive; real clustered data prunes far more. For large lists, pack_fastscan_blocks and Query.estimate_fastscan can replace estimate_bitset in the first pass.

C++: provisional API

include/rabitqlib/core/provisional/quantize.hpp adds two functions in rabitqlib::quant::provisional. They are not part of quant::v1 and may change in any minor release.

• combined_code(residual, dim, total_bits, code, t_const) returns the unsigned codes of one residual that quant::v1::encode and the indexes store, as uint8_t (1–8 bits) or uint16_t (1–9 bits). t_const > 0 selects the fast-mode scale.
• code_factors(residual, centroid, dim, code, total_bits, metric) returns the factors of a caller-supplied code. This is the only way to get factors for a prefix (code >> ex_bits), which the two-stage estimate needs.

The header includes no Eigen, detail/ or SIMD headers and reuses quant::v1::Factors. The implementation (src/core/provisional.cpp) wraps the existing rabitq_impl functions. Inputs are validated with the same checks and messages as v1, and outputs are written only after every check passes.

Those checks moved from src/core/api.cpp into a private header, src/core/validation.hpp, which both APIs now share. v1 behaviour and messages are unchanged.

The namespace name was agreed in the planning discussion, pending confirmation in this review.

Python: rabitqlib.quant

Every call works in the rotated, padded working space: float32 arrays of shape (n, d), with d a positive multiple of 32.

Layer	API
Rotation	Rotator(dim, *, method, padded_dim, seed), rotate, to_bytes/from_bytes, copy and pickle; padded_dim(dim)
1. Encoding	quantize, split (prefix and extra bits), code_factors, pack_codes/unpack_codes, FACTORS_DTYPE
2. Estimation	Query(q, *, metric, centroid) with set_centroid, estimate_bitset, estimate_codes, refine, refine_packed, copy; ESTIMATE_DTYPE; query_terms and estimate_from_ip for inner products from the caller's own kernel
3. FastScan	pack_fastscan_blocks and Query.estimate_fastscan

• metric is required wherever factors or query terms are computed. Factors for the wrong metric would otherwise give wrong distances without an error.
• An estimate has three fields: distance, error (the half-width of the library's probabilistic interval, ε₀ = 1.9) and refine_input, the prefix inner product that refinement reads.
• Scoring calls accept indices= to read selected rows in place.
• pack_codes(nbits=1) writes LSB-first bitsets, equal to numpy.packbits(..., bitorder="little"). The FastScan packer reverses the bit order to the MSB-first layout its kernels read.
• quant.py re-exports the submodule _rabitqlib._quant and defines the version constants. quant.pyi types it. __init__.py and __init__.pyi are unchanged.

Contracts at the boundary

• Validation. Shapes, dtypes and values are checked; violations raise ValueError, and non-integer options raise TypeError. Non-finite or oversized coordinates are rejected as the indexes reject them.
• Output buffers. Every call that returns an array accepts out= with the exact dtype and shape. out must not overlap any input, exact aliasing included. If a call raises, a caller's out is left unchanged.
• Threads. The GIL is released during computation.
  • Build-side calls (rotate, quantize, code_factors, pack_*, unpack_codes) default to num_threads=0, meaning the available CPU count capped by OMP_NUM_THREADS (the 0.8.1 rule the indexes use), and give the same results for any thread count.
  • Scoring calls default to 1 thread.
  • quant_bindings.cpp is compiled with -ffp-contract=off on GCC and Clang. GCC otherwise contracts the inlined estimator into FMA separately in the serial and OpenMP loops (by default on aarch64, or with -mfma), so a row's distance could differ in the last bit with the thread count. The SIMD kernels are in other files and keep their own flags.
• Sharing. Rotator is immutable. Scoring calls on one Query may run concurrently; a set_centroid that overlaps one raises RuntimeError. This reuses IndexAccess from bindings_common.hpp, which gains a constructor for a custom busy message.
• Copies. Query supports copy, copy.copy, copy.deepcopy and pickling. A copy never shares storage with the original, including the lazily built FastScan table.

Byte formats

docs/formats.md specifies every layout the module reads or writes. A layout change is a new version, never an edit of an existing one.

• Raw arrays have no header. The module exports a version constant for each: PACKED_CODES_VERSION (pack_codes, including the 1-bit bitset), FACTORS_VERSION (the three record dtypes) and FASTSCAN_BLOCKS_VERSION (pack_fastscan_blocks, in the C++ BatchDataMap<float> layout). All are 1.

• Blobs (Rotator.to_bytes()) start with a 40-byte common header, followed by the kind's own fields and then the body.

  • Header fields: magic RABQBLOB, kind, version, header and body sizes, flags, and a CRC-32 as zlib.crc32 computes it.
  • Readers validate sizes before allocating anything and reject unknown kinds, versions, flags and non-zero reserved fields with FormatError.
  • The rotator body is exactly what Rotator<float>::save() writes.

• Golden files in tests/python/fixtures/quant/ freeze every version:

  • two rotator blobs;
  • packed codes for 1–8 bits, with a 32-coordinate tail;
  • FastScan blocks with a padded last block, checked against an independent NumPy implementation of the spec before freezing.

  Their README documents how each was produced. They are never regenerated.

While the module is provisional, a release may stop reading an old version, but it raises FormatError instead of returning wrong results. The promise that every released version stays readable comes with stabilisation.

Docs

• docs/docs/rabitq/quant_python.md, linked as "Python building blocks" in the site nav, covers:
  • the working space and rotation;
  • encoding;
  • two-stage estimation, with a candidate rule that is sound under the error interval;
  • FastScan and the reference formula;
  • threads, copies and byte formats.
• docs/docs/releases/unreleased.md adds the Python and C++ release notes.
• docs/development.md adds the provisional header, the module and docs/formats.md to the code map.

Known limitations and follow-ups

• 1-bit quantize speed. quantize(nbits=1) validates each row twice: the binding checks the whole array, then combined_code checks each row. A review measured 1.6 GB/s per thread on a Xeon Gold 6418H, about 4× the cost of computing the codes alone. A batch combined_code would remove the second pass. This should be decided before the API is stabilised; it doesn't block this PR.
• refine vs refine_packed. They agree only within float32 rounding, because the packed path uses the dispatched SIMD kernel.
• FastScan CPU support. estimate_fastscan needs AVX2/FMA, AVX-512 or NEON, and raises RuntimeError otherwise. Its estimates match estimate_bitset only within the error of the quantized query table.
• No IvfIndex parity test. Nothing yet checks the module end-to-end against an index's own encoding.

Testing

• C++:

  • ProvisionalQuantizeTest (tests/unit/rabitqlib/quantization/provisional_quantize_test.cpp) checks:
    • parity with quant::v1::encode and the detail functions;
    • a null centroid meaning zero;
    • zero and underflowing residuals;
    • rejected arguments, with no output written.
  • Release build with RABITQ_ENABLE_NATIVE_OPTIMIZATION=OFF: ctest --no-tests=error passes 494 of 494. One test is skipped, as on main.
  • A core-only build with OpenMP off passes 190 of 190.

• Python (tests/python/test_quant_*.py, tests/python/typing/public_api.py):

  • parity with the C++ encoder and estimator, across both metrics and every bit width;
  • the golden byte formats;
  • error types and messages;
  • out= overlap and staging;
  • thread-count independence;
  • copy and pickle independence;
  • the busy guard;
  • estimate accuracy against true distances;
  • stub and runtime signature agreement.

  The 121 quant tests pass. The full suite passes 1129 tests, with 4 skipped, and the mypy tests (tests/test_mypy.py) pass.

• Static checks:

  • ./scripts/check-tidy.sh build-tidy passes, as do ./scripts/check-format.sh, ./scripts/check-python.sh and git diff --check.
  • ./scripts/check-includes.sh reports 21 failures, none in files this PR adds. All are in code from main: the failures in bindings_common.hpp are CPython calls already on main, and the newest, a missing <algorithm> for std::max in ivf_bindings.cpp, came with 0.8.2 in a file this PR doesn't touch.

• Docs: python -m mkdocs build --strict --config-file docs/mkdocs.yml is clean, every code block in the guide runs in order, and so does the example above.

• Candidate rule (a one-off script, not in the suite): on 10,000 vectors in 10 clusters with 50 queries, distance ± error misses the true distance 5.5% of the time, as ε₀ = 1.9 predicts. The guide's rule reaches the recall of refining every vector.

• Not run locally: Pyright, which needs a newer Node than the test machine has; and the Windows (MSVC), macOS and ARM64 CI jobs.

[gouyt13]

Thanks for this PR. The module design and the tests made it easy to review in depth. While reviewing I found a few issues and fixed them locally as four commits on top of ea2070b. Each commit message explains the reason for its change, and none of them reverts an earlier one. The patches follow in the next four comments. Apply them with:

git checkout quant-python
# Save each patch below as 000N-*.patch, then:
git am 0001-*.patch 0002-*.patch 0003-*.patch 0004-*.patch

Please review them like any other change. If you disagree with something, say so in this thread.

The four commits

1. fix(core): reject a code that does not quantize its residual at the divisor. code_factors checked <residual, centered code> > 0 with its own sequential sum. The factor derivation, however, divides by a separately dispatched dot product, which was guarded only by an assert. The two sums round differently. So for a caller-supplied code near zero, the check could pass while the divisor was not positive: debug builds aborted, and release builds returned wrong factors. The check now runs on the divisor itself. Codes produced by the library never reach it, and NaN inputs to the legacy encoders behave as before.
2. refactor(quant): share the v1 query terms and the unpacked inner product. The bindings had their own copies of quant::v1's query-term and unpacked-code arithmetic. Both now live once in quant::detail, and v1 and the bindings both use them. I checked that the output of v1 and of every rabitqlib.quant call is bitwise unchanged on fixed inputs.
3. perf(quant): run row loops in ranges and check inputs in parallel passes. index_detail::parallel_for hands out one index at a time, which is too fine-grained for rows that cost nanoseconds. Every row loop was slower on 4 threads than on 1. Rows now go out in ranges. The input checks run in a parallel pass before anything is written, and the reported row is deterministic. This also replaces the per-worker scratch in pack_fastscan_blocks, which could be indexed past its end when the two thread-count lookups disagreed.
4. fix(quant): raise OverflowError for estimates that are not finite. Scoring calls returned inf or NaN where quant::v1::compute raises. code_factors overflow errors now name their row.

The benchmarks are in the commit messages of 3 and 4.

Design decisions on the maintainer side

These change documented behavior, so I am stating them explicitly:

• out= contract. A call that raises leaves out= unchanged, except after OverflowError or MemoryError, when its contents are unspecified. Detecting overflow before writing would need a staging copy on every call that passes out=, which is the cost out= is meant to avoid. Input errors are still raised before anything is written.
• OverflowError is part of the error model. It covers factors, query terms, distances and errors that are not representable in float32, as in quant::v1. The guide now documents it.
• v1 shares numerical helpers with this layer. A comment on the helpers says that changing them changes v1 and requires the v1 parity check.

Considered and left unchanged

• Per-row allocations in code_factors. I tried a caller-scratch overload. It made no measurable difference (41.8 ms vs 41.5 ms on 200k rows), so I dropped it rather than grow the provisional API.
• -ffp-contract=off under clang-cl. It is consistent with how CMakeLists.txt already treats every Clang as GNU-like (for example -march=native), so I did not change it.

Open question: split

split is now the only call that returns arrays but takes neither out= nor num_threads=. Commit 4 documents it as the exception to the out= rule. Could you consider whether it should accept both, like the other calls? If it should, that belongs in this PR, since the API is provisional and has not been released yet. If you think it should not, please add a sentence to the guide explaining why.

Please update the PR description

It still says "If a call raises, a caller's out is left unchanged" and that every array-returning call accepts out=. After commit 4, the guide states the narrower contract and that split takes no out=.

[gouyt13]

Patch 1 of 4: fix(core): reject a code that does not quantize its

0001-fix-core-reject-a-code-that-does-not-quantize-its-re.patch

From 0e603e8c1c6b98f2f8bc8bf028c60d44034a29d1 Mon Sep 17 00:00:00 2001
From: gouyt13clear <yutong003@e.ntu.edu.sg>
Date: Tue, 3 Nov 2026 16:28:09 +0800
Subject: [PATCH 1/4] fix(core): reject a code that does not quantize its
 residual at the divisor

The provisional code_factors accepts caller-supplied codes, so the factor
derivation can now meet a code whose centered inner product with the residual is
not positive. code_factors checked this with its own sequential sum, but
code_factor_terms divides by a separately dispatched dot product guarded only by
assert(). The two sums round differently, so near zero the check could pass while
the divisor was not positive: debug builds aborted the process and release builds
returned wrong factors.

code_factor_terms now throws std::invalid_argument when the divisor itself is not
positive, and code_factors drops its duplicate sum. The library's own codes match
the residual's signs, so every term of the divisor is nonnegative and they never
reach the new error. A NaN still passes through, as the per-vector encoders do not
check that inputs are finite (quantizer.md); a test pins that behavior.
---
 include/rabitqlib/core/detail/rabitq_impl.hpp | 10 ++-
 src/core/provisional.cpp                      | 13 +---
 .../rabitqlib/quantization/factor_test.cpp    | 68 +++++++++++++++++++
 3 files changed, 79 insertions(+), 12 deletions(-)

diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp
index efc0374..8c145b4 100644
--- a/include/rabitqlib/core/detail/rabitq_impl.hpp
+++ b/include/rabitqlib/core/detail/rabitq_impl.hpp
@@ -2,7 +2,6 @@
 
 #include <algorithm>
 #include <array>
-#include <cassert>
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
@@ -69,7 +68,14 @@ inline CodeFactorTerms<T> code_factor_terms(
     const T l2_norm = std::sqrt(l2_sqr);
     const T ip_resi_xucb = dot_product<T>(residual, centered, dim);
     const T ip_cent_xucb = dot_product<T>(centroid, centered, dim);
-    assert(ip_resi_xucb > 0);
+    // The library's own codes match the residual's signs, so every term is
+    // nonnegative; only a caller-supplied code can make the divisor nonpositive. A NaN
+    // passes: the per-vector encoders do not check that inputs are finite.
+    if (ip_resi_xucb <= 0) {
+        throw std::invalid_argument(
+            "code does not quantize residual: <residual, centered code> must be positive"
+        );
+    }
 
     // Cauchy-Schwarz makes this nonnegative; collinear vectors can round below zero.
     const T normalized_error =
diff --git a/src/core/provisional.cpp b/src/core/provisional.cpp
index 64a37f0..419c1dd 100644
--- a/src/core/provisional.cpp
+++ b/src/core/provisional.cpp
@@ -86,27 +86,20 @@ v1::Factors code_factors_impl(
     if (code == nullptr) {
         throw std::invalid_argument("code must not be null");
     }
-    const float squared_norm = validation::checked_residual_squared_norm(residual, dim);
+    (void)validation::checked_residual_squared_norm(residual, dim);
 
     // The detail function takes int codes and their offset-binary constant.
     const auto maximum = static_cast<int>((1U << total_bits) - 1);
     const float cb = -static_cast<float>(maximum) / 2.0F;
     std::vector<int> widened(dim);
-    float centered_ip = 0;
     for (size_t i = 0; i < dim; ++i) {
         widened[i] = static_cast<int>(code[i]);
         if (widened[i] > maximum) {
             throw std::invalid_argument("Unpacked code exceeds the configured bit width");
         }
-        centered_ip += residual[i] * (static_cast<float>(widened[i]) + cb);
-    }
-    // The factor derivation divides by this inner product; only a code that quantizes
-    // this residual (its own code or a prefix of it) makes it positive.
-    if (squared_norm != 0 && !(centered_ip > 0)) {
-        throw std::invalid_argument(
-            "code does not quantize residual: <residual, centered code> must be positive"
-        );
     }
+    // The factor derivation rejects a code that does not quantize this residual (its
+    // own code or a prefix of it), testing the inner product it divides by.
 
     std::vector<float> zero;
     if (centroid == nullptr) {
diff --git a/tests/unit/rabitqlib/quantization/factor_test.cpp b/tests/unit/rabitqlib/quantization/factor_test.cpp
index a297e49..184146b 100644
--- a/tests/unit/rabitqlib/quantization/factor_test.cpp
+++ b/tests/unit/rabitqlib/quantization/factor_test.cpp
@@ -89,6 +89,74 @@ TEST(RabitqFactorTest, CollinearMetricFactorsFloat) { check_collinear_factors<fl
 
 TEST(RabitqFactorTest, CollinearMetricFactorsDouble) { check_collinear_factors<double>(); }
 
+template <typename T>
+void check_rejected_foreign_code() {
+    constexpr size_t kDim = 64;
+    std::vector<T> residual(kDim), centroid(kDim, T{0});
+    std::vector<int> codes(kDim);
+    T add, rescale, error;
+    for (const auto metric : {METRIC_L2, METRIC_IP}) {
+        // A code of the opposite signs, then one orthogonal to the residual: the
+        // factor derivation divides by <residual, centered code>, which must be > 0.
+        for (const bool orthogonal : {false, true}) {
+            for (size_t i = 0; i < kDim; ++i) {
+                residual[i] = orthogonal && i % 2 == 1 ? T{-1} : T{1};
+                codes[i] = orthogonal ? 1 : 0;
+            }
+            try {
+                rabitq_impl::xy_bits::code_factors(
+                    residual.data(),
+                    centroid.data(),
+                    kDim,
+                    codes.data(),
+                    -0.5F,
+                    add,
+                    rescale,
+                    error,
+                    metric
+                );
+                FAIL() << "Expected rejection of a code that does not quantize residual";
+            } catch (const std::invalid_argument& exception) {
+                EXPECT_STREQ(
+                    exception.what(),
+                    "code does not quantize residual: <residual, centered code> must be "
+                    "positive"
+                );
+            }
+        }
+    }
+}
+
+TEST(RabitqFactorTest, NanInputIsNotCheckedByTheEncoders) {
+    // The per-vector encoders do not check that inputs are finite (quantizer.md); the
+    // rejection of a nonpositive divisor must not turn a NaN into an exception.
+    constexpr size_t kDim = 64;
+    std::vector<float> data(kDim, 0.5F), centroid(kDim, 0), residual(kDim), centered(kDim);
+    data[7] = std::numeric_limits<float>::quiet_NaN();
+    std::vector<int> codes(kDim);
+    float add, rescale, error;
+    rabitq_impl::one_bit::one_bit_code_with_factor_into(
+        data.data(),
+        centroid.data(),
+        kDim,
+        codes.data(),
+        add,
+        rescale,
+        error,
+        METRIC_L2,
+        residual.data(),
+        centered.data()
+    );
+    EXPECT_TRUE(std::isnan(add));
+    EXPECT_TRUE(std::isnan(rescale));
+}
+
+TEST(RabitqFactorTest, ForeignCodeIsRejectedFloat) { check_rejected_foreign_code<float>(); }
+
+TEST(RabitqFactorTest, ForeignCodeIsRejectedDouble) {
+    check_rejected_foreign_code<double>();
+}
+
 TEST(RabitqFactorTest, InvalidMetricMessagesArePreserved) {
     constexpr size_t kDim = 64;
     const auto invalid_metric = static_cast<MetricType>(2);
-- 
2.34.1

[gouyt13]

Patch 2 of 4: refactor(quant): share the v1 query terms and unpacked

0002-refactor-quant-share-the-v1-query-terms-and-unpacked.patch

From 7d05a71ed27ed5e57ef14d3b990be3f76aedf7cd Mon Sep 17 00:00:00 2001
From: gouyt13clear <yutong003@e.ntu.edu.sg>
Date: Tue, 3 Nov 2026 16:31:01 +0800
Subject: [PATCH 2/4] refactor(quant): share the v1 query terms and unpacked
 inner product

The bindings computed the query terms and the inner product of unpacked codes with
their own copies of quant::v1's arithmetic, so a later change to v1 would not
reach rabitqlib.quant and the two would drift apart silently.

The arithmetic now lives once in quant::detail (core/encoder.hpp, next to
full_est_dist_from_ip): query_corrections gives g_add and g_k1xsumq, and
unpacked_ip sums <query, code> in order and reports the first code above the bit
width. quant::v1::prepare_query and compute use them with unchanged arithmetic
order, checks and messages, and the bindings use them for compute_query_terms and
the unpacked scoring paths. A comment on the helpers states that a change to them
changes v1 and needs the v1 parity check against a pre-change build.

v1 encode, prepare_query and compute output is bitwise identical before and after
(2016 values over dims 64/96/128/130, total_bits 1-9, both metrics, with and
without a centroid, exact and fast quantization, both code layouts), and so are
the outputs of every rabitqlib.quant call on fixed inputs. The bindings compile
without FP contraction and api.cpp does not, so Python distances agree with v1's
within float32 rounding, as the comment on estimate_of says.
---
 include/rabitqlib/core/encoder.hpp | 46 ++++++++++++++++++++++++++
 python_bindings/quant_bindings.cpp | 52 +++++++++++++++++-------------
 src/core/api.cpp                   | 28 +++++++++-------
 3 files changed, 92 insertions(+), 34 deletions(-)

diff --git a/include/rabitqlib/core/encoder.hpp b/include/rabitqlib/core/encoder.hpp
index aa04205..2009aa8 100644
--- a/include/rabitqlib/core/encoder.hpp
+++ b/include/rabitqlib/core/encoder.hpp
@@ -2,6 +2,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <numeric>
 #include <optional>
 #include <vector>
 
@@ -11,6 +12,7 @@
 #include "rabitqlib/core/detail/rabitq_impl.hpp"
 #include "rabitqlib/core/metric.hpp"
 #include "rabitqlib/core/types.hpp"  // IWYU pragma: export
+#include "rabitqlib/utils/space.hpp"
 
 namespace rabitqlib::quant {
 
@@ -311,6 +313,50 @@ inline T full_est_dist_from_ip(
     return f_add + g_add +
            (f_rescale * (ip + g_k1xsumq * static_cast<float>((1 << bits) - 1)));
 }
+
+// query_corrections and unpacked_ip back the stable quant::v1 prepare_query and compute
+// as well as rabitqlib.quant. A change to either changes v1 results: validate v1
+// output against a saved pre-change build ("Stable core API changes" in
+// docs/development.md), not only the provisional layer's tests.
+
+// The query terms of full_est_dist_from_ip for a rotated query and a centroid (nullptr
+// means zero): g_add is |q - c|^2 for L2 and -<q, c> for IP; g_k1xsumq is
+// -0.5 * sum(q). Callers check that both are finite.
+struct QueryCorrections {
+    float g_add;
+    float g_k1xsumq;
+};
+
+inline QueryCorrections query_corrections(
+    const float* query, const float* centroid, size_t dim, MetricType metric
+) {
+    QueryCorrections corrections{};
+    corrections.g_k1xsumq = std::accumulate(query, query + dim, 0.0F) * -0.5F;
+    corrections.g_add =
+        metric == METRIC_L2
+            ? (centroid == nullptr ? l2norm_sqr(query, dim)
+                                   : euclidean_sqr(query, centroid, dim))
+            : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, dim));
+    return corrections;
+}
+
+// <query, code> for unpacked codes, summed in order: the reference inner product of
+// quant::v1::compute. code_at(i) returns the i-th code value. Returns dim, or the index
+// of the first code above maximum, in which case ip is incomplete.
+template <typename CodeAt>
+inline size_t unpacked_ip(
+    const float* query, size_t dim, uint32_t maximum, CodeAt code_at, float& ip
+) {
+    ip = 0;
+    for (size_t i = 0; i < dim; ++i) {
+        const uint32_t value = code_at(i);
+        if (value > maximum) {
+            return i;
+        }
+        ip += query[i] * static_cast<float>(value);
+    }
+    return dim;
+}
 }  // namespace detail
 
 template <typename T, typename TP>
diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp
index 4f4de1a..f135f53 100644
--- a/python_bindings/quant_bindings.cpp
+++ b/python_bindings/quant_bindings.cpp
@@ -18,7 +18,6 @@
 #include <initializer_list>
 #include <limits>
 #include <memory>
-#include <numeric>
 #include <optional>
 #include <random>
 #include <stdexcept>
@@ -193,6 +192,21 @@ void check_code_values(
     }
 }
 
+// <query, code> of row i for unpacked codes already checked against their bit width.
+template <typename TP>
+float unpacked_row_ip(const float* query, const RowView<TP>& codes, size_t i) {
+    const TP* row = codes.row(i);
+    float ip = 0;
+    (void)quant::detail::unpacked_ip(
+        query,
+        codes.cols,
+        std::numeric_limits<uint32_t>::max(),
+        [row](size_t at) { return row[at]; },
+        ip
+    );
+    return ip;
+}
+
 // A float32 vector copied into owned storage and checked like index inputs. Its length
 // must be expected_d, or with expected_d == 0 a positive multiple of 32.
 std::vector<float> owned_vector(py::handle value, const char* name, size_t expected_d) {
@@ -644,18 +658,19 @@ struct QueryTerms {
     float g_k1xsumq = 0;
 };
 
+// g_add and g_k1xsumq are quant::v1::prepare_query's; g_error is |q - c|.
 QueryTerms compute_query_terms(
     const float* query, const float* centroid, size_t d, MetricType metric
 ) {
+    const auto corrections = quant::detail::query_corrections(query, centroid, d, metric);
     const float distance_sqr =
-        centroid == nullptr ? l2norm_sqr(query, d) : euclidean_sqr(query, centroid, d);
+        metric == METRIC_L2 ? corrections.g_add
+                            : (centroid == nullptr ? l2norm_sqr(query, d)
+                                                   : euclidean_sqr(query, centroid, d));
     QueryTerms terms;
-    terms.g_add = metric == METRIC_L2
-                      ? distance_sqr
-                      : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, d));
+    terms.g_add = corrections.g_add;
     terms.g_error = std::sqrt(distance_sqr);
-    terms.g_k1xsumq =
-        std::accumulate(query, query + d, 0.0F) * fastscan::detail::code_center(1);
+    terms.g_k1xsumq = corrections.g_k1xsumq;
     if (!std::isfinite(terms.g_add) || !std::isfinite(terms.g_error) ||
         !std::isfinite(terms.g_k1xsumq)) {
         throw std::overflow_error("Query terms are not representable in float32");
@@ -663,6 +678,10 @@ QueryTerms compute_query_terms(
     return terms;
 }
 
+// The estimate of one row by the reference formula. Unpacked codes reach ip through
+// quant::detail::unpacked_ip, the inner product of quant::v1::compute. The arithmetic
+// is v1's, but this file is compiled without FP contraction and api.cpp is not, so
+// distances agree with v1's within float32 rounding rather than bitwise.
 inline Estimate estimate_of(
     float ip, size_t bits, const quant::v1::Factors& f, const QueryTerms& terms
 ) {
@@ -674,17 +693,6 @@ inline Estimate estimate_of(
         ip};
 }
 
-// <q, code> for unpacked codes, summed in order as quant::v1::compute does. This is
-// the portable reference path, not a SIMD kernel.
-template <typename TP>
-inline float unpacked_ip(const float* query, const TP* code, size_t d) {
-    float ip = 0;
-    for (size_t j = 0; j < d; ++j) {
-        ip += query[j] * static_cast<float>(code[j]);
-    }
-    return ip;
-}
-
 py::object query_terms(
     py::handle rotated_query, const std::string& metric, const py::object& centroid
 ) {
@@ -1037,7 +1045,7 @@ class QuantQuery {
                         0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) {
                             const size_t i = selected.row(k);
                             estimates[k] = estimate_of(
-                                unpacked_ip(query_.data(), view.row(i), d()),
+                                unpacked_row_ip(query_.data(), view, i),
                                 nbits,
                                 records[i],
                                 terms_
@@ -1121,9 +1129,9 @@ class QuantQuery {
                 index_detail::parallel_for(
                     0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) {
                         const size_t i = selected.row(k);
-                        const float ip_ex =
-                            Packed ? ip_func(query_.data(), codes.row(i), d())
-                                   : unpacked_ip(query_.data(), codes.row(i), d());
+                        const float ip_ex = Packed
+                                                ? ip_func(query_.data(), codes.row(i), d())
+                                                : unpacked_row_ip(query_.data(), codes, i);
                         const quant::v1::Factors& f = records[i];
                         distances[k] = quant::detail::full_est_dist_from_ip(
                             (scale * estimates[k].refine_input) + ip_ex,
diff --git a/src/core/api.cpp b/src/core/api.cpp
index 1af0c56..7efac9b 100644
--- a/src/core/api.cpp
+++ b/src/core/api.cpp
@@ -6,7 +6,6 @@
 #include <cstdint>
 #include <limits>
 #include <memory_resource>
-#include <numeric>
 #include <stdexcept>
 #include <utility>
 
@@ -165,11 +164,15 @@ uint16_t read_u16(const uint8_t* bytes) {
     return static_cast<uint16_t>(bytes[0] | (uint16_t{bytes[1]} << 8U));
 }
 
+[[noreturn]] void reject_unpacked_code() {
+    throw std::invalid_argument("Unpacked code exceeds the configured bit width");
+}
+
 // The i-th UnpackedU16 code, rejected if it does not fit the configured bit width.
 uint16_t checked_unpacked_code(const uint8_t* bytes, size_t i, uint16_t maximum) {
     const uint16_t value = read_u16(bytes + 2 * i);
     if (value > maximum) {
-        throw std::invalid_argument("Unpacked code exceeds the configured bit width");
+        reject_unpacked_code();
     }
     return value;
 }
@@ -362,12 +365,11 @@ QueryState prepare_query(
     state.total_bits_ = parameters.total_bits;
     state.metric_ = parameters.metric;
     state.extra_ip_ = select_excode_ipfunc(parameters.total_bits - 1);
-    state.g_k1xsumq_ = std::accumulate(query, query + parameters.dim, 0.0F) * -0.5F;
-    state.g_add_ =
-        parameters.metric == METRIC_L2
-            ? (centroid == nullptr ? l2norm_sqr(query, parameters.dim)
-                                   : euclidean_sqr(query, centroid, parameters.dim))
-            : (centroid == nullptr ? 0.0F : -dot_product(query, centroid, parameters.dim));
+    const auto corrections = quant::detail::query_corrections(
+        query, centroid, parameters.dim, parameters.metric
+    );
+    state.g_k1xsumq_ = corrections.g_k1xsumq;
+    state.g_add_ = corrections.g_add;
     if (!std::isfinite(state.g_add_) || !std::isfinite(state.g_k1xsumq_)) {
         throw std::overflow_error("Query factors are not representable in float32");
     }
@@ -392,10 +394,12 @@ float QueryState::compute_checked(
     float ip = 0;
     if (code.layout == CodeLayout::UnpackedU16) {
         const auto* bytes = static_cast<const uint8_t*>(code.codes.data);
-        const auto maximum = static_cast<uint16_t>((1U << parameters.total_bits) - 1);
-        for (size_t i = 0; i < parameters.dim; ++i) {
-            ip += query_wrapper.query_[i] *
-                  static_cast<float>(checked_unpacked_code(bytes, i, maximum));
+        const auto maximum = static_cast<uint32_t>((1U << parameters.total_bits) - 1);
+        const auto code_at = [bytes](size_t i) { return read_u16(bytes + (2 * i)); };
+        if (quant::detail::unpacked_ip(
+                query_wrapper.query_, parameters.dim, maximum, code_at, ip
+            ) != parameters.dim) {
+            reject_unpacked_code();
         }
     } else {
         const size_t ex_bits = parameters.total_bits - 1;
-- 
2.34.1

[gouyt13]

Patch 3 of 4: perf(quant): run row loops in ranges and check inputs in

0003-perf-quant-run-row-loops-in-ranges-and-check-inputs-.patch

From b608478d9aa66bdb0dbbe1c37bc3f788ae7857ce Mon Sep 17 00:00:00 2001
From: gouyt13clear <yutong003@e.ntu.edu.sg>
Date: Tue, 3 Nov 2026 16:36:17 +0800
Subject: [PATCH 3/4] perf(quant): run row loops in ranges and check inputs in
 parallel passes

index_detail::parallel_for hands out one index at a time from a shared counter.
That suits index queries, but a row of rabitqlib.quant costs nanoseconds to
microseconds, so the threads mostly contended for the counter: every row loop was
slower on 4 threads than on 1 (estimate_from_ip on 200k rows: 18.8 ms against
1.6 ms).

Row loops now go out in ranges, about eight per worker: at least 64 rows for the
cheap loops (scoring, packing, input checks) and at least one row for rotate,
quantize and code_factors, whose rows cost enough that a small batch should still
use every worker.

The input checks that ran serially before the parallel loop (code values, finite
factors, ip, refine_input, FastScan block factors) now run in a parallel pass of
their own, still before anything is written, so a rejected call leaves out=
unchanged. Factors of contiguous rows are tested once per range, and code values
with a vectorizable maximum. When several rows fail a check, the error reported is
the first failing row's in call order (the lowest row, or the earliest entry of
indices=), whatever the thread count; the guide says so. code_factors reports its
first rejected row without recomputing it.

pack_fastscan_blocks allocates its sign-byte scratch per range. It used to index a
per-worker buffer with the worker number of a separately resolved thread count,
which could run past the buffer if the two counts differed.

Outputs are unchanged: every rabitqlib.quant call returns bitwise identical arrays
on fixed inputs.

Benchmark: synthetic standard-normal data, 200k rows, d=128 (codes at nbits 1 and
4), plus 100 rows at d=4096; Xeon Gold 6418H, GCC 14, Release, runtime dispatch
(AVX-512); median of 3 interleaved rounds of 7 timed calls each, pinned to 8 idle
cores of one NUMA node on a shared host; PR head -> this commit, ms:

  call                     1 thread         8 threads
  estimate_bitset          2.40 -> 2.33     24.68 -> 1.00
  estimate_codes (nbits 4) 27.48 -> 22.13   34.80 -> 3.43
  refine (ex_bits 3)       25.79 -> 22.28   34.97 -> 3.40
  estimate_from_ip         1.61 -> 1.27     20.68 -> 1.47
  estimate_fastscan        1.65 -> 1.80     2.57 -> 0.96
  pack_codes (nbits 1)     15.88 -> 9.67    24.63 -> 2.02
  code_factors (nbits 4)   67.54 -> 57.08   45.72 -> 30.65
  rotate                   102.20 -> 96.27  86.18 -> 36.85
  quantize, 100 x 4096, 9b 297.87 -> 299.82 46.42 -> 48.09
---
 docs/docs/rabitq/quant_python.md   |   2 +-
 python_bindings/quant_bindings.cpp | 496 +++++++++++++++++++----------
 tests/python/test_quant_layer1.py  |  22 ++
 tests/python/test_quant_layer2.py  | 103 ++++++
 tests/python/test_quant_layer3.py  |  19 ++
 5 files changed, 479 insertions(+), 163 deletions(-)

diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md
index bc29c10..c5fc0ee 100644
--- a/docs/docs/rabitq/quant_python.md
+++ b/docs/docs/rabitq/quant_python.md
@@ -115,7 +115,7 @@ If your `ip` is itself approximate, `error` covers only the code's quantization
 - The GIL is released during computation. `Rotator` is immutable and safe to share. Scoring calls on one `Query` may run concurrently; a `set_centroid` during one raises `RuntimeError`. Use `qry.copy()` for one query per thread.
 - `Query` supports `copy.copy`, `copy.deepcopy` and pickling of its query vector, metric and centroid; a copy never shares storage with the original. Queries have no `to_bytes`, because they are not a stored format.
 - Every call that returns an array accepts `out=` with the exact dtype and shape. It must not share memory with any input, or the call raises `ValueError`. If a call raises, a caller's `out=` is left unchanged.
-- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them.
+- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. Inputs are checked before anything is computed; within each check, the error reported is that of the first failing row in call order (the lowest row, or the earliest entry of `indices=`), whatever the thread count.
 
 ## Byte formats
 
diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp
index f135f53..22db653 100644
--- a/python_bindings/quant_bindings.cpp
+++ b/python_bindings/quant_bindings.cpp
@@ -15,9 +15,11 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <exception>
 #include <initializer_list>
 #include <limits>
 #include <memory>
+#include <mutex>
 #include <optional>
 #include <random>
 #include <stdexcept>
@@ -172,23 +174,38 @@ Selection selection(const py::object& indices, size_t rows) {
 }
 
 template <typename TP>
-void check_code_values(
-    const RowView<TP>& codes, const Selection& selected, size_t nbits, const char* name
+[[noreturn]] void reject_code_value(
+    const RowView<TP>& codes, size_t i, size_t j, size_t nbits, const char* name
 ) {
+    throw std::invalid_argument(
+        std::string(name) + " values must be below 2**" + std::to_string(nbits) + " = " +
+        std::to_string(size_t{1} << nbits) + "; got " + std::to_string(codes.row(i)[j]) +
+        " at (" + std::to_string(i) + ", " + std::to_string(j) + ")"
+    );
+}
+
+template <typename TP>
+void check_code_row(const RowView<TP>& codes, size_t i, size_t nbits, const char* name) {
     const auto maximum = static_cast<TP>((1U << nbits) - 1);
-    for (size_t k = 0; k < selected.count; ++k) {
-        const size_t i = selected.row(k);
-        const TP* row = codes.row(i);
-        for (size_t j = 0; j < codes.cols; ++j) {
-            if (row[j] > maximum) {
-                throw std::invalid_argument(
-                    std::string(name) + " values must be below 2**" +
-                    std::to_string(nbits) + " = " + std::to_string(size_t{1} << nbits) +
-                    "; got " + std::to_string(row[j]) + " at (" + std::to_string(i) + ", " +
-                    std::to_string(j) + ")"
-                );
-            }
-        }
+    const TP* row = codes.row(i);
+    // A vectorizable maximum first; the position is searched only for a rejected row.
+    TP largest = 0;
+    for (size_t j = 0; j < codes.cols; ++j) {
+        largest = std::max(largest, row[j]);
+    }
+    if (largest > maximum) {
+        reject_code_value(
+            codes,
+            i,
+            static_cast<size_t>(
+                std::find_if(
+                    row, row + codes.cols, [&](TP value) { return value > maximum; }
+                ) -
+                row
+            ),
+            nbits,
+            name
+        );
     }
 }
 
@@ -207,6 +224,12 @@ float unpacked_row_ip(const float* query, const RowView<TP>& codes, size_t i) {
     return ip;
 }
 
+void check_finite_factor(const quant::v1::Factors& f) {
+    if (!std::isfinite(f.add) || !std::isfinite(f.rescale) || !std::isfinite(f.error)) {
+        throw std::invalid_argument("factors must contain only finite values");
+    }
+}
+
 // A float32 vector copied into owned storage and checked like index inputs. Its length
 // must be expected_d, or with expected_d == 0 a positive multiple of 32.
 std::vector<float> owned_vector(py::handle value, const char* name, size_t expected_d) {
@@ -401,16 +424,137 @@ Records<quant::v1::Factors> factors_input(py::handle value, size_t rows, const c
     return records;
 }
 
-// Factors are read only for the selected rows; call with the GIL released.
-void check_finite_factors(
-    const Records<quant::v1::Factors>& factors, const Selection& selected
-) {
-    for (size_t k = 0; k < selected.count; ++k) {
-        const quant::v1::Factors& f = factors[selected.row(k)];
-        if (!std::isfinite(f.add) || !std::isfinite(f.rescale) || !std::isfinite(f.error)) {
-            throw std::invalid_argument("factors must contain only finite values");
+// --- Parallel rows ----------------------------------------------------------------
+
+// parallel_for hands out one index at a time from a shared counter, which costs more
+// than a cheap row. Rows go out in ranges instead, about eight per worker and at least
+// min_rows long: kCheapRows for rows of up to a few hundred nanoseconds, such as
+// scoring, packing and checks, and kCostlyRows for rotate, quantize and code_factors,
+// whose rows cost enough that a small batch should still use every worker.
+constexpr size_t kCheapRows = 64;
+constexpr size_t kCostlyRows = 1;
+
+// Runs task(begin, end) on consecutive ranges covering [0, rows). Call with the GIL
+// released.
+template <typename Task>
+void parallel_rows(size_t rows, size_t num_threads, size_t min_rows, Task task) {
+    if (rows == 0) {
+        return;
+    }
+    const size_t workers = resolve_num_threads(num_threads);
+    const size_t task_rows = std::max(min_rows, div_round_up(rows, workers * 8));
+    index_detail::parallel_for(
+        0, div_round_up(rows, task_rows), workers, [&](size_t t, size_t /*worker*/) {
+            const size_t begin = t * task_rows;
+            task(begin, std::min(rows, begin + task_rows));
+        }
+    );
+}
+
+// The exception of the lowest failing row of a parallel loop, so the error a call
+// raises does not depend on the thread schedule. Rows after a recorded failure are
+// skipped.
+class FirstFailure {
+   public:
+    explicit FirstFailure(size_t rows) : row_(rows) {}
+
+    // A stale value only delays the skip; record() decides under the lock.
+    [[nodiscard]] bool after(size_t row) const {
+        return row > row_.load(std::memory_order_relaxed);
+    }
+
+    // Call from the catch block of the failing row.
+    void record(size_t row) {
+        const std::lock_guard<std::mutex> lock(mutex_);
+        if (row < row_.load()) {
+            row_.store(row);
+            error_ = std::current_exception();
         }
     }
+
+    void rethrow() const {
+        if (error_) {
+            std::rethrow_exception(error_);
+        }
+    }
+
+   private:
+    std::atomic<size_t> row_;
+    std::mutex mutex_;
+    std::exception_ptr error_;
+};
+
+// Runs row(i) for every i in [0, rows) and rethrows the lowest failing row's exception.
+// range(begin, end) first handles a whole range without per-row checks, so its loops
+// can vectorize, and returns false if any row in it fails; only such a range runs
+// through row(i), whose checks find that row. range must not throw.
+template <typename Range, typename Row>
+void for_each_range(
+    size_t rows, size_t num_threads, size_t min_rows, Range range, Row row
+) {
+    FirstFailure failure(rows);
+    parallel_rows(rows, num_threads, min_rows, [&](size_t begin, size_t end) {
+        if (failure.after(begin) || range(begin, end)) {
+            return;
+        }
+        for (size_t i = begin; i < end && !failure.after(i); ++i) {
+            try {
+                row(i);
+            } catch (...) {
+                failure.record(i);
+                return;
+            }
+        }
+    });
+    failure.rethrow();
+}
+
+template <typename Row>
+void for_each_row(size_t rows, size_t num_threads, size_t min_rows, Row row) {
+    for_each_range(
+        rows,
+        num_threads,
+        min_rows,
+        [](size_t /*begin*/, size_t /*end*/) { return false; },
+        row
+    );
+}
+
+template <typename TP>
+void check_code_values(
+    const RowView<TP>& codes, size_t nbits, const char* name, size_t num_threads
+) {
+    for_each_row(codes.rows, num_threads, kCheapRows, [&](size_t i) {
+        check_code_row(codes, i, nbits, name);
+    });
+}
+
+// Whether every float of records [begin, end) is finite.
+template <typename T>
+bool all_finite_records(const T* records, size_t begin, size_t end) {
+    static_assert(sizeof(T) == 3 * sizeof(float));
+    return all_finite(reinterpret_cast<const float*>(records + begin), 3 * (end - begin));
+}
+
+// A vectorized test that the selected rows [begin, end) have finite factors, for every
+// row selected; false means only that the rows must be checked one by one.
+bool finite_factor_range(
+    const Records<quant::v1::Factors>& factors,
+    const Selection& selected,
+    size_t begin,
+    size_t end
+) {
+    return selected.indices == nullptr && all_finite_records(factors.data, begin, end);
+}
+
+// Writes value_at(k) to out[k] for k in [0, count), for rows already validated.
+template <typename T, typename ValueAt>
+void write_rows(T* out, size_t count, size_t num_threads, ValueAt value_at) {
+    parallel_rows(count, num_threads, kCheapRows, [&](size_t begin, size_t end) {
+        for (size_t k = begin; k < end; ++k) {
+            out[k] = value_at(k);
+        }
+    });
 }
 
 // --- Layer 1 --------------------------------------------------------------------
@@ -455,13 +599,13 @@ py::array quantize(
         {
             py::gil_scoped_release release;
             check_coordinates(input.data(), rows * d, d, "residuals");
-            index_detail::parallel_for(
-                0, rows, num_threads, [&](size_t i, size_t /*thread*/) {
+            parallel_rows(rows, num_threads, kCostlyRows, [&](size_t begin, size_t end) {
+                for (size_t i = begin; i < end; ++i) {
                     quant::provisional::combined_code(
                         input.data() + (i * d), d, nbits, view.row(i), t_const
                     );
                 }
-            );
+            });
         }
         return view.owner;
     };
@@ -479,10 +623,9 @@ py::tuple split_rows(const RowView<TP>& codes, size_t nbits, size_t prefix_bits)
     auto ex = new_rows<uint8_t>(codes.rows, codes.cols);
     uint8_t* prefix_data = prefix.mutable_data();
     uint8_t* ex_data = ex.mutable_data();
-    const Selection every_row = all_rows(codes.rows);
     {
         py::gil_scoped_release release;
-        check_code_values(codes, every_row, nbits, "codes");
+        check_code_values(codes, nbits, "codes", 1);
         for (size_t i = 0; i < codes.rows; ++i) {
             const TP* row = codes.row(i);
             for (size_t j = 0; j < codes.cols; ++j) {
@@ -526,39 +669,34 @@ py::array code_factors_rows(
         rows,
         {{"residuals", residuals}, {"codes", codes.owner}}
     );
+    const bool caller_out = !out.is_none();
     {
         py::gil_scoped_release release;
         check_coordinates(residuals.data(), rows * d, d, "residuals");
-        // A caller's out= is written only after every row succeeds.
-        std::vector<quant::v1::Factors> staged(out.is_none() ? 0 : rows);
-        quant::v1::Factors* factors = out.is_none() ? data : staged.data();
-        // code_factors checks each row's code values and that the code quantizes the
-        // residual; the binding adds the row number to its message.
-        const auto factors_of = [&](size_t i) {
-            return quant::provisional::code_factors(
-                residuals.data() + (i * d), centroid.data(), d, codes.row(i), nbits, metric
-            );
-        };
-        // Report the lowest rejected row, whatever the thread schedule.
-        std::atomic<size_t> first_rejected{rows};
-        index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) {
-            try {
-                factors[i] = factors_of(i);
-            } catch (const std::invalid_argument&) {
-                size_t current = first_rejected.load();
-                while (i < current && !first_rejected.compare_exchange_weak(current, i)) {}
-            }
-        });
-        if (const size_t i = first_rejected.load(); i < rows) {
+        // code_factors checks while it computes, so a caller's out= is written only after
+        // every row succeeds.
+        std::vector<quant::v1::Factors> staged(caller_out ? rows : 0);
+        quant::v1::Factors* factors = caller_out ? staged.data() : data;
+        for_each_row(rows, num_threads, kCostlyRows, [&](size_t i) {
+            // code_factors checks the row's code values and that the code quantizes the
+            // residual; the binding adds the row number to its message.
+            const auto row = [i] {
+                return "codes row " + std::to_string(i) + " for residuals row " +
+                       std::to_string(i) + ": ";
+            };
             try {
-                (void)factors_of(i);
-            } catch (const std::invalid_argument& error) {
-                throw std::invalid_argument(
-                    "codes row " + std::to_string(i) + " for residuals row " +
-                    std::to_string(i) + ": " + error.what()
+                factors[i] = quant::provisional::code_factors(
+                    residuals.data() + (i * d),
+                    centroid.data(),
+                    d,
+                    codes.row(i),
+                    nbits,
+                    metric
                 );
+            } catch (const std::invalid_argument& error) {
+                throw std::invalid_argument(row() + error.what());
             }
-        }
+        });
         std::copy(staged.begin(), staged.end(), data);
     }
     return result;
@@ -606,17 +744,17 @@ py::array pack_codes(
     const size_t d = checked_working_dim(view.cols, view.owner, "codes");
     const auto packed =
         output_rows<uint8_t>(out, view.rows, d * nbits / 8, {{"codes", view.owner}});
-    const Selection every_row = all_rows(view.rows);
     {
         py::gil_scoped_release release;
-        check_code_values(view, every_row, nbits, "codes");
-        index_detail::parallel_for(
-            0, view.rows, num_threads, [&](size_t i, size_t /*thread*/) {
+        // Checked in a pass of its own, so a rejected call leaves out= unchanged.
+        check_code_values(view, nbits, "codes", num_threads);
+        parallel_rows(view.rows, num_threads, kCheapRows, [&](size_t begin, size_t end) {
+            for (size_t i = begin; i < end; ++i) {
                 quant::rabitq_impl::ex_bits::packing_rabitqplus_code(
                     view.row(i), packed.row(i), d, nbits
                 );
             }
-        );
+        });
     }
     return packed.owner;
 }
@@ -638,13 +776,13 @@ py::array unpack_codes(
     const auto codes = output_rows<uint8_t>(out, view.rows, d, {{"packed", view.owner}});
     {
         py::gil_scoped_release release;
-        index_detail::parallel_for(
-            0, view.rows, num_threads, [&](size_t i, size_t /*thread*/) {
+        parallel_rows(view.rows, num_threads, kCheapRows, [&](size_t begin, size_t end) {
+            for (size_t i = begin; i < end; ++i) {
                 quant::rabitq_impl::ex_bits::unpacking_rabitqplus_code(
                     view.row(i), codes.row(i), d, nbits
                 );
             }
-        );
+        });
     }
     return codes.owner;
 }
@@ -745,12 +883,23 @@ py::array estimate_from_ip(
     const Selection every_row = all_rows(rows);
     {
         py::gil_scoped_release release;
-        if (!all_finite(ips.data(), rows)) {
-            throw std::invalid_argument("ip must contain only finite values");
-        }
-        check_finite_factors(records, every_row);
-        index_detail::parallel_for(0, rows, num_threads, [&](size_t i, size_t /*thread*/) {
-            estimates[i] = estimate_of(ips.data()[i], nbits, records[i], query);
+        for_each_range(
+            rows,
+            num_threads,
+            kCheapRows,
+            [&](size_t begin, size_t end) {
+                return all_finite(ips.data() + begin, end - begin) &&
+                       finite_factor_range(records, every_row, begin, end);
+            },
+            [&](size_t i) {
+                if (!std::isfinite(ips.data()[i])) {
+                    throw std::invalid_argument("ip must contain only finite values");
+                }
+                check_finite_factor(records[i]);
+            }
+        );
+        write_rows(estimates, rows, num_threads, [&](size_t i) {
+            return estimate_of(ips.data()[i], nbits, records[i], query);
         });
     }
     return result;
@@ -797,18 +946,23 @@ py::array pack_fastscan_blocks(
     const Selection every_row = all_rows(n);
     {
         py::gil_scoped_release release;
-        check_finite_factors(records, every_row);
-        // One block of sign bytes per worker, in the order the FastScan packer reads:
-        // the first coordinate of each byte in its highest bit.
-        const size_t workers =
-            std::max<size_t>(1, std::min(resolve_num_threads(num_threads), num_blocks));
-        std::vector<uint8_t> scratch(workers * fastscan::kBatchSize * codes.cols);
-        index_detail::parallel_for(
-            0, num_blocks, num_threads, [&](size_t b, size_t thread) {
+        // Checked in a pass of its own, so a rejected call leaves out= unchanged.
+        for_each_range(
+            n,
+            num_threads,
+            kCheapRows,
+            [&](size_t begin, size_t end) {
+                return finite_factor_range(records, every_row, begin, end);
+            },
+            [&](size_t i) { check_finite_factor(records[i]); }
+        );
+        parallel_rows(num_blocks, num_threads, kCheapRows, [&](size_t first, size_t last) {
+            // One block of sign bytes in the order the FastScan packer reads: the first
+            // coordinate of each byte in its highest bit.
+            std::vector<uint8_t> msb_first(fastscan::kBatchSize * codes.cols);
+            for (size_t b = first; b < last; ++b) {
                 const size_t begin = b * fastscan::kBatchSize;
                 const size_t lanes = std::min(fastscan::kBatchSize, n - begin);
-                uint8_t* msb_first =
-                    scratch.data() + (thread * fastscan::kBatchSize * codes.cols);
                 for (size_t lane = 0; lane < lanes; ++lane) {
                     const uint8_t* row = codes.row(begin + lane);
                     for (size_t byte = 0; byte < codes.cols; ++byte) {
@@ -816,7 +970,7 @@ py::array pack_fastscan_blocks(
                     }
                 }
                 uint8_t* block = blocks.row(b);
-                fastscan::pack_codes(d, msb_first, lanes, block);
+                fastscan::pack_codes(d, msb_first.data(), lanes, block);
                 BatchDataMap<float> map(reinterpret_cast<char*>(block), d);
                 for (size_t lane = 0; lane < fastscan::kBatchSize; ++lane) {
                     const quant::v1::Factors f =
@@ -826,7 +980,7 @@ py::array pack_fastscan_blocks(
                     map.f_error()[lane] = f.error;
                 }
             }
-        );
+        });
     }
     return blocks.owner;
 }
@@ -1002,15 +1156,21 @@ class QuantQuery {
             const auto ip_func = select_excode_ipfunc(1);
             {
                 py::gil_scoped_release release;
-                check_finite_factors(records, selected);
-                index_detail::parallel_for(
-                    0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) {
-                        const size_t i = selected.row(k);
-                        estimates[k] = estimate_of(
-                            ip_func(query_.data(), codes.row(i), d()), 1, records[i], terms_
-                        );
-                    }
+                for_each_range(
+                    selected.count,
+                    num_threads,
+                    kCheapRows,
+                    [&](size_t begin, size_t end) {
+                        return finite_factor_range(records, selected, begin, end);
+                    },
+                    [&](size_t k) { check_finite_factor(records[selected.row(k)]); }
                 );
+                write_rows(estimates, selected.count, num_threads, [&](size_t k) {
+                    const size_t i = selected.row(k);
+                    return estimate_of(
+                        ip_func(query_.data(), codes.row(i), d()), 1, records[i], terms_
+                    );
+                });
             }
             return result;
         });
@@ -1039,19 +1199,20 @@ class QuantQuery {
                 );
                 {
                     py::gil_scoped_release release;
-                    check_code_values(view, selected, nbits, "codes");
-                    check_finite_factors(records, selected);
-                    index_detail::parallel_for(
-                        0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) {
-                            const size_t i = selected.row(k);
-                            estimates[k] = estimate_of(
-                                unpacked_row_ip(query_.data(), view, i),
-                                nbits,
-                                records[i],
-                                terms_
-                            );
-                        }
-                    );
+                    for_each_row(selected.count, num_threads, kCheapRows, [&](size_t k) {
+                        const size_t i = selected.row(k);
+                        check_code_row(view, i, nbits, "codes");
+                        check_finite_factor(records[i]);
+                    });
+                    write_rows(estimates, selected.count, num_threads, [&](size_t k) {
+                        const size_t i = selected.row(k);
+                        return estimate_of(
+                            unpacked_row_ip(query_.data(), view, i),
+                            nbits,
+                            records[i],
+                            terms_
+                        );
+                    });
                 }
                 return result;
             };
@@ -1115,34 +1276,32 @@ class QuantQuery {
             const auto scale = static_cast<float>(1U << ex_bits);
             {
                 py::gil_scoped_release release;
-                if constexpr (!Packed) {
-                    check_code_values(codes, selected, ex_bits, name);
-                }
-                check_finite_factors(records, selected);
-                for (size_t k = 0; k < estimates.rows; ++k) {
+                for_each_row(selected.count, num_threads, kCheapRows, [&](size_t k) {
+                    const size_t i = selected.row(k);
+                    if constexpr (!Packed) {
+                        check_code_row(codes, i, ex_bits, name);
+                    }
+                    check_finite_factor(records[i]);
                     if (!std::isfinite(estimates[k].refine_input)) {
                         throw std::invalid_argument(
                             "first['refine_input'] must contain only finite values"
                         );
                     }
-                }
-                index_detail::parallel_for(
-                    0, selected.count, num_threads, [&](size_t k, size_t /*thread*/) {
-                        const size_t i = selected.row(k);
-                        const float ip_ex = Packed
-                                                ? ip_func(query_.data(), codes.row(i), d())
-                                                : unpacked_row_ip(query_.data(), codes, i);
-                        const quant::v1::Factors& f = records[i];
-                        distances[k] = quant::detail::full_est_dist_from_ip(
-                            (scale * estimates[k].refine_input) + ip_ex,
-                            total_bits,
-                            f.add,
-                            f.rescale,
-                            terms_.g_add,
-                            terms_.g_k1xsumq
-                        );
-                    }
-                );
+                });
+                write_rows(distances, selected.count, num_threads, [&](size_t k) {
+                    const size_t i = selected.row(k);
+                    const float ip_ex = Packed ? ip_func(query_.data(), codes.row(i), d())
+                                               : unpacked_row_ip(query_.data(), codes, i);
+                    const quant::v1::Factors& f = records[i];
+                    return quant::detail::full_est_dist_from_ip(
+                        (scale * estimates[k].refine_input) + ip_ex,
+                        total_bits,
+                        f.add,
+                        f.rescale,
+                        terms_.g_add,
+                        terms_.g_k1xsumq
+                    );
+                });
             }
             return result;
         });
@@ -1168,51 +1327,64 @@ class QuantQuery {
             const SplitBatchQuery<float>& query = batch_query();
             {
                 py::gil_scoped_release release;
+                const auto lanes_of = [&](size_t b) {
+                    return std::min(
+                        fastscan::kBatchSize, num_points - (b * fastscan::kBatchSize)
+                    );
+                };
                 // Hand-built blocks get the finite-factor check of the other calls.
-                for (size_t b = 0; b < num_blocks; ++b) {
+                const auto finite_factors = [&](size_t b) {
                     const ConstBatchDataMap<float> map(
                         reinterpret_cast<const char*>(view.row(b)), d()
                     );
                     std::array<float, fastscan::kBatchSize> values{};
-                    const size_t lanes = std::min(
-                        fastscan::kBatchSize, num_points - (b * fastscan::kBatchSize)
-                    );
                     for (const auto& field :
                          {map.f_add(), map.f_rescale(), map.f_error()}) {
                         field.copy_to(values.data(), values.size());
-                        if (!all_finite(values.data(), lanes)) {
-                            throw std::invalid_argument(
-                                "blocks must contain only finite factors; block " +
-                                std::to_string(b) + " does not"
-                            );
+                        if (!all_finite(values.data(), lanes_of(b))) {
+                            return false;
                         }
                     }
-                }
-                index_detail::parallel_for(
-                    0, num_blocks, num_threads, [&](size_t b, size_t /*thread*/) {
-                        std::array<float, fastscan::kBatchSize> distances{};
-                        std::array<float, fastscan::kBatchSize> lower_bounds{};
-                        std::array<float, fastscan::kBatchSize> ips{};
-                        std::array<float, fastscan::kBatchSize> errors{};
-                        const auto* block = reinterpret_cast<const char*>(view.row(b));
-                        split_batch_estdist(
-                            block,
-                            query,
-                            d(),
-                            distances.data(),
-                            lower_bounds.data(),
-                            ips.data(),
-                            true
+                    return true;
+                };
+                for_each_row(num_blocks, num_threads, kCheapRows, [&](size_t b) {
+                    if (!finite_factors(b)) {
+                        throw std::invalid_argument(
+                            "blocks must contain only finite factors; block " +
+                            std::to_string(b) + " does not"
                         );
-                        ConstBatchDataMap<float>(block, d())
-                            .f_error()
-                            .copy_to(errors.data(), errors.size());
-                        const size_t begin = b * fastscan::kBatchSize;
-                        const size_t lanes =
-                            std::min(fastscan::kBatchSize, num_points - begin);
-                        for (size_t lane = 0; lane < lanes; ++lane) {
-                            estimates[begin + lane] = {
-                                distances[lane], errors[lane] * terms_.g_error, ips[lane]};
+                    }
+                });
+                // Writes the estimates of block b.
+                const auto score_block = [&](size_t b) {
+                    const auto* block = reinterpret_cast<const char*>(view.row(b));
+                    std::array<float, fastscan::kBatchSize> distances{};
+                    std::array<float, fastscan::kBatchSize> lower_bounds{};
+                    std::array<float, fastscan::kBatchSize> ips{};
+                    std::array<float, fastscan::kBatchSize> errors{};
+                    ConstBatchDataMap<float>(block, d())
+                        .f_error()
+                        .copy_to(errors.data(), errors.size());
+                    split_batch_estdist(
+                        block,
+                        query,
+                        d(),
+                        distances.data(),
+                        lower_bounds.data(),
+                        ips.data(),
+                        true
+                    );
+                    const size_t begin = b * fastscan::kBatchSize;
+                    const size_t lanes = lanes_of(b);
+                    for (size_t lane = 0; lane < lanes; ++lane) {
+                        estimates[begin + lane] = {
+                            distances[lane], errors[lane] * terms_.g_error, ips[lane]};
+                    }
+                };
+                parallel_rows(
+                    num_blocks, num_threads, kCheapRows, [&](size_t first, size_t last) {
+                        for (size_t b = first; b < last; ++b) {
+                            score_block(b);
                         }
                     }
                 );
@@ -1412,11 +1584,11 @@ class QuantRotator {
         {
             py::gil_scoped_release release;
             check_coordinates(input.data(), rows * dim_, dim_, "x");
-            index_detail::parallel_for(
-                0, rows, num_threads, [&](size_t i, size_t /*thread*/) {
+            parallel_rows(rows, num_threads, kCostlyRows, [&](size_t begin, size_t end) {
+                for (size_t i = begin; i < end; ++i) {
                     rotator_->rotate(input.data() + (i * dim_), view.row(i));
                 }
-            );
+            });
         }
         return view.owner;
     }
diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py
index 178319c..e2b51f5 100644
--- a/tests/python/test_quant_layer1.py
+++ b/tests/python/test_quant_layer1.py
@@ -296,3 +296,25 @@ def test_code_factors_leaves_out_unchanged_on_rejection():
     with pytest.raises(ValueError, match="codes row 4 for residuals row 4"):
         quant.code_factors(x, codes, nbits=4, metric="l2", out=out)
     assert not out.view(np.float32).any()
+
+
+def test_rejected_codes_leave_pack_out_unchanged():
+    # The codes are checked in a pass of their own, before out= is written.
+    codes = quant.quantize(residuals(300, seed=11), nbits=4)
+    codes[200, 7] = 16
+    out = np.zeros((300, D * 4 // 8), np.uint8)
+    with pytest.raises(ValueError, match=r"got 16 at \(200, 7\)$"):
+        quant.pack_codes(codes, nbits=4, num_threads=4, out=out)
+    assert not out.any()
+
+
+def test_code_factors_names_the_lowest_rejected_row_across_threads():
+    # 1000 rows span several parallel ranges; the lowest rejected row is reported.
+    x = residuals(1000, seed=10)
+    codes = quant.quantize(x, nbits=4)
+    codes[[700, 300]] = 15 - codes[[700, 300]]
+    for threads in (1, 4):
+        with pytest.raises(
+            ValueError, match="^codes row 300 for residuals row 300: code does not"
+        ):
+            quant.code_factors(x, codes, nbits=4, metric="l2", num_threads=threads)
diff --git a/tests/python/test_quant_layer2.py b/tests/python/test_quant_layer2.py
index 0c94e7f..c34add8 100644
--- a/tests/python/test_quant_layer2.py
+++ b/tests/python/test_quant_layer2.py
@@ -414,6 +414,109 @@ def test_out_overlapping_an_input_is_rejected():
         )
 
 
+def test_rejections_name_the_lowest_failing_row_across_threads():
+    _, c, residuals, codes, factors, q = setup("l2", 4, n=1000)
+    query = quant.Query(q, metric="l2", centroid=c)
+    bad = codes.copy()
+    bad[700, 5] = 16
+    bad[300, 9] = 16
+    prefix, _ = quant.split(codes, nbits=4, prefix_bits=1)
+    bitsets = quant.pack_codes(prefix, nbits=1)
+    broken = quant.code_factors(residuals, prefix, nbits=1, metric="l2", centroid=c)
+    broken["error"][[800, 500]] = np.nan
+    out = np.zeros(1000, quant.ESTIMATE_DTYPE)
+    for threads in (1, 4):
+        with pytest.raises(
+            ValueError,
+            match=r"^codes values must be below 2\*\*4 = 16; got 16 at \(300, 9\)$",
+        ):
+            query.estimate_codes(bad, factors, nbits=4, num_threads=threads, out=out)
+        # The vectorized range check fails and the row checks name the cause.
+        with pytest.raises(
+            ValueError, match="^factors must contain only finite values$"
+        ):
+            query.estimate_bitset(bitsets, broken, num_threads=threads, out=out)
+        # A rejected input leaves out= unchanged.
+        assert not out.view(np.float32).any()
+    # With indices=, the first failing entry is reported, not the lowest row.
+    with pytest.raises(ValueError, match=r"got 16 at \(700, 5\)$"):
+        query.estimate_codes(bad, factors, nbits=4, indices=[700, 300])
+
+
+def test_rejected_inputs_leave_out_unchanged():
+    # Inputs are checked in a pass of their own, before out= is written; the failing
+    # rows lie past the first parallel range.
+    _, c, _, codes, factors, q = setup("l2", 4, n=300)
+    query = quant.Query(q, metric="l2", centroid=c)
+    _, ex = quant.split(codes, nbits=4, prefix_bits=1)
+    first = query.estimate_codes(codes, factors, nbits=4)
+    terms = quant.query_terms(q, metric="l2", centroid=c)
+    broken = factors.copy()
+    broken["add"][200] = np.nan
+    bad_ex = ex.copy()
+    bad_ex[100, 3] = 8
+    bad_first = first.copy()
+    bad_first["refine_input"][50] = np.nan
+    bad_ip = first["refine_input"].copy()
+    bad_ip[150] = np.inf
+    distances = np.zeros(300, np.float32)
+    estimates = np.zeros(300, quant.ESTIMATE_DTYPE)
+    finite = "factors must contain only finite values"
+    cases = [
+        (
+            lambda: query.refine(
+                bad_ex, factors, first, ex_bits=3, num_threads=4, out=distances
+            ),
+            r"ex values must be below 2\*\*3 = 8; got 8 at \(100, 3\)",
+        ),
+        (
+            lambda: query.refine(
+                ex, broken, first, ex_bits=3, num_threads=4, out=distances
+            ),
+            finite,
+        ),
+        (
+            lambda: query.refine(
+                ex, factors, bad_first, ex_bits=3, num_threads=4, out=distances
+            ),
+            r"first\['refine_input'\] must contain only finite values",
+        ),
+        (
+            lambda: query.refine_packed(
+                quant.pack_codes(ex, nbits=3),
+                broken,
+                first,
+                ex_bits=3,
+                num_threads=4,
+                out=distances,
+            ),
+            finite,
+        ),
+        (
+            lambda: quant.estimate_from_ip(
+                bad_ip, factors, terms, nbits=4, num_threads=4, out=estimates
+            ),
+            "ip must contain only finite values",
+        ),
+        (
+            lambda: quant.estimate_from_ip(
+                first["refine_input"],
+                broken,
+                terms,
+                nbits=4,
+                num_threads=4,
+                out=estimates,
+            ),
+            finite,
+        ),
+    ]
+    for call, message in cases:
+        with pytest.raises(ValueError, match=message):
+            call()
+        assert not distances.any()
+        assert not estimates.view(np.float32).any()
+
+
 def test_corrupt_pickle_state_is_rejected():
     query = quant.Query.__new__(quant.Query)
     with pytest.raises(ValueError, match="Invalid Query pickle state"):
diff --git a/tests/python/test_quant_layer3.py b/tests/python/test_quant_layer3.py
index 133fd5c..8083570 100644
--- a/tests/python/test_quant_layer3.py
+++ b/tests/python/test_quant_layer3.py
@@ -190,6 +190,25 @@ def test_hand_built_blocks_with_non_finite_factors_are_rejected():
     query.estimate_fastscan(blocks, num_points=40)
 
 
+def test_rejected_inputs_leave_out_unchanged():
+    # Inputs are checked in a pass of their own, before out= is written.
+    c, bitsets, factors, q = setup("l2", 300)
+    query = quant.Query(q, metric="l2", centroid=c)
+    broken = factors.copy()
+    broken["rescale"][250] = np.nan
+    packed = np.zeros((10, 768), np.uint8)
+    with pytest.raises(ValueError, match="factors must contain only finite values"):
+        quant.pack_fastscan_blocks(bitsets, broken, num_threads=4, out=packed)
+    assert not packed.any()
+    blocks = quant.pack_fastscan_blocks(bitsets, factors)
+    add = 96 * 4  # Lane 0 of the add factors.
+    blocks[7, add : add + 4] = np.frombuffer(np.float32(np.nan).tobytes(), np.uint8)
+    estimates = np.zeros(300, quant.ESTIMATE_DTYPE)
+    with pytest.raises(ValueError, match="block 7 does not"):
+        query.estimate_fastscan(blocks, num_points=300, num_threads=4, out=estimates)
+    assert not estimates.view(np.float32).any()
+
+
 def test_out_overlapping_an_input_is_rejected():
     _, bitsets, factors, _ = setup("l2", 40)
     shared = np.zeros(2 * 768, np.uint8)
-- 
2.34.1

[gouyt13]

Patch 4 of 4: fix(quant): raise OverflowError for estimates that are

0004-fix-quant-raise-OverflowError-for-estimates-that-are.patch

From 9db54901ebf44a176c5b13fb2b16f3bf37eba473 Mon Sep 17 00:00:00 2001
From: gouyt13clear <yutong003@e.ntu.edu.sg>
Date: Tue, 3 Nov 2026 16:40:33 +0800
Subject: [PATCH 4/4] fix(quant): raise OverflowError for estimates that are
 not finite

The scoring calls returned an inf or NaN distance when finite but extreme factors
overflowed, where quant::v1::compute raises std::overflow_error. NumPy sorts NaN
last, so a ranking silently dropped those vectors. estimate_bitset,
estimate_codes, estimate_from_ip, estimate_fastscan, refine and refine_packed now
raise OverflowError for a distance that is not finite, with v1's message, and for
an estimated error that is not finite. code_factors now adds the row number to its
OverflowError, as it already did to its ValueError.

Each range is written and then tested at once, and only a failing range is
recomputed row by row to report the row, so the check adds no measurable cost
(same benchmark as the previous commit, 1 thread, previous commit -> this one, ms:
estimate_bitset 2.33 -> 2.41, estimate_codes 22.13 -> 22.41, estimate_from_ip
1.27 -> 1.17, estimate_fastscan 1.80 -> 1.56; differences are within run-to-run
noise).

Overflow is found after the values are written, so the out= contract becomes: a
call that raises leaves out= unchanged, except after OverflowError or MemoryError,
when its contents are unspecified. Finding overflow before writing would need a
staging copy on every call that passes out=, the cost out= exists to avoid. Input
errors are still raised before anything is written. The guide states this,
documents which values raise OverflowError, and notes that split, which returns a
tuple of new arrays, takes no out=.
---
 docs/docs/rabitq/quant_python.md   |  4 +-
 python_bindings/quant_bindings.cpp | 72 +++++++++++++++++++++++++-----
 tests/python/test_quant_layer1.py  | 21 +++++++++
 tests/python/test_quant_layer2.py  | 58 ++++++++++++++++++++++++
 4 files changed, 142 insertions(+), 13 deletions(-)

diff --git a/docs/docs/rabitq/quant_python.md b/docs/docs/rabitq/quant_python.md
index c5fc0ee..82576c9 100644
--- a/docs/docs/rabitq/quant_python.md
+++ b/docs/docs/rabitq/quant_python.md
@@ -114,8 +114,8 @@ If your `ip` is itself approximate, `error` covers only the code's quantization
 - Build-side calls (`Rotator.rotate`, `quantize`, `code_factors`, `pack_codes`, `unpack_codes`, `pack_fastscan_blocks`) default to `num_threads=0`, the available CPU count capped by `OMP_NUM_THREADS`, as for the indexes (see [threading](../index/usage.md#threading-and-file-paths)); their results do not depend on the thread count. Scoring calls (`estimate_*`, `refine*`, `estimate_from_ip`) default to 1.
 - The GIL is released during computation. `Rotator` is immutable and safe to share. Scoring calls on one `Query` may run concurrently; a `set_centroid` during one raises `RuntimeError`. Use `qry.copy()` for one query per thread.
 - `Query` supports `copy.copy`, `copy.deepcopy` and pickling of its query vector, metric and centroid; a copy never shares storage with the original. Queries have no `to_bytes`, because they are not a stored format.
-- Every call that returns an array accepts `out=` with the exact dtype and shape. It must not share memory with any input, or the call raises `ValueError`. If a call raises, a caller's `out=` is left unchanged.
-- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. Inputs are checked before anything is computed; within each check, the error reported is that of the first failing row in call order (the lowest row, or the earliest entry of `indices=`), whatever the thread count.
+- Every call that returns an array, except `split`, accepts `out=` with the exact dtype and shape. It must not share memory with any input, or the call raises `ValueError`. A call that raises leaves a caller's `out=` unchanged, except after `OverflowError` or `MemoryError`, when its contents are unspecified.
+- Inputs are validated: wrong shapes, dtypes or values raise `ValueError`, and non-finite or oversized coordinates are rejected as the indexes reject them. Factors and query terms that are not representable in float32 raise `OverflowError`, and so does an estimated distance, as in `quant::v1`, or an estimated error. Inputs are checked before anything is computed; within each check, the error reported is that of the first failing row in call order (the lowest row, or the earliest entry of `indices=`), whatever the thread count.
 
 ## Byte formats
 
diff --git a/python_bindings/quant_bindings.cpp b/python_bindings/quant_bindings.cpp
index 22db653..1294ec9 100644
--- a/python_bindings/quant_bindings.cpp
+++ b/python_bindings/quant_bindings.cpp
@@ -530,6 +530,10 @@ void check_code_values(
 }
 
 // Whether every float of records [begin, end) is finite.
+bool all_finite_records(const float* values, size_t begin, size_t end) {
+    return all_finite(values + begin, end - begin);
+}
+
 template <typename T>
 bool all_finite_records(const T* records, size_t begin, size_t end) {
     static_assert(sizeof(T) == 3 * sizeof(float));
@@ -547,14 +551,40 @@ bool finite_factor_range(
     return selected.indices == nullptr && all_finite_records(factors.data, begin, end);
 }
 
-// Writes value_at(k) to out[k] for k in [0, count), for rows already validated.
+// quant::v1::compute rejects the distances that are not finite; so do the scoring calls.
+float checked(float distance) {
+    if (!std::isfinite(distance)) {
+        throw std::overflow_error("Estimated distance is not representable in float32");
+    }
+    return distance;
+}
+
+Estimate checked(const Estimate& estimate) {
+    (void)checked(estimate.distance);
+    if (!std::isfinite(estimate.error)) {
+        throw std::overflow_error("Estimated error is not representable in float32");
+    }
+    return estimate;
+}
+
+// Writes value_at(k) to out[k] for k in [0, count), for rows already validated, and
+// raises OverflowError for the lowest row whose value is not finite. Ranges are
+// written without per-row checks and tested at once; only a failing range is
+// recomputed row by row. out may then be partly written.
 template <typename T, typename ValueAt>
 void write_rows(T* out, size_t count, size_t num_threads, ValueAt value_at) {
-    parallel_rows(count, num_threads, kCheapRows, [&](size_t begin, size_t end) {
-        for (size_t k = begin; k < end; ++k) {
-            out[k] = value_at(k);
-        }
-    });
+    for_each_range(
+        count,
+        num_threads,
+        kCheapRows,
+        [&](size_t begin, size_t end) {
+            for (size_t k = begin; k < end; ++k) {
+                out[k] = value_at(k);
+            }
+            return all_finite_records(out, begin, end);
+        },
+        [&](size_t k) { out[k] = checked(value_at(k)); }
+    );
 }
 
 // --- Layer 1 --------------------------------------------------------------------
@@ -679,7 +709,8 @@ py::array code_factors_rows(
         quant::v1::Factors* factors = caller_out ? staged.data() : data;
         for_each_row(rows, num_threads, kCostlyRows, [&](size_t i) {
             // code_factors checks the row's code values and that the code quantizes the
-            // residual; the binding adds the row number to its message.
+            // residual, and rejects factors that overflow float32; the binding adds the
+            // row number to its message.
             const auto row = [i] {
                 return "codes row " + std::to_string(i) + " for residuals row " +
                        std::to_string(i) + ": ";
@@ -695,6 +726,8 @@ py::array code_factors_rows(
                 );
             } catch (const std::invalid_argument& error) {
                 throw std::invalid_argument(row() + error.what());
+            } catch (const std::overflow_error& error) {
+                throw std::overflow_error(row() + error.what());
             }
         });
         std::copy(staged.begin(), staged.end(), data);
@@ -820,6 +853,7 @@ QueryTerms compute_query_terms(
 // quant::detail::unpacked_ip, the inner product of quant::v1::compute. The arithmetic
 // is v1's, but this file is compiled without FP contraction and api.cpp is not, so
 // distances agree with v1's within float32 rounding rather than bitwise.
+// It is not checked here: write_rows rejects an estimate that is not finite.
 inline Estimate estimate_of(
     float ip, size_t bits, const quant::v1::Factors& f, const QueryTerms& terms
 ) {
@@ -1355,7 +1389,7 @@ class QuantQuery {
                         );
                     }
                 });
-                // Writes the estimates of block b.
+                // Whether block b writes finite estimates.
                 const auto score_block = [&](size_t b) {
                     const auto* block = reinterpret_cast<const char*>(view.row(b));
                     std::array<float, fastscan::kBatchSize> distances{};
@@ -1380,11 +1414,27 @@ class QuantQuery {
                         estimates[begin + lane] = {
                             distances[lane], errors[lane] * terms_.g_error, ips[lane]};
                     }
+                    return all_finite_records(estimates, begin, begin + lanes);
                 };
-                parallel_rows(
-                    num_blocks, num_threads, kCheapRows, [&](size_t first, size_t last) {
+                for_each_range(
+                    num_blocks,
+                    num_threads,
+                    kCheapRows,
+                    [&](size_t first, size_t last) {
                         for (size_t b = first; b < last; ++b) {
-                            score_block(b);
+                            if (!score_block(b)) {
+                                return false;
+                            }
+                        }
+                        return true;
+                    },
+                    [&](size_t b) {
+                        if (score_block(b)) {
+                            return;
+                        }
+                        const size_t begin = b * fastscan::kBatchSize;
+                        for (size_t lane = 0; lane < lanes_of(b); ++lane) {
+                            (void)checked(estimates[begin + lane]);
                         }
                     }
                 );
diff --git a/tests/python/test_quant_layer1.py b/tests/python/test_quant_layer1.py
index e2b51f5..a83fc88 100644
--- a/tests/python/test_quant_layer1.py
+++ b/tests/python/test_quant_layer1.py
@@ -318,3 +318,24 @@ def test_code_factors_names_the_lowest_rejected_row_across_threads():
             ValueError, match="^codes row 300 for residuals row 300: code does not"
         ):
             quant.code_factors(x, codes, nbits=4, metric="l2", num_threads=threads)
+
+
+@pytest.mark.parametrize("metric", ["l2", "ip"])
+def test_code_factors_reports_the_overflowing_row(metric):
+    # With every code 1, <residual, centered code> is half the residual's sum. The
+    # +-2**25 at coordinates 0 and 64 share an accumulator slot in every dot-product
+    # backend and cancel exactly, and coordinate 65 keeps its tiny value apart, so the
+    # sum is a positive 1e-30 and the rescale factor overflows float32.
+    x = residuals(6, seed=9)
+    codes = quant.quantize(x, nbits=1)
+    for row in (3, 5):
+        x[row] = 0
+        x[row, [0, 64, 65]] = [2.0**25, -(2.0**25), 1e-30]
+        codes[row] = 1
+    out = np.zeros(6, quant.FACTORS_DTYPE)
+    with pytest.raises(
+        OverflowError,
+        match="^codes row 3 for residuals row 3: Quantization factors are not "
+        "representable in float32$",
+    ):
+        quant.code_factors(x, codes, nbits=1, metric=metric, num_threads=4, out=out)
diff --git a/tests/python/test_quant_layer2.py b/tests/python/test_quant_layer2.py
index c34add8..6c0f541 100644
--- a/tests/python/test_quant_layer2.py
+++ b/tests/python/test_quant_layer2.py
@@ -414,6 +414,64 @@ def test_out_overlapping_an_input_is_rejected():
         )
 
 
+@pytest.mark.parametrize(
+    "field,message",
+    [
+        ("rescale", "Estimated distance is not representable in float32"),
+        ("error", "Estimated error is not representable in float32"),
+    ],
+)
+def test_estimates_that_overflow_are_rejected(field, message):
+    # With q all ones and every code 1, each centered inner product is 48 and |q| is
+    # sqrt(96), so a factor of 3e38 overflows that row's distance or error.
+    n, d = 300, 96
+    q = np.ones(d, np.float32)
+    query = quant.Query(q, metric="l2")
+    codes = np.ones((n, d), np.uint8)
+    bitsets = quant.pack_codes(codes, nbits=1)
+    factors = np.zeros(n, quant.FACTORS_DTYPE)
+    factors[field][[250, 120]] = 3e38
+    terms = quant.query_terms(q, metric="l2")
+    blocks = quant.pack_fastscan_blocks(bitsets, factors)
+    calls = [
+        lambda out: query.estimate_bitset(bitsets, factors, num_threads=4, out=out),
+        lambda out: query.estimate_codes(
+            codes, factors, nbits=1, num_threads=4, out=out
+        ),
+        lambda out: quant.estimate_from_ip(
+            np.full(n, d, np.float32), factors, terms, nbits=1, num_threads=4, out=out
+        ),
+        lambda out: query.estimate_fastscan(
+            blocks, num_points=n, num_threads=4, out=out
+        ),
+    ]
+    for call in calls:
+        # After OverflowError the contents of out= are unspecified.
+        with pytest.raises(OverflowError, match=f"^{message}$"):
+            call(np.zeros(n, quant.ESTIMATE_DTYPE))
+
+
+@pytest.mark.parametrize("packed", [False, True])
+def test_refined_distances_that_overflow_are_rejected(packed):
+    # Prefix and extra bits all 1 give total codes 3, centered at 1.5: the inner product
+    # with q all ones is 144, so a rescale of 3e38 overflows the distance.
+    n, d = 300, 96
+    query = quant.Query(np.ones(d, np.float32), metric="l2")
+    ex = np.ones((n, d), np.uint8)
+    factors = np.zeros(n, quant.FACTORS_DTYPE)
+    factors["rescale"][[250, 120]] = 3e38
+    first = np.zeros(n, quant.ESTIMATE_DTYPE)
+    first["refine_input"] = d
+    refine = query.refine_packed if packed else query.refine
+    codes = quant.pack_codes(ex, nbits=1) if packed else ex
+    with pytest.raises(
+        OverflowError, match="^Estimated distance is not representable in float32$"
+    ):
+        refine(
+            codes, factors, first, ex_bits=1, num_threads=4, out=np.zeros(n, np.float32)
+        )
+
+
 def test_rejections_name_the_lowest_failing_row_across_threads():
     _, c, residuals, codes, factors, q = setup("l2", 4, n=1000)
     query = quant.Query(q, metric="l2", centroid=c)
-- 
2.34.1

[TamimEhsan]

@gouyt13 patched with the 4 commits. Also, now split is also aligned with the other functions and takes an out and a num_threads parameter for the result.