outputs) {
+ QueryBatchRows rows = QUERY_BATCH_ROWS.get();
+ int count = inputs.size();
+ rows.ensure(count);
+ try {
+ for (int i = 0; i < count; i++) {
+ rows.input[i] = inputs.get(i);
+ rows.output[i] = outputs.get(i);
+ }
+ target.projectBatch(rows.input, rows.output, count);
+ } finally {
+ Arrays.fill(rows.input, 0, count, null);
+ Arrays.fill(rows.output, 0, count, null);
+ }
+ }
+ };
+ int budget = Math.max(1, Math.min(Runtime.getRuntime().availableProcessors(),
+ PhysicalCoreExecutor.getPhysicalCoreCount()));
+ queryBatches = CallerBatchProcessor.opportunistic(
+ budget, queryBatchSize, operation);
+ retained += queryBatches.ramBytesUsed() + RamUsageEstimator.shallowSizeOf(operation);
+ }
+ queryProjectionHeapBytes = retained;
+ queryProjection = ready; // Publish the representation and its heap charge together.
+ }
+ }
+ }
+ return ready;
+ }
+
+ /** Projects one vector with the original FP32 representation, including encoding/training. */
+ public void project(float[] query, float[] out) {
+ projection.project(query, out);
+ }
+
+ /** Projects base vectors in FP32; query precision selection does not apply. */
+ public void projectBatch(float[][] vectors, float[][] out, int count) {
+ projection.projectBatch(vectors, out, count);
+ }
+
+ /** Number of rows in A (quantizedDim = d). */
+ public final int rows;
+
+ /** Number of columns in A (originalDim = D). */
+ public final int cols;
+
+ public StiefelTransform(RealMatrix W) {
+ this.W = Objects.requireNonNull(W, "W");
+ String batchSize = System.getProperty("jvector.ash.queryBatchSize", "8");
+ try { queryBatchSize = Integer.parseInt(batchSize); }
+ catch (NumberFormatException invalid) {
+ throw new IllegalArgumentException("jvector.ash.queryBatchSize must be a positive integer: " + batchSize, invalid);
+ }
+ if (queryBatchSize < 1 || queryBatchSize == Integer.MAX_VALUE)
+ throw new IllegalArgumentException("jvector.ash.queryBatchSize must be between 1 and " + (Integer.MAX_VALUE - 1));
+
+ // Dimensions: W is [D × d]
+ final int originalDim = W.getRowDimension(); // D
+ final int quantizedDim = W.getColumnDimension(); // d
+
+ this.rows = quantizedDim;
+ this.cols = originalDim;
+
+ // Explicitly materialize A = Wᵀ into row-major arrays. Models persist A in FP32, so AData holds
+ // those FP32 values: a trained model and its reloaded copy then derive identical landmark terms.
+ this.AData = new double[quantizedDim][originalDim];
+ for (int i = 0; i < quantizedDim; i++) {
+ final double[] row = this.AData[i];
+ for (int j = 0; j < originalDim; j++) {
+ // A[i][j] = W[j][i]
+ row[j] = (float) W.getEntry(j, i);
+ }
+ }
+
+ this.AFloat = new float[quantizedDim][originalDim];
+ for (int i = 0; i < quantizedDim; i++) {
+ final double[] src = this.AData[i];
+ final float[] dst = this.AFloat[i];
+ for (int j = 0; j < originalDim; j++) {
+ dst[j] = (float) src[j];
+ }
+ }
+ this.projection = VectorizationProvider.getInstance().getVectorUtilSupport()
+ .prepareAshProjection(AFloat);
+
+ // Hard invariants: these are programmer-contract checks, not user input checks.
+ if (AFloat.length != rows) {
+ throw new AssertionError("AFloat rows " + AFloat.length + " != quantizedDim " + rows);
+ }
+ if (rows > 0 && AFloat[0].length != cols) {
+ throw new AssertionError("AFloat cols " + AFloat[0].length + " != originalDim " + cols);
+ }
+ if (AData.length != rows) {
+ throw new AssertionError("AData rows " + AData.length + " != quantizedDim " + rows);
+ }
+ if (rows > 0 && AData[0].length != cols) {
+ throw new AssertionError("AData cols " + AData[0].length + " != originalDim " + cols);
+ }
+ heapBytes = RamUsageEstimator.shallowSizeOfInstance(StiefelTransform.class)
+ + estimateMatrixBytes(W) + estimateArray2DBytes(AData) + estimateArray2DBytes(AFloat)
+ + RamUsageEstimator.shallowSizeOf(projection);
+ }
+ }
+
+ /**
+ * RANDOM (untrained) Stiefel initialization.
+ *
+ * Samples a Gaussian matrix G ∈ R^{D×d} and computes its polar factor
+ * W = U·Vᵀ via a thin SVD (G = U·S·Vᵀ). The resulting W has orthonormal
+ * columns (WᵀW = I) and can be used as an orthonormal projection matrix
+ * (e.g., y = Wᵀx maps R^D → R^d).
+ *
+ * @param D dimensionality of the input vectors
+ * @param d dimensionality of the projected space
+ * @param rng random number generator used to sample G
+ * @return a {@link StiefelTransform} containing W ∈ R^{D×d} with orthonormal columns
+ * @throws IllegalArgumentException if {@code d <= 0 or d > D}
+ * @throws RuntimeException if the native SVD fails
+ */
+ public static StiefelTransform runWithoutTraining(int D, int d, Random rng) {
+ try (var blas = BlasThreads.singleUnlessLimited()) {
+ return randomTransform(D, d, rng);
+ }
+ }
+
+ private static StiefelTransform randomTransform(int D, int d, Random rng) {
+ if (d <= 0 || d > D) {
+ throw new IllegalArgumentException("Invalid d=" + d + " for D=" + D);
+ }
+
+ // gData is interpreted as a [D x d] column-major matrix by LAPACK.
+ double[] gData = new double[D * d];
+ for (int i = 0; i < gData.length; i++) {
+ gData[i] = rng.nextGaussian();
+ }
+
+ double[] wData = orthogonalize(gData, D, d);
+ RealMatrix W = new Array2DRowRealMatrix(deserializeColumnMajor(wData, D, d), false);
+ return new StiefelTransform(W);
+ }
+
+ @Override
+ @Deprecated
+ public CompressedVectors createCompressedVectors(Object[] compressedVectors) {
+ throw new UnsupportedOperationException("Deprecated createCompressedVectors() method not supported");
+ }
+
+
+ // An internal working-set bound, not a projection API limit. Each task owns
+ // at most this many normalized vectors and projected results, reused until
+ // its range is complete. Native code supplies the SIMD/cache micro-tiling.
+ private static final int ENCODING_BATCH_SIZE = 32;
+
+ /**
+ * Encodes all base vectors in bounded FP32 projection batches. The supplied
+ * pool controls parallelism; native calls execute on the calling worker.
+ * Missing vectors retain the existing all-zero encoded representation.
+ */
+ @Override
+ public ASHVectors encodeAll(RandomAccessVectorValues ravv, ForkJoinPool simdExecutor) {
+ Objects.requireNonNull(ravv, "ravv");
+ Objects.requireNonNull(simdExecutor, "simdExecutor");
+ if (ravv.dimension() != originalDimension)
+ throw new IllegalArgumentException("ASH encoding vector dimension mismatch");
+ final int n = ravv.size();
+ final int parallelism = Math.max(1, simdExecutor.getParallelism());
+ final int chunkSize = (int) (((long) n + parallelism - 1) / parallelism);
+ // Encoded straight into the compact store; no per-vector objects are retained.
+ final ASHCodes out = ASHCodes.allocate(n, quantizedDim, bitsPerDimension);
+ // Fill while each encoded vector is hot; never reread the raw input for setup.
+ final float[] symmetricOffsets = new float[n];
+ final var ravvSupplier = ravv.threadLocalSupplier();
+ var tasks = new java.util.ArrayList>(parallelism);
+ Throwable failure = null;
+ try {
+ for (int start = 0; start < n;) {
+ final int first = start;
+ final int end = (int) Math.min((long) n, (long) first + chunkSize);
+ tasks.add(simdExecutor.submit(() -> encodeRange(ravvSupplier.get(), out, symmetricOffsets, first, end)));
+ start = end;
+ }
+ } catch (RuntimeException | Error error) {
+ failure = error;
+ }
+ // Await every submitted task, including when the pool rejected a later
+ // submission. The caller may close its source when this method exits.
+ for (var task : tasks) {
+ try { task.join(); }
+ catch (RuntimeException | Error error) {
+ if (failure == null) failure = error;
+ else if (failure != error) failure.addSuppressed(error);
+ }
+ }
+ if (failure instanceof Error) throw (Error) failure;
+ if (failure != null) throw (RuntimeException) failure;
+ return new ASHVectors(this, out, symmetricOffsets);
+ }
+
+ private void encodeRange(RandomAccessVectorValues source, ASHCodes out,
+ float[] symmetricOffsets, int start, int end) {
+ int capacity = Math.min(ENCODING_BATCH_SIZE, end - start);
+ float[][] inputs = new float[capacity][originalDimension];
+ float[][] projected = new float[capacity][quantizedDim];
+ float[] residualNorms = new float[capacity];
+ int[] ordinals = new int[capacity];
+ double[] absNorm = new double[quantizedDim];
+ // Per-batch scratch, copied into the store once complete. Never modified: the code for missing vectors.
+ QuantizedVector[] scratch = new QuantizedVector[capacity];
+ for (int k = 0; k < capacity; k++) scratch[k] = QuantizedVector.createEmpty(quantizedDim, bitsPerDimension);
+ final QuantizedVector missing = QuantizedVector.createEmpty(quantizedDim, bitsPerDimension);
+ final float missingOffset = symmetricOffsets == null ? 0f : centeredOffset(missing);
+ int count = 0;
+ try (var cursor = VectorAccess.openRange(source, start, end)) {
+ int expected = start;
+ while (cursor.next()) {
+ int i = cursor.ordinal();
+ if (i != expected++) throw new IllegalStateException("Unexpected encoding ordinal");
+ VectorFloat vector = cursor.vector();
+ if (vector == null) {
+ out.set(i, missing);
+ if (symmetricOffsets != null) symmetricOffsets[i] = missingOffset;
+ continue;
+ }
+ QuantizedVector dest = scratch[count];
+ float residualNorm = prepareEncoding(vector, dest);
+ normalizeResidual(vector, landmarks[dest.landmark & 0xFF], residualNorm, inputs[count]);
+ residualNorms[count] = residualNorm;
+ ordinals[count] = i;
+ if (++count == capacity) {
+ encodeBatch(inputs, projected, residualNorms, ordinals, count, absNorm, scratch, out, symmetricOffsets);
+ count = 0;
+ }
+ }
+ if (expected != end) throw new IllegalStateException("Incomplete encoding range");
+ }
+ if (count != 0) encodeBatch(inputs, projected, residualNorms, ordinals, count, absNorm, scratch, out, symmetricOffsets);
+ }
+
+ private void encodeBatch(float[][] inputs, float[][] projected, float[] residualNorms,
+ int[] ordinals, int count, double[] absNorm, QuantizedVector[] scratch,
+ ASHCodes out, float[] symmetricOffsets) {
+ stiefelTransform.projectBatch(inputs, projected, count);
+ for (int i = 0; i < count; i++) {
+ QuantizedVector dest = scratch[i];
+ // Quantizers OR bits and nibbles into place, so reused scratch starts from zero.
+ Arrays.fill(dest.binaryVector, 0L);
+ Arrays.fill(dest.extraBits, (byte) 0);
+ float codeNorm = ProjectedBinaryQuantizer.quantizeProjected(projected[i], bitsPerDimension,
+ dest.binaryVector, dest.extraBits, absNorm);
+ float correction = finishEncoding(dest, residualNorms[i], codeNorm);
+ out.set(ordinals[i], dest);
+ if (symmetricOffsets != null) symmetricOffsets[ordinals[i]] = correction;
+ }
+ }
+
+ private static void normalizeResidual(VectorFloat vector, VectorFloat mean,
+ float residualNorm, float[] normalized) {
+ if (vector.length() != normalized.length || mean.length() != normalized.length)
+ throw new IllegalArgumentException("ASH encoding vector dimension mismatch");
+ float invNorm = residualNorm > 0f ? 1f / residualNorm : 0f;
+ if (vector instanceof FloatArray
+ && mean instanceof FloatArray) {
+ // Resolve heap backing once, including MemorySegment views. The
+ // simple array loop avoids per-component virtual/scope checks and
+ // permits JVM SIMD without changing the FP32 operation order.
+ float[] values = ((FloatArray) vector).array();
+ float[] center = ((FloatArray) mean).array();
+ for (int d = 0; d < normalized.length; d++)
+ normalized[d] = (values[d] - center[d]) * invNorm;
+ } else {
+ for (int d = 0; d < normalized.length; d++)
+ normalized[d] = (vector.get(d) - mean.get(d)) * invNorm;
+ }
+ }
+
+ /**
+ * Encodes the input vector using ASH.
+ * @return the header and encoded vector payload
+ */
+ @Override
+ public QuantizedVector encode(VectorFloat vector) {
+ var qv = QuantizedVector.createEmpty(this.quantizedDim, this.bitsPerDimension);
+ encodeTo(vector, qv);
+ return qv;
+ }
+
+ @Override
+ public void encodeTo(VectorFloat vector, QuantizedVector dest) {
+ float residualNorm = prepareEncoding(vector, dest);
+ float codeNorm = quantizeVector(vector, residualNorm, dest.landmark, dest);
+ finishEncoding(dest, residualNorm, codeNorm);
+ }
+
+ /** Sets landmark and raw offset; returns the residual norm needed by both paths. */
+ private float prepareEncoding(VectorFloat vector, QuantizedVector dest) {
+ Objects.requireNonNull(vector, "vector");
+ Objects.requireNonNull(dest, "dest");
+ if (vector.length() != originalDimension)
+ throw new IllegalArgumentException("ASH encoding vector dimension mismatch");
+ if (dest.binaryVector == null || dest.binaryVector.length < QuantizedVector.signWordsForDims(quantizedDim, bitsPerDimension)
+ || dest.extraBits == null || dest.extraBits.length < QuantizedVector.bodyBytesForDims(quantizedDim, bitsPerDimension))
+ throw new IllegalArgumentException("ASH encoding destination too short");
+ // Landmark assignment: nearest centroid by L2 distance.
+ byte landmark = 0;
+ float bestDist = Float.POSITIVE_INFINITY;
+
+ for (int c = 0; c < landmarkCount; c++) {
+ float dist = VectorUtil.squareL2Distance(vector, landmarks[c]);
+ if (dist < bestDist) {
+ bestDist = dist;
+ landmark = (byte) c;
+ }
+ }
+
+ final int c = landmark & 0xFF;
+ final VectorFloat mu = landmarks[c];
+
+ // Compute and ||x_i − μ_i*||_2.
+ final float dotXMu = VectorUtil.dotProduct(vector, mu);
+ final float sqDist = VectorUtil.squareL2Distance(vector, mu);
+ final float residualNorm = (float) Math.sqrt(sqDist);
+
+ assert quantizedDim > 0;
+
+ final float rawOffset = dotXMu - landmarkNormSq[c];
+ // ||x||^2 = ||x - mu||^2 + 2 - ||mu||^2, from the terms already computed.
+ recordEncodedSquaredNorm(sqDist + 2f * dotXMu - landmarkNormSq[c]);
+
+ // Header fields needed before quantization. The final offset may be
+ // adjusted after quantization for C++-style fast-scan projection codes.
+ dest.offset = rawOffset;
+ dest.landmark = landmark;
+ assert Float.isFinite(rawOffset) : "offset is not finite";
+
+ return residualNorm;
+ }
+
+ /** Applies the existing code-dependent header adjustment and FP16 storage rounding. */
+ /** Returns the symmetric correction using the same dot product and stored half headers. */
+ private float finishEncoding(QuantizedVector dest, float residualNorm, float codeNorm) {
+ final int c = dest.landmark & 0xFF;
+ final float rawOffset = dest.offset;
+ final float scale = codeNorm > 0f ? residualNorm / codeNorm : 0f;
+
+ // Projection-mode offset: offset = - ||μ||² - scale * , so query-time
+ // scoring can use Aq instead of A(q - μ). A 1-bit code is its ±1 signs.
+ final float landmarkDot = landmarkCodeDot(c, dest.binaryVector, 0, dest.extraBits, 0);
+ dest.scale = roundToFloat16(scale);
+ dest.offset = roundToFloat16(rawOffset - scale * landmarkDot);
+ return dest.offset + dest.scale * landmarkDot;
+ }
+
+ // ---------------------------------------------------------------------
+ // Training of projection matrix
+ // ---------------------------------------------------------------------
+
+ /** Sampled, normalized residuals in the column-major layout consumed by LAPACK. */
+ private static final class TrainingData {
+ final double[] xCol;
+ final int rows;
+ final int columns;
+
+ TrainingData(double[] xCol, int rows, int columns) {
+ this.xCol = xCol;
+ this.rows = rows;
+ this.columns = columns;
+ }
+
+ TrainingData(double[][] rows) {
+ this(serializeColumnMajor(rows), rows.length, rows[0].length);
+ }
+ }
+
+ private static final class SingleLandmarkTraining {
+ final VectorFloat landmark;
+ final TrainingData trainingData;
+
+ SingleLandmarkTraining(VectorFloat landmark, TrainingData trainingData) {
+ this.landmark = landmark;
+ this.trainingData = trainingData;
+ }
+ }
+
+ private static int trainingSampleSize(int count, int dimensions) {
+ return (int) Math.min((long) count, Math.max(1L, (long) dimensions * ITQ_TRAINING_FACTOR));
+ }
+
+ /**
+ * Computes the C=1 landmark and ITQ inputs from the same deterministic
+ * sample. Only sampled source vectors are read; normalized sample values
+ * are retained until the mean is known. One reusable vector protects
+ * borrowed source storage. Double sums follow sorted ordinal order, and
+ * normalization retains the provider's FP32 dot/scale operations.
+ */
+ private static SingleLandmarkTraining prepareSingleLandmarkTraining(
+ RandomAccessVectorValues source, int dimensions, Random rng, boolean learn) {
+ final int count = source.size();
+ if (count <= 0 || dimensions <= 0) {
+ throw new IllegalArgumentException("ASH training requires positive size and dimension");
+ }
+ final int sampleSize = trainingSampleSize(count, dimensions);
+ // Validate the flat-array representation before allocating ordinal or
+ // normalization buffers.
+ long sampleElements = (long) sampleSize * dimensions;
+ if (learn && sampleElements > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException("ASH training sample exceeds the Java array size limit: "
+ + sampleSize + " rows x " + dimensions + " dimensions");
+ }
+ // The existing PCA driver also requires a full D-by-D right basis,
+ // including null-space directions when the training matrix is wide.
+ if (learn && (long) dimensions * dimensions > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException("ASH PCA basis exceeds the Java array size limit: "
+ + dimensions + " x " + dimensions);
+ }
+ // RANDOM projections keep their existing Gaussian RNG sequence. ITQ
+ // consumes exactly the same reservoir-sampling draws as before.
+ final Random sampleRandom = learn ? rng : new Random(TRAINING_SEED);
+ final int[] ordinals = reservoirSampleOrdinals(count, sampleSize, sampleRandom);
+ Arrays.sort(ordinals);
+ final float[] sample = learn ? new float[(int) sampleElements] : null;
+ final double[] accumulator = new double[dimensions];
+ final VectorFloat normalized = vectorTypeSupport.createFloatVector(dimensions);
+ final float[] values = normalized instanceof FloatArray ? ((FloatArray) normalized).array() : null;
+ for (int row = 0; row < sampleSize; row++) {
+ int ordinal = ordinals[row];
+ VectorFloat input = source.getVector(ordinal);
+ if (input == null) {
+ throw new NullPointerException("Null ASH training vector at ordinal " + ordinal);
+ }
+ if (input.length() != dimensions) {
+ throw new IllegalArgumentException("ASH training vector at ordinal " + ordinal
+ + " has dimension " + input.length() + "; expected " + dimensions);
+ }
+ normalized.copyFrom(input, 0, 0, dimensions);
+ // Dimensions are already checked. Call the provider directly so invalid
+ // norms reach argument validation even when VectorUtil assertions are enabled.
+ float norm = (float) Math.sqrt(vectorUtilSupport.dotProduct(normalized, normalized));
+ if (!Float.isFinite(norm)) {
+ throw new IllegalArgumentException("ASH training vector at ordinal " + ordinal
+ + " has a non-finite FP32 L2 norm");
+ }
+ if (norm > 1e-10f) {
+ VectorUtil.scale(normalized, 1.0f / norm);
+ }
+ if (values != null) {
+ for (int d = 0; d < dimensions; d++) accumulator[d] += values[d];
+ } else {
+ for (int d = 0; d < dimensions; d++) accumulator[d] += normalized.get(d);
+ }
+ if (learn) {
+ int offset = row * dimensions;
+ if (values != null) {
+ System.arraycopy(values, 0, sample, offset, dimensions);
+ } else {
+ for (int d = 0; d < dimensions; d++) sample[offset + d] = normalized.get(d);
+ }
+ }
+ }
+ VectorFloat mean = vectorTypeSupport.createFloatVector(dimensions);
+ double invCount = 1.0 / sampleSize;
+ for (int d = 0; d < dimensions; d++) mean.set(d, (float) (accumulator[d] * invCount));
+ if (!learn) return new SingleLandmarkTraining(mean, null);
+ return new SingleLandmarkTraining(mean, centerTrainingSample(sample, sampleSize, dimensions, mean));
+ }
+
+ /**
+ * Centers/normalizes contiguous sample rows in FP64, then writes small tiles
+ * into the final column-major matrix. No full-size double row matrix or
+ * strided, element-at-a-time whole-matrix conversion is needed. Tile width
+ * bounds scratch only; it does not change arithmetic or sample boundaries.
+ */
+ private static TrainingData centerTrainingSample(float[] sample, int rows, int dimensions,
+ VectorFloat mean) {
+ final int tileRows = Math.min(TRAINING_TRANSPOSE_TILE_ROWS, rows);
+ double[] tile = new double[Math.multiplyExact(tileRows, dimensions)];
+ double[] columnMajor = new double[Math.multiplyExact(rows, dimensions)];
+ float[] center = new float[dimensions];
+ for (int d = 0; d < dimensions; d++) center[d] = mean.get(d);
+ for (int first = 0; first < rows; first += tileRows) {
+ int rowsInTile = Math.min(tileRows, rows - first);
+ for (int row = 0; row < rowsInTile; row++) {
+ int sampleOffset = (first + row) * dimensions;
+ int tileOffset = row * dimensions;
+ double norm2 = 0.0;
+ for (int d = 0; d < dimensions; d++) {
+ double value = (double) sample[sampleOffset + d] - (double) center[d];
+ tile[tileOffset + d] = value;
+ norm2 += value * value;
+ }
+ double inverse = 1.0 / Math.max(Math.sqrt(norm2), 1e-20);
+ for (int d = 0; d < dimensions; d++) tile[tileOffset + d] *= inverse;
+ }
+ for (int d = 0; d < dimensions; d++) {
+ int columnOffset = d * rows + first;
+ for (int row = 0; row < rowsInTile; row++) {
+ columnMajor[columnOffset + row] = tile[row * dimensions + d];
+ }
+ }
+ }
+ return new TrainingData(columnMajor, rows, dimensions);
+ }
+
+ /**
+ * Builds normalized centered residuals for training:
+ * - Choose N_train = min(N, D * training_factor).
+ * - Sample ordinals uniformly at random (deterministic RNG).
+ * - For each sampled ordinal i:
+ * c = argmin ||x_i - mu_c||^2 (same rule as encoding)
+ * r = x_i - mu_c
+ * r_hat = r / ||r||
+ * store r_hat into xTrainNorm
+ */
+ private static TrainingData prepareTrainingData(
+ VectorFloat[] points,
+ VectorFloat[] landmarks,
+ int originalDim,
+ Random rng
+ ) {
+ final int N = points.length;
+
+ final int nTrain = trainingSampleSize(N, originalDim);
+
+ final int[] sample = reservoirSampleOrdinals(N, nTrain, rng);
+ Arrays.sort(sample); // deterministic order for reproducibility
+
+ final double[][] xTrainNorm = new double[nTrain][originalDim];
+
+ int row = 0;
+ int samplePos = 0;
+
+ for (int i = 0; i < N && samplePos < nTrain; i++) {
+ if (sample[samplePos] != i) {
+ continue;
+ }
+ samplePos++;
+
+ final VectorFloat x = points[i];
+
+ // Nearest landmark assignment (same semantics as encodeTo).
+ byte best = 0;
+ float bestDist = Float.POSITIVE_INFINITY;
+ for (int c = 0; c < landmarks.length; c++) {
+ float dist = VectorUtil.squareL2Distance(x, landmarks[c]);
+ if (dist < bestDist) {
+ bestDist = dist;
+ best = (byte) c;
+ }
+ }
+
+ final VectorFloat mu = landmarks[best & 0xFF];
+
+ // Center and normalize into double precision row.
+ final double[] out = xTrainNorm[row++];
+
+ double ss = 0.0;
+ for (int d = 0; d < originalDim; d++) {
+ double v = (double) x.get(d) - (double) mu.get(d);
+ out[d] = v;
+ ss += v * v;
+ }
+
+ // Normalize residuals while guarding against zero residuals.
+ double inv = 1.0 / Math.max(Math.sqrt(ss), 1e-20);
+ for (int d = 0; d < originalDim; d++) {
+ out[d] *= inv;
+ }
+ }
+
+ if (row != nTrain) {
+ throw new IllegalStateException("Training sample fill mismatch: got " + row + ", expected " + nTrain);
+ }
+
+ return new TrainingData(xTrainNorm);
+ }
+
+ /**
+ * Reservoir sample k ordinals uniformly from [0, N).
+ * Deterministic given rng.
+ */
+ private static int[] reservoirSampleOrdinals(int N, int k, Random rng) {
+ if (k < 0 || k > N) throw new IllegalArgumentException("k out of range: " + k + " for N=" + N);
+ int[] r = new int[k];
+ for (int i = 0; i < k; i++) r[i] = i;
+ for (int i = k; i < N; i++) {
+ int j = rng.nextInt(i + 1);
+ if (j < k) r[j] = i;
+ }
+ return r;
+ }
+
+ /**
+ * ITQ / learned projection trainer.
+ *
+ * Mirrors the current learned_projection path:
+ * - PCA basis P from checked FP64 covariance PCA, with direct SVD fallback
+ * - Xld = X @ P
+ * - iterative Procrustes:
+ * R = polar(M)
+ * Xtr = Xld @ R
+ * Xenc = encodeTrainingDirections(Xtr, bitsPerDimension)
+ * M = Xld^T @ Xenc
+ * - W = P @ R
+ */
+ private static StiefelTransform runItqTrainer(
+ TrainingData trainingData, // [N x D] normalized residuals, column-major
+ int latentDim, // projected dimensions, d
+ int bitsPerDimension,
+ Random rng,
+ int nTrainingIterations
+ ) {
+ LOG.info("ASH ITQ training started");
+
+ validateBitsPerDimension(bitsPerDimension);
+
+ final int N = trainingData.rows;
+ final int D = trainingData.columns;
+ final int d = latentDim;
+
+ if (d <= 0 || d > D) {
+ throw new IllegalArgumentException("Invalid latentDim=" + d + " for D=" + D);
+ }
+
+ // X is [N x D], flat column-major for BLAS/LAPACK.
+ double[] xCol = trainingData.xCol;
+
+ LOG.debug("ASH ITQ: PCA started");
+ double[] pCol = computePcaBasis(xCol, N, D, d);
+ LOG.debug("ASH ITQ: PCA finished");
+
+ // Xld = X @ P, shape [N x d].
+ double[] xldCol = nativeMultiply(xCol, N, D, pCol, d);
+
+ // Initial random update matrix M, shape [d x d].
+ double[] mCol = gaussianMatrix(d, d, rng);
+
+ // Each projected vector is contiguous. Quantization replaces it in place:
+ // the loss uses the update matrix, so unquantized rows need not be retained.
+ double[] projectedRows = new double[Math.multiplyExact(N, d)];
+
+ LOG.debug("ASH ITQ: training iterations started");
+ long startTime = System.nanoTime();
+ // Repeated polar steps need FP32 accuracy; the accumulated update,
+ // projection products and loss retain FP64. Final W uses a FP64 polar
+ // factor so the stored transform keeps its orthogonality contract.
+ ASHFloatOrthogonalizer orthogonalizer = new ASHFloatOrthogonalizer(d, d);
+ double[] rCol = new double[Math.multiplyExact(d, d)];
+ TrainingDirectionEncoder encoder = new TrainingDirectionEncoder(N, d, bitsPerDimension);
+
+ double[] recentLosses = new double[3];
+ for (int epoch = 0; epoch < nTrainingIterations; epoch++) {
+ orthogonalizer.orthogonalize(mCol, rCol);
+
+ projectTrainingRows(xldCol, N, d, rCol, projectedRows);
+ encoder.encodeInPlace(projectedRows);
+
+ // ITQ update: M = Xld^T @ Xenc. Reuse the initial Gaussian buffer.
+ updateTrainingMatrix(xldCol, N, d, projectedRows, mCol);
+
+ // C++-style cheap loss: trace(R^T M) / N.
+ ProjectionTrainingLoss loss = computeTrainingLossFromUpdateMatrix(rCol, mCol, N);
+ printTrainingLoss(epoch, loss);
+ if (epoch >= 8 && recentLosses[(epoch - 3) % 3] - loss.loss <= 0.01 * loss.loss) {
+ LOG.info("ASH ITQ training converged at epoch {}", epoch);
+ break;
+ }
+ recentLosses[epoch % 3] = loss.loss;
+ }
+
+ // Recompute from the final update, in FP64, rather than retaining the
+ // previous epoch's rotation or widening a FP32 polar factor.
+ new ASHOrthogonalizer(d, d).orthogonalize(mCol, rCol);
+
+ double loopTime = (System.nanoTime() - startTime) / 1e9;
+ if (LOG.isInfoEnabled()) {
+ LOG.info("ASH ITQ training finished in {} s", String.format(java.util.Locale.ROOT, "%.1f", loopTime));
+ }
+
+ // W = P @ R, shape [D x d].
+ double[] wCol = nativeMultiply(pCol, D, d, rCol, d);
+
+ RealMatrix W = new Array2DRowRealMatrix(deserializeColumnMajor(wCol, D, d), false);
+ return new StiefelTransform(W);
+ }
+
+ private static final class ProjectionTrainingLoss {
+ final double ip;
+ final double loss;
+
+ ProjectionTrainingLoss(double ip, double loss) {
+ this.ip = ip;
+ this.loss = loss;
+ }
+ }
+
+ /**
+ * Java equivalent of C++ compute_training_loss_from_update_matrix(R, M).
+ *
+ * R and M are [d x d] column-major.
+ */
+ private static ProjectionTrainingLoss computeTrainingLossFromUpdateMatrix(
+ double[] rCol,
+ double[] mCol,
+ int rows
+ ) {
+ double ipSum = 0.0;
+ for (int i = 0; i < rCol.length; i++) {
+ ipSum += rCol[i] * mCol[i];
+ }
+
+ double ip = ipSum / rows;
+ return new ProjectionTrainingLoss(ip, Math.max(0.0, 1.0 - ip));
+ }
+
+ private static void printTrainingLoss(int epoch, ProjectionTrainingLoss loss) {
+ if (LOG.isDebugEnabled()) {
+ LOG.debug(String.format(java.util.Locale.ROOT,
+ "ASH ITQ: epoch %2d inner-product loss %.5f, loss %.5f", epoch, loss.ip, loss.loss));
+ }
+ }
+
+ private static final float[] K_TIGHT_START = {
+ 0.0f,
+ 0.15f,
+ 0.20f,
+ 0.52f,
+ 0.59f,
+ 0.71f,
+ 0.75f,
+ 0.77f,
+ 0.81f
+ };
+
+ private static final double K_EPS = 1e-5;
+
+ private static final ThreadLocal SCALING_WORKSPACE =
+ ThreadLocal.withInitial(ScalingWorkspace::new);
+
+ private static final class ScalingWorkspace {
+ float[] invOAbs = new float[0];
+ int[] curOBar = new int[0];
+ long[] events = new long[0];
+
+ void ensureCapacity(int d, int maxEvents) {
+ if (invOAbs.length < d) {
+ invOAbs = new float[d];
+ }
+ if (curOBar.length < d) {
+ curOBar = new int[d];
+ }
+ if (events.length < maxEvents) {
+ events = new long[maxEvents];
+ }
+ }
+ }
+
+ /**
+ * Java port of ash::compute_optimal_scaling_factor.
+ *
+ * The input array is double because the Java trainer stores training matrices
+ * as double, but this method intentionally performs the search in float
+ * precision to match the C++ implementation.
+ */
+ private static double computeOptimalScalingFactor(
+ double[] oAbs,
+ int d,
+ int bitsPerDimension
+ ) {
+ final int exBits = bitsPerDimension - 1;
+ if (exBits < 1 || exBits > 8) {
+ throw new IllegalArgumentException("exBits must be in [1,8], got " + exBits);
+ }
+ if (d <= 0 || d > oAbs.length) {
+ throw new IllegalArgumentException("Invalid d=" + d + " for oAbs.length=" + oAbs.length);
+ }
+
+ final int kNEnum = 10;
+ final int maxCode = (1 << exBits) - 1;
+
+ float maxO = 0.0f;
+ for (int i = 0; i < d; i++) {
+ float v = (float) oAbs[i];
+ if (v > maxO) {
+ maxO = v;
+ }
+ }
+
+ if (!(maxO > 0.0f)) {
+ return 0.0;
+ }
+
+ final float tEnd = (float) (maxCode + kNEnum) / maxO;
+ final float tStart = tEnd * K_TIGHT_START[exBits];
+
+ final int maxEvents = Math.multiplyExact(d, maxCode + 1);
+ ScalingWorkspace ws = SCALING_WORKSPACE.get();
+ ws.ensureCapacity(d, maxEvents);
+
+ float sqrDenominator = (float) d * 0.25f;
+ float numerator = 0.0f;
+
+ for (int i = 0; i < d; i++) {
+ float oi = (float) oAbs[i];
+
+ ws.invOAbs[i] = 1.0f / oi;
+
+ int cur = (int) ((double) (tStart * oi) + K_EPS);
+ ws.curOBar[i] = cur;
+
+ sqrDenominator += (float) (cur * cur + cur);
+ numerator += (cur + 0.5f) * oi;
+ }
+
+ float invSqrtDenom = 1.0f / (float) Math.sqrt(sqrDenominator);
+
+ int eventCount = 0;
+ for (int i = 0; i < d; i++) {
+ final float invO = ws.invOAbs[i];
+ final int first = ws.curOBar[i] + 1;
+
+ // Match C++: always consider the first event, then only generate
+ // follow-on events while k <= maxCode.
+ float tn = (float) first * invO;
+ if (tn >= tEnd) {
+ continue;
+ }
+
+ ws.events[eventCount++] = packScalingEvent(tn, i);
+
+ for (int k = first + 1; k <= maxCode; k++) {
+ tn = (float) k * invO;
+ if (tn >= tEnd) {
+ break;
+ }
+
+ ws.events[eventCount++] = packScalingEvent(tn, i);
+ }
+ }
+
+ Arrays.sort(ws.events, 0, eventCount);
+
+ float maxIp = 0.0f;
+ float t = 0.0f;
+
+ for (int e = 0; e < eventCount; e++) {
+ long event = ws.events[e];
+ int id = scalingEventId(event);
+
+ ws.curOBar[id]++;
+ int update = ws.curOBar[id];
+
+ float delta = 2.0f * update;
+ sqrDenominator += delta;
+ numerator += (float) oAbs[id];
+
+ float oldDenom = sqrDenominator - delta;
+ invSqrtDenom *=
+ (1.0f - 0.5f * delta / (oldDenom + delta * 0.5f));
+
+ float curIp = numerator * invSqrtDenom;
+ if (curIp > maxIp) {
+ maxIp = curIp;
+ t = scalingEventT(event);
+ }
+ }
+
+ return t;
+ }
+
+ /**
+ * Packs event sort key as:
+ * high 32 bits: positive float t bits
+ * low 32 bits: dimension id
+ *
+ * For positive finite t, raw float bits sort in the same order as numeric t.
+ * The low bits reproduce the C++ tie-breaker by id.
+ */
+ private static long packScalingEvent(float t, int id) {
+ return ((long) Float.floatToRawIntBits(t) << 32)
+ | (id & 0xFFFF_FFFFL);
+ }
+
+ private static float scalingEventT(long event) {
+ return Float.intBitsToFloat((int) (event >>> 32));
+ }
+
+ private static int scalingEventId(long event) {
+ return (int) event;
+ }
+
+ /**
+ * Samples a Gaussian matrix in flat column-major storage.
+ *
+ * Values are assigned in row-major logical order.
+ */
+ private static double[] gaussianMatrix(int rows, int cols, Random rng) {
+ double[] data = new double[rows * cols];
+ for (int r = 0; r < rows; r++) {
+ int dst = r;
+ for (int c = 0; c < cols; c++, dst += rows) {
+ data[dst] = rng.nextGaussian();
+ }
+ }
+ return data;
+ }
+
+ /** Converts a row-major double[][] to a flat column-major array for BLAS/LAPACK. */
+ private static double[] serializeColumnMajor(double[][] matrix) {
+ int rows = matrix.length;
+ int cols = matrix[0].length;
+ double[] data = new double[rows * cols];
+
+ for (int r = 0; r < rows; r++) {
+ double[] row = matrix[r];
+ int dst = r;
+ for (int c = 0; c < cols; c++, dst += rows) {
+ data[dst] = row[c];
+ }
+ }
+ return data;
+ }
+
+ /** Converts a flat column-major array back into a row-major double[][] for Commons Math. */
+ private static double[][] deserializeColumnMajor(double[] data, int rows, int cols) {
+ double[][] matrix = new double[rows][cols];
+ for (int r = 0; r < rows; r++) {
+ double[] row = matrix[r];
+ int src = r;
+ for (int c = 0; c < cols; c++, src += rows) {
+ row[c] = data[src];
+ }
+ }
+ return matrix;
+ }
+
+ /**
+ * Extracts M = V[:, :d] from VT, where VT is [dim x dim] in flat column-major layout.
+ * Returns M as [dim x d], also in flat column-major layout.
+ */
+ private static double[] extractPcaBasis(double[] vtCol, int dim, int d) {
+ double[] mCol = new double[dim * d];
+
+ for (int col = 0; col < d; col++) {
+ int dstOffset = col * dim;
+ for (int row = 0; row < dim; row++) {
+ // M(row, col) = V(row, col) = VT(col, row)
+ mCol[dstOffset + row] = vtCol[col + row * dim];
+ }
+ }
+
+ return mCol;
+ }
+
+ /** Native DGEMM for flat column-major matrices: C = A @ B. */
+ private static double[] nativeMultiply(double[] aCol, int m, int k, double[] bCol, int n) {
+ double[] cCol = new double[m * n];
+
+ BLAS.getInstance().dgemm(
+ "N", "N",
+ m, n, k,
+ 1.0,
+ aCol, m,
+ bCol, k,
+ 0.0,
+ cCol, m
+ );
+
+ return cCol;
+ }
+
+ /**
+ * Produces contiguous projected vectors without a separate transpose pass.
+ * Xld is [rows x dims] column-major; output rows are its product with R,
+ * stored as [dims x rows] column-major: output = R^T * Xld^T.
+ */
+ private static void projectTrainingRows(double[] xldCol, int rows, int dims,
+ double[] rCol, double[] outRows) {
+ checkTrainingProduct(xldCol, rows, dims, rCol, outRows);
+ if (outRows == xldCol || outRows == rCol) {
+ throw new IllegalArgumentException("ITQ projection output must not alias an input");
+ }
+ BLAS.getInstance().dgemm("T", "T", dims, rows, dims, 1.0,
+ rCol, dims, xldCol, rows, 0.0, outRows, dims);
+ }
+
+ /**
+ * Recomputes M = Xld^T * Xenc into an existing square matrix. Encoded rows
+ * are viewed as [dims x rows] column-major, hence the second transpose.
+ */
+ private static void updateTrainingMatrix(double[] xldCol, int rows, int dims,
+ double[] encodedRows, double[] mCol) {
+ checkTrainingProduct(xldCol, rows, dims, mCol, encodedRows);
+ if (mCol == xldCol || mCol == encodedRows) {
+ throw new IllegalArgumentException("ITQ update output must not alias an input");
+ }
+ BLAS.getInstance().dgemm("T", "T", dims, dims, rows, 1.0,
+ xldCol, rows, encodedRows, dims, 0.0, mCol, dims);
+ }
+
+ private static void checkTrainingProduct(double[] xldCol, int rows, int dims,
+ double[] square, double[] rowMajor) {
+ Objects.requireNonNull(xldCol, "xldCol");
+ Objects.requireNonNull(square, "square");
+ Objects.requireNonNull(rowMajor, "rowMajor");
+ if (rows <= 0 || dims <= 0 || (long) rows * dims > Integer.MAX_VALUE
+ || (long) dims * dims > Integer.MAX_VALUE
+ || xldCol.length != rows * dims || rowMajor.length != rows * dims
+ || square.length != dims * dims) {
+ throw new IllegalArgumentException("Invalid ITQ matrix dimensions");
+ }
+ }
+
+ /**
+ * Computes all right singular vectors with the divide-and-conquer LAPACK driver.
+ * Input X is [rows x cols] in column-major order; the returned VT is [cols x cols].
+ *
+ * For tall matrices, overwrite the scratch input with the unused left singular
+ * vectors instead of allocating another rows-by-cols array. Wide matrices need
+ * the full VT, including its null-space basis, because the requested projection
+ * dimension can exceed the number of training rows.
+ *
+ * Check the singular-value energy before accepting the basis. A native backend
+ * can return success with an inaccurate decomposition; its status alone does not
+ * establish that the PCA input was decomposed correctly.
+ */
+ private static double[] computeNativeVt(double[] xCol, int rows, int cols) {
+ if (rows <= 0 || cols <= 0 || xCol.length != Math.multiplyExact(rows, cols)) {
+ throw new IllegalArgumentException("Invalid PCA matrix dimensions");
+ }
+ double inputEnergy = 0.0;
+ for (double value : xCol) {
+ if (!Double.isFinite(value)) {
+ throw new IllegalArgumentException("PCA training input contains a non-finite value");
+ }
+ inputEnergy += value * value;
+ }
+
+ int k = Math.min(rows, cols);
+ boolean tall = rows >= cols;
+ String job = tall ? "O" : "A";
+ int ldu = tall ? 1 : rows;
+ double[] u = new double[tall ? 1 : Math.multiplyExact(rows, rows)];
+ double[] s = new double[k];
+ double[] vt = new double[Math.multiplyExact(cols, cols)];
+ int[] iwork = new int[Math.multiplyExact(8, k)];
+ intW info = new intW(0);
+ double[] workQuery = new double[1];
+ // LAPACK's workspace query does not modify input matrix values.
+ LAPACK.getInstance().dgesdd(job, rows, cols, xCol, rows, s, u, ldu,
+ vt, cols, workQuery, -1, iwork, info);
+ if (info.val != 0) {
+ throw new IllegalStateException("PCA SVD workspace query failed with info=" + info.val);
+ }
+ if (!Double.isFinite(workQuery[0]) || workQuery[0] < 1 || workQuery[0] > Integer.MAX_VALUE) {
+ throw new IllegalStateException("Invalid PCA SVD workspace size: " + workQuery[0]);
+ }
+ double[] work = new double[(int) Math.ceil(workQuery[0])];
+ double[] a = xCol.clone(); // The actual decomposition overwrites its input.
+ LAPACK.getInstance().dgesdd(job, rows, cols, a, rows, s, u, ldu,
+ vt, cols, work, work.length, iwork, info);
+ if (info.val != 0) {
+ throw new IllegalStateException("PCA SVD failed with info=" + info.val);
+ }
+
+ double singularEnergy = 0.0;
+ for (double value : s) {
+ if (!Double.isFinite(value) || value < 0.0) {
+ throw new IllegalStateException("PCA SVD returned an invalid singular value");
+ }
+ singularEnergy += value * value;
+ }
+ // Training residuals are normalized; allow accumulation roundoff while rejecting
+ // materially incorrect decompositions, including ones reported as successful.
+ double tolerance = 1e-8 * Math.max(1.0, inputEnergy);
+ if (!Double.isFinite(singularEnergy) || Math.abs(singularEnergy - inputEnergy) > tolerance) {
+ throw new IllegalStateException("PCA SVD failed its energy check: input=" + inputEnergy
+ + ", singular values=" + singularEnergy);
+ }
+ return vt;
+ }
+
+ /** Prefer checked covariance PCA; preserve direct SVD for uncertain subspaces. */
+ private static double[] computePcaBasis(double[] xCol, int rows, int dims, int latent) {
+ double[] basis = ASHCovariancePCA.projection(xCol, rows, dims, latent);
+ if (basis != null) {
+ LOG.debug("ASH ITQ: PCA used FP64 covariance eigendecomposition");
+ return basis;
+ }
+ LOG.debug("ASH ITQ: PCA used guarded direct SVD");
+ return extractPcaBasis(computeNativeVt(xCol, rows, dims), dims, latent);
+ }
+
+ /** Single-use convenience path for random projection initialization. */
+ private static double[] orthogonalize(double[] mCol, int rows, int cols) {
+ ASHOrthogonalizer workspace = new ASHOrthogonalizer(rows, cols);
+ double[] result = new double[Math.multiplyExact(rows, cols)];
+ workspace.orthogonalize(mCol, result);
+ return result;
+ }
+
+ /**
+ * Quantizes contiguous training vectors in place, retaining the existing
+ * FP64 norms and FP32 scaling-factor search. Scratch belongs to logical
+ * workers and is reused across epochs, regardless of which pool threads
+ * execute them. One encoder belongs to one trainer and is not thread-safe.
+ */
+ private static final class TrainingDirectionEncoder {
+ private final int rows;
+ private final int dims;
+ private final int elements;
+ private final int bits;
+ private final DirectionScratch[] workers;
+
+ TrainingDirectionEncoder(int rows, int dims, int bits) {
+ validateBitsPerDimension(bits);
+ if (rows <= 0 || dims <= 0 || (long) rows * dims > Integer.MAX_VALUE) {
+ throw new IllegalArgumentException("Invalid ITQ direction dimensions");
+ }
+ this.rows = rows;
+ this.dims = dims;
+ this.elements = rows * dims;
+ this.bits = bits;
+ int count = bits == 1 ? 0
+ : Math.min(rows, PhysicalCoreExecutor.getPhysicalCoreCount());
+ workers = new DirectionScratch[count];
+ for (int i = 0; i < count; i++) workers[i] = new DirectionScratch(dims);
+ }
+
+ void encodeInPlace(double[] vectors) {
+ Objects.requireNonNull(vectors, "vectors");
+ if (vectors.length != elements) {
+ throw new IllegalArgumentException("Invalid ITQ direction buffer length");
+ }
+ if (bits == 1) {
+ double inverse = 1.0 / Math.sqrt(dims);
+ for (int i = 0; i < elements; i++) {
+ vectors[i] = vectors[i] >= 0.0 ? inverse : -inverse;
+ }
+ return;
+ }
+ int maxCode = (1 << (bits - 1)) - 1;
+ if (workers.length == 1) {
+ encodeRows(vectors, workers[0], 0, rows, maxCode);
+ } else {
+ // Share the library's bounded physical-core pool. A trainer's
+ // logical workers own their scratch and disjoint output rows.
+ PhysicalCoreExecutor.instance().execute(() ->
+ IntStream.range(0, workers.length).parallel().forEach(worker -> {
+ int first = (int) ((long) worker * rows / workers.length);
+ int end = (int) ((long) (worker + 1) * rows / workers.length);
+ encodeRows(vectors, workers[worker], first, end, maxCode);
+ }));
+ }
+ }
+
+ private void encodeRows(double[] vectors, DirectionScratch scratch, int first, int end, int maxCode) {
+ double[] absNorm = scratch.absNorm;
+ double[] mag = scratch.mag;
+ for (int r = first; r < end; r++) {
+ int offset = r * dims;
+ double norm2 = 0.0;
+ for (int j = 0; j < dims; j++) {
+ double value = vectors[offset + j];
+ norm2 += value * value;
+ }
+ double inverse = 1.0 / Math.sqrt(Math.max(norm2, 1e-20));
+ for (int j = 0; j < dims; j++) {
+ absNorm[j] = Math.abs(vectors[offset + j]) * inverse;
+ }
+ double scale = computeOptimalScalingFactor(absNorm, dims, bits);
+ double encodedNorm2 = 0.0;
+ for (int j = 0; j < dims; j++) {
+ int q = Math.min((int) (scale * absNorm[j] + K_EPS), maxCode);
+ double magnitude = q + 0.5;
+ mag[j] = magnitude;
+ encodedNorm2 += magnitude * magnitude;
+ }
+ double encodedInverse = 1.0 / Math.sqrt(Math.max(encodedNorm2, 1e-20));
+ for (int j = 0; j < dims; j++) {
+ double signed = vectors[offset + j] < 0.0 ? -mag[j] : mag[j];
+ vectors[offset + j] = signed * encodedInverse;
+ }
+ }
+ }
+ }
+
+ private static final class DirectionScratch {
+ final double[] absNorm;
+ final double[] mag;
+
+ DirectionScratch(int dims) {
+ absNorm = new double[dims];
+ mag = new double[dims];
+ }
+ }
+
+ private static void writeStiefelTransform(IndexWriter out, StiefelTransform st) throws IOException {
+ // Serialize A = W^T in float precision as a [d][D] row-major matrix.
+ // This keeps the format compact and matches the encoding/scoring fast paths.
+ out.writeInt(st.rows);
+ out.writeInt(st.cols);
+
+ for (int i = 0; i < st.rows; i++) {
+ float[] row = st.AFloat[i];
+ if (row.length != st.cols) {
+ throw new IOException("Invalid StiefelTransform AFloat row length: " + row.length + " != " + st.cols);
+ }
+ for (int j = 0; j < st.cols; j++) {
+ out.writeFloat(row[j]);
+ }
+ }
+ }
+
+ private static StiefelTransform readStiefelTransform(RandomAccessReader in, int expectedRows, int expectedCols) throws IOException {
+ int rows = in.readInt();
+ int cols = in.readInt();
+
+ if (rows != expectedRows) {
+ throw new IOException("StiefelTransform rows mismatch: expected " + expectedRows + ", got " + rows);
+ }
+ if (cols != expectedCols) {
+ throw new IOException("StiefelTransform cols mismatch: expected " + expectedCols + ", got " + cols);
+ }
+
+ // Read A = W^T as [rows][cols] (row-major) in float precision.
+ float[][] aFloat = new float[rows][cols];
+ for (int i = 0; i < rows; i++) {
+ float[] row = aFloat[i];
+ for (int j = 0; j < cols; j++) {
+ row[j] = in.readFloat();
+ }
+ }
+
+ // Reconstruct W [D x d] so we can build StiefelTransform(W).
+ // W[j][i] = A[i][j]
+ double[][] wData = new double[cols][rows];
+ for (int i = 0; i < rows; i++) {
+ float[] aRow = aFloat[i];
+ for (int j = 0; j < cols; j++) {
+ wData[j][i] = aRow[j];
+ }
+ }
+
+ RealMatrix W = new Array2DRowRealMatrix(wData, false);
+ return new StiefelTransform(W);
+ }
+
+ @Override
+ public int compressorSize() {
+ // Must exactly match write(out, version)
+ int size = 0;
+ size += Integer.BYTES; // MAGIC
+ size += Integer.BYTES; // version
+ size += Integer.BYTES; // originalDimension
+ size += Integer.BYTES; // encodedBits
+ size += Integer.BYTES; // quantizedDim
+ size += Integer.BYTES; // bitsPerDimension
+ size += Integer.BYTES; // optimizer
+ size += Integer.BYTES; // landmarkCount
+ for (int c = 0; c < landmarkCount; c++) {
+ size += Integer.BYTES; // landmark dimension
+ size += Float.BYTES * landmarks[c].length();
+ }
+ size += Integer.BYTES; // stiefel rows
+ size += Integer.BYTES; // stiefel cols
+ size += Float.BYTES * quantizedDim * originalDimension;
+ size += Float.BYTES; // largest encoded vector length
+ return size;
+ }
+
+ private static float roundToFloat16(float value) {
+ return halfBitsToFloat(floatToHalfBits(value));
+ }
+
+ static void writeFloat16(IndexWriter out, float value) throws IOException {
+ int bits = floatToHalfBits(value);
+
+ // Explicit big-endian byte order, matching Java DataOutput-style primitives.
+ out.writeByte((byte) ((bits >>> 8) & 0xFF));
+ out.writeByte((byte) (bits & 0xFF));
+ }
+
+ static float readFloat16(RandomAccessReader in) throws IOException {
+ int hi = readUnsignedByte(in);
+ int lo = readUnsignedByte(in);
+ return halfBitsToFloat((hi << 8) | lo);
+ }
+
+ private static int readUnsignedByte(RandomAccessReader in) throws IOException {
+ if (in instanceof io.github.jbellis.jvector.disk.ByteReadable) {
+ return ((io.github.jbellis.jvector.disk.ByteReadable) in).readByte() & 0xFF;
+ }
+
+ byte[] one = ONE_BYTE.get();
+ in.readFully(one);
+ return one[0] & 0xFF;
+ }
+
+ private static final ThreadLocal ONE_BYTE =
+ ThreadLocal.withInitial(() -> new byte[1]);
+
+ /**
+ * IEEE-754 binary16 round-to-nearest-even conversion.
+ */
+ private static int floatToHalfBits(float value) {
+ int bits = Float.floatToRawIntBits(value);
+ int sign = (bits >>> 16) & 0x8000;
+ int abs = bits & 0x7FFF_FFFF;
+
+ // NaN / infinity
+ if (abs >= 0x7F80_0000) {
+ if ((abs & 0x007F_FFFF) == 0) {
+ return sign | 0x7C00;
+ }
+ return sign | 0x7E00; // canonical quiet NaN
+ }
+
+ int exp = ((abs >>> 23) & 0xFF) - 127 + 15;
+ int mant = abs & 0x007F_FFFF;
+
+ // Overflow to infinity.
+ if (exp >= 31) {
+ return sign | 0x7C00;
+ }
+
+ // Subnormal / underflow.
+ if (exp <= 0) {
+ if (exp < -10) {
+ return sign;
+ }
+
+ mant |= 0x0080_0000;
+ int shift = 14 - exp;
+ int halfMant = mant >>> shift;
+
+ int roundBit = 1 << (shift - 1);
+ int remainder = mant & (roundBit - 1);
+
+ if ((mant & roundBit) != 0 && (remainder != 0 || (halfMant & 1) != 0)) {
+ halfMant++;
+ }
+
+ return sign | halfMant;
+ }
+
+ // Normalized.
+ int halfMant = mant >>> 13;
+ int roundBit = 0x0000_1000;
+ int remainder = mant & (roundBit - 1);
+
+ if ((mant & roundBit) != 0 && (remainder != 0 || (halfMant & 1) != 0)) {
+ halfMant++;
+
+ if (halfMant == 0x0400) {
+ halfMant = 0;
+ exp++;
+
+ if (exp >= 31) {
+ return sign | 0x7C00;
+ }
+ }
+ }
+
+ return sign | (exp << 10) | halfMant;
+ }
+
+ private static float halfBitsToFloat(int half) {
+ int h = half & 0xFFFF;
+ int sign = (h & 0x8000) << 16;
+ int exp = (h >>> 10) & 0x1F;
+ int mant = h & 0x03FF;
+
+ if (exp == 0) {
+ if (mant == 0) {
+ return Float.intBitsToFloat(sign);
+ }
+
+ // Normalize half subnormal.
+ while ((mant & 0x0400) == 0) {
+ mant <<= 1;
+ exp--;
+ }
+
+ exp++;
+ mant &= 0x03FF;
+ } else if (exp == 31) {
+ return Float.intBitsToFloat(sign | 0x7F80_0000 | (mant << 13));
+ }
+
+ int floatExp = exp + (127 - 15);
+ int bits = sign | (floatExp << 23) | (mant << 13);
+ return Float.intBitsToFloat(bits);
+ }
+
+ @Override
+ public int compressedVectorSize() {
+ return QuantizedVector.serializedSizeBytes(quantizedDim, bitsPerDimension);
+ }
+
+ @Override
+ public double reconstructionError(VectorFloat vector) {
+ throw new UnsupportedOperationException("Reconstruction error not defined for ASH");
+ }
+
+ /**
+ * Retained model storage, including provider-owned native weights and lazy query weights.
+ * Like PQ's compressor accounting, this excludes worker scratch and training temporaries.
+ * Immutable Java backing is measured once; querying this estimate does not scan matrices.
+ */
+ @Override
+ public long ramBytesUsed() {
+ long bytes = heapBytes + stiefelTransform.projection.ramBytesUsed();
+ var queryProjection = stiefelTransform.queryProjection;
+ if (queryProjection != null) {
+ bytes += stiefelTransform.queryProjectionHeapBytes;
+ if (queryProjection != stiefelTransform.projection) bytes += queryProjection.ramBytesUsed();
+ }
+ return bytes;
+ }
+
+ static long estimateArray2DBytes(double[][] a) {
+ if (a == null) return 0L;
+ long bytes = RamUsageEstimator.shallowSizeOf(a);
+ for (double[] row : a) bytes = Math.addExact(bytes, RamUsageEstimator.shallowSizeOf(row));
+ return bytes;
+ }
+
+ static long estimateArray2DBytes(float[][] a) {
+ if (a == null) return 0L;
+ long bytes = RamUsageEstimator.shallowSizeOf(a);
+ for (float[] row : a) bytes = Math.addExact(bytes, RamUsageEstimator.shallowSizeOf(row));
+ return bytes;
+ }
+
+ private static long estimateMatrixBytes(RealMatrix matrix) {
+ long bytes = RamUsageEstimator.shallowSizeOf(matrix);
+ if (matrix instanceof Array2DRowRealMatrix) {
+ return bytes + estimateArray2DBytes(((Array2DRowRealMatrix) matrix).getDataRef());
+ }
+ int rows = matrix.getRowDimension(), cols = matrix.getColumnDimension();
+ if (matrix instanceof org.apache.commons.math3.linear.BlockRealMatrix) {
+ int width = org.apache.commons.math3.linear.BlockRealMatrix.BLOCK_SIZE;
+ long blocks = ((rows - 1L) / width + 1) * ((cols - 1L) / width + 1);
+ bytes += RamUsageEstimator.alignObjectSize(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER
+ + blocks * RamUsageEstimator.NUM_BYTES_OBJECT_REF);
+ // Tail blocks have their own alignment padding; aggregating their
+ // payloads before alignment undercounts some 16-byte-aligned layouts.
+ for (long firstRow = 0; firstRow < rows; firstRow += width) {
+ long blockRows = Math.min(width, rows - firstRow);
+ for (long firstColumn = 0; firstColumn < cols; firstColumn += width) {
+ long blockColumns = Math.min(width, cols - firstColumn);
+ bytes = Math.addExact(bytes, RamUsageEstimator.alignObjectSize(
+ RamUsageEstimator.NUM_BYTES_ARRAY_HEADER
+ + blockRows * blockColumns * Double.BYTES));
+ }
+ }
+ return bytes;
+ }
+ // Other RealMatrix implementations retain the existing dense estimate.
+ // Avoid getData(), which can allocate a full matrix copy just to measure it.
+ return bytes + RamUsageEstimator.alignObjectSize(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER
+ + (long) rows * RamUsageEstimator.NUM_BYTES_OBJECT_REF)
+ + rows * RamUsageEstimator.alignObjectSize(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER
+ + (long) cols * Double.BYTES);
+ }
+
+ // Exact, element-wise comparison of VectorFloat contents.
+ private static boolean vectorEquals(VectorFloat a, VectorFloat b) {
+ if (a == b) return true;
+ if (a == null || b == null) return false;
+ if (a.length() != b.length()) return false;
+
+ for (int i = 0; i < a.length(); i++) {
+ if (Float.floatToIntBits(a.get(i)) != Float.floatToIntBits(b.get(i))) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ // Hash based on the raw bit representation of each float.
+ private static int vectorHash(VectorFloat v) {
+ if (v == null) return 0;
+ int h = 1;
+ for (int i = 0; i < v.length(); i++) {
+ h = 31 * h + Integer.hashCode(Float.floatToIntBits(v.get(i)));
+ }
+ return h;
+ }
+
+ // Exact comparison of two RealMatrix instances by dimensions and entries.
+ private static boolean matrixEquals(RealMatrix a, RealMatrix b) {
+ if (a == b) return true;
+ if (a == null || b == null) return false;
+ if (a.getRowDimension() != b.getRowDimension()
+ || a.getColumnDimension() != b.getColumnDimension()) {
+ return false;
+ }
+
+ for (int i = 0; i < a.getRowDimension(); i++) {
+ for (int j = 0; j < a.getColumnDimension(); j++) {
+ if (Double.doubleToLongBits(a.getEntry(i, j))
+ != Double.doubleToLongBits(b.getEntry(i, j))) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+
+ // Hash code derived from matrix dimensions and raw double bits.
+ private static int matrixHash(RealMatrix m) {
+ if (m == null) return 0;
+ int h = 1;
+ h = 31 * h + m.getRowDimension();
+ h = 31 * h + m.getColumnDimension();
+
+ for (int i = 0; i < m.getRowDimension(); i++) {
+ for (int j = 0; j < m.getColumnDimension(); j++) {
+ long bits = Double.doubleToLongBits(m.getEntry(i, j));
+ h = 31 * h + (int) (bits ^ (bits >>> 32));
+ }
+ }
+ return h;
+ }
+
+ /**
+ * Compares this ASH compressor to another for encoding equivalence.
+ */
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+
+ AsymmetricHashing that = (AsymmetricHashing) o;
+ return originalDimension == that.originalDimension
+ && encodedBits == that.encodedBits
+ && bodyBits == that.bodyBits
+ && quantizedDim == that.quantizedDim
+ && bitsPerDimension == that.bitsPerDimension
+ && optimizer == that.optimizer
+ && landmarksEqual(this.landmarks, that.landmarks)
+ && matrixEquals(
+ this.stiefelTransform == null ? null : this.stiefelTransform.W,
+ that.stiefelTransform == null ? null : that.stiefelTransform.W
+ );
+ }
+
+ private static boolean landmarksEqual(VectorFloat[] a, VectorFloat[] b) {
+ if (a == b) return true;
+ if (a == null || b == null) return false;
+ if (a.length != b.length) return false;
+ for (int i = 0; i < a.length; i++) {
+ if (!vectorEquals(a[i], b[i])) return false;
+ }
+ return true;
+ }
+
+ /**
+ * Hash code consistent with {@link #equals(Object)}.
+ */
+ @Override
+ public int hashCode() {
+ int result = Integer.hashCode(originalDimension);
+ result = 31 * result + Integer.hashCode(encodedBits);
+ result = 31 * result + Integer.hashCode(bodyBits);
+ result = 31 * result + Integer.hashCode(quantizedDim);
+ result = 31 * result + Integer.hashCode(bitsPerDimension);
+ result = 31 * result + Integer.hashCode(optimizer);
+ result = 31 * result + landmarksHash(landmarks);
+ result = 31 * result + matrixHash(
+ stiefelTransform == null ? null : stiefelTransform.W
+ );
+ return result;
+ }
+
+ private static int landmarksHash(VectorFloat[] a) {
+ if (a == null) return 0;
+ int h = 1;
+ for (VectorFloat v : a) {
+ h = 31 * h + vectorHash(v);
+ }
+ return h;
+ }
+
+ @Override
+ public String toString() {
+ String optimizerName;
+ switch (optimizer) {
+ case RANDOM:
+ optimizerName = "RANDOM";
+ break;
+ case ITQ:
+ optimizerName = "ITQ";
+ break;
+ default:
+ optimizerName = "UNKNOWN(" + optimizer + ")";
+ }
+
+ return String.format(
+ "AsymmetricHashing[origDim=%d, encodedBits=%d, bodyBits=%d, quantizedDim=%d, bitsPerDimension=%d, optimizer=%s, landmarks=%d]",
+ originalDimension,
+ encodedBits,
+ bodyBits,
+ quantizedDim,
+ bitsPerDimension,
+ optimizerName,
+ landmarkCount
+ );
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHBlockPacker.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHBlockPacker.java
new file mode 100644
index 000000000..17c5fe9c1
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHBlockPacker.java
@@ -0,0 +1,103 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.quantization;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+import java.nio.ByteOrder;
+
+/** Packs canonical ASH codes into the existing fused neighborhood layout. */
+public final class FusedASHBlockPacker {
+ private static final VarHandle LONGS =
+ MethodHandles.byteArrayViewVarHandle(long[].class, ByteOrder.LITTLE_ENDIAN);
+ private static final ThreadLocal WORK = ThreadLocal.withInitial(() -> new long[8]);
+
+ private FusedASHBlockPacker() {}
+
+ /**
+ * Packs up to {@code blockSize} vectors from {@code vectors[first]} into a
+ * zero-initialized block. Supported block sizes are 8, 16 and 32.
+ * Header encoding and lane order are identical to packQuantizedVector.
+ */
+ public static void pack(byte[] destination, int offset,
+ AsymmetricHashing.QuantizedVector[] vectors, int first, int count,
+ int dimension, int bits, int blockSize) {
+ if (bits == 1 || FusedASHLayout.usesFieldBlocks(bits)) {
+ for (int lane = 0; lane < count; lane++) {
+ FusedASHLayout.packQuantizedVector(destination, offset, lane,
+ vectors[first + lane], dimension, bits, blockSize);
+ }
+ return;
+ }
+ int groups = FusedASHLayout.codeGroups(dimension, bits);
+ int fullBytes = groups / 2;
+ for (int lane = 0; lane < count; lane++) {
+ byte[] code = vectors[first + lane].extraBits;
+ if (code == null || code.length < (groups + 1) / 2) {
+ throw new IllegalArgumentException("Missing or truncated ASH code");
+ }
+ }
+
+ long[] work = WORK.get();
+ int column = 0;
+ // Transpose an 8-neighbor by 8-byte tile using three swap stages.
+ // Explicit little-endian loads/stores preserve the format on either endianness.
+ for (; column + 8 <= fullBytes; column += 8) {
+ for (int laneBase = 0; laneBase < blockSize; laneBase += 8) {
+ for (int i = 0; i < 8; i++) {
+ work[i] = laneBase + i < count
+ ? (long) LONGS.get(vectors[first + laneBase + i].extraBits, column)
+ : 0;
+ }
+ for (int step = 1; step <= 4; step *= 2) {
+ int shift = step * 8;
+ long mask = step == 1 ? 0x00ff00ff00ff00ffL
+ : step == 2 ? 0x0000ffff0000ffffL : 0x00000000ffffffffL;
+ for (int tile = 0; tile < 8; tile += 2 * step) {
+ for (int j = 0; j < step; j++) {
+ long a = work[tile + j];
+ long b = work[tile + j + step];
+ work[tile + j] = (a & mask) | ((b & mask) << shift);
+ work[tile + j + step] = ((a >>> shift) & mask) | (b & ~mask);
+ }
+ }
+ }
+ for (int i = 0; i < 8; i++) {
+ LONGS.set(destination, offset + (column + i) * blockSize + laneBase, work[i]);
+ }
+ }
+ }
+ for (; column < fullBytes; column++) {
+ for (int lane = 0; lane < count; lane++) {
+ destination[offset + column * blockSize + lane] = vectors[first + lane].extraBits[column];
+ }
+ }
+ // An odd group count leaves a high padding nibble; preserve it.
+ if ((groups & 1) != 0) {
+ for (int lane = 0; lane < count; lane++) {
+ int at = offset + fullBytes * blockSize + lane;
+ destination[at] = (byte) ((destination[at] & 0xf0)
+ | (vectors[first + lane].extraBits[fullBytes] & 15));
+ }
+ }
+ for (int lane = 0; lane < count; lane++) {
+ var vector = vectors[first + lane];
+ FusedASHLayout.writeLaneHeader(destination, offset, lane, dimension, bits, blockSize,
+ vector.scale, vector.offset, vector.landmark);
+ }
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHDecoder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHDecoder.java
new file mode 100644
index 000000000..e1c860948
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHDecoder.java
@@ -0,0 +1,393 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.quantization;
+
+import io.github.jbellis.jvector.graph.disk.feature.FusedFeature;
+import io.github.jbellis.jvector.graph.similarity.ScoreFunction;
+import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import org.agrona.collections.Int2ObjectHashMap;
+
+import java.util.Objects;
+
+/**
+ * Scores FusedASH neighborhood blocks.
+ *
+ * The fused layout is intentionally block-oriented: {@link #enableSimilarityToNeighbors(int)}
+ * reads and scores the whole packed neighborhood for one origin node, and
+ * {@link #similarityToNeighbor(int, int)} becomes a cheap array lookup. This preserves the
+ * existing fused-feature search API while avoiding one-neighbor-at-a-time decode work in the
+ * hot graph-expansion loop.
+ *
+ * Each decoder and its mutable scratch/cache are confined to one graph view/search thread.
+ * A decoder borrowing view-owned query scratch is invalidated by the next query on that view.
+ * Standalone constructors own independent preparation state.
+ *
+ * Scalar and SIMD execution share the same packed layout. The property
+ * {@code jvector.ash.blockKernel=auto|scalar|simd} selects execution once per decoder;
+ * forced SIMD fails at construction if the backend cannot provide it.
+ *
+ * FusedASH has two primary scoring paths:
+ *
+ *
+ * - L0 neighborhood scoring uses fused packed ASH blocks via
+ * {@link #enableSimilarityToNeighbors(int)} and
+ * {@link #similarityToNeighbor(int, int)}.
+ * - Upper-layer/direct node scoring uses canonical per-node ASH source
+ * features via {@link #similarityTo(int)}. Upper layers are not fused.
+ *
+ *
+ * FusedASH lane headers are expected to contain projection-adjusted offsets:
+ *
+ *
+ * offset = <x, μ> - ||μ||² - scale * <Aμ, code>
+ *
+ *
+ * Given that encode-time adjustment, the raw dot-product estimate is:
+ *
+ *
+ * rawDotProduct = scale * <Aq, code> + <q, μ> + offset
+ * similarity = (1 + rawDotProduct) / 2 (unclamped, the exact DOT_PRODUCT scale)
+ *
+ */
+public final class FusedASHDecoder implements ScoreFunction.AdaptiveApproximateScoreFunction {
+
+ /**
+ * Source of fused neighborhood bytes. The implementation is provided by the graph feature,
+ * typically by reading the FusedASH inline payload for {@code origin} into {@code dest}.
+ */
+ public interface PackedNeighborhoods {
+ /** Replaces the complete serialized payload for {@code origin} in caller-owned storage. */
+ void readInto(int origin, byte[] dest);
+
+ /** Returns the scored capacity, including padding, of each neighborhood. */
+ int maxDegree();
+
+ /** Returns the complete serialized byte count for the configured layout. */
+ int featureSize();
+ }
+
+ /**
+ * Converts an in-memory upper-layer/source feature into the canonical ASH vector
+ * used for primary direct-node ASH scoring via {@link #similarityTo(int)}.
+ */
+ @FunctionalInterface
+ public interface InlineSourceAccessor {
+ AsymmetricHashing.QuantizedVector vector(FusedFeature.InlineSource source);
+ }
+
+ /**
+ * Reusable preparation and scoring arrays for one graph view and one ASH model.
+ * Starting a new query invalidates that view's previous decoder; use a separate
+ * view for overlapping searches. Standalone constructors retain independent state.
+ */
+ public static final class QueryScratch {
+ private final AsymmetricHashing ash;
+ private final ASHScorer scorer;
+ private final ASHScorer.QueryPrecompute prepared;
+ private final float[] input;
+ private final float[] lut;
+ private final byte[] neighborhoodBytes;
+ private final float[] scores;
+ private final int blockSize;
+ private long generation;
+
+ public QueryScratch(AsymmetricHashing ash, int maxDegree, int blockSize) {
+ this.ash = Objects.requireNonNull(ash, "ash");
+ int size = FusedASHLayout.featureSize(maxDegree, ash.quantizedDim, ash.bitsPerDimension, blockSize);
+ this.blockSize = blockSize;
+ this.scorer = new ASHScorer(ash);
+ this.prepared = scorer.newQueryPrecompute();
+ this.input = new float[ash.originalDimension];
+ this.lut = new float[Math.multiplyExact(FusedASHLayout.codeGroups(ash.quantizedDim, ash.bitsPerDimension), 16)];
+ this.neighborhoodBytes = new byte[size];
+ this.scores = new float[maxDegree];
+ }
+ }
+
+ private final QueryScratch queryScratch;
+ private final long generation;
+ private final ASHQueryPreparer.State externalPrepared;
+ private final long preparedGeneration;
+
+ private final PackedNeighborhoods packedNeighborhoods;
+ private final Int2ObjectHashMap hierarchyCachedFeatures;
+ private final InlineSourceAccessor inlineSourceAccessor;
+ private final ASHScorer.ASHScoreFunction sourceScorer;
+ private final float dotProductBound;
+
+ private final int maxDegree;
+
+ private final byte[] neighborhoodScratch;
+ private final float[] neighborScores;
+
+ private final ASHNeighborhoodScorer neighborhoodScorer;
+
+ private int origin = -1;
+
+ /**
+ * Creates a query-local decoder and shares one projection between L0 and upper-layer scoring.
+ *
+ * @param packedNeighborhoods source of complete fused payloads with a fixed layout
+ * @param ash compressor that encoded the graph's ASH features
+ * @param hierarchyCachedFeatures canonical features available for direct/upper-layer scoring
+ * @param inlineSourceAccessor reads a canonical ASH vector from a cached feature
+ * @param query original-dimensional query, normalized by the caller when appropriate
+ * @param reusableNeighborhoodBytes decoder-owned scratch, at least the complete feature size
+ * @param reusableScores decoder-owned scratch, at least the source's maximum degree
+ * @param blockSize serialized block width: 8, 16, or 32
+ * @param similarityFunction must be DOT_PRODUCT
+ * @throws IllegalArgumentException if dimensions, payload size, or scratch capacity disagree
+ * @throws UnsupportedOperationException for another similarity function
+ */
+ public FusedASHDecoder(
+ PackedNeighborhoods packedNeighborhoods,
+ AsymmetricHashing ash,
+ Int2ObjectHashMap hierarchyCachedFeatures,
+ InlineSourceAccessor inlineSourceAccessor,
+ VectorFloat query,
+ byte[] reusableNeighborhoodBytes,
+ float[] reusableScores,
+ int blockSize,
+ VectorSimilarityFunction similarityFunction) {
+
+ this(packedNeighborhoods, ash, hierarchyCachedFeatures, inlineSourceAccessor, query,
+ reusableNeighborhoodBytes, reusableScores, blockSize, similarityFunction, null, null);
+ }
+
+ private FusedASHDecoder(
+ PackedNeighborhoods packedNeighborhoods, AsymmetricHashing ash,
+ Int2ObjectHashMap hierarchyCachedFeatures,
+ InlineSourceAccessor inlineSourceAccessor, VectorFloat query,
+ byte[] reusableNeighborhoodBytes, float[] reusableScores, int blockSize,
+ VectorSimilarityFunction similarityFunction, QueryScratch scratch,
+ ASHQueryPreparer.State externalPrepared) {
+
+ if (similarityFunction != VectorSimilarityFunction.DOT_PRODUCT) {
+ throw new UnsupportedOperationException("FusedASH supports DOT_PRODUCT only");
+ }
+
+ this.packedNeighborhoods = Objects.requireNonNull(packedNeighborhoods, "packedNeighborhoods");
+ Objects.requireNonNull(ash, "ash");
+ this.hierarchyCachedFeatures = Objects.requireNonNull(
+ hierarchyCachedFeatures,
+ "hierarchyCachedFeatures"
+ );
+ this.inlineSourceAccessor = Objects.requireNonNull(
+ inlineSourceAccessor,
+ "inlineSourceAccessor"
+ );
+ Objects.requireNonNull(query, "query");
+ if (query.length() != ash.originalDimension) {
+ throw new IllegalArgumentException("ASH query dimension must be " + ash.originalDimension);
+ }
+
+ FusedASHLayout.validateBitsPerDimension(ash.bitsPerDimension);
+ FusedASHLayout.validateBlockSize(blockSize);
+
+ this.maxDegree = packedNeighborhoods.maxDegree();
+ if (maxDegree <= 0) {
+ throw new IllegalArgumentException("maxDegree must be > 0");
+ }
+
+ int quantizedDim = ash.quantizedDim;
+ int bitsPerDimension = ash.bitsPerDimension;
+
+ int expectedFeatureSize = FusedASHLayout.featureSize(
+ maxDegree,
+ quantizedDim,
+ bitsPerDimension,
+ blockSize
+ );
+ if (packedNeighborhoods.featureSize() != expectedFeatureSize) {
+ throw new IllegalArgumentException(
+ "FusedASH feature size mismatch: expected " + expectedFeatureSize +
+ ", got " + packedNeighborhoods.featureSize());
+ }
+ if (reusableNeighborhoodBytes.length < expectedFeatureSize) {
+ throw new IllegalArgumentException("reusableNeighborhoodBytes is smaller than FusedASH feature size");
+ }
+ if (reusableScores.length < maxDegree) {
+ throw new IllegalArgumentException("reusableScores is smaller than maxDegree");
+ }
+
+ this.neighborhoodScratch = reusableNeighborhoodBytes;
+ this.neighborScores = reusableScores;
+
+ // Upper-layer source scoring and L0 block scoring use the same Aq and
+ // values. Prepare them once, then build each scoring form.
+ if (scratch != null && (scratch.ash != ash || scratch.blockSize != blockSize
+ || scratch.scores.length != maxDegree)) {
+ throw new IllegalArgumentException("Query scratch belongs to a different ASH model or graph layout");
+ }
+ this.queryScratch = scratch;
+ this.externalPrepared = externalPrepared;
+ this.preparedGeneration = externalPrepared == null ? 0 : externalPrepared.generation;
+ var scorer = scratch == null ? new ASHScorer(ash) : scratch.scorer;
+ ASHScorer.QueryPrecompute preparedQuery;
+ if (externalPrepared != null) {
+ externalPrepared.check(preparedGeneration);
+ this.generation = scratch == null ? 0 : ++scratch.generation;
+ preparedQuery = externalPrepared.prepared;
+ } else if (scratch == null) {
+ this.generation = 0;
+ preparedQuery = scorer.precomputeQuery(query);
+ } else {
+ // Invalidate before projection: failure must not leave a usable old decoder
+ // whose query buffers have been partially overwritten.
+ this.generation = ++scratch.generation;
+ preparedQuery = scratch.prepared;
+ scorer.prepareQuery(query, scratch.input, preparedQuery);
+ }
+ this.neighborhoodScorer = new ASHNeighborhoodScorer(ash, preparedQuery, maxDegree, blockSize,
+ externalPrepared == null ? (scratch == null ? null : scratch.lut) : externalPrepared.lut,
+ externalPrepared != null);
+ this.sourceScorer = scorer.scoreFunctionFor(preparedQuery, neighborhoodScorer.queryLut());
+ this.dotProductBound = scorer.dotProductBound(preparedQuery);
+ }
+
+ /** The query's length times the largest encoded vector length, taken when this decoder was prepared. */
+ @Override
+ public float dotProductBound() {
+ return dotProductBound;
+ }
+
+ public static FusedASHDecoder newDecoder(
+ PackedNeighborhoods packedNeighborhoods,
+ AsymmetricHashing ash,
+ Int2ObjectHashMap hierarchyCachedFeatures,
+ InlineSourceAccessor inlineSourceAccessor,
+ VectorFloat query,
+ byte[] reusableNeighborhoodBytes,
+ float[] reusableScores,
+ int blockSize,
+ VectorSimilarityFunction similarityFunction) {
+ return new FusedASHDecoder(
+ packedNeighborhoods,
+ ash,
+ hierarchyCachedFeatures,
+ inlineSourceAccessor,
+ query,
+ reusableNeighborhoodBytes,
+ reusableScores,
+ blockSize,
+ similarityFunction);
+ }
+
+ /** Creates a query decoder using scratch owned by its graph view. */
+ public static FusedASHDecoder newDecoder(
+ PackedNeighborhoods neighborhoods, AsymmetricHashing ash,
+ Int2ObjectHashMap hierarchyFeatures,
+ InlineSourceAccessor sourceAccessor, VectorFloat query, QueryScratch scratch,
+ int blockSize, VectorSimilarityFunction similarityFunction) {
+ Objects.requireNonNull(scratch, "scratch");
+ return new FusedASHDecoder(neighborhoods, ash, hierarchyFeatures, sourceAccessor, query,
+ scratch.neighborhoodBytes, scratch.scores, blockSize, similarityFunction, scratch, null);
+ }
+
+ /** Binds caller-owned preparation; no projection or LUT construction is repeated. */
+ public static FusedASHDecoder newDecoder(
+ PackedNeighborhoods neighborhoods, AsymmetricHashing ash,
+ Int2ObjectHashMap hierarchyFeatures,
+ InlineSourceAccessor sourceAccessor,
+ io.github.jbellis.jvector.graph.similarity.PreparedQuery preparedQuery,
+ QueryScratch scratch, int blockSize) {
+ Objects.requireNonNull(scratch, "scratch");
+ ASHQueryPreparer.State prepared = ASHQueryPreparer.state(ash, preparedQuery, true);
+ return new FusedASHDecoder(neighborhoods, ash, hierarchyFeatures, sourceAccessor,
+ prepared.queryVector(), scratch.neighborhoodBytes, scratch.scores, blockSize,
+ VectorSimilarityFunction.DOT_PRODUCT, scratch, prepared);
+ }
+
+ private void checkQueryLifetime() {
+ if (externalPrepared != null) externalPrepared.check(preparedGeneration);
+ if (queryScratch != null && generation != queryScratch.generation) {
+ throw new IllegalStateException("FusedASH decoder was invalidated by a new query on the same graph view");
+ }
+ }
+
+ @Override
+ public boolean supportsSimilarityToNeighbors() {
+ return true;
+ }
+
+ @Override
+ public void enableSimilarityToNeighbors(int origin) {
+ checkQueryLifetime();
+ if (this.origin == origin) {
+ return;
+ }
+
+ // Invalidate before I/O: a failed read must not leave partially replaced scores cached.
+ this.origin = -1;
+ packedNeighborhoods.readInto(origin, neighborhoodScratch);
+ scoreNeighborhood();
+ this.origin = origin;
+ }
+
+ /**
+ * Scores a directly-addressed node using its canonical ASH source feature.
+ *
+ * This method is used for upper-layer graph traversal and other direct-node
+ * primary scoring paths. It must not call the configured reranker. Missing
+ * source features indicate an invalid fused graph/source-feature layout.
+ */
+ @Override
+ public float similarityTo(int node2) {
+ checkQueryLifetime();
+ FusedFeature.InlineSource source = hierarchyCachedFeatures.get(node2);
+ if (source == null) {
+ throw new IllegalStateException(
+ "Missing FusedASH source feature for direct ASH scoring of node " + node2 +
+ ". FusedASH requires canonical ASH source features for every upper-layer node " +
+ "that can be scored directly."
+ );
+ }
+
+ AsymmetricHashing.QuantizedVector vector = inlineSourceAccessor.vector(source);
+ if (vector == null) {
+ throw new IllegalStateException(
+ "FusedASH source accessor returned null for node " + node2
+ );
+ }
+
+ return ASHScorer.toSimilarity(sourceScorer.similarityTo(vector));
+ }
+
+ @Override
+ public float similarityToNeighbor(int origin, int neighborIndex) {
+ checkQueryLifetime();
+ if (this.origin != origin) {
+ throw new IllegalArgumentException(
+ "origin must be the same as the origin used to enable similarityToNeighbor");
+ }
+ if (neighborIndex < 0 || neighborIndex >= maxDegree) {
+ throw new IllegalArgumentException("neighborIndex out of range: " + neighborIndex);
+ }
+ return neighborScores[neighborIndex];
+ }
+
+ /** Uses the same prepared scorer as the isolated benchmark; graph I/O stays in this decoder. */
+ private void scoreNeighborhood() {
+ neighborhoodScorer.score(neighborhoodScratch, neighborScores);
+ }
+
+ @Override
+ public String toString() {
+ return "FusedASHDecoder[" + neighborhoodScorer + "]";
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHLayout.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHLayout.java
new file mode 100644
index 000000000..e24c758fa
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedASHLayout.java
@@ -0,0 +1,613 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.quantization;
+
+/**
+ * Layout and scalar reference helpers for FusedASH neighborhood blocks.
+ *
+ * The canonical ASH vector format remains owned by {@link AsymmetricHashing}.
+ * FusedASH stores each origin node's neighbor codes in a neighborhood-local block
+ * layout optimized for scoring the whole adjacency list at once. The projection
+ * code itself follows the C++ fast-scan nibble model for bits-per-dimension
+ * values {@code 1}, {@code 2}, and {@code 4}:
+ *
+ *
+ * - 1 bit/dim: four projected dimensions per nibble
+ * - 2 bits/dim: two projected dimensions per nibble
+ * - 4 bits/dim: one projected dimension per nibble
+ *
+ *
+ * Within a fused block, bytes are stored by code-group pair first and lane
+ * second:
+ *
+ *
+ * byteIndex = blockOffset + (group / 2) * blockSize + lane
+ *
+ *
+ * The low nibble stores even groups and the high nibble stores odd groups.
+ * This is deliberately graph-neighborhood-oriented: it supports block sizes
+ * {@code 8}, {@code 16}, and {@code 32} without copying the IVF-specific B32
+ * addressing scheme.
+ *
+ * Offset-binary codes (3 and 5 to 9 bits per dimension) use field blocks instead, scored with an
+ * INT8 query: each unit of eight dimensions holds every lane's {@code b} code bytes contiguously,
+ *
+ *
+ * byteIndex = blockOffset + (unit * blockSize + lane) * b + k, k in [0, b)
+ *
+ *
+ * so one 64-bit word per lane holds that unit's eight fields. Both layouts are followed by the
+ * same per-lane headers.
+ *
+ * Block and neighborhood byte-size calculations use checked arithmetic. Shapes whose
+ * serialized storage cannot fit a Java array fail with {@link ArithmeticException}.
+ */
+public final class FusedASHLayout {
+ private FusedASHLayout() {}
+
+ public static final int SCALE_BYTES = Short.BYTES;
+ public static final int OFFSET_BYTES = Short.BYTES;
+ public static final int LANDMARK_BYTES = Byte.BYTES;
+ public static final int HEADER_BYTES_PER_LANE = AsymmetricHashing.HEADER_BYTES;
+
+ /**
+ * Every ASH width is supported. Each lane's canonical code is a bit stream scanned in 4-bit
+ * groups: 1-bit signs and 2/4-bit sign-magnitude fields fill whole groups, and the offset-binary
+ * fields of the other widths are linear in their bits, so any group of 4 bits has a 16-entry
+ * query table.
+ */
+ public static boolean supportsBitsPerDimension(int bitsPerDimension) {
+ return bitsPerDimension >= 1 && bitsPerDimension <= 9;
+ }
+
+ public static void validateBitsPerDimension(int bitsPerDimension) {
+ if (!supportsBitsPerDimension(bitsPerDimension)) {
+ throw new IllegalArgumentException(
+ "FusedASH supports bitsPerDimension 1 through 9; got " + bitsPerDimension);
+ }
+ }
+
+ public static void validateBlockSize(int blockSize) {
+ if (blockSize != 8 && blockSize != 16 && blockSize != 32) {
+ throw new IllegalArgumentException(
+ "FusedASH blockSize must be one of {8,16,32}; got " + blockSize);
+ }
+ }
+
+ public static int chooseBlockSize(int maxDegree) {
+ if (maxDegree <= 0) {
+ throw new IllegalArgumentException("maxDegree must be > 0");
+ }
+ if (maxDegree <= 8) return 8;
+ if (maxDegree <= 16) return 16;
+ return 32;
+ }
+
+ /** Dimensions per 4-bit group for the widths whose fields fill whole groups: 1, 2 and 4 bits. */
+ public static int projectionDimsPerNibble(int bitsPerDimension) {
+ if (bitsPerDimension != 1 && bitsPerDimension != 2 && bitsPerDimension != 4) {
+ throw new IllegalArgumentException("Whole-nibble fields require 1, 2 or 4 bits; got " + bitsPerDimension);
+ }
+ return 4 / bitsPerDimension;
+ }
+
+ /** 4-bit groups in one canonical code: {@code ceil(quantizedDim * bits / 4)}. */
+ public static int codeGroups(int quantizedDim, int bitsPerDimension) {
+ if (quantizedDim <= 0) {
+ throw new IllegalArgumentException("quantizedDim must be > 0");
+ }
+ validateBitsPerDimension(bitsPerDimension);
+ return Math.toIntExact(((long) quantizedDim * bitsPerDimension + 3) / 4);
+ }
+
+ public static int canonicalCodeBytes(int quantizedDim, int bitsPerDimension) {
+ return (codeGroups(quantizedDim, bitsPerDimension) - 1) / 2 + 1;
+ }
+
+ /**
+ * Whether the graph neighborhood body uses field blocks: offset-binary codes (3 and 5 to 9 bits) are stored
+ * as {@code [unit][lane][bits bytes]}, where a unit is eight dimensions, so each lane's eight fields are
+ * contiguous for INT8 dot products. Other widths store byte {@code p} of every lane's code in row {@code p}.
+ */
+ public static boolean usesFieldBlocks(int bitsPerDimension) {
+ return AsymmetricHashing.usesOffsetBinaryCode(bitsPerDimension);
+ }
+
+ /** Eight-dimension units in a field-block code. */
+ public static int fieldUnits(int quantizedDim) {
+ if (quantizedDim <= 0) {
+ throw new IllegalArgumentException("quantizedDim must be > 0");
+ }
+ return (quantizedDim - 1) / 8 + 1;
+ }
+
+ /** Body bytes of one graph neighborhood block, in this width's layout. */
+ public static int blockBodyBytes(int quantizedDim, int bitsPerDimension, int blockSize) {
+ validateBlockSize(blockSize);
+ if (usesFieldBlocks(bitsPerDimension)) {
+ return Math.multiplyExact(Math.multiplyExact(fieldUnits(quantizedDim), bitsPerDimension), blockSize);
+ }
+ return nibbleBlockBodyBytes(quantizedDim, bitsPerDimension, blockSize);
+ }
+
+ /** Body bytes of one block whose rows hold byte {@code p} of every lane's code (the LUT layout). */
+ public static int nibbleBlockBodyBytes(int quantizedDim, int bitsPerDimension, int blockSize) {
+ validateBlockSize(blockSize);
+ return Math.multiplyExact(canonicalCodeBytes(quantizedDim, bitsPerDimension), blockSize);
+ }
+
+ /**
+ * Copies one lane's canonical code out of a graph neighborhood block into {@code dest}, which must hold
+ * {@link #canonicalCodeBytes} bytes. For 1-bit codes these are the packed sign nibbles.
+ */
+ public static void readLaneCode(byte[] block, int blockOffset, int lane, int quantizedDim, int bitsPerDimension,
+ int blockSize, byte[] dest) {
+ validateLane(lane, blockSize);
+ int bytes = canonicalCodeBytes(quantizedDim, bitsPerDimension);
+ if (usesFieldBlocks(bitsPerDimension)) {
+ int b = bitsPerDimension;
+ for (int i = 0; i < bytes; i++) {
+ int unit = i / b;
+ dest[i] = block[blockOffset + (unit * blockSize + lane) * b + i % b];
+ }
+ return;
+ }
+ for (int i = 0; i < bytes; i++) dest[i] = block[blockOffset + i * blockSize + lane];
+ }
+
+ public static int blockHeaderBytes(int blockSize) {
+ validateBlockSize(blockSize);
+ return HEADER_BYTES_PER_LANE * blockSize;
+ }
+
+ public static int blockBytes(int quantizedDim, int bitsPerDimension, int blockSize) {
+ return Math.addExact(blockBodyBytes(quantizedDim, bitsPerDimension, blockSize), blockHeaderBytes(blockSize));
+ }
+
+ public static int blocksPerNode(int maxDegree, int blockSize) {
+ if (maxDegree <= 0) {
+ throw new IllegalArgumentException("maxDegree must be > 0");
+ }
+ validateBlockSize(blockSize);
+ return (maxDegree - 1) / blockSize + 1;
+ }
+
+ public static int featureSize(int maxDegree, int quantizedDim, int bitsPerDimension, int blockSize) {
+ return Math.multiplyExact(blocksPerNode(maxDegree, blockSize), blockBytes(quantizedDim, bitsPerDimension, blockSize));
+ }
+
+ public static int blockOffset(int blockIndex, int quantizedDim, int bitsPerDimension, int blockSize) {
+ if (blockIndex < 0) {
+ throw new IllegalArgumentException("blockIndex must be nonnegative");
+ }
+ return Math.multiplyExact(blockIndex, blockBytes(quantizedDim, bitsPerDimension, blockSize));
+ }
+
+ public static int scaleOffset(int blockOffset, int quantizedDim, int bitsPerDimension, int blockSize) {
+ return blockOffset + blockBodyBytes(quantizedDim, bitsPerDimension, blockSize);
+ }
+
+ public static int offsetOffset(int blockOffset, int quantizedDim, int bitsPerDimension, int blockSize) {
+ return scaleOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + SCALE_BYTES * blockSize;
+ }
+
+ public static int landmarkOffset(int blockOffset, int quantizedDim, int bitsPerDimension, int blockSize) {
+ return offsetOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + OFFSET_BYTES * blockSize;
+ }
+
+ public static int getPackedNibble(byte[] block, int blockOffset, int lane, int group, int blockSize) {
+ validateLane(lane, blockSize);
+ int byteIndex = blockOffset + (group >>> 1) * blockSize + lane;
+ int value = block[byteIndex] & 0xFF;
+ return ((group & 1) == 0) ? (value & 0x0F) : ((value >>> 4) & 0x0F);
+ }
+
+ public static void setPackedNibble(byte[] block, int blockOffset, int lane, int group, int blockSize, int nibble) {
+ validateLane(lane, blockSize);
+ int byteIndex = blockOffset + (group >>> 1) * blockSize + lane;
+ int old = block[byteIndex] & 0xFF;
+ int n = nibble & 0x0F;
+ block[byteIndex] = (byte) (((group & 1) == 0)
+ ? ((old & 0xF0) | n)
+ : ((old & 0x0F) | (n << 4)));
+ }
+
+ public static int flatNibble(byte[] code, int group) {
+ int b = code[group >>> 1] & 0xFF;
+ return ((group & 1) == 0) ? (b & 0x0F) : ((b >>> 4) & 0x0F);
+ }
+
+ public static void setFlatNibble(byte[] code, int group, int nibble) {
+ int byteIndex = group >>> 1;
+ int old = code[byteIndex] & 0xFF;
+ int n = nibble & 0x0F;
+ code[byteIndex] = (byte) (((group & 1) == 0)
+ ? ((old & 0xF0) | n)
+ : ((old & 0x0F) | (n << 4)));
+ }
+
+ /** Value of a 1-bit sign or a 2/4-bit sign-magnitude field. */
+ public static float decodeProjectionComponent(int field, int bitsPerDimension) {
+ projectionDimsPerNibble(bitsPerDimension);
+ if (bitsPerDimension == 1) {
+ return field != 0 ? 1.0f : -1.0f;
+ }
+
+ int exBits = bitsPerDimension - 1;
+ int magMask = (1 << exBits) - 1;
+ boolean positive = ((field >>> exBits) & 1) != 0;
+ float mag = (field & magMask) + 0.5f;
+ return positive ? mag : -mag;
+ }
+
+ public static int signNibbleFromWords(long[] signWords, int group, int quantizedDim) {
+ int baseDim = group << 2;
+ int nibble = 0;
+ for (int slot = 0; slot < 4; slot++) {
+ int dim = baseDim + slot;
+ if (dim >= quantizedDim) break;
+ if (((signWords[dim >>> 6] >>> (dim & 63)) & 1L) != 0L) {
+ nibble |= 1 << slot;
+ }
+ }
+ return nibble;
+ }
+
+ public static void packQuantizedVector(
+ byte[] block,
+ int blockOffset,
+ int lane,
+ AsymmetricHashing.QuantizedVector vector,
+ int quantizedDim,
+ int bitsPerDimension,
+ int blockSize) {
+ validateBitsPerDimension(bitsPerDimension);
+ validateLane(lane, blockSize);
+
+ int groups = codeGroups(quantizedDim, bitsPerDimension);
+ if (bitsPerDimension == 1) {
+ for (int group = 0; group < groups; group++) {
+ setPackedNibble(block, blockOffset, lane, group, blockSize,
+ signNibbleFromWords(vector.binaryVector, group, quantizedDim));
+ }
+ } else if (usesFieldBlocks(bitsPerDimension)) {
+ byte[] code = vector.extraBits;
+ int bytes = canonicalCodeBytes(quantizedDim, bitsPerDimension);
+ if (code == null || code.length < bytes) {
+ throw new IllegalArgumentException("ASH projection code is missing or too short");
+ }
+ // Bytes past the canonical code, in the last unit, are zero padding.
+ int b = bitsPerDimension, units = fieldUnits(quantizedDim);
+ for (int unit = 0; unit < units; unit++) {
+ int destination = blockOffset + (unit * blockSize + lane) * b;
+ for (int k = 0; k < b; k++) {
+ int i = unit * b + k;
+ block[destination + k] = i < bytes ? code[i] : 0;
+ }
+ }
+ } else {
+ byte[] code = vector.extraBits;
+ int bytes = canonicalCodeBytes(quantizedDim, bitsPerDimension);
+ if (code == null || code.length < bytes) {
+ throw new IllegalArgumentException("ASH projection code is missing or too short");
+ }
+ int fullBytes = groups >>> 1;
+ int destination = blockOffset + lane;
+ for (int i = 0; i < fullBytes; i++, destination += blockSize) {
+ block[destination] = code[i];
+ }
+ if ((groups & 1) != 0) {
+ block[destination] = (byte) ((block[destination] & 0xF0) | (code[fullBytes] & 0x0F));
+ }
+ }
+
+ writeLaneHeader(block, blockOffset, lane, quantizedDim, bitsPerDimension, blockSize,
+ vector.scale, vector.offset, vector.landmark);
+ }
+
+ /**
+ * Packs one canonical ASH vector into the graph's fused neighborhood representation.
+ *
+ * Fused scoring uses {@code Aq} for every bit width. Canonical 1/2/4-bit offsets already
+ * contain the landmark adjustment {@code -scale * dot(A(mu), code)}, so the header is copied as
+ * stored. Benchmark fixtures call this to use the graph writer's exact bytes.
+ *
+ * @param block destination containing a complete serialized block
+ * @param blockOffset byte offset of that block
+ * @param lane destination lane within the block
+ * @param vector canonical quantized vector, whose landmark belongs to {@code ash}
+ * @param ash compressor that owns the vector's projection and landmarks
+ * @param blockSize serialized block width: 8, 16, or 32
+ * @throws IllegalArgumentException if the landmark or block shape is invalid
+ */
+ public static void packForFused(byte[] block, int blockOffset, int lane,
+ AsymmetricHashing.QuantizedVector vector,
+ AsymmetricHashing ash, int blockSize) {
+ int landmark = vector.landmark & 0xFF;
+ if (landmark >= ash.landmarkCount) {
+ throw new IllegalArgumentException("Invalid ASH landmark: " + landmark);
+ }
+ packQuantizedVector(block, blockOffset, lane, vector,
+ ash.quantizedDim, ash.bitsPerDimension, blockSize);
+ }
+
+ public static void writeLaneHeader(
+ byte[] block,
+ int blockOffset,
+ int lane,
+ int quantizedDim,
+ int bitsPerDimension,
+ int blockSize,
+ float scale,
+ float offset,
+ byte landmark) {
+ validateLane(lane, blockSize);
+ writeFloat16(block, scaleOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + lane * SCALE_BYTES, scale);
+ writeFloat16(block, offsetOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + lane * OFFSET_BYTES, offset);
+ block[landmarkOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + lane] = landmark;
+ }
+
+ public static float readScale(byte[] block, int blockOffset, int lane, int quantizedDim, int bitsPerDimension, int blockSize) {
+ validateLane(lane, blockSize);
+ return readFloat16(block, scaleOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + lane * SCALE_BYTES);
+ }
+
+ public static float readOffset(byte[] block, int blockOffset, int lane, int quantizedDim, int bitsPerDimension, int blockSize) {
+ validateLane(lane, blockSize);
+ return readFloat16(block, offsetOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + lane * OFFSET_BYTES);
+ }
+
+ public static int readLandmark(byte[] block, int blockOffset, int lane, int quantizedDim, int bitsPerDimension, int blockSize) {
+ validateLane(lane, blockSize);
+ return block[landmarkOffset(blockOffset, quantizedDim, bitsPerDimension, blockSize) + lane] & 0xFF;
+ }
+
+ public static void buildQueryLut(float[] qProj, int quantizedDim, int bitsPerDimension, float[] lut) {
+ if (AsymmetricHashing.usesOffsetBinaryCode(bitsPerDimension)) {
+ buildOffsetBinaryLut(qProj, quantizedDim, bitsPerDimension, lut);
+ return;
+ }
+ validateBitsPerDimension(bitsPerDimension);
+ int groupDims = projectionDimsPerNibble(bitsPerDimension);
+ int groups = codeGroups(quantizedDim, bitsPerDimension);
+ if (qProj.length < quantizedDim) {
+ throw new IllegalArgumentException("qProj length is smaller than quantizedDim");
+ }
+ if (lut.length < groups * 16) {
+ throw new IllegalArgumentException("lut length is smaller than groups * 16");
+ }
+
+ // Every entry is still 0 + q[d0]*c0 + q[d1]*c1 + ..., added in slot order,
+ // so the table is bit-identical to the per-entry definition. Each product
+ // is formed once per dimension instead of once per entry, and the field
+ // decode is a constant table rather than a branch per term.
+ int complete = quantizedDim / groupDims;
+ if (bitsPerDimension == 2) {
+ for (int group = 0; group < complete; group++) lut2(qProj, group << 1, lut, group << 4);
+ } else if (bitsPerDimension == 1) {
+ for (int group = 0; group < complete; group++) lut1(qProj, group << 2, lut, group << 4);
+ } else {
+ for (int group = 0; group < complete; group++) {
+ float q = qProj[group];
+ int base = group << 4;
+ for (int field = 0; field < 16; field++) lut[base + field] = 0f + q * DECODE_4BIT[field];
+ }
+ }
+ // A partial final group adds only the dimensions that exist.
+ for (int group = complete; group < groups; group++) {
+ int dimBase = group * groupDims;
+ int lutBase = group << 4;
+ for (int nibble = 0; nibble < 16; nibble++) {
+ float sum = 0f;
+ for (int slot = 0; slot < groupDims; slot++) {
+ int dim = dimBase + slot;
+ if (dim >= quantizedDim) break;
+ int mask = (1 << bitsPerDimension) - 1;
+ int field = (nibble >>> (slot * bitsPerDimension)) & mask;
+ sum += qProj[dim] * decodeProjectionComponent(field, bitsPerDimension);
+ }
+ lut[lutBase + nibble] = sum;
+ }
+ }
+ }
+
+ /**
+ * Table for offset-binary codes. Group {@code k} covers code bits {@code 4k..4k+3}; bit {@code i} is
+ * bit {@code i % b} of dimension {@code i / b} and adds {@code 2^(i % b) * q[i / b]}. Each dimension's
+ * centering term {@code -(2^b - 1)/2 * q[j]} is added to every entry of the group holding its bit 0,
+ * so each lookup sum includes it exactly once.
+ */
+ private static void buildOffsetBinaryLut(float[] qProj, int quantizedDim, int bits, float[] lut) {
+ int groups = codeGroups(quantizedDim, bits);
+ if (qProj.length < quantizedDim) {
+ throw new IllegalArgumentException("qProj length is smaller than quantizedDim");
+ }
+ if (lut.length < groups * 16) {
+ throw new IllegalArgumentException("lut length is smaller than groups * 16");
+ }
+ float center = ((1 << bits) - 1) * 0.5f;
+ float[] weight = new float[4];
+ for (int group = 0; group < groups; group++) {
+ float base = 0f;
+ for (int slot = 0; slot < 4; slot++) {
+ int bit = 4 * group + slot, dim = bit / bits, power = bit % bits;
+ weight[slot] = dim < quantizedDim ? qProj[dim] * (float) (1 << power) : 0f;
+ if (dim < quantizedDim && power == 0) base -= center * qProj[dim];
+ }
+ int at = group << 4;
+ lut[at] = base;
+ for (int n = 1; n < 16; n++) lut[at + n] = lut[at + (n & (n - 1))] + weight[Integer.numberOfTrailingZeros(n)];
+ }
+ }
+
+ private static final float[] DECODE_4BIT = new float[16];
+ static {
+ for (int field = 0; field < 16; field++) DECODE_4BIT[field] = decodeProjectionComponent(field, 4);
+ }
+
+ // Two 2-bit fields per nibble: entry = (0 + a[n & 3]) + b[n >>> 2].
+ // Field values -1.5, -0.5, 0.5, 1.5 (sign bit high, magnitude low).
+ private static void lut2(float[] qProj, int dim, float[] lut, int base) {
+ float q0 = qProj[dim], q1 = qProj[dim + 1];
+ float a0 = 0f + q0 * -0.5f, a1 = 0f + q0 * -1.5f, a2 = 0f + q0 * 0.5f, a3 = 0f + q0 * 1.5f;
+ float b0 = q1 * -0.5f, b1 = q1 * -1.5f, b2 = q1 * 0.5f, b3 = q1 * 1.5f;
+ lut[base] = a0 + b0; lut[base + 1] = a1 + b0; lut[base + 2] = a2 + b0; lut[base + 3] = a3 + b0;
+ lut[base + 4] = a0 + b1; lut[base + 5] = a1 + b1; lut[base + 6] = a2 + b1; lut[base + 7] = a3 + b1;
+ lut[base + 8] = a0 + b2; lut[base + 9] = a1 + b2; lut[base + 10] = a2 + b2; lut[base + 11] = a3 + b2;
+ lut[base + 12] = a0 + b3; lut[base + 13] = a1 + b3; lut[base + 14] = a2 + b3; lut[base + 15] = a3 + b3;
+ }
+
+ // Four sign bits per nibble: entry = (((0 + ±q0) + ±q1) + ±q2) + ±q3.
+ private static void lut1(float[] qProj, int dim, float[] lut, int base) {
+ float p0 = qProj[dim] * 1.0f, n0 = qProj[dim] * -1.0f;
+ float p1 = qProj[dim + 1] * 1.0f, n1 = qProj[dim + 1] * -1.0f;
+ float p2 = qProj[dim + 2] * 1.0f, n2 = qProj[dim + 2] * -1.0f;
+ float p3 = qProj[dim + 3] * 1.0f, n3 = qProj[dim + 3] * -1.0f;
+ for (int nibble = 0; nibble < 16; nibble++) {
+ float sum = 0f + ((nibble & 1) != 0 ? p0 : n0);
+ sum += (nibble & 2) != 0 ? p1 : n1;
+ sum += (nibble & 4) != 0 ? p2 : n2;
+ sum += (nibble & 8) != 0 ? p3 : n3;
+ lut[base + nibble] = sum;
+ }
+ }
+
+ public static float scoreLane(
+ byte[] block,
+ int blockOffset,
+ int lane,
+ int quantizedDim,
+ int bitsPerDimension,
+ int blockSize,
+ float[] lut,
+ float[] dotQMuByLandmark) {
+ validateLane(lane, blockSize);
+ int groups = codeGroups(quantizedDim, bitsPerDimension);
+ float ip = 0f;
+ for (int group = 0; group < groups; group++) {
+ int nibble = getPackedNibble(block, blockOffset, lane, group, blockSize);
+ ip += lut[(group << 4) + nibble];
+ }
+
+ int landmark = readLandmark(block, blockOffset, lane, quantizedDim, bitsPerDimension, blockSize);
+ float dotQMu = landmark < dotQMuByLandmark.length ? dotQMuByLandmark[landmark] : 0f;
+ return readScale(block, blockOffset, lane, quantizedDim, bitsPerDimension, blockSize) * ip
+ + dotQMu
+ + readOffset(block, blockOffset, lane, quantizedDim, bitsPerDimension, blockSize);
+ }
+
+ private static void validateLane(int lane, int blockSize) {
+ validateBlockSize(blockSize);
+ if (lane < 0 || lane >= blockSize) {
+ throw new IllegalArgumentException("lane out of range: " + lane + " for blockSize=" + blockSize);
+ }
+ }
+
+ private static void writeFloat16(byte[] out, int offset, float value) {
+ int bits = floatToHalfBits(value);
+ out[offset] = (byte) ((bits >>> 8) & 0xFF);
+ out[offset + 1] = (byte) (bits & 0xFF);
+ }
+
+ private static float readFloat16(byte[] in, int offset) {
+ int hi = in[offset] & 0xFF;
+ int lo = in[offset + 1] & 0xFF;
+ return halfBitsToFloat((hi << 8) | lo);
+ }
+
+ /** IEEE-754 binary16 round-to-nearest-even conversion. */
+ private static int floatToHalfBits(float value) {
+ int bits = Float.floatToRawIntBits(value);
+ int sign = (bits >>> 16) & 0x8000;
+ int abs = bits & 0x7FFF_FFFF;
+
+ if (abs >= 0x7F80_0000) {
+ if ((abs & 0x007F_FFFF) == 0) {
+ return sign | 0x7C00;
+ }
+ return sign | 0x7E00;
+ }
+
+ int exp = ((abs >>> 23) & 0xFF) - 127 + 15;
+ int mant = abs & 0x007F_FFFF;
+
+ if (exp >= 31) {
+ return sign | 0x7C00;
+ }
+
+ if (exp <= 0) {
+ if (exp < -10) {
+ return sign;
+ }
+
+ mant |= 0x0080_0000;
+ int shift = 14 - exp;
+ int halfMant = mant >>> shift;
+
+ int roundBit = 1 << (shift - 1);
+ int remainder = mant & (roundBit - 1);
+
+ if ((mant & roundBit) != 0 && (remainder != 0 || (halfMant & 1) != 0)) {
+ halfMant++;
+ }
+
+ return sign | halfMant;
+ }
+
+ int halfMant = mant >>> 13;
+ int roundBit = 0x0000_1000;
+ int remainder = mant & (roundBit - 1);
+
+ if ((mant & roundBit) != 0 && (remainder != 0 || (halfMant & 1) != 0)) {
+ halfMant++;
+ if (halfMant == 0x0400) {
+ halfMant = 0;
+ exp++;
+ if (exp >= 31) {
+ return sign | 0x7C00;
+ }
+ }
+ }
+
+ return sign | (exp << 10) | halfMant;
+ }
+
+ private static float halfBitsToFloat(int half) {
+ int h = half & 0xFFFF;
+ int sign = (h & 0x8000) << 16;
+ int exp = (h >>> 10) & 0x1F;
+ int mant = h & 0x03FF;
+
+ if (exp == 0) {
+ if (mant == 0) {
+ return Float.intBitsToFloat(sign);
+ }
+
+ while ((mant & 0x0400) == 0) {
+ mant <<= 1;
+ exp--;
+ }
+ exp++;
+ mant &= 0x03FF;
+ } else if (exp == 31) {
+ return Float.intBitsToFloat(sign | 0x7F80_0000 | (mant << 13));
+ }
+
+ int floatExp = exp + (127 - 15);
+ return Float.intBitsToFloat(sign | (floatExp << 23) | (mant << 13));
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedPQDecoder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedPQDecoder.java
index b52d0ca66..282da9db1 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedPQDecoder.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/FusedPQDecoder.java
@@ -124,6 +124,17 @@ public DotProductDecoder(FusedPQ.PackedNeighbors neighbors, ProductQuantization
protected float distanceToScore(float distance) {
return (distance + 1) / 2;
}
+
+ /** PQ scores the centered query against centered reconstructions, so their lengths bound the score. */
+ @Override
+ public float dotProductBound() {
+ var center = pq.globalCentroid;
+ double squared = VectorUtil.dotProduct(query, query);
+ if (center != null) {
+ squared += VectorUtil.dotProduct(center, center) - 2.0 * VectorUtil.dotProduct(query, center);
+ }
+ return (float) Math.sqrt(Math.max(0.0, squared)) * pq.maxEncodedNorm();
+ }
}
static class EuclideanDecoder extends FusedPQDecoder {
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ImmutablePQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ImmutablePQVectors.java
index 8df05b675..e8afe756a 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ImmutablePQVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ImmutablePQVectors.java
@@ -37,11 +37,24 @@ public class ImmutablePQVectors extends PQVectors {
* @param vectorsPerChunk the number of vectors per chunk
*/
public ImmutablePQVectors(ProductQuantization pq, ByteSequence[] compressedDataChunks, int vectorCount, int vectorsPerChunk) {
+ this(pq, compressedDataChunks, vectorCount, vectorsPerChunk, true);
+ }
+
+ /**
+ * @param recordCodes whether {@code pq} must count these codes' lengths in its
+ * {@link ProductQuantization#maxEncodedNorm}: true unless it encoded them, or they were saved
+ * with it
+ */
+ ImmutablePQVectors(ProductQuantization pq, ByteSequence[] compressedDataChunks, int vectorCount,
+ int vectorsPerChunk, boolean recordCodes) {
super(pq);
this.compressedDataChunks = compressedDataChunks;
this.vectorCount = vectorCount;
this.vectorsPerChunk = vectorsPerChunk;
this.codebookPartialSumsMap = new ConcurrentHashMap<>();
+ if (recordCodes) {
+ for (int i = 0; i < vectorCount; i++) pq.recordCode(getChunk(i), getOffsetInChunk(i));
+ }
}
@Override
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java
index 29f0e884a..dc05709b8 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/KMeansPlusPlusClusterer.java
@@ -37,8 +37,17 @@
* A KMeans++ implementation for float vectors. Optimizes to use SIMD vector instructions if available.
*/
public class KMeansPlusPlusClusterer {
+ private static final org.slf4j.Logger LOG = org.slf4j.LoggerFactory.getLogger(KMeansPlusPlusClusterer.class);
+
private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+ // Observability flags (disabled by default)
+ private static final boolean VERBOSE_KMEANS =
+ Boolean.getBoolean("jvector.kmeans.verbose");
+
+ private static final boolean VERBOSE_KMEANS_TIMING =
+ Boolean.getBoolean("jvector.kmeans.timing");
+
public static final float UNWEIGHTED = -1.0f;
// number of centroids to compute
@@ -130,18 +139,70 @@ static float computeParallelCostMultiplier(double threshold, int dimensions) {
* @return a VectorFloat of cluster centroids.
*/
public VectorFloat cluster(int unweightedIterations, int anisotropicIterations) {
- // Always cluster unweighted first, it is significantly faster
+
+ final int n = points.length;
+ final int threshold = Math.max(1, (int) (0.01 * n));
+
+ // ------------------------------------------------------------
+ // Unweighted (isotropic) k-means phase
+ // ------------------------------------------------------------
for (int i = 0; i < unweightedIterations; i++) {
+
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
int changedCount = clusterOnceUnweighted();
- if (changedCount <= 0.01 * points.length) {
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ if (VERBOSE_KMEANS) {
+ LOG.info(String.format("KMeans unweighted iter %d: %d/%d reassigned (%.3f%%)",
+ i,
+ changedCount,
+ n,
+ 100.0 * changedCount / n)
+ );
+ }
+
+ if (VERBOSE_KMEANS_TIMING) {
+ LOG.info(String.format(" time = %.3f ms",
+ (t1 - t0) * 1e-6)
+ );
+ }
+
+ if (changedCount <= threshold) {
+ if (VERBOSE_KMEANS) {
+ LOG.info("KMeans unweighted converged early");
+ }
break;
}
}
- // Optionally, refine using anisotropic clustering
+ // ------------------------------------------------------------
+ // Anisotropic refinement phase
+ // ------------------------------------------------------------
for (int i = 0; i < anisotropicIterations; i++) {
+
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
int changedCount = clusterOnceAnisotropic();
- if (changedCount <= 0.01 * points.length) {
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ if (VERBOSE_KMEANS) {
+ LOG.info(String.format("KMeans anisotropic iter %d: %d/%d reassigned (%.3f%%)",
+ i,
+ changedCount,
+ n,
+ 100.0 * changedCount / n)
+ );
+ }
+
+ if (VERBOSE_KMEANS_TIMING) {
+ LOG.info(String.format(" time = %.3f ms",
+ (t1 - t0) * 1e-6)
+ );
+ }
+
+ if (changedCount <= threshold) {
+ if (VERBOSE_KMEANS) {
+ LOG.info("KMeans anisotropic converged early");
+ }
break;
}
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MiniBatchKMeansClusterer.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MiniBatchKMeansClusterer.java
new file mode 100644
index 000000000..ccf1d42dd
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MiniBatchKMeansClusterer.java
@@ -0,0 +1,419 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.quantization;
+
+import io.github.jbellis.jvector.vector.Matrix;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import io.github.jbellis.jvector.vector.types.VectorTypeSupport;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static io.github.jbellis.jvector.util.MathUtil.square;
+import static io.github.jbellis.jvector.vector.VectorUtil.*;
+import static java.lang.Math.max;
+
+/**
+ * A MiniBatch KMeans implementation for float vectors.
+ * Drop-in replacement for KMeansPlusPlusClusterer, optimized for large datasets (10M+).
+ * * Includes:
+ * 1. Mini-Batch Stochastic Updates (fast convergence).
+ * 2. K-Means++ Subsampling Initialization (fast startup).
+ */
+public class MiniBatchKMeansClusterer {
+ private static final VectorTypeSupport vectorTypeSupport = VectorizationProvider.getInstance().getVectorTypeSupport();
+
+ private static final boolean VERBOSE_KMEANS = Boolean.getBoolean("jvector.kmeans.verbose");
+ private static final boolean VERBOSE_KMEANS_TIMING = Boolean.getBoolean("jvector.kmeans.timing");
+
+ public static final float UNWEIGHTED = -1.0f;
+
+ private final int k;
+ private final VectorFloat[] points;
+ private final int[] assignments;
+ private final VectorFloat centroids;
+ private final float anisotropicThreshold;
+
+ // MiniBatch specific
+ private final int batchSize;
+ private final int[] centroidCounts;
+
+ /**
+ * @param batchSize Recommended: 1024 or 2048 for large datasets
+ */
+ public MiniBatchKMeansClusterer(VectorFloat[] points, int k, int batchSize) {
+ this(points, chooseInitialCentroids(points, k), UNWEIGHTED, batchSize);
+ }
+
+ public MiniBatchKMeansClusterer(VectorFloat[] points, int k, float anisotropicThreshold, int batchSize) {
+ this(points, chooseInitialCentroids(points, k), anisotropicThreshold, batchSize);
+ }
+
+ public MiniBatchKMeansClusterer(VectorFloat[] points, VectorFloat centroids, float anisotropicThreshold, int batchSize) {
+ if (Float.isNaN(anisotropicThreshold) || anisotropicThreshold < -1.0 || anisotropicThreshold >= 1.0) {
+ throw new IllegalArgumentException("Valid range for anisotropic threshold T is -1.0 <= t < 1.0");
+ }
+
+ this.points = points;
+ this.k = centroids.length() / points[0].length();
+ this.centroids = centroids.copy();
+ this.anisotropicThreshold = anisotropicThreshold;
+ this.batchSize = Math.min(batchSize, points.length);
+
+ this.centroidCounts = new int[k];
+ this.assignments = new int[points.length];
+ Arrays.fill(this.assignments, -1);
+ }
+
+ static float computeParallelCostMultiplier(double threshold, int dimensions) {
+ assert Double.isFinite(threshold) : "threshold=" + threshold;
+ double parallelCost = threshold * threshold;
+ double perpendicularCost = (1 - parallelCost) / (dimensions - 1);
+ return (float) max(1.0, (parallelCost / perpendicularCost));
+ }
+
+ /**
+ * Performs clustering.
+ * @param unweightedIterations If small ({@code < 100}), interpreted as EPOCHS (full passes).
+ * If large, interpreted as TOTAL BATCH UPDATES.
+ * @param anisotropicIterations Number of full-batch anisotropic refinement steps.
+ */
+ public VectorFloat cluster(int unweightedIterations, int anisotropicIterations) {
+
+ // ------------------------------------------------------------
+ // Mini-Batch Unweighted (isotropic) k-means phase
+ // ------------------------------------------------------------
+ int totalBatches = unweightedIterations;
+ // Interpret small iteration counts as Epochs to ensure sufficient data coverage
+ if (unweightedIterations < 100) {
+ int batchesPerEpoch = Math.max(1, points.length / batchSize);
+ totalBatches = unweightedIterations * batchesPerEpoch;
+ }
+
+ long tTotal0 = System.nanoTime();
+
+ for (int i = 0; i < totalBatches; i++) {
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ clusterOnceMiniBatch();
+
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+ if (VERBOSE_KMEANS_TIMING && i % 100 == 0) {
+ System.out.printf("MiniBatch iter %d time = %.3f ms%n", i, (t1 - t0) * 1e-6);
+ }
+ }
+
+ if (VERBOSE_KMEANS) {
+ System.out.printf("MiniBatch unweighted phase finished. Total batches: %d. Total time: %.3f s%n",
+ totalBatches, (System.nanoTime() - tTotal0) * 1e-9);
+ }
+
+ // ------------------------------------------------------------
+ // Anisotropic refinement phase (Full Batch)
+ // ------------------------------------------------------------
+ if (anisotropicIterations > 0) {
+ // MiniBatch leaves assignments in flux; we must perform a full assignment pass
+ // so the anisotropic phase has valid global clusters.
+ assignAllPoints();
+
+ final int threshold = Math.max(1, (int) (0.01 * points.length));
+
+ for (int i = 0; i < anisotropicIterations; i++) {
+ long t0 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+ int changedCount = clusterOnceAnisotropic();
+ long t1 = VERBOSE_KMEANS_TIMING ? System.nanoTime() : 0L;
+
+ if (VERBOSE_KMEANS) {
+ System.out.printf("KMeans anisotropic iter %d: %d/%d reassigned (%.3f%%)%n",
+ i, changedCount, points.length, 100.0 * changedCount / points.length);
+ }
+
+ if (VERBOSE_KMEANS_TIMING) {
+ System.out.printf(" time = %.3f ms%n", (t1 - t0) * 1e-6);
+ }
+
+ if (changedCount <= threshold) {
+ if (VERBOSE_KMEANS) System.out.println("KMeans anisotropic converged early");
+ break;
+ }
+ }
+ }
+
+ return centroids;
+ }
+
+ /**
+ * Performs a single MiniBatch update step:
+ * 1. Sample batch
+ * 2. Find nearest centroids
+ * 3. Update centroids immediately using learning rate
+ */
+ public void clusterOnceMiniBatch() {
+ var random = ThreadLocalRandom.current();
+
+ // 1. Sample indices
+ int[] batchIndices = new int[batchSize];
+ for (int i = 0; i < batchSize; i++) {
+ batchIndices[i] = random.nextInt(points.length);
+ }
+
+ // 2. Compute assignments for the batch
+ int[] batchAssignments = new int[batchSize];
+ for(int i = 0; i < batchSize; i++) {
+ batchAssignments[i] = getNearestCluster(points[batchIndices[i]]);
+ }
+
+ // 3. Update Centroids (Online)
+ int dim = centroids.length() / k;
+ VectorFloat tempDiff = vectorTypeSupport.createFloatVector(dim);
+
+ for (int i = 0; i < batchSize; i++) {
+ int pointIdx = batchIndices[i];
+ int clusterIdx = batchAssignments[i];
+ VectorFloat point = points[pointIdx];
+
+ centroidCounts[clusterIdx]++;
+ float learningRate = 1.0f / centroidCounts[clusterIdx];
+
+ // tempDiff = X
+ tempDiff.copyFrom(point, 0, 0, dim);
+ // tempDiff = X - C_old
+ subFromOffset(tempDiff, centroids, clusterIdx * dim);
+
+ // tempDiff = lr * (X - C_old)
+ scale(tempDiff, learningRate);
+
+ // C_new = C_old + tempDiff
+ // Pass the offset explicitly as the second argument
+ addToOffset(centroids, clusterIdx * dim, tempDiff);
+ }
+ }
+
+ // --- Helpers ---
+
+ private void assignAllPoints() {
+ for (int i = 0; i < points.length; i++) {
+ assignments[i] = getNearestCluster(points[i]);
+ }
+ }
+
+ public int clusterOnceAnisotropic() {
+ updateCentroidsAnisotropic();
+ return updateAssignedPointsAnisotropic();
+ }
+
+ // addToOffset: Adds a SMALL vector (source) into a LARGE vector (dest) at a specific offset.
+ private static void addToOffset(VectorFloat dest, int destOffset, VectorFloat source) {
+ // Fix: Iterate over source.length() (Small Temp Vector)
+ for (int i = 0; i < source.length(); i++) {
+ dest.set(destOffset + i, dest.get(destOffset + i) + source.get(i));
+ }
+ }
+
+ // subFromOffset: Subtracts a slice of a LARGE vector (source) from a SMALL vector (dest).
+ // (This logic was actually correct, but verifying strictly ensures safety)
+ private static void subFromOffset(VectorFloat dest, VectorFloat source, int sourceOffset) {
+ for (int i = 0; i < dest.length(); i++) {
+ dest.set(i, dest.get(i) - source.get(sourceOffset + i));
+ }
+ }
+
+ // --- Optimized Initialization Logic ---
+
+ /**
+ * Chooses the initial centroids for clustering using K-Means++.
+ * If the dataset is large, uses a random subsample to accelerate initialization.
+ */
+ private static VectorFloat chooseInitialCentroids(VectorFloat[] points, int k) {
+ if (k <= 0) throw new IllegalArgumentException("Number of clusters must be positive.");
+ if (k > points.length) throw new IllegalArgumentException("K cannot exceed N");
+
+ // Heuristic: If N is massive, run KMeans++ on a subset only.
+ // Cap at 50,000 points or 10*k to ensure statistical representation without performance penalty.
+ int maxInitSamples = Math.max(k * 10, 50_000);
+
+ if (points.length <= maxInitSamples) {
+ return runStandardKMeansPlusPlus(points, k);
+ } else {
+ if (VERBOSE_KMEANS) {
+ System.out.printf("Subsampling %d points for K-Means++ initialization...%n", maxInitSamples);
+ }
+ var random = ThreadLocalRandom.current();
+ VectorFloat[] subset = new VectorFloat[maxInitSamples];
+ for (int i = 0; i < maxInitSamples; i++) {
+ subset[i] = points[random.nextInt(points.length)];
+ }
+ return runStandardKMeansPlusPlus(subset, k);
+ }
+ }
+
+ /**
+ * Standard KMeans++ implementation.
+ * When called by chooseInitialCentroids, 'points' may be the full set or a subset.
+ */
+ private static VectorFloat runStandardKMeansPlusPlus(VectorFloat[] points, int k) {
+ var random = ThreadLocalRandom.current();
+ VectorFloat centroids = vectorTypeSupport.createFloatVector(k * points[0].length());
+
+ float[] distances = new float[points.length];
+ Arrays.fill(distances, Float.MAX_VALUE);
+ VectorFloat distancesVector = vectorTypeSupport.createFloatVector(distances);
+
+ VectorFloat firstCentroid = points[random.nextInt(points.length)];
+ centroids.copyFrom(firstCentroid, 0, 0, firstCentroid.length());
+
+ VectorFloat newDistancesVector = vectorTypeSupport.createFloatVector(points.length);
+ for (int i = 0; i < points.length; i++) {
+ newDistancesVector.set(i, squareL2Distance(points[i], firstCentroid));
+ }
+ VectorUtil.minInPlace(distancesVector, newDistancesVector);
+
+ for (int i = 1; i < k; i++) {
+ float totalDistance = VectorUtil.sum(distancesVector);
+ float r = random.nextFloat() * totalDistance;
+ int selectedIdx = -1;
+
+ for (int j = 0; j < points.length; j++) {
+ r -= distancesVector.get(j);
+ if (r < 1e-6) {
+ selectedIdx = j;
+ break;
+ }
+ }
+ if (selectedIdx == -1) selectedIdx = random.nextInt(points.length);
+
+ VectorFloat nextCentroid = points[selectedIdx];
+ centroids.copyFrom(nextCentroid, 0, i * nextCentroid.length(), nextCentroid.length());
+
+ for (int j = 0; j < points.length; j++) {
+ newDistancesVector.set(j, squareL2Distance(points[j], nextCentroid));
+ }
+ VectorUtil.minInPlace(distancesVector, newDistancesVector);
+ }
+ return centroids;
+ }
+
+ // --- Standard Distance and Update Logic ---
+
+ private int getNearestCluster(VectorFloat point) {
+ float minDistance = Float.MAX_VALUE;
+ int nearestCluster = 0;
+ int dim = point.length();
+
+ for (int i = 0; i < k; i++) {
+ float distance = squareL2Distance(point, 0, centroids, i * dim, dim);
+ if (distance < minDistance) {
+ minDistance = distance;
+ nearestCluster = i;
+ }
+ }
+ return nearestCluster;
+ }
+
+ private int updateAssignedPointsAnisotropic() {
+ float pcm = computeParallelCostMultiplier(anisotropicThreshold, points[0].length());
+ int dim = points[0].length();
+
+ float[] cNormSquared = new float[k];
+ for (int i = 0; i < k; i++) {
+ cNormSquared[i] = dotProduct(centroids, i * dim, centroids, i * dim, dim);
+ }
+
+ int changedCount = 0;
+ for (int i = 0; i < points.length; i++) {
+ var x = points[i];
+ var xNormSquared = dotProduct(x, x);
+
+ int index = assignments[i];
+ float minDist = Float.MAX_VALUE;
+ for (int j = 0; j < k; j++) {
+ float dist = weightedDistance(x, j, pcm, cNormSquared[j], xNormSquared);
+ if (dist < minDist) {
+ minDist = dist;
+ index = j;
+ }
+ }
+
+ if (index != assignments[i]) {
+ changedCount++;
+ assignments[i] = index;
+ }
+ }
+ return changedCount;
+ }
+
+ private float weightedDistance(VectorFloat x, int centroid, float parallelCostMultiplier, float cNormSquared, float xNormSquared) {
+ float cDotX = VectorUtil.dotProduct(centroids, centroid * x.length(), x, 0, x.length());
+ float parallelErrorSubtotal = cDotX - xNormSquared;
+ float residualSquaredNorm = cNormSquared - 2 * cDotX + xNormSquared;
+ float parallelError = square(parallelErrorSubtotal);
+ float perpendicularError = residualSquaredNorm - parallelError;
+ return parallelCostMultiplier * parallelError + perpendicularError;
+ }
+
+ private void updateCentroidsAnisotropic() {
+ int dimensions = points[0].length();
+ float pcm = computeParallelCostMultiplier(anisotropicThreshold, dimensions);
+ float orthogonalCostMultiplier = 1.0f / pcm;
+
+ var pointsByCluster = new HashMap>();
+ for (int i = 0; i < assignments.length; i++) {
+ pointsByCluster.computeIfAbsent(assignments[i], __ -> new ArrayList<>()).add(i);
+ }
+
+ for (int i = 0; i < k; i++) {
+ var L = pointsByCluster.getOrDefault(i, List.of());
+ if (L.isEmpty()) {
+ initializeCentroidToRandomPoint(i);
+ continue;
+ }
+
+ var mean = vectorTypeSupport.createFloatVector(dimensions);
+ var outerProdSums = new Matrix(dimensions, dimensions);
+ for (int j : L) {
+ var point = points[j];
+ // Offset is 0 because 'mean' is a fresh vector of the same dimension as 'point'
+ addToOffset(mean, 0, point);
+ float denom = dotProduct(point, point);
+ if (denom > 0) {
+ var op = Matrix.outerProduct(point, point);
+ op.scale(1.0f / denom);
+ outerProdSums.addInPlace(op);
+ }
+ }
+ outerProdSums.scale((1 - orthogonalCostMultiplier) / L.size());
+ scale(mean, 1.0f / L.size());
+
+ for (int j = 0; j < dimensions; j++) {
+ outerProdSums.addTo(j, j, orthogonalCostMultiplier);
+ }
+
+ var invertedMatrix = outerProdSums.invert();
+ centroids.copyFrom(invertedMatrix.multiply(mean), 0, i * dimensions, dimensions);
+ }
+ }
+
+ private void initializeCentroidToRandomPoint(int i) {
+ var random = ThreadLocalRandom.current();
+ centroids.copyFrom(points[random.nextInt(points.length)], 0, i * points[0].length(), points[0].length());
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutablePQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutablePQVectors.java
index b47d2b9ce..c4ceb75d4 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutablePQVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/MutablePQVectors.java
@@ -63,6 +63,8 @@ public void setZero(int ordinal) {
// increase count first so get() works
vectorCount.updateAndGet(current -> max(current, ordinal + 1));
get(ordinal).zero();
+ // A zero code reconstructs to each subspace's first centroid, which searches can reach.
+ pq.recordCode(get(ordinal), 0);
}
private synchronized void ensureChunkCapacity(int ordinal) {
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java
index 6ebc11701..151446070 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/NVQuantization.java
@@ -158,7 +158,10 @@ public static NVQuantization compute(RandomAccessVectorValues ravv, int nSubVect
/**
* Creates a NVQuantization instance by using the global mean and computing the data structures used to divide each
- * vector into subvectors.
+ * vector into subvectors. The mean can be collected during any earlier input scan,
+ * for example with {@link io.github.jbellis.jvector.graph.VectorMean}, independently
+ * of the primary quantizer or loader. This method does not read input vectors.
+ * Use {@link #compute} when no precomputed mean is available.
*
* @param scaledGlobalMean the mean of the graph (its average vector)
* @param nSubVectors number of subvectors
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQDecoder.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQDecoder.java
index ced0124fd..73cf40c28 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQDecoder.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQDecoder.java
@@ -59,8 +59,22 @@ protected float decodedSimilarity(ByteSequence encoded, int offset, int lengt
}
static class DotProductDecoder extends CachingDecoder {
+ private final VectorFloat query;
+
public DotProductDecoder(PQVectors cv, VectorFloat query) {
super(cv, query, VectorSimilarityFunction.DOT_PRODUCT);
+ this.query = query;
+ }
+
+ /** PQ scores the centered query against centered reconstructions, so their lengths bound the score. */
+ @Override
+ public float dotProductBound() {
+ var center = cv.pq.globalCentroid;
+ double squared = VectorUtil.dotProduct(query, query);
+ if (center != null) {
+ squared += VectorUtil.dotProduct(center, center) - 2.0 * VectorUtil.dotProduct(query, center);
+ }
+ return (float) Math.sqrt(Math.max(0.0, squared)) * cv.maxReconstructedNorm();
}
@Override
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java
index 5e5f172dc..0581d77cc 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/PQVectors.java
@@ -52,6 +52,15 @@ protected PQVectors(ProductQuantization pq) {
this.pq = pq;
}
+ /**
+ * The largest length of any stored vector's reconstruction, relative to the global centroid if any, or more:
+ * the PQ model's {@link ProductQuantization#maxEncodedNorm}, which counts every code it encoded or holds. With the
+ * centered query's length it bounds every dot product PQ scores. NaN when unknown.
+ */
+ public float maxReconstructedNorm() {
+ return pq.maxEncodedNorm();
+ }
+
public static ImmutablePQVectors load(RandomAccessReader in) throws IOException {
// pq codebooks
var pq = ProductQuantization.load(in);
@@ -72,7 +81,7 @@ public static ImmutablePQVectors load(RandomAccessReader in) throws IOException
chunks[layout.fullSizeChunks] = vectorTypeSupport.readByteSequence(in, layout.lastChunkBytes);
}
- return new ImmutablePQVectors(pq, chunks, vectorCount, layout.fullChunkVectors);
+ return new ImmutablePQVectors(pq, chunks, vectorCount, layout.fullChunkVectors, false);
}
public static PQVectors load(RandomAccessReader in, long offset) throws IOException {
@@ -146,7 +155,7 @@ private static ImmutablePQVectors encodeAndBuild(ProductQuantization pq, int vec
else
VectorAccess.forEach(ravv, vectorCount, ordinalsMapping, simdExecutor, encode);
- return new ImmutablePQVectors(pq, chunks, vectorCount, layout.fullChunkVectors);
+ return new ImmutablePQVectors(pq, chunks, vectorCount, layout.fullChunkVectors, false);
}
@Override
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java
index ba96d8f2f..bc80c310f 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/quantization/ProductQuantization.java
@@ -72,6 +72,11 @@ public class ProductQuantization implements VectorCompressor>, A
final int[][] subvectorSizesAndOffsets;
final float anisotropicThreshold; // parallel cost multiplier
private final float[][] centroidNormsSquared; // precomputed norms of the centroids, for encoding
+ // The largest length of any reconstruction this model has encoded or holds codes for. Unknown for a model loaded
+ // from a format that did not store it, since it cannot know the lengths of the codes encoded before.
+ private final io.github.jbellis.jvector.util.MaxVectorLength maxEncoded = new io.github.jbellis.jvector.util.MaxVectorLength();
+ /** The first on-disk format version that stores {@link #maxEncodedNorm}. */
+ static final int MAX_NORM_VERSION = 7;
private final ThreadLocal> partialSums; // for dot product, euclidean, and cosine partials
private final AtomicReference> partialSquaredMagnitudes; // for cosine partials
@@ -439,6 +444,31 @@ public void encodeTo(VectorFloat vector, ByteSequence dest) {
encodeAnisotropic(vector, dest);
else
encodeUnweighted(vector, dest);
+ recordCode(dest, 0);
+ }
+
+ /** The squared length of the reconstruction of the code at {@code offset}, relative to the global centroid. */
+ float reconstructedSquaredNorm(ByteSequence code, int offset) {
+ float sum = 0;
+ for (int m = 0; m < M; m++) sum += centroidNormsSquared[m][Byte.toUnsignedInt(code.get(offset + m))];
+ return sum;
+ }
+
+ /** Counts the code at {@code offset} in {@link #maxEncodedNorm}. */
+ void recordCode(ByteSequence code, int offset) {
+ maxEncoded.recordSquared(reconstructedSquaredNorm(code, offset));
+ }
+
+ /**
+ * The largest length of the reconstruction, relative to the global centroid if any, of any vector this model has
+ * encoded or holds codes for, including vectors encoded before it was saved and loaded. With the centered query's
+ * length it bounds every dot product PQ scores, which adaptive stopping needs to read PQ scores as distances.
+ * NaN until the model encodes or holds a code; positive infinity when unknown, for a model loaded from a format
+ * older than version {@value #MAX_NORM_VERSION}, which did not store it. In both cases dot-product searches stop
+ * by the standard rule.
+ */
+ public float maxEncodedNorm() {
+ return maxEncoded.get();
}
/**
@@ -608,6 +638,10 @@ public void write(IndexWriter out, int version) throws IOException
assert codebook.length() == clusterCount * subvectorSizesAndOffsets[i][0];
vectorTypeSupport.writeFloatVector(out, codebook);
}
+
+ if (version >= MAX_NORM_VERSION) {
+ out.writeFloat(maxEncodedNorm());
+ }
}
/**
@@ -641,7 +675,12 @@ public VectorFloat createCodebookPartialSums(VectorSimilarityFunction vectorS
@Override
public int compressorSize() {
- int size = 0;
+ return compressorSize(OnDiskGraphIndex.CURRENT_VERSION);
+ }
+
+ /** The bytes {@link #write} writes at {@code version}, which must be 3 or later. */
+ public int compressorSize(int version) {
+ int size = version >= MAX_NORM_VERSION ? Float.BYTES : 0; // largest encoded length
size += Integer.BYTES; // MAGIC
size += Integer.BYTES; // STORAGE_VERSION
size += Integer.BYTES; // globalCentroidLength
@@ -701,7 +740,14 @@ public static ProductQuantization load(RandomAccessReader in) throws IOException
codebooks[m] = codebook;
}
- return new ProductQuantization(codebooks, clusters, subvectorSizes, globalCentroid, anisotropicThreshold);
+ var pq = new ProductQuantization(codebooks, clusters, subvectorSizes, globalCentroid, anisotropicThreshold);
+ if (version >= MAX_NORM_VERSION) {
+ pq.maxEncoded.restore(in.readFloat());
+ } else {
+ // Older formats did not store the largest length, so codes encoded with this model have unknown lengths.
+ pq.maxEncoded.markUnknown();
+ }
+ return pq;
}
@Override
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/BatchingExecutor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/BatchingExecutor.java
new file mode 100644
index 000000000..fa0fa0952
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/BatchingExecutor.java
@@ -0,0 +1,394 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.github.jbellis.jvector.util;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicLongArray;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * Bounded asynchronous execution with optional preparation of compatible requests
+ * in batches. Batches contain only requests already waiting; there is no fill timer.
+ * Prepared work takes priority, and the preparing worker executes its first request
+ * while other workers may execute the rest.
+ *
+ * The supplied executor must schedule asynchronously, reject explicitly when it
+ * cannot accept work, and remain operational until this service closes. Direct
+ * executors, caller-runs rejection policies and silent-discard policies are not
+ * supported. At most {@code workerCount} reusable drain tasks are submitted at a
+ * time. Each exits when the queues are empty, releasing shared execution capacity.
+ * This class creates no threads and does not shut down the supplied executor.
+ *
+ *
Requests and their buffers remain caller-owned. An accepted request must not
+ * be modified or submitted again until its completion callback. Rejection at the
+ * admission limit leaves ownership with the caller. Accepted work completes once,
+ * including preparation or dispatch failure. Completion may run on a submitting
+ * thread if the executor rejects dispatch; successful work never runs inline.
+ * Queue arrays, worker runnables and batch lists are allocated at construction.
+ *
+ *
Close rejects new submissions and waits for accepted work and callbacks.
+ * Do not close from a callback or block a host executor thread needed to drain this
+ * service. The host must not cancel or discard already accepted drain tasks.
+ */
+public final class BatchingExecutor implements AutoCloseable {
+ /** Callbacks must be thread-safe. Completion must return promptly. */
+ public interface Operation {
+ /**
+ * Immutable identity key for compatible preparation, or null to prepare
+ * alone. Called on the submitting thread, outside internal locks.
+ */
+ Object batchKey(T request);
+ /** Borrowed read-only list, valid only during this call; never retain it. */
+ void prepareBatch(List requests) throws Exception;
+ /** Executes one prepared request, without a batch-wide completion barrier. */
+ void execute(T request) throws Exception;
+ /**
+ * Called exactly once for accepted work. The request may be recycled here.
+ * Exceptions are retained for inspection and close, without poisoning work.
+ */
+ void complete(T request, Throwable failure);
+ }
+
+ // Keep the entry on host threads but clear its value on exit. Removing it
+ // would allocate a fresh ThreadLocalMap entry for every isolated drain task.
+ private static final ThreadLocal> CURRENT = new ThreadLocal<>();
+
+ private final ReentrantLock lock = new ReentrantLock();
+ private final Condition drained = lock.newCondition();
+ private final Executor executor;
+ private final Operation operation;
+ private final int capacity, batchSize, workerLimit;
+ private final Object[] pending, pendingKeys, ready;
+ private final Throwable[] pendingFailures;
+ private int pendingHead, pendingTail, pendingSize;
+ private int pendingGroups, tailGroupSize;
+ private int readyHead, readyTail, readySize;
+ private final ArrayDeque idleWorkers;
+ private final List workers;
+ private int activeWorkers, busyWorkers, callbacks;
+ private volatile int outstanding;
+ private volatile boolean closed;
+ private final AtomicLong callbackFailures = new AtomicLong();
+ private final AtomicReference firstCallbackFailure = new AtomicReference<>();
+
+ /**
+ * @param workerCount maximum concurrent preparation/execution drain tasks
+ * @param capacity maximum requests queued, preparing or executing; excludes callbacks
+ * @param batchSize maximum currently waiting compatible requests per preparation
+ * @param executor caller-owned asynchronous execution capacity
+ * @param operation request preparation, execution and completion operations
+ */
+ public BatchingExecutor(int workerCount, int capacity, int batchSize,
+ Executor executor, Operation operation) {
+ if (workerCount <= 0 || capacity <= 0 || batchSize <= 0
+ || batchSize > capacity || batchSize == Integer.MAX_VALUE)
+ throw new IllegalArgumentException("Workers, capacity and batch size must be positive; "
+ + "batch size must fit its histogram and not exceed capacity");
+ this.executor = Objects.requireNonNull(executor, "executor");
+ this.operation = Objects.requireNonNull(operation, "operation");
+ if (executor instanceof ThreadPoolExecutor) {
+ var policy = ((ThreadPoolExecutor) executor).getRejectedExecutionHandler();
+ if (policy instanceof ThreadPoolExecutor.DiscardPolicy
+ || policy instanceof ThreadPoolExecutor.DiscardOldestPolicy
+ || policy instanceof ThreadPoolExecutor.CallerRunsPolicy)
+ throw new IllegalArgumentException("Executor must reject explicitly without discarding or running inline");
+ }
+ this.capacity = capacity;
+ this.batchSize = batchSize;
+ this.workerLimit = Math.min(workerCount, capacity);
+ pending = new Object[capacity];
+ pendingKeys = new Object[capacity];
+ pendingFailures = new Throwable[capacity];
+ ready = new Object[capacity];
+ idleWorkers = new ArrayDeque<>(workerLimit);
+ workers = new ArrayList<>(workerLimit);
+ for (int i = 0; i < workerLimit; i++) {
+ Worker worker = new Worker();
+ workers.add(worker);
+ idleWorkers.addLast(worker);
+ }
+ }
+
+ /**
+ * Accepts without waiting for capacity. Returns false at the admission limit;
+ * throws after close. A true return transfers ownership until completion.
+ */
+ public boolean trySubmit(T request) {
+ Objects.requireNonNull(request, "request");
+ if (closed) throw new RejectedExecutionException("Batching executor is closed");
+ if (outstanding == capacity) return false;
+ Object key = null;
+ Throwable failure = null;
+ try { key = operation.batchKey(request); }
+ catch (Throwable error) { failure = error; }
+ Worker launches;
+ lock.lock();
+ try {
+ if (closed) throw new RejectedExecutionException("Batching executor is closed");
+ if (outstanding == capacity) return false;
+ int previousTail = pendingTail == 0 ? capacity - 1 : pendingTail - 1;
+ if (pendingSize != 0 && failure == null && key != null
+ && pendingFailures[previousTail] == null && pendingKeys[previousTail] == key
+ && tailGroupSize < batchSize) {
+ tailGroupSize++;
+ } else {
+ pendingGroups++;
+ tailGroupSize = 1;
+ }
+ pending[pendingTail] = request;
+ pendingKeys[pendingTail] = key;
+ pendingFailures[pendingTail] = failure;
+ pendingTail = advance(pendingTail);
+ pendingSize++;
+ outstanding++;
+ launches = reserveWorkers();
+ } finally { lock.unlock(); }
+ launch(launches);
+ return true;
+ }
+
+ private int advance(int position) { return position == capacity - 1 ? 0 : position + 1; }
+
+ /** Called under lock; links reusable tasks rather than allocating a launch list. */
+ private Worker reserveWorkers() {
+ // Track compatible groups at admission rather than scanning the queue or
+ // treating unrelated models as one batch and underutilizing the budget.
+ int target = (int) Math.min(workerLimit, (long) busyWorkers + readySize + pendingGroups);
+ Worker first = null, last = null;
+ while (activeWorkers < target) {
+ Worker worker = idleWorkers.removeFirst();
+ activeWorkers++;
+ if (last == null) first = worker;
+ else last.nextLaunch = worker;
+ last = worker;
+ }
+ return first;
+ }
+
+ /** Executor calls and operation callbacks are always outside lock. */
+ private void launch(Worker first) {
+ while (first != null) {
+ Worker worker = first;
+ first = worker.nextLaunch;
+ worker.nextLaunch = null;
+ worker.launchingThread.set(Thread.currentThread());
+ try { executor.execute(worker); }
+ catch (RuntimeException | Error failure) { dispatchFailed(worker, failure); }
+ finally {
+ // A host worker may later dequeue a task that it submitted.
+ // Only execution on the submitting call stack is inline. CAS
+ // cannot clear the marker of a reused task launched elsewhere.
+ worker.launchingThread.compareAndSet(Thread.currentThread(), null);
+ }
+ }
+ }
+
+ private void dispatchFailed(Worker worker, Throwable failure) {
+ // Allocation is confined to dispatch failure, never ordinary submission.
+ List failed = null;
+ lock.lock();
+ try {
+ worker.launchingThread.set(null);
+ activeWorkers--;
+ idleWorkers.addLast(worker);
+ if (activeWorkers == 0) {
+ failed = new ArrayList<>(pendingSize + readySize);
+ while (readySize != 0) failed.add(takeReady());
+ while (pendingSize != 0) failed.add(takePending());
+ pendingGroups = 0;
+ }
+ signalDrained();
+ } finally { lock.unlock(); }
+ if (failed != null) for (T request : failed) finish(request, failure);
+ // Another accepted drain task, if present, will process the queues.
+ }
+
+ @SuppressWarnings("unchecked")
+ private T takePending() {
+ T value = (T) pending[pendingHead];
+ pending[pendingHead] = null;
+ pendingKeys[pendingHead] = null;
+ pendingFailures[pendingHead] = null;
+ pendingHead = advance(pendingHead);
+ pendingSize--;
+ if (pendingSize == 0) tailGroupSize = 0;
+ return value;
+ }
+
+ @SuppressWarnings("unchecked")
+ private T takeReady() {
+ T value = (T) ready[readyHead];
+ ready[readyHead] = null;
+ readyHead = advance(readyHead);
+ readySize--;
+ return value;
+ }
+
+ private final class Worker implements Runnable {
+ final ArrayList requests = new ArrayList<>(batchSize);
+ final List borrowedRequests = Collections.unmodifiableList(requests);
+ final AtomicLongArray histogram = new AtomicLongArray(batchSize + 1);
+ Worker nextLaunch;
+ final AtomicReference launchingThread = new AtomicReference<>();
+ boolean busy, alreadyPrepared;
+ Throwable preparationFailure;
+
+ /** Acquires the next work, or atomically releases this task's worker lease. */
+ boolean takeWork() {
+ lock.lock();
+ try {
+ if (busy) { busyWorkers--; busy = false; }
+ requests.clear();
+ preparationFailure = null;
+ if (readySize != 0) {
+ requests.add(takeReady());
+ alreadyPrepared = true;
+ } else if (pendingSize != 0) {
+ Object key = pendingKeys[pendingHead];
+ preparationFailure = pendingFailures[pendingHead];
+ requests.add(takePending());
+ alreadyPrepared = false;
+ while (preparationFailure == null && key != null
+ && requests.size() < batchSize && pendingSize != 0
+ && pendingKeys[pendingHead] == key && pendingFailures[pendingHead] == null)
+ requests.add(takePending());
+ pendingGroups--;
+ } else {
+ activeWorkers--;
+ idleWorkers.addLast(this);
+ signalDrained();
+ return false;
+ }
+ busyWorkers++;
+ busy = true;
+ return true;
+ } finally { lock.unlock(); }
+ }
+
+ @Override
+ public void run() {
+ if (Thread.currentThread() == launchingThread.get())
+ throw new RejectedExecutionException("Executor must schedule work asynchronously");
+ launchingThread.set(null);
+ BatchingExecutor previous = CURRENT.get();
+ CURRENT.set(BatchingExecutor.this);
+ try {
+ while (takeWork()) {
+ Throwable failure = preparationFailure;
+ if (!alreadyPrepared && failure == null) {
+ try {
+ operation.prepareBatch(borrowedRequests);
+ histogram.incrementAndGet(requests.size());
+ } catch (Throwable error) { failure = error; }
+ }
+ if (failure != null) {
+ for (T request : requests) finish(request, failure);
+ continue;
+ }
+ if (requests.size() > 1) {
+ Worker launches;
+ lock.lock();
+ try {
+ for (int i = 1; i < requests.size(); i++) {
+ ready[readyTail] = requests.get(i);
+ readyTail = advance(readyTail);
+ readySize++;
+ }
+ launches = reserveWorkers();
+ } finally { lock.unlock(); }
+ launch(launches);
+ }
+ T first = requests.get(0);
+ try { operation.execute(first); }
+ catch (Throwable error) { failure = error; }
+ finish(first, failure);
+ }
+ } finally {
+ CURRENT.set(previous); // Do not retain the service/model on host threads.
+ }
+ }
+ }
+
+ private void finish(T request, Throwable failure) {
+ lock.lock();
+ try {
+ // Release admission before the callback may recycle the request.
+ callbacks++;
+ outstanding--;
+ } finally { lock.unlock(); }
+ BatchingExecutor previous = CURRENT.get();
+ boolean entering = previous != this;
+ if (entering) CURRENT.set(this);
+ try { operation.complete(request, failure); }
+ catch (Throwable error) {
+ firstCallbackFailure.compareAndSet(null, error);
+ callbackFailures.incrementAndGet();
+ } finally {
+ if (entering) CURRENT.set(previous);
+ lock.lock();
+ try { callbacks--; signalDrained(); }
+ finally { lock.unlock(); }
+ }
+ }
+
+ private void signalDrained() {
+ if (closed && outstanding == 0 && activeWorkers == 0 && callbacks == 0)
+ drained.signalAll();
+ }
+
+ /** Atomic snapshot; exact after all accepted requests have completed. */
+ public long[] batchHistogram() {
+ long[] result = new long[batchSize + 1];
+ for (Worker worker : workers)
+ for (int i = 1; i < result.length; i++) result[i] += worker.histogram.get(i);
+ return result;
+ }
+
+ /** Accepted work, excluding completion callbacks. */
+ public int outstanding() { return outstanding; }
+
+ /** Number of completion callbacks that threw, without blocking later submissions. */
+ public long completionFailureCount() { return callbackFailures.get(); }
+
+ /** First callback exception for monitoring, or null when no callback failed. */
+ public Throwable firstCompletionFailure() { return firstCallbackFailure.get(); }
+
+ /** Drains accepted work, preserves interrupts, then reports any callback failure. */
+ @Override
+ public void close() {
+ if (CURRENT.get() == this)
+ throw new IllegalStateException("Cannot close a batching executor from its operation or callback");
+ lock.lock();
+ try {
+ closed = true;
+ while (outstanding != 0 || activeWorkers != 0 || callbacks != 0)
+ drained.awaitUninterruptibly();
+ } finally { lock.unlock(); }
+ Throwable failure = firstCallbackFailure.get();
+ if (failure != null)
+ throw new IllegalStateException(completionFailureCount() + " completion callbacks failed", failure);
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/BlasThreads.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/BlasThreads.java
new file mode 100644
index 000000000..144bf16a4
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/BlasThreads.java
@@ -0,0 +1,76 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.github.jbellis.jvector.util;
+
+import io.github.jbellis.jvector.annotations.Experimental;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+
+/**
+ * Sets how many threads native BLAS and LAPACK may use for calls made on the current thread.
+ *
+ * Native BLAS libraries such as MKL and OpenBLAS otherwise start one thread per CPU on their
+ * first call, regardless of the host's own executors. JVector's BLAS calls therefore run inside a
+ * scope: one thread by default, or the size of the executor a caller passes to the operation.
+ * The native vectorization provider applies the limit through the library's per-thread setting
+ * (MKL {@code MKL_Set_Num_Threads_Local}, OpenBLAS {@code openblas_set_num_threads_local}), or,
+ * for OpenBLAS releases without it, the process-wide setting. Other providers leave BLAS threading
+ * to the library's own configuration, such as {@code MKL_NUM_THREADS}.
+ *
+ * Scopes nest: closing one restores the count that was in effect when it opened.
+ */
+@Experimental
+public final class BlasThreads {
+ private static final ThreadLocal ACTIVE = new ThreadLocal<>();
+ private static final Scope NONE = () -> { };
+
+ private BlasThreads() {
+ }
+
+ /** Restores the previous BLAS thread count when closed. */
+ public interface Scope extends AutoCloseable {
+ @Override
+ void close();
+ }
+
+ /**
+ * Uses {@code threads} BLAS threads for native calls on this thread until the scope closes.
+ *
+ * @throws IllegalArgumentException if {@code threads} is not positive
+ */
+ public static Scope limit(int threads) {
+ if (threads < 1) throw new IllegalArgumentException("BLAS thread count must be positive, got " + threads);
+ var support = VectorizationProvider.getInstance().getVectorUtilSupport();
+ Integer outer = ACTIVE.get();
+ int previous = support.setBlasThreadsLocal(threads);
+ ACTIVE.set(threads);
+ return () -> {
+ if (previous >= 0) support.setBlasThreadsLocal(previous);
+ if (outer == null) ACTIVE.remove();
+ else ACTIVE.set(outer);
+ };
+ }
+
+ /** One BLAS thread, unless an enclosing {@link #limit} scope on this thread already sets the count. */
+ public static Scope singleUnlessLimited() {
+ return ACTIVE.get() != null ? NONE : limit(1);
+ }
+
+ /** The count set by the innermost open scope on this thread, or 0 if none is open. */
+ public static int current() {
+ Integer active = ACTIVE.get();
+ return active == null ? 0 : active;
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/CallerBatchProcessor.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/CallerBatchProcessor.java
new file mode 100644
index 000000000..e0bc62351
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/CallerBatchProcessor.java
@@ -0,0 +1,472 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.util;
+
+import java.lang.invoke.MethodHandles;
+import java.lang.invoke.VarHandle;
+import java.util.AbstractList;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.locks.LockSupport;
+
+/**
+ * Synchronously combines work from concurrent callers, using only their threads.
+ * The strict constructor permits at most {@code maxLeaders} operations at once;
+ * {@link #opportunistic} additionally supplies a direct pass-through budget.
+ * An available
+ * leader starts immediately; it never waits for a batch to fill. Waiting requests
+ * are taken in FIFO order by a caller that acquires free capacity. Capacity is
+ * never reserved for a parked caller. Reused request records are published
+ * without a blocking queue lock; only batch removal needs brief exclusion.
+ *
+ * Input and output belong to the submitting caller but must be accessible to
+ * other threads until {@link #process} returns. The operation must not retain
+ * them or reenter this processor. A batch finishes before any of its callers
+ * returns; subsequent application work is entirely outside this processor.
+ * Interrupts are preserved, but do not abandon an input that another caller may
+ * already be using. A failed operation fails its batch and still hands off work.
+ *
+ *
Request slots, batch lists and callbacks are reused. This class creates no
+ * executor, worker, completion future, timer, or per-operation request object.
+ * A slot is allocated on each participating thread's first use of this instance.
+ */
+public final class CallerBatchProcessor implements Accountable {
+ /** An operation that writes caller-owned outputs before returning. */
+ public interface Operation {
+ /** Handles a single input without wrapping it in a batch. */
+ void apply(I input, O output) throws Exception;
+
+ /** Borrowed read-only lists; neither the lists nor their rows may escape. */
+ void applyBatch(List inputs, List outputs) throws Exception;
+ }
+
+ private static final int WAITING = 0, ASSIGNED = 1, DONE = 2;
+ private static final int CLAIM_BUSY = 0, CLAIM_SKIPPED = 1, CLAIM_READY = 2;
+ private static final VarHandle INCOMING, DRAINING;
+ static {
+ try {
+ var lookup = MethodHandles.lookup();
+ INCOMING = lookup.findVarHandle(CallerBatchProcessor.class, "incoming", Request.class);
+ DRAINING = lookup.findVarHandle(CallerBatchProcessor.class, "draining", int.class);
+ } catch (ReflectiveOperationException failure) {
+ throw new ExceptionInInitializerError(failure);
+ }
+ }
+ private final int maxBatchSize;
+ private final Operation operation;
+ private final Context[] contexts;
+ private final DirectSlot[] directSlots;
+ private final ThreadLocal> local = ThreadLocal.withInitial(Request::new);
+ // Producers publish their reusable records with one successful CAS. A
+ // consumer reverses a detached chain once, preserving FIFO publication order.
+ // ready and incoming are disjoint; only the drain owner modifies ready.
+ private volatile Request incoming, ready;
+ private volatile int draining;
+ private final long retainedBytes;
+
+ /**
+ * @param maxLeaders maximum concurrent operations, without creating threads
+ * @param maxBatchSize maximum existing requests claimed by one leader
+ * @param operation computation; all its methods must support concurrent calls
+ */
+ public CallerBatchProcessor(int maxLeaders, int maxBatchSize, Operation operation) {
+ this(maxLeaders, maxBatchSize, 0, operation);
+ }
+
+ /**
+ * Creates an opportunistic processor for a caller-owned CPU budget. Up to
+ * {@code parallelism} overlapping calls pass through directly while no work
+ * is queued or executing as a multi-input batch. Only overlap beyond that budget starts
+ * combining; it does not throttle ordinary calls to the combiner limit.
+ *
+ * The fixed internal limit is one combiner per four budgeted CPUs, with
+ * a minimum of one. No thread is created. Direct calls are outside that
+ * combiner limit; total concurrent operations cannot exceed the direct budget
+ * plus the combiner limit. Emptying the queue and completing its multi-input
+ * batches restores
+ * the direct path. An overflowing singleton never suppresses free direct
+ * capacity on its own. There are no timing thresholds,
+ * learned decisions, dataset-dependent rules or batch-fill delays.
+ *
+ * @param parallelism CPU capacity available to the callers, positive
+ * @param maxBatchSize maximum ready inputs passed to one batch operation
+ * @param operation operation whose methods support concurrent calls
+ */
+ public static CallerBatchProcessor opportunistic(
+ int parallelism, int maxBatchSize, Operation operation) {
+ if (parallelism < 1) throw new IllegalArgumentException("Parallelism must be positive");
+ return new CallerBatchProcessor<>(Math.max(1, parallelism / 4), maxBatchSize, parallelism, operation);
+ }
+
+ @SuppressWarnings("unchecked")
+ private CallerBatchProcessor(int maxLeaders, int maxBatchSize, int directParallelism,
+ Operation operation) {
+ if (maxLeaders < 1 || maxBatchSize < 1 || maxBatchSize == Integer.MAX_VALUE)
+ throw new IllegalArgumentException("Leader and batch limits must be positive");
+ this.maxBatchSize = maxBatchSize;
+ this.operation = Objects.requireNonNull(operation);
+ directSlots = new DirectSlot[directParallelism];
+ for (int i = 0; i < directSlots.length; i++) directSlots[i] = new DirectSlot(i);
+ contexts = (Context[]) new Context[maxLeaders];
+ for (int i = 0; i < maxLeaders; i++) contexts[i] = new Context<>(maxBatchSize, i);
+ long bytes = RamUsageEstimator.shallowSizeOf(this)
+ + RamUsageEstimator.shallowSizeOf(local) + RamUsageEstimator.shallowSizeOf(contexts)
+ + RamUsageEstimator.shallowSizeOf(directSlots);
+ for (DirectSlot slot : directSlots) bytes = Math.addExact(bytes, RamUsageEstimator.shallowSizeOf(slot));
+ for (Context context : contexts) {
+ bytes = Math.addExact(bytes, RamUsageEstimator.shallowSizeOf(context)
+ + listBytes(context.requests, maxBatchSize)
+ + RamUsageEstimator.shallowSizeOf(context.inputs)
+ + RamUsageEstimator.shallowSizeOf(context.outputs)
+ + RamUsageEstimator.shallowSizeOf(context.histogram));
+ }
+ retainedBytes = bytes;
+ }
+
+ private static long listBytes(ArrayList list, int capacity) {
+ return RamUsageEstimator.shallowSizeOf(list)
+ + RamUsageEstimator.alignObjectSize(RamUsageEstimator.NUM_BYTES_ARRAY_HEADER
+ + (long) capacity * RamUsageEstimator.NUM_BYTES_OBJECT_REF);
+ }
+
+ /** Owned coordinator storage; the operation and caller-owned request slots are excluded. */
+ @Override
+ public long ramBytesUsed() { return retainedBytes; }
+
+ private static final class Request {
+ final Thread thread = Thread.currentThread();
+ volatile int state;
+ volatile boolean parked;
+ boolean inUse;
+ int preferred = (int) (thread.getId() ^ (thread.getId() >>> 32)) & Integer.MAX_VALUE;
+ int preferredDirect = preferred;
+ I input;
+ O output;
+ Throwable failure;
+ Request next;
+ }
+
+ private static final class DirectSlot {
+ private static final VarHandle CLAIMED, COUNT;
+ static {
+ try {
+ var lookup = MethodHandles.lookup();
+ CLAIMED = lookup.findVarHandle(DirectSlot.class, "claimed", int.class);
+ COUNT = lookup.findVarHandle(DirectSlot.class, "calls", long.class);
+ } catch (ReflectiveOperationException failure) {
+ throw new ExceptionInInitializerError(failure);
+ }
+ }
+ final int index;
+ volatile int claimed;
+ private long calls;
+ DirectSlot(int index) { this.index = index; }
+ void record() { COUNT.setOpaque(this, (long) COUNT.getOpaque(this) + 1); }
+ }
+
+ // Static slots/contexts do not retain the processor (and its model) in an
+ // idle thread's ThreadLocal value after the application releases the model.
+ private static final class Context {
+ private static final VarHandle CLAIMED;
+ private static final VarHandle COUNT = MethodHandles.arrayElementVarHandle(long[].class);
+ static {
+ try {
+ CLAIMED = MethodHandles.lookup().findVarHandle(Context.class, "claimed", int.class);
+ } catch (ReflectiveOperationException failure) {
+ throw new ExceptionInInitializerError(failure);
+ }
+ }
+ volatile int claimed;
+ final int slot;
+ final long[] histogram;
+ final ArrayList> requests;
+ final List inputs;
+ final List outputs;
+ Context(int capacity, int slot) {
+ this.slot = slot;
+ histogram = new long[capacity + 1];
+ requests = new ArrayList<>(capacity);
+ inputs = new AbstractList<>() {
+ public int size() { return requests.size(); }
+ public I get(int i) { return requests.get(i).input; }
+ };
+ outputs = new AbstractList<>() {
+ public int size() { return requests.size(); }
+ public O get(int i) { return requests.get(i).output; }
+ };
+ }
+
+ // Only the owning leader writes these counts. Opaque access makes the
+ // diagnostic snapshot non-tearing without a shared atomic increment.
+ void record(int size) {
+ COUNT.setOpaque(histogram, size, (long) COUNT.getOpaque(histogram, size) + 1);
+ }
+ }
+
+ /** Executes synchronously, preserving interruption and caller buffer lifetime. */
+ public void process(I input, O output) {
+ Objects.requireNonNull(input); Objects.requireNonNull(output);
+ Request request = local.get();
+ if (request.inUse) throw new IllegalStateException("Reentrant batch processing is unsupported");
+ request.inUse = true;
+ boolean interrupted = Thread.interrupted();
+ try {
+ if (directSlots.length != 0 && !hasPending() && !isBatching()) {
+ DirectSlot slot = acquireDirect(request.preferredDirect % directSlots.length);
+ if (slot != null) {
+ try {
+ // Honor older requests published while capacity was acquired.
+ // A singleton overflow must not suppress free direct
+ // capacity. Real batches retain sharing until completion.
+ if (!hasPending() && !isBatching()) {
+ request.preferredDirect = slot.index;
+ slot.record();
+ operation.apply(input, output);
+ return;
+ }
+ } finally { DirectSlot.CLAIMED.setVolatile(slot, 0); }
+ }
+ }
+ // An idle caller executes directly. Only saturated callers enter the
+ // FIFO; no monitor, list mutation or request publication is needed.
+ if (!hasPending()) {
+ Context context = acquire(request.preferred % contexts.length);
+ if (context != null) {
+ if (!hasPending()) {
+ request.preferred = context.slot;
+ context.record(1);
+ try { operation.apply(input, output); }
+ finally { release(context); }
+ return;
+ }
+ // A waiter arrived while we acquired capacity; give it priority.
+ release(context);
+ }
+ }
+ request.input = input;
+ request.output = output;
+ request.failure = null;
+ request.state = WAITING;
+ publish(request);
+ for (;;) {
+ int state = request.state; // Acquire publishes outputs and failures.
+ if (state == DONE) break;
+ if (state == WAITING) {
+ Context context = acquire(request.preferred % contexts.length);
+ if (context != null) {
+ request.preferred = context.slot;
+ int claim = take(context, request);
+ if (claim == CLAIM_READY) { run(context); continue; }
+ if (claim == CLAIM_SKIPPED) { release(context); continue; }
+ // The drain owner will signal pending work after release.
+ // Publish this slot first, then recheck the drain gate to
+ // cover an unlock that observed the slot still occupied.
+ free(context);
+ if ((int) DRAINING.getVolatile(this) == 0) continue;
+ }
+ }
+ request.parked = true;
+ // Completion publishes DONE before checking parked. Rechecking
+ // after this volatile store ensures an early completion either
+ // prevents parking or supplies the unpark permit.
+ int latest = request.state;
+ if (latest != DONE && (latest != WAITING
+ || (int) DRAINING.getVolatile(this) != 0 || !hasFreeCapacity())) {
+ LockSupport.park(this);
+ if (Thread.interrupted()) interrupted = true;
+ }
+ request.parked = false;
+ }
+ if (request.failure instanceof Error) throw (Error) request.failure;
+ if (request.failure instanceof RuntimeException) throw (RuntimeException) request.failure;
+ if (request.failure != null) throw new RuntimeException("Batch operation failed", request.failure);
+ } catch (RuntimeException | Error failure) {
+ throw failure;
+ } catch (Exception failure) {
+ if (failure instanceof InterruptedException) interrupted = true;
+ throw new RuntimeException("Batch operation failed", failure);
+ } finally {
+ request.input = null;
+ request.output = null;
+ request.failure = null;
+ request.parked = false;
+ request.inUse = false;
+ if (interrupted) Thread.currentThread().interrupt();
+ }
+ }
+
+ private DirectSlot acquireDirect(int preferred) {
+ int index = preferred;
+ do {
+ DirectSlot slot = directSlots[index];
+ if ((int) DirectSlot.CLAIMED.getVolatile(slot) == 0
+ && DirectSlot.CLAIMED.compareAndSet(slot, 0, 1)) return slot;
+ if (++index == directSlots.length) index = 0;
+ } while (index != preferred);
+ return null;
+ }
+
+ private boolean isBatching() {
+ for (Context context : contexts)
+ if ((int) Context.CLAIMED.getVolatile(context) > 1) return true;
+ return false;
+ }
+
+ // Each slot has its own admission word. Callers start at their last slot,
+ // avoiding a global CAS counter and repeated competition for the first slot.
+ private Context acquire(int preferred) {
+ int slot = preferred;
+ do {
+ Context context = contexts[slot];
+ if ((int) Context.CLAIMED.getVolatile(context) == 0
+ && Context.CLAIMED.compareAndSet(context, 0, 1)) return context;
+ if (++slot == contexts.length) slot = 0;
+ } while (slot != preferred);
+ return null;
+ }
+
+ private void free(Context context) {
+ // Volatile publication orders the subsequent waiter check after freeing
+ // capacity, closing the enqueue-versus-release lost-wakeup race.
+ Context.CLAIMED.setVolatile(context, 0);
+ }
+
+ private boolean hasFreeCapacity() {
+ for (Context context : contexts)
+ if ((int) Context.CLAIMED.getVolatile(context) == 0) return true;
+ return false;
+ }
+
+ private boolean hasPending() { return ready != null || incoming != null; }
+
+ private void publish(Request request) {
+ Request previous;
+ do {
+ previous = incoming;
+ request.next = previous;
+ } while (!INCOMING.compareAndSet(this, previous, request));
+ }
+
+ private void signalPending() {
+ Request next = ready;
+ if (next == null) next = incoming;
+ // Any pending caller can act as combiner; take() still serves FIFO.
+ // A concurrent completion can make this a harmless spurious wakeup.
+ if (next != null) LockSupport.unpark(next.thread);
+ }
+
+ private void release(Context context) {
+ free(context);
+ if (hasPending()) signalPending();
+ }
+
+ private void releaseDrain() {
+ // A failed drainer frees its slot before checking this publication. The
+ // two checks ensure it either retries or receives a wakeup, never sleeps
+ // with free capacity and no owner responsible for further progress.
+ DRAINING.setVolatile(this, 0);
+ if (hasPending() && hasFreeCapacity()) signalPending();
+ }
+
+ @SuppressWarnings("unchecked")
+ private int take(Context context, Request owner) {
+ if (!DRAINING.compareAndSet(this, 0, 1)) return CLAIM_BUSY;
+ try {
+ if (owner.state != WAITING) return CLAIM_SKIPPED;
+ Request first = ready;
+ if (first == null) {
+ Request stack = (Request) INCOMING.getAndSet(this, null);
+ while (stack != null) {
+ Request next = stack.next;
+ stack.next = first;
+ first = stack;
+ stack = next;
+ }
+ }
+ if (first == null) return CLAIM_SKIPPED;
+ do {
+ Request next = first.next;
+ first.next = null;
+ first.state = ASSIGNED;
+ context.requests.add(first);
+ first = next;
+ } while (first != null && context.requests.size() < maxBatchSize);
+ ready = first;
+ return CLAIM_READY;
+ } finally { releaseDrain(); }
+ }
+
+ private void run(Context context) {
+ int size = context.requests.size();
+ // Admission distinguishes a true shared computation from a singleton
+ // overflow. Zero is free; one is reserved/singleton; larger values are
+ // executing batch widths. All occupied states retain the same limit.
+ if (size > 1) Context.CLAIMED.setVolatile(context, size);
+ context.record(size);
+ Throwable failure = null;
+ try {
+ if (context.requests.size() == 1) {
+ Request only = context.requests.get(0);
+ operation.apply(only.input, only.output);
+ } else operation.applyBatch(context.inputs, context.outputs);
+ } catch (Throwable error) {
+ failure = error;
+ if (error instanceof InterruptedException) Thread.currentThread().interrupt();
+ } finally {
+ // This context is exclusively owned until release. Publish results
+ // and wake callers without holding the drain gate.
+ for (int i = 0; i < context.requests.size(); i++) {
+ Request request = context.requests.get(i);
+ request.failure = failure;
+ request.state = DONE; // Release before followers access their output.
+ if (request.parked) LockSupport.unpark(request.thread);
+ }
+ context.requests.clear();
+ release(context);
+ }
+ }
+
+ /**
+ * Diagnostic snapshot: index is batch size, value is batch count. Concurrent
+ * calls may observe an in-progress operation; totals are exact when idle.
+ */
+ public long[] batchHistogram() {
+ long[] snapshot = new long[maxBatchSize + 1];
+ for (DirectSlot slot : directSlots) snapshot[1] += (long) DirectSlot.COUNT.getOpaque(slot);
+ for (Context context : contexts)
+ for (int i = 1; i < snapshot.length; i++)
+ snapshot[i] += (long) Context.COUNT.getOpaque(context.histogram, i);
+ return snapshot;
+ }
+
+ /**
+ * Diagnostic snapshot of waiting requests, excluding executing batches. The
+ * short drain exclusion is used only by this explicit inspection method.
+ */
+ public int waitingCount() {
+ while (!DRAINING.compareAndSet(this, 0, 1)) Thread.yield();
+ try {
+ int count = 0;
+ for (Request r = ready; r != null; r = r.next) count = Math.incrementExact(count);
+ // Producers can prepend records, but cannot change this chain while
+ // the drain gate is held. This is a snapshot at the incoming read.
+ for (Request r = incoming; r != null; r = r.next) count = Math.incrementExact(count);
+ return count;
+ } finally { releaseDrain(); }
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/MaxVectorLength.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/MaxVectorLength.java
new file mode 100644
index 000000000..228564e58
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/MaxVectorLength.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.github.jbellis.jvector.util;
+
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * The largest length of a set of vectors, recorded as they are added, from many threads at once. Dot-product searches
+ * that stop adaptively bound their scores with it. It is in one of three states:
+ *
+ *
+ * - none recorded yet: {@link #get} is NaN, and the first recorded length sets it;
+ * - known: {@link #get} is the largest length recorded, and longer lengths raise it;
+ * - unknown: {@link #get} is positive infinity, and it stays so. Use it when vectors exist whose lengths were not
+ * recorded, such as vectors encoded by a model loaded from a format that did not store their largest length; a
+ * length recorded later could not bound them.
+ *
+ *
+ * {@link #get} and {@link #restore} round-trip all three states through one float, for storage.
+ */
+public final class MaxVectorLength {
+ private static final int NONE = -1;
+ private static final int UNKNOWN = Float.floatToIntBits(Float.POSITIVE_INFINITY);
+
+ // The largest squared length as float bits: nonnegative floats order like their bits, NONE orders below every
+ // length and UNKNOWN above, so one compare-and-set loop keeps the maximum in every state.
+ private final AtomicInteger squaredBits = new AtomicInteger(NONE);
+
+ /** Records a vector's length. */
+ public void record(VectorFloat vector) {
+ recordSquared(VectorUtil.dotProduct(vector, vector));
+ }
+
+ /** Records a squared length; a NaN length makes the largest length unknown. */
+ public void recordSquared(float squaredLength) {
+ int bits = Float.isNaN(squaredLength) ? UNKNOWN : Float.floatToIntBits(Math.max(0f, squaredLength));
+ int current;
+ while (bits > (current = squaredBits.get())) {
+ if (squaredBits.compareAndSet(current, bits)) return;
+ }
+ }
+
+ /** Makes the largest length unknown for good. */
+ public void markUnknown() {
+ squaredBits.set(UNKNOWN);
+ }
+
+ /** The largest length recorded: NaN if none was, positive infinity if unknown. */
+ public float get() {
+ int bits = squaredBits.get();
+ return bits == NONE ? Float.NaN : (float) Math.sqrt(Float.intBitsToFloat(bits));
+ }
+
+ /** Sets the state from a value {@link #get} returned, as when loading a saved one. */
+ public void restore(float stored) {
+ squaredBits.set(Float.isNaN(stored) ? NONE : Float.floatToIntBits(Math.abs(stored) * Math.abs(stored)));
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/util/PagedIntMap.java b/jvector-base/src/main/java/io/github/jbellis/jvector/util/PagedIntMap.java
new file mode 100644
index 000000000..5bdbdcbc7
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/util/PagedIntMap.java
@@ -0,0 +1,149 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.github.jbellis.jvector.util;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReferenceArray;
+
+/**
+ * A growable map of dense-ish int keys to values, like {@link DenseIntMap}, whose updates take no lock.
+ *
+ * Values live in fixed-size pages. A page, once installed, is never copied or replaced, so a slot
+ * has one location for the life of the map. Growing installs new pages and, rarely, replaces the
+ * spine (the array of page references) with a larger copy that refers to the same page objects.
+ * Because no slot ever moves, a compare-and-set can never be lost to a concurrent resize, which is
+ * the race {@link DenseIntMap}'s read-write lock exists to prevent. {@link #compareAndPut} is
+ * therefore a single slot CAS, and readers take no locks.
+ *
+ * Publication: pages are installed with a volatile store into the spine, and a replaced spine is
+ * published through a volatile field after it holds every existing page. A thread that observes a
+ * value written into a page therefore also observes that page through the spine.
+ *
+ * As with {@link DenseIntMap}, space is allocated for all keys from 0 to the highest page in use,
+ * and gap keys map to null.
+ */
+public final class PagedIntMap implements IntMap {
+ static final int PAGE_BITS = 12;
+ static final int PAGE_SIZE = 1 << PAGE_BITS;
+ private static final int PAGE_MASK = PAGE_SIZE - 1;
+
+ /** Spine of pages. Entries are installed once, under {@link #growLock}, and never change. */
+ private volatile AtomicReferenceArray> spine;
+ private final Object growLock = new Object();
+ private final AtomicInteger size = new AtomicInteger();
+
+ public PagedIntMap(int initialCapacity) {
+ if (initialCapacity < 0) throw new IllegalArgumentException("initialCapacity must be nonnegative");
+ int pages = Math.max(1, (int) ((initialCapacity + (long) PAGE_SIZE - 1) >>> PAGE_BITS));
+ var s = new AtomicReferenceArray>(pages);
+ for (int p = 0; p < pages; p++) s.set(p, new AtomicReferenceArray<>(PAGE_SIZE));
+ spine = s;
+ }
+
+ @Override
+ public boolean compareAndPut(int key, T existing, T value) {
+ if (value == null) {
+ throw new IllegalArgumentException("compareAndPut() value cannot be null -- use remove() instead");
+ }
+ if (key < 0) throw new IndexOutOfBoundsException("negative key " + key);
+ boolean success = pageForWrite(key).compareAndSet(key & PAGE_MASK, existing, value);
+ if (success && existing == null) size.incrementAndGet();
+ return success;
+ }
+
+ @Override
+ public int size() {
+ return size.get();
+ }
+
+ @Override
+ public T get(int key) {
+ var page = pageOrNull(key);
+ return page == null ? null : page.get(key & PAGE_MASK);
+ }
+
+ @Override
+ public T remove(int key) {
+ var page = pageOrNull(key);
+ if (page == null) return null;
+ int slot = key & PAGE_MASK;
+ T old = page.get(slot);
+ if (old == null) return null;
+ if (page.compareAndSet(slot, old, null)) {
+ size.decrementAndGet();
+ return old;
+ }
+ return null;
+ }
+
+ @Override
+ public boolean containsKey(int key) {
+ return get(key) != null;
+ }
+
+ @Override
+ public void forEach(IntBiConsumer consumer) {
+ var s = spine;
+ for (int p = 0; p < s.length(); p++) {
+ var page = s.get(p);
+ if (page == null) continue;
+ int base = p << PAGE_BITS;
+ for (int i = 0; i < PAGE_SIZE; i++) {
+ T value = page.get(i);
+ if (value != null) consumer.consume(base + i, value);
+ }
+ }
+ }
+
+ private AtomicReferenceArray pageOrNull(int key) {
+ if (key < 0) throw new IndexOutOfBoundsException("negative key " + key);
+ int p = key >>> PAGE_BITS;
+ var s = spine;
+ return p < s.length() ? s.get(p) : null;
+ }
+
+ private AtomicReferenceArray pageForWrite(int key) {
+ int p = key >>> PAGE_BITS;
+ var s = spine;
+ if (p < s.length()) {
+ var page = s.get(p);
+ if (page != null) return page;
+ }
+ return installPage(p);
+ }
+
+ /** Slow path, once per new page: install it, growing the spine geometrically if needed. */
+ private AtomicReferenceArray installPage(int p) {
+ synchronized (growLock) {
+ var s = spine;
+ if (p >= s.length()) {
+ // Keys up to Integer.MAX_VALUE need (MAX_VALUE >>> PAGE_BITS) + 1 pages.
+ int newLength = (int) Math.min((Integer.MAX_VALUE >>> PAGE_BITS) + 1L,
+ Math.max(p + 1L, s.length() + (s.length() >> 1) + 1L));
+ var grown = new AtomicReferenceArray>(newLength);
+ for (int i = 0; i < s.length(); i++) grown.set(i, s.get(i));
+ spine = grown; // volatile publish after copying every existing page
+ s = grown;
+ }
+ var page = s.get(p);
+ if (page == null) {
+ page = new AtomicReferenceArray<>(PAGE_SIZE);
+ s.set(p, page);
+ }
+ return page;
+ }
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ASHLutScoring.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ASHLutScoring.java
new file mode 100644
index 000000000..27a63597c
--- /dev/null
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/ASHLutScoring.java
@@ -0,0 +1,82 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.vector;
+
+import java.util.Objects;
+
+/** Portable reference and bounds contract for interleaved ASH nibble blocks. */
+public final class ASHLutScoring {
+ private ASHLutScoring() {}
+
+ /** Checks a canonical packed code against a 16-entry table for each nibble. */
+ public static void checkCanonicalBounds(byte[] code, int groups, float[] lut) {
+ checkCanonicalBounds(code, 0, groups, lut);
+ }
+
+ /** As {@link #checkCanonicalBounds(byte[], int, float[])} for the code starting at {@code codeOffset}. */
+ public static void checkCanonicalBounds(byte[] code, int codeOffset, int groups, float[] lut) {
+ Objects.requireNonNull(code, "code");
+ Objects.requireNonNull(lut, "lut");
+ if (groups <= 0) throw new IllegalArgumentException("groups must be positive");
+ if (codeOffset < 0 || codeOffset + (groups + 1L) / 2 > code.length || (long) groups * 16 > lut.length) {
+ throw new IndexOutOfBoundsException("Canonical ASH code or LUT is too small");
+ }
+ }
+
+ /** Scalar reference; the unused high nibble of an odd final group is ignored. */
+ public static float dotCanonical(byte[] code, int groups, float[] lut) {
+ return dotCanonical(code, 0, groups, lut);
+ }
+
+ /** As {@link #dotCanonical(byte[], int, float[])} for the code starting at {@code codeOffset}. */
+ public static float dotCanonical(byte[] code, int codeOffset, int groups, float[] lut) {
+ checkCanonicalBounds(code, codeOffset, groups, lut);
+ float sum = 0;
+ for (int group = 0; group < groups; group++) {
+ int nibble = (code[codeOffset + (group >>> 1)] >>> ((group & 1) * 4)) & 15;
+ sum += lut[group * 16 + nibble];
+ }
+ return sum;
+ }
+
+ public static void checkBounds(byte[] codes, int offset, int groups, int stride,
+ int lane, int count, float[] lut, float[] out, int outOffset) {
+ if (groups < 0 || stride <= 0) {
+ throw new IllegalArgumentException("groups must be nonnegative and stride must be positive");
+ }
+ Objects.checkFromIndexSize(lane, count, stride);
+ Objects.checkFromIndexSize(outOffset, count, out.length);
+ Objects.checkFromIndexSize(0, Math.multiplyExact(groups, 16), lut.length);
+ int bytes = Math.multiplyExact(groups / 2 + groups % 2, stride);
+ Objects.checkFromIndexSize(offset, bytes, codes.length);
+ }
+
+ /** Returns projection dot products only; callers apply scale, offset and landmark terms. */
+ public static void score(byte[] codes, int offset, int groups, int stride,
+ int lane, int count, float[] lut, float[] out, int outOffset) {
+ checkBounds(codes, offset, groups, stride, lane, count, lut, out, outOffset);
+ for (int i = 0; i < count; i++) {
+ float sum = 0;
+ for (int group = 0; group < groups; group++) {
+ int packed = codes[offset + (group >>> 1) * stride + lane + i] & 255;
+ int nibble = (packed >>> ((group & 1) * 4)) & 15;
+ sum += lut[group * 16 + nibble];
+ }
+ out[outOffset + i] = sum;
+ }
+ }
+}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java
index 5843dc5f6..92b12ab92 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/DefaultVectorUtilSupport.java
@@ -547,4 +547,6 @@ public float nvqUniformLoss(VectorFloat vector, float minValue, float maxValu
return squaredSum;
}
+
+
}
diff --git a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java
index 9d663d10c..d3c6ec76c 100644
--- a/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java
+++ b/jvector-base/src/main/java/io/github/jbellis/jvector/vector/VectorUtilSupport.java
@@ -35,6 +35,200 @@
*/
public interface VectorUtilSupport {
+ /** Effective ASH kernel tuning for benchmark diagnostics. */
+ default String ashKernelDescription() { return "portable scalar ASH kernels"; }
+
+ /** Query factories select the tuned entry point once, outside the scoring hot loop. */
+ default boolean usesAshProjectionTuning() { return false; }
+
+ /** Explicitly tuned projection kernel; the ordinary entry point retains its compact default loop. */
+ default float ashProjectionDotTuned(float[] query, byte[] code, int dimensions, int bitsPerDimension) {
+ return ashProjectionDotTuned(query, code, 0, dimensions, bitsPerDimension);
+ }
+
+ /** As {@link #ashProjectionDotTuned(float[], byte[], int, int)} for the code starting at {@code codeOffset}. */
+ default float ashProjectionDotTuned(float[] query, byte[] code, int codeOffset, int dimensions, int bitsPerDimension) {
+ return ashProjectionDot(query, code, codeOffset, dimensions, bitsPerDimension);
+ }
+
+ /** Whether ASH nibble lookup and accumulation are implemented with SIMD. */
+ default boolean supportsAshLutScoring() { return false; }
+
+ /** Exact integer lookup for encoded C=1 query pairs. LUT stride is 32 (two repeated 16-entry tables). */
+ default void ashSymmetricLutScore(byte[] codes, int groups, int stride, int lane, int count,
+ short[] lut, float[] out, int outOffset) {
+ for (int i = 0; i < count; i++) {
+ long sum = 0;
+ for (int g = 0; g < groups; g++) {
+ int nibble = (codes[(g / 2) * stride + lane + i] >>> (4 * (g % 2))) & 15;
+ sum += lut[g * 32 + nibble];
+ }
+ out[outOffset + i] = sum * .25f;
+ }
+ }
+
+ /** Packed-to-packed integer dot products for symmetric ASH. */
+ default boolean supportsAshSymmetricScoring() { return false; }
+
+ default float ashSymmetricDot(long[] aBits, byte[] aCodes, long[] bBits, byte[] bCodes,
+ int dimensions, int bits) {
+ throw new UnsupportedOperationException("Symmetric ASH SIMD is unavailable on this backend");
+ }
+
+ /** Whether {@link #ashSymmetricDotAt} is implemented (2- and 4-bit codes at array offsets). */
+ default boolean supportsAshSymmetricDotAt() { return false; }
+
+ /**
+ * Symmetric 2/4-bit packed dot product of codes stored at offsets inside larger arrays.
+ * Same result as {@link #ashSymmetricDot} on the corresponding standalone code arrays.
+ */
+ default float ashSymmetricDotAt(byte[] aCodes, int aOffset, byte[] bCodes, int bOffset,
+ int dimensions, int bits) {
+ throw new UnsupportedOperationException("Offset symmetric ASH SIMD is unavailable on this backend");
+ }
+
+ /** Whether {@link #ashSymmetricDotRows} scores all rows in one call rather than one call per row. */
+ default boolean supportsAshSymmetricDotRows() { return false; }
+
+ /**
+ * Symmetric 2/4-bit dot products of one code against {@code count} codes stored {@code stride} bytes apart from
+ * the start of {@code rows}: {@code out[i]} equals {@link #ashSymmetricDotAt} of the query and row {@code i}.
+ * Requires {@link #supportsAshSymmetricDotAt()}.
+ */
+ default void ashSymmetricDotRows(byte[] query, int queryOffset, byte[] rows, int stride, int count,
+ int dimensions, int bits, float[] out) {
+ for (int i = 0; i < count; i++) out[i] = ashSymmetricDotAt(query, queryOffset, rows, i * stride, dimensions, bits);
+ }
+
+ /**
+ * Symmetric 1-bit dot product, {@code dimensions - 2 * mismatches}, of sign words stored at offsets
+ * inside larger arrays. Same result as {@link #ashSymmetricDot} with {@code bits == 1}.
+ * Available whenever {@link #supportsAshSymmetricScoring()} is.
+ */
+ default float ashSymmetricBinaryDotAt(long[] aBits, int aOffset, long[] bBits, int bOffset, int dimensions) {
+ throw new UnsupportedOperationException("Symmetric ASH SIMD is unavailable on this backend");
+ }
+
+ /** Whether canonical 2/4-bit projection codes have a single-vector SIMD kernel. */
+ default boolean supportsAshProjectionScoring() { return false; }
+
+ /** Whether canonical ASH codes can use an already prepared nibble LUT in SIMD. */
+ default boolean supportsAshCanonicalLutScoring() { return false; }
+
+ /** Whether canonical LUT scoring is implemented by the native backend. */
+ default boolean supportsNativeAshCanonicalLutScoring() { return false; }
+
+ /** Fail before scoring when the caller explicitly requires a native ASH single scorer. */
+ default void requireNativeAshSingleScoring(int bitsPerDimension) {
+ throw new IllegalStateException("Native ASH single scoring is unavailable for b=" + bitsPerDimension);
+ }
+
+ /** Projection dot only; callers apply the vector's scale, offset and landmark term. */
+ default float ashCanonicalLutDot(byte[] code, int groups, float[] lut) {
+ return ashCanonicalLutDot(code, 0, groups, lut);
+ }
+
+ /** As {@link #ashCanonicalLutDot(byte[], int, float[])} for the code starting at {@code codeOffset}. */
+ default float ashCanonicalLutDot(byte[] code, int codeOffset, int groups, float[] lut) {
+ return ASHLutScoring.dotCanonical(code, codeOffset, groups, lut);
+ }
+
+ /** Dot product of a projected query with a canonical signed-magnitude 2/4-bit code. */
+ default float ashProjectionDot(float[] query, byte[] code, int dimensions, int bitsPerDimension) {
+ return ashProjectionDot(query, code, 0, dimensions, bitsPerDimension);
+ }
+
+ /** As {@link #ashProjectionDot(float[], byte[], int, int)} for the code starting at {@code codeOffset}. */
+ default float ashProjectionDot(float[] query, byte[] code, int codeOffset, int dimensions, int bitsPerDimension) {
+ if (bitsPerDimension != 2 && bitsPerDimension != 4) {
+ throw new IllegalArgumentException("Projection scoring requires 2 or 4 bits per dimension");
+ }
+ java.util.Objects.checkFromIndexSize(0, dimensions, query.length);
+ int perByte = 8 / bitsPerDimension;
+ java.util.Objects.checkFromIndexSize(codeOffset, dimensions / perByte + (dimensions % perByte == 0 ? 0 : 1), code.length);
+ int sign = 1 << (bitsPerDimension - 1);
+ float sum = 0;
+ for (int i = 0; i < dimensions; i++) {
+ int field = (code[codeOffset + i / perByte] & 255) >>> ((i % perByte) * bitsPerDimension);
+ float magnitude = (field & (sign - 1)) + 0.5f;
+ sum += query[i] * ((field & sign) == 0 ? -magnitude : magnitude);
+ }
+ return sum;
+ }
+
+ /**
+ * Scores contiguous lanes in a byte-interleaved ASH block using a 16-entry LUT per nibble.
+ * Even groups occupy low nibbles; odd groups occupy high nibbles. Each byte row has
+ * {@code stride} lanes. Writes projection dot products, without per-vector headers.
+ */
+ default void ashLutScore(byte[] codes, int offset, int groups, int stride,
+ int lane, int count, float[] lut, float[] out, int outOffset) {
+ ASHLutScoring.score(codes, offset, groups, stride, lane, count, lut, out, outOffset);
+ }
+
+ /**
+ * Whether this provider implements complete FusedASH neighborhood scoring for one landmark.
+ * Query factories use this capability independently of the standalone LUT primitive.
+ *
+ * @return {@code true} when {@link #ashScoreNeighborhood} is implemented
+ */
+ default boolean supportsAshNeighborhoodScoring() { return false; }
+
+ /**
+ * Scores one complete C=1 FusedASH neighborhood, including all serialized blocks and headers.
+ * The result is {@code (1 + rawDotProduct) / 2}, unclamped, matching graph DOT_PRODUCT similarity.
+ *
+ * Inputs are read-only for the call; output storage must be distinct from {@code lut}.
+ * Calls may share immutable inputs across threads when their output arrays are separate.
+ * On failure, output contents are unspecified. Implementations validate sizes before native access.
+ *
+ * @param codes complete packed neighborhood, with full serialized storage for its final block
+ * @param groups positive nibble group count, including a possible partial final group
+ * @param blockSize serialized block width: 8, 16, or 32
+ * @param capacity positive number of lanes to score, including graph padding
+ * @param lut query LUT with at least {@code 16 * groups} floats
+ * @param landmarkDot query dot product with the single landmark
+ * @param out receives {@code capacity} similarities; later elements are left unchanged
+ * @throws UnsupportedOperationException if this provider has no implementation
+ * @throws IllegalArgumentException for an invalid shape or aliased LUT/output
+ * @throws IndexOutOfBoundsException if an array is too small
+ * @throws IllegalStateException if a scored lane contains a nonzero landmark ID
+ */
+ default void ashScoreNeighborhood(byte[] codes, int groups, int blockSize, int capacity,
+ float[] lut, float landmarkDot, float[] out) {
+ throw new UnsupportedOperationException("Native ASH neighborhood scoring is unavailable");
+ }
+
+ /** Whether {@link #ashScoreNeighborhoodInt8} is implemented. */
+ default boolean supportsAshInt8NeighborhoodScoring() { return false; }
+
+ /**
+ * Scores one complete C=1 FusedASH neighborhood of offset-binary codes (3 and 5 to 9 bits) with an
+ * INT8 query. Each block's body is {@code [unit][lane][bits bytes]} for {@code units} units of
+ * eight dimensions, followed by the usual headers. Per lane,
+ * {@code dot = ( - center) * invScale}, and the result is
+ * {@code (1 + scale * dot + landmarkDot + offset) / 2}, unclamped.
+ *
+ * @param codes complete packed neighborhood, with full serialized storage for its final block
+ * @param bits bits per dimension: 3, 5, 6, 7, 8 or 9
+ * @param units positive number of eight-dimension units
+ * @param blockSize serialized block width: 8, 16, or 32
+ * @param capacity positive number of lanes to score, including graph padding
+ * @param q INT8 query, {@code units * 8} entries
+ * @param invScale reciprocal of the query quantization scale
+ * @param center {@code (2^bits - 1)/2 * sum(q)}
+ * @param landmarkDot query dot product with the single landmark
+ * @param out receives {@code capacity} similarities; later elements are left unchanged
+ * @throws UnsupportedOperationException if this provider has no implementation
+ * @throws IllegalArgumentException for an invalid shape
+ * @throws IndexOutOfBoundsException if an array is too small
+ * @throws IllegalStateException if a scored lane contains a nonzero landmark ID
+ */
+ default void ashScoreNeighborhoodInt8(byte[] codes, int bits, int units, int blockSize, int capacity,
+ byte[] q, float invScale, float center, float landmarkDot, float[] out) {
+ throw new UnsupportedOperationException("Native INT8 ASH neighborhood scoring is unavailable");
+ }
+
/** Calculates the dot product of the given float arrays. */
float dotProduct(VectorFloat a, VectorFloat b);
@@ -343,6 +537,214 @@ default long nvqCosine8bitPacked(VectorFloat vector, ByteSequence bytes, f
*/
float nvqUniformLoss(VectorFloat vector, float minValue, float maxValue, int nBits);
+ // ------------------------------------------------------------------
+ // ASH kernels (default scalar implementations)
+ // ------------------------------------------------------------------
+
+ /**
+ * @return true if this backend provides a SIMD implementation of the masked-load kernel as described in the ASH paper.
+ */
+ default boolean supportsAshMaskedLoad() {
+ return false;
+ }
+
+ /** Immutable provider-specific projection, safe to share between query threads. */
+ interface AshProjection {
+ /** Writes the full matrix-vector product; input and output must not alias. */
+ void project(float[] query, float[] out);
+
+ /**
+ * Projects the first {@code count} vectors, each stored in its own row.
+ * All input rows must have the model's original dimension and output rows
+ * its projected dimension. Output rows must be distinct and must not alias
+ * any input row. Rows beyond count are ignored; count zero is a no-op.
+ * Providers may batch the multiply; the default retains their single-vector
+ * arithmetic without allocating temporary arrays.
+ */
+ default void projectBatch(float[][] inputs, float[][] outputs, int count) {
+ validateBatch(inputs, outputs, count);
+ for (int i = 0; i < count; i++) project(inputs[i], outputs[i]);
+ }
+
+ /**
+ * Whether this view has a genuine multi-vector projection implementation.
+ * Query dispatch uses this capability to avoid coordinating providers whose
+ * batch method simply repeats singleton calls. This is a capability flag,
+ * not a performance promise, and may become false after a native fallback.
+ */
+ default boolean supportsBatchProjection() { return false; }
+
+ /** Validates a batch before any output is changed. */
+ static void validateBatch(float[][] inputs, float[][] outputs, int count) {
+ java.util.Objects.requireNonNull(inputs, "inputs");
+ java.util.Objects.requireNonNull(outputs, "outputs");
+ if (count < 0 || count > inputs.length || count > outputs.length)
+ throw new IllegalArgumentException("Invalid ASH projection batch count: " + count);
+ for (int i = 0; i < count; i++) {
+ java.util.Objects.requireNonNull(inputs[i], "input row");
+ java.util.Objects.requireNonNull(outputs[i], "output row");
+ if (inputs[i].length != inputs[0].length || outputs[i].length != outputs[0].length)
+ throw new IllegalArgumentException("Ragged ASH projection batch");
+ }
+ if (count <= 32) {
+ // The normal encoding batch needs no temporary collection. For larger
+ // API requests use linear validation rather than an unbounded scan.
+ for (int i = 0; i < count; i++)
+ for (int j = 0; j < count; j++)
+ if (outputs[i] == inputs[j] || (j < i && outputs[i] == outputs[j]))
+ throw new IllegalArgumentException("Aliased ASH projection batch");
+ } else {
+ var used = new java.util.IdentityHashMap(count);
+ for (int i = 0; i < count; i++) used.put(inputs[i], Boolean.TRUE);
+ for (int i = 0; i < count; i++)
+ if (used.put(outputs[i], Boolean.TRUE) != null)
+ throw new IllegalArgumentException("Aliased ASH projection batch");
+ }
+ }
+
+ /**
+ * Returns an immutable view for query scoring, which may use lower operand
+ * precision. Encoding/training must continue to use this original object.
+ * The caller caches the returned view for the model's lifetime. Providers
+ * without a query-specific representation return this same object.
+ */
+ default AshProjection forQueryScoring() {
+ String mode = System.getProperty("jvector.ash.queryProjection", "auto");
+ switch (mode) {
+ case "auto":
+ case "fp32": return this;
+ case "fp16": throw new UnsupportedOperationException(
+ "This provider does not support FP16 ASH query projection");
+ case "int8": throw new UnsupportedOperationException(
+ "This provider does not support INT8 ASH query projection");
+ default: throw new IllegalArgumentException(
+ "jvector.ash.queryProjection must be auto, fp32, fp16 or int8: " + mode);
+ }
+ }
+
+ /** Operand precision of this prepared view; reporting must not initialize another view. */
+ default String operandPrecision() { return "fp32"; }
+
+ /** Additional backing storage, excluding the caller-owned matrix. */
+ default long ramBytesUsed() { return 0L; }
+ }
+
+ /**
+ * Prepares an immutable [projected dimension][original dimension] matrix once.
+ * The caller must not mutate its rows after preparation. Providers may cache a
+ * packed representation; each project call still prepares exactly one query.
+ */
+ default AshProjection prepareAshProjection(float[][] matrix) {
+ java.util.Objects.requireNonNull(matrix, "matrix");
+ if (matrix.length == 0 || matrix[0].length == 0)
+ throw new IllegalArgumentException("ASH projection dimensions must be positive");
+ int columns = matrix[0].length;
+ for (float[] row : matrix) {
+ if (row.length != columns) throw new IllegalArgumentException("Ragged ASH projection");
+ }
+ return (query, out) -> {
+ if (query.length != columns || out.length != matrix.length || query == out)
+ throw new IllegalArgumentException("ASH projection shape or aliasing mismatch");
+ for (int row = 0; row < matrix.length; row++) out[row] = ashDotRow(matrix[row], query);
+ };
+ }
+
+ /**
+ * Dot product between a matrix row and a dense vector.
+ * Used by ASH projection: {@code y_i = }.
+ *
+ * Contract: Arow.length == x.length.
+ * Default is scalar; SIMD backends may override.
+ */
+ default float ashDotRow(float[] Arow, float[] x) {
+ assert Arow.length == x.length : "Arow.length != x.length";
+ float acc = 0.0f;
+ for (int i = 0; i < x.length; i++) {
+ acc += Arow[i] * x[i];
+ }
+ return acc;
+ }
+
+ /**
+ * Sums {@code tildeQ[qOffset + i]} over the set bits {@code i < d} of one vector's packed sign words.
+ * The default is portable scalar code; SIMD backends override it.
+ */
+ default float ashMaskedAdd_512(float[] tildeQ,
+ int qOffset,
+ long[] allPackedVectors,
+ int packedBase,
+ int d,
+ int words) {
+ float sum = 0.0f;
+ int baseDim = 0;
+
+ for (int w = 0; w < words; w++) {
+ long word = allPackedVectors[packedBase + w];
+
+ // Standard bit-twiddling iteration (iterate only set bits)
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+
+ // Bounds check required for the last word
+ if (idx < d) {
+ sum += tildeQ[qOffset + idx];
+ }
+
+ // Clear the lowest set bit
+ word &= (word - 1);
+ }
+
+ baseDim += 64;
+ }
+ return sum;
+ }
+
+ /**
+ * Computes maskedAdd for a block slice of vectors stored in packed block-column-major form.
+ * For each lane in [0, blockLen), writes:
+ * {@code outMaskedAdd[lane] = }
+ * packedBits layout:
+ * packedBits[blockWordBase + w*blockSize + laneIndex] is 64-bit word w for that lane.
+ * Default implementation is scalar bit-walk with a register accumulator per lane.
+ * SIMD backends may override.
+ */
+ default void ashMaskedAddBlockAllWords(
+ float[] tildeQ,
+ int d,
+ long[] packedBits,
+ int packedBase, // blockId * nWords * blockSize
+ int nWords,
+ int blockSize,
+ int laneStart,
+ int blockLen,
+ float[] outMaskedAdd
+ ) {
+ assert laneStart >= 0;
+ assert blockLen >= 0;
+ assert laneStart + blockLen <= blockSize;
+ assert outMaskedAdd.length >= blockLen : "outMaskedAdd too small";
+
+ for (int lane = 0; lane < blockLen; lane++) {
+ float maskedAdd = 0f; // register accumulator (critical!)
+ int laneIndex = laneStart + lane;
+
+ for (int w = 0; w < nWords; w++) {
+ long word = packedBits[packedBase + (w * blockSize) + laneIndex];
+ int baseDim = w * 64;
+
+ while (word != 0L) {
+ int bit = Long.numberOfTrailingZeros(word);
+ int idx = baseDim + bit;
+ if (idx < d) maskedAdd += tildeQ[idx];
+ word &= (word - 1);
+ }
+ }
+
+ outMaskedAdd[lane] = maskedAdd;
+ }
+ }
+
/**
* Optionally performs the existing coarse/fine NVQ growth-rate search in one
* call. Returns NaN when unsupported; the caller then uses its Java search.
@@ -352,4 +754,13 @@ default float tryNvqGrowthRate(VectorFloat vector, float minValue, float maxV
return Float.NaN;
}
+ /**
+ * Sets the native BLAS thread count for calls made on the current thread, as
+ * {@link io.github.jbellis.jvector.util.BlasThreads} requests. Returns the setting to restore
+ * afterward, or -1 when this provider cannot control BLAS threads.
+ */
+ default int setBlasThreadsLocal(int threads) {
+ return -1;
+ }
+
}
diff --git a/jvector-examples/README.md b/jvector-examples/README.md
index 46990445b..1cce8a61f 100644
--- a/jvector-examples/README.md
+++ b/jvector-examples/README.md
@@ -2,6 +2,60 @@
JVector comes with the following sample programs to try:
+### Hello ASH world
+
+Two self-contained introductions build and query a fused ASH disk index over the
+public `e5-small-v2-100k` dataset (downloaded and cached automatically):
+
+- [HelloASHWorldInt8](src/main/java/io/github/jbellis/jvector/example/HelloASHWorldInt8.java):
+ INT8 query projection for Intel Ice Lake and other supported INT8 processors.
+- [HelloASHWorldFp16](src/main/java/io/github/jbellis/jvector/example/HelloASHWorldFp16.java):
+ native FP16 query projection for Intel Sapphire Rapids and newer FP16 processors.
+
+Both teach the same six steps: load vectors, train/encode ASH, construct the graph,
+write fused ASH neighborhoods with NVQ reranking vectors, reopen and search
+synchronously, then submit individual queries through the queued API.
+They use ASH(b=2,d=192) for 384-dimensional inputs. Stored-code precision (`b`)
+is independent of query-projection precision; training and encoding remain FP32.
+NVQ's mean is collected during ASH encoding, avoiding another input scan.
+
+Build the examples with the native provider and oneDNN support described in the
+[native build guide](../jvector-native/src/main/native/README.md#recommended-onednn-build-and-package).
+From the repository root, using JDK 22 or newer:
+
+```bash
+mvn package -DskipTests
+java --add-modules=jdk.incubator.vector --enable-native-access=ALL-UNNAMED -Xmx8g \
+ -cp jvector-examples/target/jvector-examples-4.1.0-SNAPSHOT-jar-with-dependencies.jar \
+ io.github.jbellis.jvector.example.HelloASHWorldInt8
+java --add-modules=jdk.incubator.vector --enable-native-access=ALL-UNNAMED -Xmx8g \
+ -cp jvector-examples/target/jvector-examples-4.1.0-SNAPSHOT-jar-with-dependencies.jar \
+ io.github.jbellis.jvector.example.HelloASHWorldFp16
+```
+
+Run the program matching your CPU; explicit unsupported precisions fail instead
+of silently changing precision. Each saves `ash-int8.graph` or `ash-fp16.graph`
+and prints top-ten vector ordinals, similarity scores and visited-node counts
+for the first held-out query using both synchronous and queued search. An optional
+argument selects a new output path;
+existing files are never overwritten. Original vectors are not needed to reopen
+the finished index.
+
+The synchronous section introduces preparation and graph search for one query.
+The queued section submits eight independent queries using two application-assigned
+workers, admission 16 and maximum preparation batch 8. The shared
+[HelloASHQueuedSearch](src/main/java/io/github/jbellis/jvector/example/HelloASHQueuedSearch.java)
+shows the host integration explicitly: compatible-query preparation, worker-owned
+searchers, completion futures, admission rejection and orderly shutdown. Callers
+submit individual requests; the queue prepares currently waiting compatible work
+together, without a fill timer. Batch sizes depend on what is waiting at that moment.
+Each request owns its input snapshot and prepared state through completion.
+
+For clarity, this small example allocates request records and preparation arrays.
+Services can reuse bounded storage as described in
+[bounded query batching and state ownership](../docs/ash-query-integration.md).
+For more hardware choices, see [projection precision](../docs/ash-projection-precision.md).
+
### Batched file-backed vectors
Two hello-world programs use public-catalog `e5-small-v2-100k`, or accept an existing
@@ -18,7 +72,7 @@ and vector ownership; the second trains and encodes with PQ and NVQ. See
source lifecycle, and integration guidance.
### SiftSmall
-A simple benchmark for the sift dataset located in the [siftsmall](./siftsmall) directory in the project root.
+A simple benchmark for the sift dataset located in the [siftsmall](../siftsmall) directory in the project root.
> `mvn compile exec:exec@sift`
@@ -29,8 +83,6 @@ the best tradeoffs between recall and throughput.
This benchmark requires `fvecs' versions of datasets from [https://github.com/erikbern/ann-benchmarks](https://github.com/erikbern/ann-benchmarks/blob/main/README.md#data-sets) to be downloaded to `dataset_cache`
directory under the project root.
-You can use [`plot_output.py`](./plot_output.py) to graph the [pareto-optimal points](https://en.wikipedia.org/wiki/Pareto_efficiency) found by `Bench`.
-
> `mvn compile exec:exec@bench`
Some sample KNN datasets for testing based on ada-002 embeddings generated on wikipedia data are available in ivec/fvec format for testing at:
diff --git a/jvector-examples/bench-profiles/README.md b/jvector-examples/bench-profiles/README.md
new file mode 100644
index 000000000..73071fd04
--- /dev/null
+++ b/jvector-examples/bench-profiles/README.md
@@ -0,0 +1,44 @@
+# Benchmark launcher and server profiles
+
+`jvector-examples/benchyaml-play.sh` runs one
+[index-parameters](../yaml-configs/index-parameters) YAML with BenchYAML in one JVM. Every launch
+first rebuilds the checkout with Maven, including the native libraries, so results always come
+from the current source. The run records the commit, with `-dirty` when tracked files have
+uncommitted changes, and the JAR checksum.
+
+```bash
+BENCH_RUN_CONFIG=jvector-examples/yaml-configs/ash-pq/run-single-jvm.yml \
+ jvector-examples/benchyaml-play.sh \
+ jvector-examples/yaml-configs/index-parameters/studies/ash-vs-pq-single-jvm/mpnet10m-2bit-ash.yml
+```
+
+The index YAML sets concurrency: `construction.workers` for training, encoding, graph insertion
+and writing, and `search.workers` for queries, each a number or `max` (the platform's physical
+cores). `search.workers` is a sweep, so `[1, 8]` queries the same index with 1 and then 8 workers.
+The launcher sets where the JVM may run, and the run YAML sets what is printed and logged.
+JVector sets native BLAS threads itself (the construction pool's size during training), so the
+launcher clears `MKL_NUM_THREADS` and the other BLAS thread variables.
+
+`BENCH_RUN_CONFIG` selects the run YAML (default `yaml-configs/run-config.yml`).
+`BENCH_DEBUG=` waits for a debugger on that port.
+
+Long construction phases print progress every 10 seconds: vectors inserted with an estimate of
+the time left, then graph cleanup and index writing. Set `-Djvector.bench.progressSeconds=0`
+to print only each phase's start and finish.
+
+## Server profiles
+
+A profile describes one machine: JDK, working directory root, dataset catalogs and, optionally,
+cores and NUMA nodes. CPU lists come from `lscpu`, one hardware thread per physical core.
+Copy [example.env](example.env) to `local/profile.env`, or keep several profiles in
+`local/.env` and pick one with `BENCH_PROFILE=`. The `local/` directory is ignored
+by git because profiles hold machine paths. Each recipe runs in its own directory under
+`BENCH_WORK_ROOT`.
+
+## IntelliJ
+
+The repository's `.run/` directory holds run configurations for experiments. Each names a recipe
+and a run YAML, and takes machine settings from your profile. To add one for another recipe,
+copy an existing file, then change its name and the recipe path in `SCRIPT_OPTIONS`; every case in
+[ash-vs-pq-single-jvm](../yaml-configs/index-parameters/studies/ash-vs-pq-single-jvm) runs with
+`run-single-jvm.yml`.
diff --git a/jvector-examples/bench-profiles/example.env b/jvector-examples/bench-profiles/example.env
new file mode 100644
index 000000000..1235deedd
--- /dev/null
+++ b/jvector-examples/bench-profiles/example.env
@@ -0,0 +1,31 @@
+# Server profile for jvector-examples/benchyaml-play.sh.
+# Copy to bench-profiles/local/profile.env (ignored by git) and edit for your machine,
+# or keep several in bench-profiles/local/.env and select one with BENCH_PROFILE=.
+
+# JDK used to build and run BenchYAML.
+BENCH_JDK=$HOME/.jdks/jdk-22.0.2+9
+
+# Parent of each experiment's working directory (indexes, models, logs). Put it on fast local
+# storage: queries read the index from here, including indexes that are not saved.
+BENCH_WORK_ROOT=/data/jvector-bench
+
+# Dataset catalogs, one per line: .
+# Repository paths are relative to the checkout. Each catalog's ${JVECTOR_BENCH_DATA_DIR}
+# is replaced by its directory. List each dataset name in only one catalog.
+BENCH_CATALOGS="
+jvector-examples/yaml-configs/ash-pq/catalog-large.yml /data/vectors/large
+jvector-examples/yaml-configs/ash-pq/catalog-1m.yml /data/vectors/1m
+"
+
+# Optional. The JVM runs on one hardware thread of each of the first BENCH_CORES physical
+# cores (default: all), with memory interleaved across BENCH_NODES (default: all NUMA nodes).
+# The index YAML's construction.workers and search.workers set how many threads work.
+#BENCH_CORES=48
+#BENCH_NODES=0,1
+
+# Optional. Where indexes that are not saved are written and reopened for queries
+# (default: each experiment's tmp directory under BENCH_WORK_ROOT). Use fast local storage.
+#BENCH_TMP_DIR=/data/jvector-bench/tmp
+
+# Optional. Maximum Java heap (default 128g).
+#BENCH_HEAP=128g
diff --git a/jvector-examples/benchyaml-play.sh b/jvector-examples/benchyaml-play.sh
new file mode 100755
index 000000000..c0f9618c4
--- /dev/null
+++ b/jvector-examples/benchyaml-play.sh
@@ -0,0 +1,141 @@
+#!/usr/bin/env bash
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Build BenchYAML from this checkout, then run one index-parameters YAML in one JVM.
+# Every launch rebuilds Java and native code, so a run always reflects the current
+# source. The YAML's construction.workers and search.workers set concurrency; this
+# script sets where the JVM may run. See jvector-examples/bench-profiles/README.md.
+#
+# benchyaml-play.sh build, then run
+# benchyaml-play.sh --build-only build and print JAR=, SOURCE_REF= and JAR_SHA256=
+#
+# BENCH_RUN_CONFIG selects the run YAML (default jvector-examples/yaml-configs/run-config.yml).
+# Machine settings come from the server profile (BENCH_PROFILE, or bench-profiles/local/profile.env).
+set -euo pipefail
+
+repo=$(git -C "$(dirname -- "$0")" rev-parse --show-toplevel)
+profiles="$repo/jvector-examples/bench-profiles"
+
+profile=${BENCH_PROFILE:-$profiles/local/profile.env}
+[[ $profile == */* ]] || profile="$profiles/local/$profile.env"
+[[ -f $profile ]] || { echo "Server profile not found: $profile. Copy bench-profiles/example.env to bench-profiles/local/profile.env and edit it." >&2; exit 2; }
+# shellcheck source=/dev/null
+source "$profile"
+jdk=${BENCH_JDK:?The server profile must set BENCH_JDK.}
+
+# Maven and native build output goes to build.log under the work root; the console gets one line.
+build() {
+ local marker start jars log
+ marker=$(mktemp)
+ start=$(date +%s)
+ log="${BENCH_WORK_ROOT:-$repo/jvector-examples/target}/build.log"
+ mkdir -p "$(dirname "$log")"
+ printf 'Building BenchYAML from %s ... ' "$repo" >&2
+ if ! (cd "$repo" && JAVA_HOME="$jdk" PATH="$jdk/bin:$PATH" JAVA_TOOL_OPTIONS= JDK_JAVA_OPTIONS= \
+ mvn -B -q package -DskipTests -pl jvector-examples -am) > "$log" 2>&1; then
+ printf 'failed. Last lines of %s:\n' "$log" >&2
+ tail -40 "$log" >&2
+ exit 1
+ fi
+ printf 'done in %s s (log: %s)\n' "$(( $(date +%s) - start ))" "$log" >&2
+ mapfile -t jars < <(find "$repo/jvector-examples/target" -maxdepth 1 -name '*-jar-with-dependencies.jar' -newer "$marker")
+ rm -f "$marker"
+ if [[ ${#jars[@]} != 1 ]]; then
+ echo "Expected the build to write one jar-with-dependencies in jvector-examples/target; found ${#jars[@]}." >&2
+ exit 1
+ fi
+ JAR=${jars[0]}
+ SOURCE_REF=$(git -C "$repo" rev-parse HEAD)
+ [[ -z $(git -C "$repo" status --porcelain --untracked-files=no) ]] || SOURCE_REF="$SOURCE_REF-dirty"
+ JAR_SHA256=$(sha256sum "$JAR" | cut -d' ' -f1)
+}
+
+if [[ ${1:-} == --build-only ]]; then
+ build
+ printf 'JAR=%q\nSOURCE_REF=%q\nJAR_SHA256=%q\n' "$JAR" "$SOURCE_REF" "$JAR_SHA256"
+ exit 0
+fi
+
+index=${1:?Pass an index-parameters YAML file.}
+[[ $index == /* ]] || index="$repo/$index"
+[[ -f $index ]] || { echo "Index config not found: $index" >&2; exit 2; }
+run_config=${BENCH_RUN_CONFIG:-jvector-examples/yaml-configs/run-config.yml}
+[[ $run_config == /* ]] || run_config="$repo/$run_config"
+[[ -f $run_config ]] || { echo "Run config not found: $run_config" >&2; exit 2; }
+work_root=${BENCH_WORK_ROOT:?The server profile must set BENCH_WORK_ROOT.}
+
+# One hardware thread on each of the first BENCH_CORES physical cores, in CPU order.
+mapfile -t physical < <(lscpu -p=CPU,CORE,SOCKET | awk -F, '!/^#/ && !seen[$3":"$2]++ {print $1}')
+cores=${BENCH_CORES:-${#physical[@]}}
+(( cores >= 1 && cores <= ${#physical[@]} )) || { echo "BENCH_CORES=$cores must be between 1 and the ${#physical[@]} physical cores." >&2; exit 2; }
+cpus=$(printf '%s\n' "${physical[@]:0:cores}" | paste -sd, -)
+# Print CPU lists as ranges, for example 0-47.
+cpu_ranges=$(printf '%s\n' "${physical[@]:0:cores}" | sort -n | awk 'NR==1{a=b=$1; next} $1==b+1{b=$1; next} {printf "%s%s,", a, (a==b?"":"-"b); a=b=$1} END{printf "%s%s", a, (a==b?"":"-"b)}')
+nodes=${BENCH_NODES:-$(lscpu -p=NODE | grep -v '^#' | sort -un | paste -sd, -)}
+
+# Each recipe gets its own working directory, so caches and logs never mix between experiments.
+# Unsaved indexes are written and reopened for queries under BENCH_TMP_DIR (java.io.tmpdir),
+# by default the experiment's tmp directory on the same storage as BENCH_WORK_ROOT.
+recipe=${index#"$repo/jvector-examples/yaml-configs/index-parameters/"}
+[[ $recipe != /* ]] || recipe="other/$(basename "$recipe")"
+work="$work_root/${recipe%.yml}"
+
+mkdir -p "$work_root"
+exec 9>"$work_root/benchyaml.lock"
+flock -n 9 || { echo "Another BenchYAML launch holds $work_root/benchyaml.lock." >&2; exit 1; }
+if pgrep -f '[.]example[.]BenchYAML' >/dev/null; then echo "A BenchYAML JVM is already running on this machine." >&2; exit 1; fi
+
+build
+
+# BenchYAML reads dataset metadata and catalogs relative to its working directory.
+catalogs="$work/jvector-examples/yaml-configs/dataset-catalogs"
+tmp=${BENCH_TMP_DIR:-$work/tmp}
+mkdir -p "$catalogs" "$tmp"
+cp "$repo/jvector-examples/yaml-configs/dataset-metadata.yml" "$repo/jvector-examples/yaml-configs/datasets.yml" "$work/jvector-examples/yaml-configs/"
+find "$catalogs" -mindepth 1 -delete
+# A recipe that does not save its index also starts without saved models, so every run
+# measures training and encoding. Recipes that reuse indexes keep both caches.
+if ! grep -Eq '^[[:space:]]+useSavedIndexIfExists:[[:space:]]*Yes' "$index"; then
+ for cache in compressor_cache index_cache; do
+ if [[ -d $work/$cache ]]; then find "$work/$cache" -mindepth 1 -delete; fi
+ done
+fi
+n=0
+while read -r catalog data; do
+ [[ -n $catalog ]] || continue
+ [[ $catalog == /* ]] || catalog="$repo/$catalog"
+ n=$((n + 1))
+ sed "s|\${JVECTOR_BENCH_DATA_DIR}|$data|g" "$catalog" > "$catalogs/$n-$(basename "${catalog%.*}").yml"
+done <<< "${BENCH_CATALOGS:?The server profile must list BENCH_CATALOGS.}"
+
+printf '\nRecipe %s\nCPUs %s (%s physical cores), memory interleaved on NUMA nodes %s\nSource %s\nJAR %s\n sha256 %s\nRun YAML %s\nWork %s\nIndexes %s\n\n' \
+ "$recipe" "$cpu_ranges" "$cores" "$nodes" "$SOURCE_REF" "$JAR" "$JAR_SHA256" "$run_config" "$work" "$tmp"
+debug=()
+[[ -z ${BENCH_DEBUG:-} ]] || debug=("-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=*:$BENCH_DEBUG")
+cd "$work"
+# JVector sets BLAS threads for its own calls (one, or the construction pool's size), so the
+# BLAS thread variables are cleared rather than set here.
+exec env -u JAVA_TOOL_OPTIONS -u JDK_JAVA_OPTIONS -u _JAVA_OPTIONS \
+ -u MKL_NUM_THREADS -u OMP_NUM_THREADS -u OPENBLAS_NUM_THREADS -u GOTO_NUM_THREADS \
+ MKL_DYNAMIC=FALSE ONEDNN_VERBOSE=0 \
+ numactl --physcpubind="$cpus" --interleave="$nodes" "$jdk/bin/java" "${debug[@]}" \
+ --add-modules=jdk.incubator.vector --enable-native-access=ALL-UNNAMED -da -Xms2g "-Xmx${BENCH_HEAP:-128g}" \
+ "-XX:ActiveProcessorCount=$cores" "-Djava.io.tmpdir=$tmp" '-Djava.util.logging.SimpleFormatter.format=%5$s%6$s%n' \
+ -Dstdout.encoding=UTF-8 -Dstderr.encoding=UTF-8 \
+ -Djvector.vectorization_provider=native -Djvector.parallel_writer=true \
+ "-Djvector.physical_core_count=$cores" \
+ -Djvector.bench.verifyConcurrentResults=true -Djvector.bench.throughputMinSeconds=3.0 \
+ "-Djvector.bench.sourceRef=$SOURCE_REF" "-Djvector.bench.runConfig=$run_config" \
+ -cp "$JAR" io.github.jbellis.jvector.example.BenchYAML "$index"
diff --git a/jvector-examples/run-ash-pq-benchmark.sh b/jvector-examples/run-ash-pq-benchmark.sh
new file mode 100755
index 000000000..6cd9a977f
--- /dev/null
+++ b/jvector-examples/run-ash-pq-benchmark.sh
@@ -0,0 +1,183 @@
+#!/usr/bin/env bash
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+set -euo pipefail
+
+if [[ $# != 5 ]]; then
+ echo "Usage: $0 " >&2
+ echo "Example: $0 mpnet-2bit-ash.yml queued 8 0-7 /data/bench/mpnet-2bit-ash" >&2
+ exit 2
+fi
+REPO_ROOT="$(git -C "$(dirname "$0")" rev-parse --show-toplevel)"
+CONFIG_NAME="$1"
+MODE="$2"
+CORES="$3"
+CPU_LIST="$4"
+CASE_DIRECTORY="$5"
+case "$MODE" in construction|direct|queued) ;; *) echo "Invalid mode: $MODE" >&2; exit 2;; esac
+if ! [[ "$CORES" =~ ^[1-9][0-9]*$ ]]; then echo "Physical core count must be positive" >&2; exit 2; fi
+for command in java numactl lscpu sha256sum awk; do
+ command -v "$command" >/dev/null || { echo "Required command not found: $command" >&2; exit 2; }
+done
+if ! [[ "$CPU_LIST" =~ ^[0-9]+(-[0-9]+)?(,[0-9]+(-[0-9]+)?)*$ ]]; then
+ echo "CPU list must contain CPU numbers or ranges, such as 0 or 0-7." >&2
+ exit 2
+fi
+lscpu --parse=CPU,CORE,SOCKET,ONLINE | awk -F, -v selection="$CPU_LIST" -v budget="$CORES" '
+ BEGIN {
+ count = split(selection, pieces, ",")
+ for (i = 1; i <= count; i++) {
+ n = split(pieces[i], bounds, "-")
+ first = bounds[1] + 0; last = n == 1 ? first : bounds[2] + 0
+ if (last < first || last - first + 1 > budget) { bad = 1; exit 2 }
+ for (cpu = first; cpu <= last; cpu++) chosen[cpu] = 1
+ }
+ }
+ !/^#/ && ($1 in chosen) {
+ if ($4 == "N") { bad = 1; exit 2 }
+ seen[$1] = 1; physical[$3 ":" $2] = 1
+ }
+ END {
+ for (cpu in chosen) wanted++
+ for (cpu in seen) found++
+ for (core in physical) cores++
+ if (bad || wanted != budget || found != budget || cores != budget) {
+ print "CPU selection must contain exactly the budgeted number of distinct online physical cores." > "/dev/stderr"
+ exit 2
+ }
+ }'
+DATA_ROOT="${JVECTOR_BENCH_DATA_DIR:-}"
+if [[ "$DATA_ROOT" != /* || ! -d "$DATA_ROOT" ]]; then
+ echo "Set JVECTOR_BENCH_DATA_DIR to the absolute local dataset directory; see the catalog template." >&2
+ exit 2
+fi
+if [[ -n "${JAVA_TOOL_OPTIONS:-}" || -n "${JDK_JAVA_OPTIONS:-}" ]]; then
+ echo "Clear JAVA_TOOL_OPTIONS/JDK_JAVA_OPTIONS so the recorded Java command is complete." >&2
+ exit 2
+fi
+git -C "$REPO_ROOT" diff --quiet
+git -C "$REPO_ROOT" diff --cached --quiet
+SOURCE_REF="$(git -C "$REPO_ROOT" rev-parse HEAD)"
+JAVA_BIN="$(command -v java)"
+CONFIG="$REPO_ROOT/jvector-examples/yaml-configs/index-parameters/studies/ash-vs-pq/$CONFIG_NAME"
+RUN_CONFIG="$REPO_ROOT/jvector-examples/yaml-configs/ash-pq/run-$MODE.yml"
+CATALOG_NAME="${JVECTOR_BENCH_CATALOG:-catalog-1m.yml}"
+case "$CATALOG_NAME" in catalog-1m.yml|catalog-large.yml) ;; *) echo "Unknown dataset catalog: $CATALOG_NAME" >&2; exit 2;; esac
+CATALOG="$REPO_ROOT/jvector-examples/yaml-configs/ash-pq/$CATALOG_NAME"
+[[ -f "$CATALOG" ]] || { echo "Dataset catalog not found: $CATALOG" >&2; exit 2; }
+[[ -f "$CONFIG" && "$CONFIG_NAME" != run-* && "$CONFIG_NAME" != catalog-* && "$CONFIG_NAME" != */* ]] || {
+ echo "Unknown index config: $CONFIG_NAME" >&2; exit 2;
+}
+
+if [[ "$MODE" != construction ]]; then
+ # The checked-in recipes put delivery under the ASH search parameters.
+ EXECUTION="$(sed -nE 's/^[[:space:]]+queryExecution:[[:space:]]*(direct|queued)[[:space:]]*(#.*)?$/\1/p' "$CONFIG")"
+ if [[ "${EXECUTION:-direct}" != "$MODE" ]]; then
+ echo "The index YAML selects ${EXECUTION:-direct}. Use its matching query report selection." >&2
+ exit 2
+ fi
+fi
+
+shopt -s nullglob
+JARS=("$REPO_ROOT"/jvector-examples/target/jvector-examples-*-jar-with-dependencies.jar)
+if [[ ${#JARS[@]} != 1 ]]; then
+ echo "Build exactly one examples assembly with Maven; see docs/ash-pq-benchmark.md." >&2
+ exit 2
+fi
+JAR="${JARS[0]}"
+mkdir -p "$CASE_DIRECTORY"
+CASE_DIRECTORY="$(cd "$CASE_DIRECTORY" && pwd)"
+CASE_ID="$(sha256sum "$CONFIG" | cut -d ' ' -f 1)"
+if [[ -f "$CASE_DIRECTORY/case-id.txt" && "$(cat "$CASE_DIRECTORY/case-id.txt")" != "$CASE_ID" ]]; then
+ echo "Case directory belongs to a different index configuration; use its matching YAML." >&2
+ exit 2
+fi
+if [[ "$MODE" == construction ]]; then
+ for cache in compressor_cache index_cache; do
+ if [[ -d "$CASE_DIRECTORY/$cache" ]] && [[ -n "$(ls -A "$CASE_DIRECTORY/$cache")" ]]; then
+ echo "Fresh construction requires empty caches; use a new case directory." >&2
+ exit 2
+ fi
+ done
+else
+ [[ -d "$CASE_DIRECTORY/index_cache" && -d "$CASE_DIRECTORY/compressor_cache" ]] || {
+ echo "Construct this case first; query runs must reuse its model and graph." >&2; exit 2;
+ }
+fi
+
+# The standard loader resolves metadata/catalogs relative to the working
+# directory. Install only the checked-in local mapping, avoiding ambiguous
+# duplicate aliases from other dataset catalogs. Numerical code is all in the JAR.
+CASE_CONFIG="$CASE_DIRECTORY/jvector-examples/yaml-configs"
+mkdir -p "$CASE_CONFIG/dataset-catalogs"
+cp "$REPO_ROOT/jvector-examples/yaml-configs/dataset-metadata.yml" "$CASE_CONFIG/"
+cp "$REPO_ROOT/jvector-examples/yaml-configs/datasets.yml" "$CASE_CONFIG/"
+cp "$CATALOG" "$CASE_CONFIG/dataset-catalogs/ash-pq-local.yml"
+printf '%s\n' "$CASE_ID" > "$CASE_DIRECTORY/case-id.txt"
+mkdir -p "$CASE_DIRECTORY/runs"
+RUN_DIRECTORY="$(mktemp -d "$CASE_DIRECTORY/runs/$MODE-c$CORES-XXXXXX")"
+cp "$CONFIG" "$RUN_DIRECTORY/index.yml"
+cp "$RUN_CONFIG" "$RUN_DIRECTORY/run.yml"
+cp "$CATALOG" "$RUN_DIRECTORY/catalog.yml"
+printf '%s\n' "$SOURCE_REF" > "$RUN_DIRECTORY/source-ref.txt"
+sha256sum "$JAR" > "$RUN_DIRECTORY/jar.sha256"
+"$JAVA_BIN" -version > "$RUN_DIRECTORY/java-version.txt" 2>&1
+lscpu --extended=CPU,CORE,SOCKET,NODE,ONLINE > "$RUN_DIRECTORY/cpu-topology.txt"
+
+# Native query kernels execute on the Java search workers. Construction has
+# its own CPU budget for training and encoding. Avoid nested query thread pools.
+NATIVE_THREADS=1
+if [[ "$MODE" == construction ]]; then NATIVE_THREADS="$CORES"; fi
+export MKL_NUM_THREADS="$NATIVE_THREADS" OMP_NUM_THREADS="$NATIVE_THREADS"
+export OPENBLAS_NUM_THREADS="$NATIVE_THREADS" GOTO_NUM_THREADS="$NATIVE_THREADS"
+export MKL_DYNAMIC=FALSE ONEDNN_VERBOSE=0
+
+NUMA_POLICY="${JVECTOR_BENCH_NUMA_POLICY:-bind:0}"
+case "$NUMA_POLICY" in
+ bind:*) MEMORY_ARGS=("--membind=${NUMA_POLICY#bind:}");;
+ interleave:*) MEMORY_ARGS=("--interleave=${NUMA_POLICY#interleave:}");;
+ *) echo "JVECTOR_BENCH_NUMA_POLICY must be bind: or interleave:" >&2; exit 2;;
+esac
+JAVA_ARGS=(--add-modules=jdk.incubator.vector --enable-native-access=ALL-UNNAMED
+ -da -Xms2g "-Xmx${JVECTOR_BENCH_HEAP:-128g}" "-XX:ActiveProcessorCount=$CORES"
+ -Djvector.vectorization_provider=native -Djvector.parallel_writer=true
+ "-Djvector.physical_core_count=$CORES" "-Djvector.bench.queryThreads=$((4 * CORES))"
+ -Djvector.bench.verifyConcurrentResults=true -Djvector.bench.throughputMinSeconds=3.0
+ "-Djvector.bench.sourceRef=$SOURCE_REF" "-Djvector.bench.runConfig=$RUN_CONFIG"
+ -cp "$JAR" io.github.jbellis.jvector.example.BenchYAML "$CONFIG")
+LAUNCH=(numactl "--physcpubind=$CPU_LIST" "${MEMORY_ARGS[@]}" "$JAVA_BIN" "${JAVA_ARGS[@]}")
+{
+ printf 'cd %q\n' "$CASE_DIRECTORY"
+ printf 'JVECTOR_BENCH_DATA_DIR=%q ' "$DATA_ROOT"
+ printf 'MKL_NUM_THREADS=%q OMP_NUM_THREADS=%q OPENBLAS_NUM_THREADS=%q GOTO_NUM_THREADS=%q MKL_DYNAMIC=FALSE ONEDNN_VERBOSE=0 ' \
+ "$MKL_NUM_THREADS" "$OMP_NUM_THREADS" "$OPENBLAS_NUM_THREADS" "$GOTO_NUM_THREADS"
+ printf '%q ' "${LAUNCH[@]}"
+ printf '\n'
+} > "$RUN_DIRECTORY/command.sh"
+echo "Artifacts: $RUN_DIRECTORY"
+cd "$CASE_DIRECTORY"
+"${LAUNCH[@]}" > "$RUN_DIRECTORY/console.log" 2>&1 &
+JVM_PID=$!
+printf '%s\n' "$JVM_PID" > "$RUN_DIRECTORY/jvm.pid"
+trap 'kill "$JVM_PID" 2>/dev/null || true; wait "$JVM_PID" 2>/dev/null || true; exit 130' INT TERM
+STATUS=0
+wait "$JVM_PID" || STATUS=$?
+trap - INT TERM
+cat "$RUN_DIRECTORY/console.log"
+printf '%s\n' "$STATUS" > "$RUN_DIRECTORY/exit-code.txt"
+if [[ "$STATUS" != 0 ]]; then exit "$STATUS"; fi
+if [[ "$MODE" != construction ]] && grep -q 'Building graph index' "$RUN_DIRECTORY/console.log"; then
+ echo "Query invocation rebuilt an index; do not treat this as a cached-index result." >&2
+ exit 1
+fi
diff --git a/jvector-examples/run-ash-pq-study.py b/jvector-examples/run-ash-pq-study.py
new file mode 100755
index 000000000..b1fcbc850
--- /dev/null
+++ b/jvector-examples/run-ash-pq-study.py
@@ -0,0 +1,334 @@
+#!/usr/bin/env python3
+"""Orchestrate the checked-in standalone BenchYAML recipe; no query loop or performance measurement code."""
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from pathlib import Path
+import argparse
+import collections
+import csv
+import hashlib
+import json
+import os
+import re
+import shutil
+import signal
+import statistics
+import sys
+import subprocess
+import time
+
+
+
+def physical_cores():
+ """One online hardware thread per physical core, in CPU order, as (cpu, node) pairs."""
+ seen, cores = set(), []
+ for line in subprocess.check_output(["lscpu", "-p=CPU,CORE,SOCKET,NODE"], text=True).splitlines():
+ if line.startswith("#"):
+ continue
+ cpu, core, socket, node = (field or "0" for field in line.split(","))
+ if (socket, core) not in seen:
+ seen.add((socket, core))
+ cores.append((int(cpu), int(node)))
+ return cores
+
+
+def cpu_ranges(cpus):
+ """Format CPU numbers as a numactl list, for example 0-7,48."""
+ ranges, cpus = [], sorted(cpus)
+ start = end = cpus[0]
+ for cpu in cpus[1:] + [None]:
+ if cpu is not None and cpu == end + 1:
+ end = cpu
+ continue
+ ranges.append(str(start) if start == end else f"{start}-{end}")
+ if cpu is not None:
+ start = end = cpu
+ return ",".join(ranges)
+
+
+# Defaults follow the published protocol on whatever machine runs it: construction on
+# one thread of every physical core with memory interleaved across all NUMA nodes, and
+# 1- and 8-core queries on the first physical cores of one node, with memory on that node.
+cores = physical_cores()
+query_node = cores[0][1]
+query_cpus = [cpu for cpu, node in cores if node == query_node]
+parser = argparse.ArgumentParser(description=__doc__)
+parser.add_argument("--output", type=Path, required=True)
+parser.add_argument("--data", type=Path, required=True)
+parser.add_argument("--build-cores", type=int, default=len(cores))
+parser.add_argument("--build-cpus", default=cpu_ranges([cpu for cpu, _ in cores]))
+parser.add_argument("--build-nodes", default=",".join(str(n) for n in sorted({node for _, node in cores})))
+parser.add_argument("--one-core-cpu", default=str(query_cpus[0]))
+parser.add_argument("--eight-core-cpus", default=cpu_ranges(query_cpus[:8]) if len(query_cpus) >= 8 else None)
+parser.add_argument("--query-node", default=str(query_node))
+parser.add_argument("--build-repetitions", type=int, default=2)
+parser.add_argument("--query-repetitions", type=int, default=1)
+parser.add_argument("--dataset-set", choices=("1m", "large"), default="1m")
+parser.add_argument("--algorithms", nargs="+", choices=("pq", "ash"), default=["pq", "ash"])
+parser.add_argument("--query-modes", nargs="+", choices=("direct", "queued"), default=["direct", "queued"])
+parser.add_argument("--plan-only", action="store_true")
+args = parser.parse_args()
+if min(args.build_cores, args.build_repetitions, args.query_repetitions) <= 0:
+ parser.error("Core and repetition counts must be positive")
+if args.eight_core_cpus is None:
+ parser.error("NUMA node %s has fewer than 8 physical cores; pass --eight-core-cpus and --query-node" % args.query_node)
+repo = Path(__file__).resolve().parent.parent
+helper = repo / "jvector-examples/run-ash-pq-benchmark.sh"
+source_ref = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=repo, text=True).strip()
+jars = list((repo / "jvector-examples/target").glob("jvector-examples-*-jar-with-dependencies.jar"))
+if len(jars) != 1:
+ parser.error("Build exactly one standalone examples assembly before starting")
+jar_sha = hashlib.sha256(jars[0].read_bytes()).hexdigest()
+output = args.output.resolve()
+data = args.data.resolve()
+if not data.is_dir():
+ parser.error("Dataset directory does not exist")
+jobs = []
+prefixes = ("openai3072", "mpnet", "cohere", "ada002") if args.dataset_set == "1m" else ("mpnet10m", "cohere10m", "cap6m")
+for group, (prefix, bits) in enumerate((p, b) for p in prefixes for b in (2, 4)):
+ for repetition in range(1, args.build_repetitions + 1):
+ algorithms = ("pq", "ash") if (group + repetition) % 2 else ("ash", "pq")
+ for algorithm in algorithms:
+ jobs.append(dict(prefix=prefix, bits=bits, algorithm=algorithm, mode="construction",
+ cores=args.build_cores, cpus=args.build_cpus, repetition=repetition,
+ numa="interleave:" + args.build_nodes))
+ for repetition in range(1, args.query_repetitions + 1):
+ algorithms = ("pq", "ash") if (group + repetition) % 2 else ("ash", "pq")
+ for cores, cpus in ((1, args.one_core_cpu), (8, args.eight_core_cpus)):
+ for algorithm in algorithms:
+ mode = "queued" if algorithm == "ash" else "direct"
+ jobs.append(dict(prefix=prefix, bits=bits, algorithm=algorithm, mode=mode,
+ cores=cores, cpus=cpus, repetition=repetition,
+ numa="bind:" + args.query_node))
+jobs = [job for job in jobs if job["algorithm"] in args.algorithms
+ and (job["mode"] == "construction" or job["mode"] in args.query_modes)]
+for job in jobs:
+ job["id"] = "{prefix}-{bits}bit-{algorithm}-{mode}-c{cores}-r{repetition}".format(**job)
+ case_repetition = job["repetition"] if job["mode"] == "construction" else 1
+ job["case"] = str(output / "cases" / f"{job['prefix']}-{job['bits']}bit-{job['algorithm']}-r{case_repetition}")
+ job["config"] = f"{job['prefix']}-{job['bits']}bit-{job['algorithm']}.yml"
+watchdog_minutes = {"construction": 180 if args.dataset_set == "large" else 30, "query": 30}
+plan = dict(source_ref=source_ref, jar_sha256=jar_sha, data=str(data), dataset_set=args.dataset_set, watchdog_minutes=watchdog_minutes, jobs=jobs,
+ construction_jobs=sum(j["mode"] == "construction" for j in jobs),
+ query_jobs=sum(j["mode"] != "construction" for j in jobs))
+if args.plan_only:
+ print(json.dumps(plan, indent=2))
+ raise SystemExit(0)
+output.mkdir(parents=True, exist_ok=True)
+
+
+def save(path, value):
+ temp = path.with_suffix(path.suffix + ".tmp")
+ temp.write_text(json.dumps(value, indent=2) + "\n")
+ temp.replace(path)
+
+
+if (output / "plan.json").exists():
+ if json.loads((output / "plan.json").read_text()) != plan:
+ raise RuntimeError("Existing campaign has different source/artifact/settings; use a new output directory")
+else:
+ save(output / "plan.json", plan)
+results_file = output / "results.json"
+results = json.loads(results_file.read_text()) if results_file.exists() else []
+complete_ids = {result["job"]["id"] for result in results}
+process = None
+
+
+def stop(signum, frame):
+ if process is not None and process.poll() is None:
+ process.send_signal(signal.SIGTERM) # The checked-in shell forwards it to its JVM.
+ try:
+ process.wait(timeout=10)
+ except subprocess.TimeoutExpired:
+ process.kill()
+ save(output / "status.json", dict(state="stopped", completed=len(results), planned=len(jobs)))
+ raise SystemExit(128 + signum)
+
+
+signal.signal(signal.SIGINT, stop)
+signal.signal(signal.SIGTERM, stop)
+
+
+def report_failure(kind, error, traceback):
+ if process is not None and process.poll() is None:
+ process.terminate()
+ try:
+ process.wait(timeout=10)
+ except subprocess.TimeoutExpired:
+ process.kill()
+ save(output / "status.json", dict(state="failed", error=str(error),
+ current=globals().get("job"), completed=len(results), planned=len(jobs)))
+ sys.__excepthook__(kind, error, traceback)
+
+
+sys.excepthook = report_failure
+
+
+def cpu_set(specification):
+ result = set()
+ for item in specification.split(","):
+ limits = item.split("-")
+ first, last = int(limits[0]), int(limits[-1])
+ result.update(range(first, last + 1))
+ return result
+
+
+def wait_idle():
+ """Read-only check between JVMs; never gates compilation inside a benchmark."""
+ while True:
+ conflicts = []
+ for directory in Path("/proc").iterdir():
+ if not directory.name.isdigit() or int(directory.name) == os.getpid():
+ continue
+ try:
+ command = (directory / "cmdline").read_bytes().replace(b"\0", b" ").decode(errors="replace")
+ if "/JetBrains/" not in command and re.search(
+ r"BenchYAML|DistancesASH|DistancesPQ|org\.junit\.runner|surefirebooter|(?:^|/)cc1plus(?: |$)", command):
+ conflicts.append(int(directory.name))
+ except (FileNotFoundError, PermissionError, ProcessLookupError):
+ pass
+ if not conflicts:
+ return
+ save(output / "status.json", dict(state="waiting_for_other_work", pids=conflicts,
+ completed=len(results), planned=len(jobs)))
+ time.sleep(5)
+
+
+def summarize():
+ groups = collections.defaultdict(list)
+ for result in results:
+ job = result["job"]
+ for row in result["rows"]:
+ if job["mode"] == "construction" and float(row["overquery"]) != 1:
+ continue # Two overquery rows describe one construction invocation.
+ key = (job["prefix"], job["bits"], job["algorithm"], job["mode"], job["cores"], row["overquery"])
+ groups[key].append(row)
+ records = []
+ for key, rows in sorted(groups.items()):
+ record = dict(zip(("dataset", "bits", "algorithm", "mode", "cores", "overquery"), key))
+ record["jvms"] = len(rows)
+ metrics = {}
+ for name in rows[0]:
+ if not name.startswith(("construction.", "search.")):
+ continue
+ if not all(row.get(name) for row in rows):
+ continue
+ values = [float(row[name]) for row in rows]
+ if any(marker in name for marker in ("p999", "p95", "p99")):
+ metrics[name] = dict(median_of_jvm_values=statistics.median(values),
+ minimum=min(values), maximum=max(values))
+ else:
+ metrics[name] = dict(mean=statistics.mean(values),
+ sample_sd=statistics.stdev(values) if len(values) > 1 else None,
+ minimum=min(values), maximum=max(values))
+ record["metrics"] = metrics
+ records.append(record)
+ save(output / "summary.json", records)
+
+
+for job in jobs:
+ if job["id"] in complete_ids:
+ continue
+ wait_idle()
+ assert subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=repo, text=True).strip() == source_ref, "Source changed during the campaign"
+ save(output / "status.json", dict(state="running", current=job,
+ completed=len(results), planned=len(jobs)))
+ case = Path(job["case"])
+ before = set((case / "runs").glob("*")) if (case / "runs").exists() else set()
+ argv = ["bash", str(helper), job["config"], job["mode"], str(job["cores"]), job["cpus"], str(case)]
+ env = os.environ.copy()
+ env.update(JVECTOR_BENCH_DATA_DIR=str(data), JVECTOR_BENCH_NUMA_POLICY=job["numa"],
+ JVECTOR_BENCH_CATALOG="catalog-large.yml" if args.dataset_set == "large" else "catalog-1m.yml")
+ log_path = output / (job["id"] + ".log")
+ print("START " + job["id"], flush=True)
+ libraries = set()
+ run_directory = None
+ with log_path.open("w") as log:
+ process = subprocess.Popen(argv, cwd=repo, env=env, stdout=log, stderr=subprocess.STDOUT)
+ started = time.monotonic()
+ while process.poll() is None:
+ limit_minutes = watchdog_minutes["construction" if job["mode"] == "construction" else "query"]
+ if time.monotonic() - started > limit_minutes * 60:
+ process.terminate()
+ process.wait(timeout=15)
+ raise RuntimeError(f"BenchYAML exceeded the {limit_minutes}-minute process watchdog; inspect {log_path}")
+ if run_directory is None:
+ created = set((case / "runs").glob("*")) - before
+ if len(created) == 1:
+ run_directory = created.pop()
+ if run_directory is not None and (run_directory / "jvm.pid").exists():
+ pid = (run_directory / "jvm.pid").read_text().strip()
+ try:
+ proc = Path("/proc") / pid
+ status = (proc / "status").read_text()
+ if "\nName:\tjava\n" in "\n" + status:
+ mask = next(line.split(":", 1)[1].strip() for line in status.splitlines()
+ if line.startswith("Cpus_allowed_list:"))
+ assert cpu_set(mask) == cpu_set(job["cpus"]), (mask, job["cpus"])
+ libraries.update(line.split(maxsplit=5)[-1] for line in (proc / "maps").read_text().splitlines()
+ if any(token in line for token in ("libjvector", "libdnnl", "libmkl", "libopenblas", "libnetlib")))
+ with (run_directory / "placement.jsonl").open("a") as placement:
+ placement.write(json.dumps(dict(cpu_mask=mask, libraries=sorted(libraries))) + "\n")
+ except (FileNotFoundError, ProcessLookupError):
+ pass
+ time.sleep(1)
+ if process.returncode:
+ save(output / "status.json", dict(state="failed", job=job, log=str(log_path), exit_code=process.returncode))
+ raise RuntimeError("Standalone BenchYAML failed: " + str(log_path))
+ created = set((case / "runs").glob("*")) - before
+ assert len(created) == 1
+ run_directory = created.pop()
+ assert (run_directory / "source-ref.txt").read_text().strip() == source_ref
+ assert (run_directory / "jar.sha256").read_text().split()[0] == jar_sha, "Artifact changed during the campaign"
+ # A new invocation contributes exactly one CSV, identified by its run mode
+ # and the source reference recorded by the normal reporting system.
+ csvs = sorted((case / "logging").rglob("experiments.csv"), key=lambda path: path.stat().st_mtime_ns)
+ assert csvs
+ rows = list(csv.DictReader(csvs[-1].open()))
+ assert len(rows) == 2 and {float(row["overquery"]) for row in rows} == {1, 2}
+ assert all("FUSED_" + job["algorithm"].upper() in row["feature_set"] and "NVQ_VECTORS" in row["feature_set"] for row in rows)
+ assert any("libjvector-" in library for library in libraries), "Native provider was not observed"
+ if job["mode"] == "construction":
+ assert all(row.get("construction.index_build_time_s") for row in rows)
+ else:
+ assert all(row["query_execution"] == job["mode"] for row in rows)
+ assert all(row.get("search.throughput.avg_qps") and row.get("search.latency.mean_ms") for row in rows)
+ reference = next(result for result in results if result["job"]["prefix"] == job["prefix"]
+ and result["job"]["bits"] == job["bits"] and result["job"]["algorithm"] == job["algorithm"]
+ and result["job"]["mode"] == "construction" and result["job"]["repetition"] == 1)
+ actual_recall = [float(row["search.accuracy.recall_at_10"]) for row in rows]
+ reference_recall = [float(row["search.accuracy.recall_at_10"]) for row in reference["rows"]]
+ if job["algorithm"] == "ash" and job["mode"] == "queued":
+ # Batch GEMM may differ from singleton GEMV by FP32 reduction rounding.
+ # Keep the actual values; an improvement is not a regression. Stop on
+ # any decline for investigation rather than silently applying a tolerance.
+ assert all(actual >= expected for actual, expected in zip(actual_recall, reference_recall)), "Execution-policy recall regression"
+ else:
+ assert actual_recall == reference_recall, "Execution-policy recall regression"
+ result = dict(job=job, rows=rows, csv=str(csvs[-1]), artifacts=str(run_directory), libraries=sorted(libraries))
+ results.append(result)
+ complete_ids.add(job["id"])
+ save(results_file, results)
+ summarize()
+ if job["mode"] == "construction" and job["repetition"] != 1:
+ # Only these extra graphs have no future use. Keep first-build indexes,
+ # raw CSVs, metadata and console logs; remove no unrelated workspace.
+ assert case.parent == output / "cases"
+ for name in ("index_cache", "compressor_cache"):
+ shutil.rmtree(case / name)
+ print("PASS " + job["id"], flush=True)
+save(output / "complete.json", dict(source_ref=source_ref, jar_sha256=jar_sha,
+ construction_jobs=plan["construction_jobs"], query_jobs=plan["query_jobs"], completed=len(results)))
+save(output / "status.json", dict(state="complete", completed=len(results), planned=len(jobs)))
+print("COMPLETE", flush=True)
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodBenchmark.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodBenchmark.java
new file mode 100644
index 000000000..b359e4964
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodBenchmark.java
@@ -0,0 +1,199 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.example;
+
+import io.github.jbellis.jvector.quantization.ASHNeighborhoodScorer;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Locale;
+import java.util.Objects;
+import java.util.Random;
+
+/**
+ * DistancesASH entry point for isolated full scans through the production neighborhood scorer.
+ * Fixtures, the independent oracle, and the pipelined workload have separate owners. The timed
+ * paths include exactly one preparation per complete scan; input loading and graph I/O are untimed.
+ */
+final class ASHNeighborhoodBenchmark {
+ private static volatile double sink;
+
+ private ASHNeighborhoodBenchmark() {}
+
+ static void run(String[] args) throws IOException {
+ String filename = System.getProperty("jvector.bench.neighborhoodFixture");
+ if (filename == null)
+ throw new IllegalArgumentException("Set -Djvector.bench.neighborhoodFixture=/path/to/fixture.ashn");
+ Path path = Path.of(filename);
+ ASHNeighborhoodFixture fixture;
+ if (Files.exists(path))
+ fixture = ASHNeighborhoodFixture.read(path);
+ else {
+ if (args.length != 2 && args.length != 3)
+ throw new IllegalArgumentException(
+ "Creating a fixture requires base.fvecs query.fvecs [ground-truth.ivecs]");
+ fixture = ASHNeighborhoodFixture.create(args[0], args[1]);
+ ASHNeighborhoodFixture.write(path, fixture);
+ }
+ checkSetting("jvector.ash.bitsPerDimension", fixture.ash.bitsPerDimension);
+ checkSetting("jvector.ash.quantizedDimensions", fixture.ash.quantizedDim);
+ checkSetting("jvector.ash.landmarkCount", 1);
+ checkSetting("jvector.bench.neighborhoodSize", fixture.capacity);
+ checkSetting("jvector.bench.blockSize", fixture.blockSize);
+ boolean scaling = Boolean.getBoolean("jvector.bench.neighborhood-scaling");
+ int threads = positive("jvector.bench.scoringThreads", 1);
+ if (!scaling && !Boolean.getBoolean("jvector.bench.recall") && threads != 1)
+ throw new IllegalArgumentException("Multiple scoring threads require neighborhood-scaling=true");
+ int queryIndex = Integer.parseInt(System.getProperty("jvector.bench.queryIndex", "0"));
+ Objects.checkIndex(queryIndex, fixture.queries.size());
+ int passes = positive("jvector.bench.scoringPasses", 5);
+ if (System.getProperty("jvector.bench.neighborhoodsPerQuery") != null) {
+ throw new IllegalArgumentException(
+ "A full scan scores every saved neighborhood; remove neighborhoodsPerQuery");
+ }
+ double seconds = seconds("jvector.bench.scoringMinSeconds", .5);
+ double warmup = seconds("jvector.bench.scoringWarmupSeconds", 1);
+ String access = System.getProperty("jvector.bench.neighborhoodAccess", "sequential");
+ int[] order = accessOrder(fixture.neighborhoods.length, access);
+ VectorFloat query = fixture.queries.get(queryIndex);
+ var backend = VectorizationProvider.getInstance().getVectorUtilSupport();
+ // A requested native comparison must fail clearly rather than time the Java fallback.
+ if ("native".equals(System.getProperty("jvector.vectorization_provider", "").trim().toLowerCase(Locale.ROOT))
+ && !"scalar".equals(
+ System.getProperty("jvector.ash.blockKernel", "auto").trim().toLowerCase(Locale.ROOT))
+ && !backend.supportsAshNeighborhoodScoring()) {
+ throw new IllegalStateException(
+ "Native ASH comparison requested, but the library lacks the neighborhood kernel");
+ }
+ System.out.printf(Locale.ROOT,
+ "Neighborhood benchmark: java=%s, C=1, singleQuery=%s, bits=%d, dimensions=%d, capacity=%d, blockSize=%d%n",
+ Runtime.version(), !Boolean.getBoolean("jvector.bench.recall"), fixture.ash.bitsPerDimension,
+ fixture.ash.quantizedDim, fixture.capacity, fixture.blockSize);
+ System.out.printf(Locale.ROOT,
+ "Fixture: sha256=%s, queries=%d, queryIndex=%d, baseVectors=%d, neighborhoods=%d, packedBytes=%d, access=%s%n",
+ ASHNeighborhoodFixture.sha256(path), fixture.queries.size(), queryIndex, fixture.baseCount,
+ order.length, (long) order.length * fixture.neighborhoods[0].length, access);
+ System.out.printf("Backend: %s; %s%n", backend.getClass().getSimpleName(), backend.ashKernelDescription());
+ var scorer = new ASHNeighborhoodScorer(fixture.ash, query, fixture.capacity, fixture.blockSize);
+ System.out.println("Scorer: " + scorer);
+ if (Boolean.getBoolean("jvector.bench.recall")) {
+ ASHNeighborhoodRecall.run(fixture, threads);
+ return;
+ }
+ ASHNeighborhoodValidator.validate(fixture);
+ if (scaling) {
+ runPipelinedScans(fixture, query, order, threads, passes, warmup, seconds);
+ return;
+ }
+ System.out.println(
+ "Full scan: query preparation once per scan, all neighborhoods, headers, similarity conversion and checksum consumption are timed. Loading, packing, graph traversal and I/O are outside timing.");
+ float[] scores = new float[fixture.capacity];
+ Runnable fullScan = () -> {
+ var current = new ASHNeighborhoodScorer(fixture.ash, query, fixture.capacity, fixture.blockSize);
+ double checksum = 0;
+ for (int n : order) {
+ current.score(fixture.neighborhoods[n], scores);
+ for (float score : scores) checksum += score;
+ }
+ sink = checksum;
+ };
+ measure(fullScan, warmup);
+ System.out.println(
+ "mode,pass,ns_per_scan,ns_per_neighborhood,base_vectors_per_second,scored_lanes_per_second,checksum");
+ for (int pass = 1; pass <= passes; pass++) {
+ double ns = measure(fullScan, seconds);
+ System.out.printf(Locale.ROOT, "full_scan,%d,%.3f,%.3f,%.3f,%.3f,%.9f%n", pass, ns, ns / order.length,
+ fixture.baseCount * 1e9 / ns, order.length * (double) fixture.capacity * 1e9 / ns, sink);
+ }
+ }
+
+ private static void runPipelinedScans(ASHNeighborhoodFixture fixture, VectorFloat query, int[] order,
+ int threads, int passes, double warmup, double seconds) {
+ System.out.printf(Locale.ROOT,
+ "Strong scaling: workers=%d, fixedNeighborhoodsPerSweep=%d, fixedScoredLanesPerSweep=%d; one shared query, disjoint neighborhood partitions%n",
+ threads, order.length, (long) order.length * fixture.capacity);
+ System.out.println(
+ "Pipelined full scans: prepare each query once while workers score preceding scans. At most two scans are in flight. Preparation, submission, completion, checksum consumption, fill and drain are timed. Persistent-worker creation is outside timing; all threads share the stated CPU affinity.");
+ try (var workload = new ASHNeighborhoodPipeline(fixture, order, threads)) {
+ workload.measure(query, warmup);
+ System.out.println(
+ "workers,pass,ns_per_scan,ns_per_neighborhood,base_vectors_per_second,scored_lanes_per_second,ns_preparation_per_scan,checksum");
+ for (int pass = 1; pass <= passes; pass++) {
+ ASHNeighborhoodPipeline.Timing timing = workload.measure(query, seconds);
+ double nanos = timing.nanosPerScan;
+ System.out.printf(Locale.ROOT, "%d,%d,%.3f,%.3f,%.3f,%.3f,%.3f,%.9f%n", threads, pass, nanos,
+ nanos / order.length, fixture.baseCount * 1e9 / nanos,
+ order.length * (double) fixture.capacity * 1e9 / nanos, timing.preparationNanosPerScan,
+ timing.checksum);
+ }
+ }
+ }
+
+ private static double measure(Runnable task, double seconds) {
+ long start = System.nanoTime(), elapsed;
+ long count = 0;
+ do {
+ task.run();
+ count++;
+ elapsed = System.nanoTime() - start;
+ } while (elapsed < seconds * 1e9);
+ return (double) elapsed / count;
+ }
+
+ static int[] accessOrder(int count, String access) {
+ if (!access.equals("sequential") && !access.equals("shuffled"))
+ throw new IllegalArgumentException("neighborhoodAccess must be sequential or shuffled");
+ int[] order = new int[count];
+ for (int i = 0; i < count; i++) order[i] = i;
+ if (access.equals("shuffled")) {
+ Random random = new Random(87342);
+ for (int i = count - 1; i > 0; i--) {
+ int j = random.nextInt(i + 1), swap = order[i];
+ order[i] = order[j];
+ order[j] = swap;
+ }
+ }
+ return order;
+ }
+
+ static int positive(String name, int fallback) {
+ int value;
+ try {
+ value = Integer.parseInt(System.getProperty(name, Integer.toString(fallback)));
+ } catch (NumberFormatException error) {
+ throw new IllegalArgumentException(name + " must be an integer", error);
+ }
+ if (value <= 0)
+ throw new IllegalArgumentException(name + " must be positive");
+ return value;
+ }
+
+ private static double seconds(String name, double fallback) {
+ double value = Double.parseDouble(System.getProperty(name, Double.toString(fallback)));
+ if (!Double.isFinite(value) || value <= 0)
+ throw new IllegalArgumentException(name + " must be finite and positive");
+ return value;
+ }
+
+ static void checkSetting(String name, int expected) {
+ String value = System.getProperty(name);
+ if (value != null && Integer.parseInt(value) != expected)
+ throw new IllegalArgumentException(name + " disagrees with saved fixture: " + expected);
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodFixture.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodFixture.java
new file mode 100644
index 000000000..2198640b0
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodFixture.java
@@ -0,0 +1,312 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.example;
+
+import io.github.jbellis.jvector.disk.ByteBufferIndexWriter;
+import io.github.jbellis.jvector.disk.ByteBufferReader;
+import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues;
+import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
+import io.github.jbellis.jvector.quantization.ASHVectors;
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
+import io.github.jbellis.jvector.quantization.FusedASHLayout;
+import io.github.jbellis.jvector.vector.VectorUtil;
+import io.github.jbellis.jvector.vector.VectorizationProvider;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import java.io.BufferedInputStream;
+import java.io.BufferedOutputStream;
+import java.io.DataInputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.BufferUnderflowException;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.security.DigestInputStream;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * Untimed fixture creation and replay for the production ASH neighborhood layout.
+ *
+ * Version 1 stores big-endian ints/floats: magic, version, capacity, block width, base count,
+ * model byte count and model, query count and original-dimensional queries, then neighborhood
+ * count, bytes per neighborhood, and packed payloads. fvecs input remains little-endian.
+ * Model metadata and lengths are checked before allocating dimensions or payload arrays.
+ *
+ * Fixture bytes and query vectors are owned by the fixture and treated as immutable during
+ * replay. Queries are normalized on creation, matching the inputs passed to the graph scorer.
+ * The final neighborhood is padded to capacity and remains part of the scored work.
+ */
+final class ASHNeighborhoodFixture {
+ private static final int MAGIC = 0x4153484e;
+ private static final int VERSION = 1;
+ private static final int MAX_FVECS_DIMENSION = 100_000;
+ private static final int FIXED_BYTES = 9 * Integer.BYTES;
+
+ final AsymmetricHashing ash;
+ final List> queries;
+ final byte[][] neighborhoods;
+ final int capacity;
+ final int blockSize;
+ final int baseCount;
+
+ ASHNeighborhoodFixture(AsymmetricHashing ash, List> queries, byte[][] neighborhoods, int capacity,
+ int blockSize, int baseCount) {
+ this.ash = Objects.requireNonNull(ash, "ash");
+ if (ash.landmarkCount != 1) {
+ throw new IllegalArgumentException("Neighborhood benchmark requires C=1");
+ }
+ FusedASHLayout.validateBitsPerDimension(ash.bitsPerDimension);
+ FusedASHLayout.validateBlockSize(blockSize);
+ if (capacity < 1 || baseCount < 1 || queries.isEmpty()) {
+ throw new IllegalArgumentException("Empty fixture");
+ }
+ if (neighborhoods.length != (baseCount - 1) / capacity + 1) {
+ throw new IllegalArgumentException("Neighborhood count mismatch");
+ }
+ int bytes = FusedASHLayout.featureSize(capacity, ash.quantizedDim, ash.bitsPerDimension, blockSize);
+ for (byte[] neighborhood : neighborhoods) {
+ if (neighborhood.length != bytes) {
+ throw new IllegalArgumentException("Payload size mismatch");
+ }
+ }
+ for (VectorFloat query : queries) {
+ if (query.length() != ash.originalDimension) {
+ throw new IllegalArgumentException("Query dimension mismatch");
+ }
+ }
+ this.queries = List.copyOf(queries);
+ this.neighborhoods = neighborhoods;
+ this.capacity = capacity;
+ this.blockSize = blockSize;
+ this.baseCount = baseCount;
+ }
+
+ static ASHNeighborhoodFixture create(String baseFile, String queryFile) throws IOException {
+ int capacity = ASHNeighborhoodBenchmark.positive("jvector.bench.neighborhoodSize", 32);
+ int blockSize =
+ ASHNeighborhoodBenchmark.positive("jvector.bench.blockSize", FusedASHLayout.chooseBlockSize(capacity));
+ FusedASHLayout.validateBlockSize(blockSize);
+ ASHNeighborhoodBenchmark.checkSetting("jvector.ash.landmarkCount", 1);
+ List> base =
+ readFvecs(baseFile, ASHNeighborhoodBenchmark.positive("jvector.bench.maxVectors", 8192));
+ List> queries =
+ readFvecs(queryFile, ASHNeighborhoodBenchmark.positive("jvector.bench.maxQueries", 8));
+ if (base.isEmpty() || queries.isEmpty()) {
+ throw new IllegalArgumentException("Empty vector input");
+ }
+ int dimensions = base.get(0).length();
+ for (VectorFloat query : queries) {
+ if (query.length() != dimensions) {
+ throw new IllegalArgumentException("Base/query dimensions differ");
+ }
+ VectorUtil.l2normalize(query);
+ }
+ int bits = ASHNeighborhoodBenchmark.positive("jvector.ash.bitsPerDimension", 4);
+ FusedASHLayout.validateBitsPerDimension(bits);
+ int projected = ASHNeighborhoodBenchmark.positive("jvector.ash.quantizedDimensions", dimensions / 2);
+ if (projected > dimensions) {
+ throw new IllegalArgumentException("Projected dimensions exceed original dimensions");
+ }
+ var values = new ListRandomAccessVectorValues(base, dimensions);
+ var ash = AsymmetricHashing.initialize(values, DistancesASH.parseOptimizerFromProperty(),
+ Math.addExact(AsymmetricHashing.HEADER_BITS, Math.multiplyExact(projected, bits)), 1, bits);
+ ASHVectors encoded = (ASHVectors) ash.encodeAll(values);
+ int count = (base.size() - 1) / capacity + 1;
+ byte[][] packed = new byte[count][FusedASHLayout.featureSize(capacity, projected, bits, blockSize)];
+ for (int i = 0; i < base.size(); i++) {
+ int neighbor = i % capacity;
+ int offset = FusedASHLayout.blockOffset(neighbor / blockSize, projected, bits, blockSize);
+ FusedASHLayout.packForFused(
+ packed[i / capacity], offset, neighbor % blockSize, encoded.get(i), ash, blockSize);
+ }
+ return new ASHNeighborhoodFixture(ash, queries, packed, capacity, blockSize, base.size());
+ }
+
+ /** Read only the requested prefix, rejecting truncated, non-finite, or inconsistent vectors. */
+ static List> readFvecs(String filename, int limit) throws IOException {
+ if (limit < 1) {
+ throw new IllegalArgumentException("Vector limit must be positive");
+ }
+ List> result = new ArrayList<>();
+ var types = VectorizationProvider.getInstance().getVectorTypeSupport();
+ try (var input = new DataInputStream(new BufferedInputStream(Files.newInputStream(Path.of(filename))))) {
+ while (result.size() < limit) {
+ int first = input.read();
+ if (first == -1) {
+ break;
+ }
+ int dimensions = first | input.readUnsignedByte() << 8 | input.readUnsignedByte() << 16
+ | input.readUnsignedByte() << 24;
+ if (dimensions < 1 || dimensions > MAX_FVECS_DIMENSION
+ || (!result.isEmpty() && dimensions != result.get(0).length())) {
+ throw new IOException("Invalid fvecs dimension: " + dimensions);
+ }
+ float[] vector = new float[dimensions];
+ for (int i = 0; i < dimensions; i++) {
+ vector[i] = Float.intBitsToFloat(Integer.reverseBytes(input.readInt()));
+ if (!Float.isFinite(vector[i])) {
+ throw new IOException("Non-finite fvecs value at row " + result.size() + ", dimension " + i);
+ }
+ }
+ result.add(types.createFloatVector(vector));
+ }
+ }
+ return result;
+ }
+
+ /** Publish a complete fixture; a failed write or existing destination never leaves a partial file. */
+ static void write(Path path, ASHNeighborhoodFixture fixture) throws IOException {
+ Path absolute = path.toAbsolutePath();
+ Path temporary = Files.createTempFile(absolute.getParent(), ".ash-neighborhood-", ".tmp");
+ try {
+ writeComplete(temporary, fixture);
+ Files.move(temporary, absolute); // No REPLACE_EXISTING: preserve existing fixture inputs.
+ } finally {
+ Files.deleteIfExists(temporary);
+ }
+ }
+
+ private static void writeComplete(Path path, ASHNeighborhoodFixture fixture) throws IOException {
+ try (var model = ByteBufferIndexWriter.create(fixture.ash.compressorSize(), false)) {
+ fixture.ash.write(model, OnDiskGraphIndex.CURRENT_VERSION);
+ ByteBuffer buffer = model.getWrittenData();
+ byte[] bytes = new byte[buffer.remaining()];
+ buffer.get(bytes);
+ try (var out = new DataOutputStream(new BufferedOutputStream(Files.newOutputStream(path)))) {
+ out.writeInt(MAGIC);
+ out.writeInt(VERSION);
+ out.writeInt(fixture.capacity);
+ out.writeInt(fixture.blockSize);
+ out.writeInt(fixture.baseCount);
+ out.writeInt(bytes.length);
+ out.write(bytes);
+ out.writeInt(fixture.queries.size());
+ for (VectorFloat query : fixture.queries) {
+ for (int d = 0; d < query.length(); d++) {
+ out.writeFloat(query.get(d));
+ }
+ }
+ out.writeInt(fixture.neighborhoods.length);
+ out.writeInt(fixture.neighborhoods[0].length);
+ for (byte[] block : fixture.neighborhoods) {
+ out.write(block);
+ }
+ }
+ }
+ }
+
+ static ASHNeighborhoodFixture read(Path path) throws IOException {
+ long fileBytes = Files.size(path);
+ try (var input = new DataInputStream(new BufferedInputStream(Files.newInputStream(path)))) {
+ if (input.readInt() != MAGIC || input.readInt() != VERSION) {
+ throw new IOException("Invalid neighborhood fixture header");
+ }
+ int capacity = input.readInt();
+ int blockSize = input.readInt();
+ int baseCount = input.readInt();
+ int modelSize = input.readInt();
+ if (capacity < 1 || baseCount < 1 || modelSize < 1 || modelSize > fileBytes - FIXED_BYTES) {
+ throw new IOException("Invalid fixture sizes");
+ }
+ FusedASHLayout.validateBlockSize(blockSize);
+ byte[] model = new byte[modelSize];
+ input.readFully(model);
+ validateModelEnvelope(model);
+ var modelBuffer = ByteBuffer.wrap(model).order(ByteOrder.BIG_ENDIAN);
+ var ash = AsymmetricHashing.load(new ByteBufferReader(modelBuffer));
+ if (modelBuffer.hasRemaining()) {
+ throw new IOException("Trailing bytes in fixture model");
+ }
+ FusedASHLayout.validateBitsPerDimension(ash.bitsPerDimension);
+ int queries = input.readInt();
+ long queryBytes = (long) queries * ash.originalDimension * Float.BYTES;
+ long payloadStart = FIXED_BYTES + (long) modelSize + queryBytes;
+ if (queries < 1 || queryBytes < 0 || payloadStart > fileBytes) {
+ throw new IOException("Invalid query count");
+ }
+ var types = VectorizationProvider.getInstance().getVectorTypeSupport();
+ List> vectors = new ArrayList<>();
+ for (int i = 0; i < queries; i++) {
+ float[] query = new float[ash.originalDimension];
+ for (int d = 0; d < query.length; d++) {
+ query[d] = input.readFloat();
+ if (!Float.isFinite(query[d])) {
+ throw new IOException("Non-finite fixture query value");
+ }
+ }
+ vectors.add(types.createFloatVector(query));
+ }
+ int count = input.readInt();
+ int bytes = input.readInt();
+ int expected = FusedASHLayout.featureSize(capacity, ash.quantizedDim, ash.bitsPerDimension, blockSize);
+ if (count != (baseCount - 1) / capacity + 1 || bytes != expected
+ || payloadStart + (long) count * bytes != fileBytes) {
+ throw new IOException("Invalid packed payload sizes or trailing bytes");
+ }
+ byte[][] blocks = new byte[count][bytes];
+ for (byte[] block : blocks) {
+ input.readFully(block);
+ }
+ return new ASHNeighborhoodFixture(ash, vectors, blocks, capacity, blockSize, baseCount);
+ } catch (IllegalArgumentException | ArithmeticException | BufferUnderflowException error) {
+ throw new IOException("Invalid ASH neighborhood fixture: " + path, error);
+ }
+ }
+
+ private static void validateModelEnvelope(byte[] model) throws IOException {
+ // Fixture V1 embeds ASH's eight-int prefix, one landmark, its row-major transform, and the largest
+ // encoded vector length. Check declared dimensions against the bytes before ASH.load allocates model arrays.
+ int fixedModelBytes = 11 * Integer.BYTES + Float.BYTES;
+ if (model.length < fixedModelBytes) {
+ throw new IOException("Truncated fixture model");
+ }
+ ByteBuffer header = ByteBuffer.wrap(model).order(ByteOrder.BIG_ENDIAN);
+ int original = header.getInt(2 * Integer.BYTES);
+ int projected = header.getInt(4 * Integer.BYTES);
+ int landmarks = header.getInt(7 * Integer.BYTES);
+ long floats = (long) original * (projected + 1L);
+ if (original < 1 || projected < 1 || projected > original || landmarks != 1
+ || floats > (model.length - fixedModelBytes) / Float.BYTES
+ || fixedModelBytes + floats * Float.BYTES != model.length) {
+ throw new IOException("Invalid C=1 fixture model dimensions");
+ }
+ }
+
+ static String sha256(Path path) throws IOException {
+ try {
+ MessageDigest digest = MessageDigest.getInstance("SHA-256");
+ try (var input = new DigestInputStream(Files.newInputStream(path), digest)) {
+ input.transferTo(OutputStream.nullOutputStream());
+ }
+ byte[] hash = digest.digest();
+ StringBuilder hex = new StringBuilder(hash.length * 2);
+ for (byte value : hash) {
+ hex.append(Character.forDigit((value >>> 4) & 0xF, 16));
+ hex.append(Character.forDigit(value & 0xF, 16));
+ }
+ return hex.toString();
+ } catch (NoSuchAlgorithmException error) {
+ throw new AssertionError("SHA-256 is required by the Java platform", error);
+ }
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodPipeline.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodPipeline.java
new file mode 100644
index 000000000..05bd1f132
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodPipeline.java
@@ -0,0 +1,208 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.example;
+
+import io.github.jbellis.jvector.quantization.ASHNeighborhoodScorer;
+import io.github.jbellis.jvector.vector.types.VectorFloat;
+import java.util.ArrayDeque;
+import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Bounded benchmark pipeline: one coordinator prepares queries while a persistent pool scores.
+ *
+ * Submit, measure and close belong to the coordinator thread. At most two scans are submitted
+ * by {@link #measure}; every scan owns its prepared query and every pool thread owns its scratch.
+ * Timing includes preparation, scheduling, score consumption, pipeline fill and complete drain.
+ * Average elapsed time per completed scan is a throughput metric, not per-query latency.
+ */
+final class ASHNeighborhoodPipeline implements AutoCloseable {
+ private static final int MAX_IN_FLIGHT = 2;
+ private static volatile double sink;
+ private final ASHNeighborhoodFixture fixture;
+ private final int[] order;
+ private final int threads;
+ private final ThreadPoolExecutor pool;
+ private final ThreadLocal output;
+ private long preparationNanos;
+
+ ASHNeighborhoodPipeline(ASHNeighborhoodFixture fixture, int[] order, int threads) {
+ if (threads < 1 || threads > order.length)
+ throw new IllegalArgumentException("Worker count must be in [1, neighborhood count]");
+ this.fixture = Objects.requireNonNull(fixture, "fixture");
+ validateOrder(order, fixture.neighborhoods.length);
+ this.order = order.clone();
+ this.threads = threads;
+ this.output = ThreadLocal.withInitial(() -> new float[fixture.capacity]);
+ var workerIds = new AtomicInteger();
+ pool = (ThreadPoolExecutor) Executors.newFixedThreadPool(threads, task -> {
+ var worker = new Thread(task, "ash-neighborhood-" + workerIds.incrementAndGet());
+ worker.setDaemon(true); // An aborted benchmark must not leave its process alive.
+ return worker;
+ });
+ pool.prestartAllCoreThreads();
+ }
+
+ /** Prepare exactly once, then publish the immutable scorer with each partition task. */
+ Scan submit(VectorFloat query) {
+ if (pool.isShutdown()) {
+ throw new IllegalStateException("ASH pipeline is closed");
+ }
+ long start = System.nanoTime();
+ var scorer = new ASHNeighborhoodScorer(fixture.ash, query, fixture.capacity, fixture.blockSize);
+ preparationNanos += System.nanoTime() - start;
+ var scan = new Scan(threads);
+ for (int worker = 0; worker < threads; worker++) {
+ final int slot = worker;
+ final int first = (int) ((long) worker * order.length / threads);
+ final int last = (int) ((long) (worker + 1) * order.length / threads);
+ var partition = new FutureTask(() -> {
+ float[] scores = output.get();
+ double checksum = 0;
+ for (int n = first; n < last; n++) {
+ scorer.score(fixture.neighborhoods[order[n]], scores);
+ for (float score : scores)
+ checksum += score;
+ }
+ scan.checksums[slot] = checksum;
+ return null;
+ }) {
+ @Override
+ protected void done() {
+ // Also complete the scan when shutdown cancels a queued partition.
+ try {
+ get(); // Already terminal: no waiting here.
+ } catch (ExecutionException error) {
+ scan.failure.compareAndSet(null, error.getCause());
+ } catch (CancellationException error) {
+ scan.failure.compareAndSet(null, error);
+ } catch (InterruptedException error) {
+ Thread.currentThread().interrupt();
+ scan.failure.compareAndSet(null, error);
+ } finally {
+ scan.done.countDown();
+ }
+ }
+ };
+ pool.execute(partition);
+ }
+ return scan;
+ }
+
+ @Override
+ public void close() {
+ pool.shutdown();
+ boolean interrupted = false;
+ try {
+ while (!pool.isTerminated()) {
+ try {
+ pool.awaitTermination(1, TimeUnit.SECONDS);
+ } catch (InterruptedException error) {
+ interrupted = true;
+ for (Runnable queued : pool.shutdownNow()) {
+ ((Future) queued).cancel(false);
+ }
+ }
+ }
+ } finally {
+ if (interrupted)
+ Thread.currentThread().interrupt();
+ }
+ }
+
+ private static void validateOrder(int[] order, int count) {
+ if (order.length != count) {
+ throw new IllegalArgumentException("A full scan must include every neighborhood exactly once");
+ }
+ boolean[] seen = new boolean[count];
+ for (int index : order) {
+ if (index < 0 || index >= count || seen[index]) {
+ throw new IllegalArgumentException("Neighborhood order must be a permutation");
+ }
+ seen[index] = true;
+ }
+ }
+
+ /** Completion state for one full scan; the latch publishes each partition's checksum. */
+ static final class Scan {
+ private final CountDownLatch done;
+ private final double[] checksums;
+ private final AtomicReference failure = new AtomicReference<>();
+
+ Scan(int threads) {
+ done = new CountDownLatch(threads);
+ checksums = new double[threads];
+ }
+
+ double await() {
+ try {
+ done.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Interrupted while scoring", e);
+ }
+ if (failure.get() != null)
+ throw new IllegalStateException("Neighborhood scoring worker failed", failure.get());
+ double sum = 0;
+ for (double checksum : checksums)
+ sum += checksum;
+ return sum;
+ }
+ }
+
+ static final class Timing {
+ final double nanosPerScan, preparationNanosPerScan, checksum;
+ Timing(double nanos, double preparation, double checksum) {
+ this.nanosPerScan = nanos;
+ this.preparationNanosPerScan = preparation;
+ this.checksum = checksum;
+ }
+ }
+
+ Timing measure(VectorFloat query, double seconds) {
+ var pending = new ArrayDeque(MAX_IN_FLIGHT);
+ long preparationStart = preparationNanos;
+ long start = System.nanoTime();
+ long completed = 0;
+ double checksum = 0;
+ do {
+ if (pending.size() == MAX_IN_FLIGHT) {
+ checksum += pending.removeFirst().await();
+ completed++;
+ }
+ // Preparation of this scan overlaps scoring of earlier submitted scans.
+ pending.addLast(submit(query));
+ } while (System.nanoTime() - start < seconds * 1e9);
+ while (!pending.isEmpty()) {
+ checksum += pending.removeFirst().await();
+ completed++;
+ }
+ long elapsed = System.nanoTime() - start;
+ sink = checksum / completed;
+ return new Timing(
+ (double) elapsed / completed, (double) (preparationNanos - preparationStart) / completed, sink);
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodRecall.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodRecall.java
new file mode 100644
index 000000000..484bd3699
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodRecall.java
@@ -0,0 +1,254 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.example;
+
+import io.github.jbellis.jvector.quantization.ASHNeighborhoodScorer;
+import java.io.BufferedInputStream;
+import java.io.DataInputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Locale;
+import java.util.PriorityQueue;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * Exact full-scan recall@10 evaluation through the production neighborhood scorer.
+ *
+ * Each query is prepared once and scores every serialized neighborhood. Only real base
+ * ordinals enter top-k; padded lanes are scored but excluded from ranking. Ties prefer the
+ * smaller original ordinal in every backend. Query parallelism changes neither the scan work
+ * nor its reduction order. This is an accuracy check, not a kernel throughput benchmark.
+ */
+final class ASHNeighborhoodRecall {
+ static final int K = 10;
+
+ private ASHNeighborhoodRecall() {}
+
+ static void run(ASHNeighborhoodFixture fixture, int threads) throws IOException {
+ String groundTruthName = System.getProperty("jvector.bench.recall.groundTruth");
+ if (groundTruthName == null) {
+ throw new IllegalArgumentException("Set jvector.bench.recall.groundTruth to the matching .ivecs file");
+ }
+ int[][] truth = readGroundTruth(Path.of(groundTruthName), fixture.queries.size(), fixture.baseCount);
+ String outputName = System.getProperty("jvector.bench.recall.output");
+ if (outputName != null && Files.exists(Path.of(outputName))) {
+ throw new IOException("Recall output already exists: " + outputName);
+ }
+ // Keep independent component validation bounded; recall itself uses every saved query.
+ var validation =
+ new ASHNeighborhoodFixture(fixture.ash, fixture.queries.subList(0, Math.min(4, fixture.queries.size())),
+ fixture.neighborhoods, fixture.capacity, fixture.blockSize, fixture.baseCount);
+ ASHNeighborhoodValidator.validate(validation);
+ System.out.printf(Locale.ROOT,
+ "Full-scan recall@10: queries=%d baseVectors=%d neighborhoods=%d queryWorkers=%d; prepare once per query; no reranking; padding excluded from top-k%n",
+ fixture.queries.size(), fixture.baseCount, fixture.neighborhoods.length, threads);
+ var pool = Executors.newFixedThreadPool(threads);
+ List results = new ArrayList<>(fixture.queries.size());
+ try {
+ var finished = new AtomicInteger();
+ List> tasks = new ArrayList<>(fixture.queries.size());
+ for (int q = 0; q < fixture.queries.size(); q++) {
+ final int query = q;
+ tasks.add(() -> {
+ Result result = scan(fixture, query, truth[query]);
+ int count = finished.incrementAndGet();
+ if (count % 100 == 0 || count == fixture.queries.size()) {
+ System.out.printf(Locale.ROOT, "Recall progress: %d/%d queries complete%n", count,
+ fixture.queries.size());
+ }
+ return result;
+ });
+ }
+ int completed = 0;
+ for (var future : pool.invokeAll(tasks)) {
+ results.add(future.get());
+ completed++;
+ }
+ if (completed != fixture.queries.size()) {
+ throw new AssertionError("Incomplete full-scan recall");
+ }
+ } catch (InterruptedException error) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException("Interrupted during full-scan recall", error);
+ } catch (ExecutionException error) {
+ throw new IllegalStateException("Full-scan recall worker failed", error.getCause());
+ } finally {
+ pool.shutdownNow();
+ boolean interrupted = false;
+ try {
+ while (!pool.isTerminated()) {
+ try {
+ pool.awaitTermination(1, TimeUnit.SECONDS);
+ } catch (InterruptedException error) {
+ interrupted = true;
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+ long hits = 0;
+ for (var result : results) {
+ hits += result.hits;
+ }
+ if (outputName != null) {
+ try (var output = Files.newBufferedWriter(
+ Path.of(outputName), StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE)) {
+ output.write("query\thits_at_10\tids_best_first\tscores_best_first\n");
+ for (var result : results) {
+ output.write(result.query + "\t" + result.hits + "\t" + Arrays.toString(result.ids) + "\t"
+ + Arrays.toString(result.scores) + "\n");
+ }
+ }
+ }
+ System.out.printf(Locale.ROOT,
+ "NEIGHBORHOOD_RECALL queries=%d baseVectors=%d hits=%d denominator=%d recall_at_10=%.9f%n",
+ results.size(), fixture.baseCount, hits, (long) K * results.size(),
+ hits / ((double) K * results.size()));
+ }
+
+ /** Scan exactly once; a worker owns the scorer, scratch, and top-k state for its query. */
+ static Result scan(ASHNeighborhoodFixture fixture, int query, int[] truth) {
+ var scorer =
+ new ASHNeighborhoodScorer(fixture.ash, fixture.queries.get(query), fixture.capacity, fixture.blockSize);
+ float[] scores = new float[fixture.capacity];
+ var top = new TopK();
+ for (int n = 0; n < fixture.neighborhoods.length; n++) {
+ if (Thread.currentThread().isInterrupted()) {
+ throw new IllegalStateException("Recall scan cancelled");
+ }
+ scorer.score(fixture.neighborhoods[n], scores);
+ int base = n * fixture.capacity;
+ int count = Math.min(fixture.capacity, fixture.baseCount - base);
+ for (int lane = 0; lane < count; lane++) {
+ float score = scores[lane];
+ if (!Float.isFinite(score)) {
+ throw new IllegalStateException(
+ "Non-finite recall score at query=" + query + " ordinal=" + (base + lane));
+ }
+ top.offer(score, base + lane);
+ }
+ }
+ Candidate[] ranked = top.ranked();
+ int[] ids = new int[K];
+ float[] bestScores = new float[K];
+ int hits = 0;
+ for (int rank = 0; rank < K; rank++) {
+ ids[rank] = ranked[rank].ordinal;
+ bestScores[rank] = ranked[rank].score;
+ for (int expected : truth) {
+ if (ids[rank] == expected) {
+ hits++;
+ break;
+ }
+ }
+ }
+ return new Result(query, ids, bestScores, hits);
+ }
+
+ /** Read little-endian ivecs; top ten must be unique, in-range original base ordinals. */
+ static int[][] readGroundTruth(Path path, int queries, int baseCount) throws IOException {
+ if (queries < 1 || baseCount < K) {
+ throw new IllegalArgumentException("Recall@10 needs queries and at least ten real base vectors");
+ }
+ int[][] result = new int[queries][K];
+ try (var input = new DataInputStream(new BufferedInputStream(Files.newInputStream(path)))) {
+ for (int q = 0; q < queries; q++) {
+ int width = Integer.reverseBytes(input.readInt());
+ if (width < K || width > 100_000) {
+ throw new IOException("Invalid ground-truth width: " + width);
+ }
+ for (int rank = 0; rank < width; rank++) {
+ int id = Integer.reverseBytes(input.readInt());
+ if (id < 0 || id >= baseCount) {
+ throw new IOException("Ground-truth ordinal outside base set: " + id);
+ }
+ if (rank < K) {
+ for (int prior = 0; prior < rank; prior++) {
+ if (result[q][prior] == id) {
+ throw new IOException("Duplicate top-ten ground-truth ordinal at query " + q);
+ }
+ }
+ result[q][rank] = id;
+ }
+ }
+ }
+ }
+ return result;
+ }
+
+ static final class Result {
+ final int query;
+ final int[] ids;
+ final float[] scores;
+ final int hits;
+
+ Result(int query, int[] ids, float[] scores, int hits) {
+ this.query = query;
+ this.ids = ids;
+ this.scores = scores;
+ this.hits = hits;
+ }
+ }
+
+ static final class Candidate {
+ final float score;
+ final int ordinal;
+
+ Candidate(float score, int ordinal) {
+ this.score = score;
+ this.ordinal = ordinal;
+ }
+ }
+
+ /** Worst retained candidate first; equal scores evict the largest ordinal first. */
+ static final class TopK {
+ private static final Comparator WORST_FIRST = (a, b) -> {
+ int comparison = Float.compare(a.score, b.score);
+ return comparison != 0 ? comparison : Integer.compare(b.ordinal, a.ordinal);
+ };
+ private final PriorityQueue heap = new PriorityQueue<>(K, WORST_FIRST);
+
+ void offer(float score, int ordinal) {
+ Candidate worst = heap.peek();
+ if (heap.size() < K) {
+ heap.add(new Candidate(score, ordinal));
+ } else if (score > worst.score || (score == worst.score && ordinal < worst.ordinal)) {
+ heap.poll();
+ heap.add(new Candidate(score, ordinal));
+ }
+ }
+
+ Candidate[] ranked() {
+ Candidate[] ranked = heap.toArray(new Candidate[0]);
+ Arrays.sort(ranked, WORST_FIRST.reversed());
+ return ranked;
+ }
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodValidator.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodValidator.java
new file mode 100644
index 000000000..796c73e7d
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ASHNeighborhoodValidator.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.example;
+
+import io.github.jbellis.jvector.quantization.ASHNeighborhoodScorer;
+import io.github.jbellis.jvector.quantization.ASHScorer;
+import io.github.jbellis.jvector.quantization.FusedASHLayout;
+import java.util.Locale;
+
+/**
+ * Untimed decoded-component oracle, independent of the scoring LUT and block-scoring implementation.
+ * It uses the model's selected query precision and checks the serialized graph headers.
+ * Projection arithmetic is covered separately by the provider's FP32 and FP16 oracle tests.
+ * Offset-binary (3 and 5 to 9 bit) neighborhoods score an INT8-rounded query, which the oracle applies too.
+ * The tolerance permits FP32 reduction-order differences, without introducing LUT quantization.
+ */
+final class ASHNeighborhoodValidator {
+ private ASHNeighborhoodValidator() {}
+
+ static void validate(ASHNeighborhoodFixture fixture) {
+ double maximum = 0;
+ int neighborhoods = Math.min(fixture.neighborhoods.length,
+ ASHNeighborhoodBenchmark.positive("jvector.bench.validationNeighborhoods", 1024));
+ for (var query : fixture.queries) {
+ var scorer = new ASHNeighborhoodScorer(fixture.ash, query, fixture.capacity, fixture.blockSize);
+ // Share the query precision, then independently decode and accumulate
+ // components in double precision instead of using the nibble LUT kernel.
+ float[] input = new float[query.length()];
+ double dotMu = 0;
+ for (int d = 0; d < input.length; d++) {
+ input[d] = query.get(d);
+ dotMu += (double) input[d] * fixture.ash.landmarks[0].get(d);
+ }
+ float[] projected = new float[fixture.ash.quantizedDim];
+ fixture.ash.stiefelTransform.projectQuery(input, projected);
+ if (FusedASHLayout.usesFieldBlocks(fixture.ash.bitsPerDimension)) {
+ // Offset-binary neighborhoods score an INT8 query: round to 127 / max|Aq| steps, as the format does.
+ float max = 0f;
+ for (float v : projected) max = Math.max(max, Math.abs(v));
+ if (max > 0f) {
+ for (int d = 0; d < projected.length; d++) {
+ projected[d] = Math.max(-127, Math.min(127, Math.round(projected[d] * (127f / max)))) * (max / 127f);
+ }
+ }
+ }
+ float[] scores = new float[fixture.capacity];
+ for (int n = 0; n < neighborhoods; n++) {
+ int index = (int) ((long) n * fixture.neighborhoods.length / neighborhoods);
+ byte[] packed = fixture.neighborhoods[index];
+ scorer.score(packed, scores);
+ for (int i = 0; i < fixture.capacity; i++) {
+ int offset = FusedASHLayout.blockOffset(
+ i / fixture.blockSize, projected.length, fixture.ash.bitsPerDimension, fixture.blockSize);
+ int lane = i % fixture.blockSize, bits = fixture.ash.bitsPerDimension;
+ double dot = 0;
+ byte[] code = new byte[FusedASHLayout.canonicalCodeBytes(projected.length, bits)];
+ FusedASHLayout.readLaneCode(packed, offset, lane, projected.length, bits, fixture.blockSize, code);
+ for (int d = 0; d < projected.length; d++) {
+ // Field d is bits [d*b, (d+1)*b) of the lane's canonical code.
+ int field = 0;
+ for (int k = 0; k < bits; k++) {
+ int bit = d * bits + k;
+ field |= ((code[bit >>> 3] >>> (bit & 7)) & 1) << k;
+ }
+ dot += (double) projected[d] * (bits == 1 || bits == 2 || bits == 4
+ ? FusedASHLayout.decodeProjectionComponent(field, bits)
+ : field - ((1 << bits) - 1) / 2.0);
+ }
+ float expected = ASHScorer.toSimilarity((float) (FusedASHLayout.readScale(packed, offset, lane,
+ projected.length, bits, fixture.blockSize)
+ * dot
+ + dotMu
+ + FusedASHLayout.readOffset(
+ packed, offset, lane, projected.length, bits, fixture.blockSize)));
+ double error = Math.abs(scores[i] - expected);
+ if (!Float.isFinite(expected) || !Float.isFinite(scores[i])
+ || error > 3e-5 + Math.abs(expected) * 3e-5)
+ throw new AssertionError("Neighborhood score mismatch at " + index + "/" + i + ": " + scores[i]
+ + " vs " + expected);
+ maximum = Math.max(maximum, error);
+ }
+ }
+ }
+ System.out.printf(Locale.ROOT,
+ "Validated %d queries x %d neighborhoods against decoded-component oracle; maxAbsError=%.9g%n",
+ fixture.queries.size(), neighborhoods, maximum);
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java
index 2d48791bc..d81ad3bb1 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java
@@ -24,6 +24,7 @@
import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets;
import io.github.jbellis.jvector.example.yaml.DatasetCollection;
import io.github.jbellis.jvector.example.yaml.MultiConfig;
+import io.github.jbellis.jvector.example.yaml.SearchParameters;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -52,9 +53,15 @@
*/
public class AutoBenchYAML {
private static final Logger logger = LoggerFactory.getLogger(AutoBenchYAML.class);
- private static final String REGRESSION_TEST_KEY = "regression-tests";
+ private static final String REGRESSION_TEST_KEY = DatasetCollection.REGRESSION_TESTS;
public static void main(String[] args) throws IOException {
+ try (var projectionSetting = new SearchParameters.AshQueryProjectionSetting()) {
+ run(args, projectionSetting);
+ }
+ }
+
+ private static void run(String[] args, SearchParameters.AshQueryProjectionSetting projectionSetting) throws IOException {
// Check for --output argument (required for this class)
String outputPath = null;
for (int i = 0; i < args.length - 1; i++) {
@@ -114,32 +121,30 @@ public static void main(String[] args) throws IOException {
// Add results from checkpoint if present
results.addAll(checkpointManager.getCompletedResults());
- // Process datasets from regex patterns
- if (!datasetNames.isEmpty()) {
- for (var datasetName : datasetNames) {
- // Skip already completed datasets
- if (checkpointManager.isDatasetCompleted(datasetName)) {
- logger.info("Skipping already completed dataset: {}", datasetName);
- continue;
- }
+ // Resolve the process-wide ASH projection before any regression builds start.
+ var configs = new LinkedHashMap();
+ for (var datasetName : datasetNames) {
+ if (checkpointManager.isDatasetCompleted(datasetName)) {
+ logger.info("Skipping already completed dataset: {}", datasetName);
+ continue;
+ }
+ try {
+ var name = datasetName.endsWith(".hdf5")
+ ? datasetName.substring(0, datasetName.length() - ".hdf5".length()) : datasetName;
+ var config = finalConfigPath == null ? MultiConfig.getDefaultConfig(name) : MultiConfig.getConfig(finalConfigPath);
+ if (finalConfigPath == null || config.dataset == null || config.dataset.isEmpty()) config.dataset = name;
+ configs.put(datasetName, config);
+ } catch (Exception e) {
+ logger.error("Exception while processing dataset {}", datasetName, e);
+ }
+ }
+ projectionSetting.apply(configs.values().stream().filter(MultiConfig::usesSchema3).collect(Collectors.toList()));
+ if (!configs.isEmpty()) {
+ for (var entry : configs.entrySet()) {
+ var datasetName = entry.getKey();
+ var config = entry.getValue();
try {
- String normalizedDatasetName = datasetName;
- if (normalizedDatasetName.endsWith(".hdf5")) {
- normalizedDatasetName = normalizedDatasetName.substring(0, normalizedDatasetName.length() - ".hdf5".length());
- }
-
- MultiConfig config;
- if (finalConfigPath != null) {
- config = MultiConfig.getConfig(finalConfigPath);
- // Override dataset name if not specified in custom config
- if (config.dataset == null || config.dataset.isEmpty()) {
- config.dataset = normalizedDatasetName;
- }
- } else {
- config = MultiConfig.getDefaultConfig(normalizedDatasetName);
- config.dataset = normalizedDatasetName;
- }
logger.info("Using configuration: {}", config);
var loading = BaseVectorLoading.forBenchmark(config);
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java
index 8eb6ee09f..769787a70 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BenchYAML.java
@@ -23,6 +23,7 @@
import io.github.jbellis.jvector.example.yaml.DatasetCollection;
import io.github.jbellis.jvector.example.yaml.MultiConfig;
import io.github.jbellis.jvector.example.yaml.RunConfig;
+import io.github.jbellis.jvector.example.yaml.SearchParameters;
import java.io.IOException;
import java.security.InvalidParameterException;
@@ -36,6 +37,12 @@
public class BenchYAML {
public static void main(String[] args) throws IOException {
+ try (var projectionSetting = new SearchParameters.AshQueryProjectionSetting()) {
+ run(args, projectionSetting);
+ }
+ }
+
+ private static void run(String[] args, SearchParameters.AshQueryProjectionSetting projectionSetting) throws IOException {
// args is one of:
// - a list of regexes, possibly needing to be split by whitespace.
// - a list of YAML files
@@ -45,17 +52,26 @@ public static void main(String[] args) throws IOException {
}
Grid.configureWriterForBenchmark();
+ // A file path is not a dataset regular expression (and may contain
+ // spaces, brackets or platform-specific path separators).
+ List configNames = Arrays.stream(args).filter(Objects::nonNull)
+ .filter(s -> s.endsWith(".yml")).collect(Collectors.toList());
+
String regex = Arrays.stream(args)
.filter(Objects::nonNull)
+ .filter(s -> !s.endsWith(".yml"))
.flatMap(s -> Arrays.stream(s.split("\\s")))
.filter(s -> !s.isEmpty())
.collect(Collectors.joining("|"));
var pattern = Pattern.compile(regex.isEmpty() ? ".*" : regex);
- System.out.println("Heap space available is " + Runtime.getRuntime().maxMemory());
+ System.out.printf(java.util.Locale.ROOT, "Maximum heap: %.1f GiB%n", Runtime.getRuntime().maxMemory() / 1073741824.0);
var datasetCollection = DatasetCollection.load();
- var datasetNames = datasetCollection.getAll().stream().filter(dn -> pattern.matcher(dn).find()).collect(Collectors.toList());
+ var datasetNames = regex.isEmpty() && !configNames.isEmpty() ? List.of()
+ : regex.isEmpty()
+ ? datasetCollection.getAllExcept(DatasetCollection.REGRESSION_TESTS)
+ : datasetCollection.getAll().stream().filter(dn -> pattern.matcher(dn).find()).collect(Collectors.toList());
List allConfigs = new ArrayList<>();
@@ -68,19 +84,13 @@ public static void main(String[] args) throws IOException {
String name = raw.endsWith(HDF5) ? raw.substring(0, raw.length() - HDF5.length()) : raw;
allConfigs.add(MultiConfig.getDefaultConfig(name));
}
-
- MultiConfig.printDefaultConfigUsageSummary();
}
// get the list of YAML files from args
- List configNames = (args == null) ? Collections.emptyList() : Arrays.stream(args)
- .filter(Objects::nonNull)
- .filter(s -> s.endsWith(".yml"))
- .collect(Collectors.toList());
if (!configNames.isEmpty()) {
for (var configName : configNames) {
- MultiConfig config = MultiConfig.getDefaultConfig(configName);
+ MultiConfig config = MultiConfig.getConfig(configName);
allConfigs.add(config);
}
}
@@ -90,12 +100,15 @@ public static void main(String[] args) throws IOException {
return;
}
+ projectionSetting.apply(allConfigs);
+
// Set up reporting for run
RunArtifacts artifacts;
try {
RunConfig runCfg = RunConfig.loadDefault();
artifacts = RunArtifacts.open(runCfg, allConfigs);
} catch (java.io.FileNotFoundException e) {
+ if (System.getProperty("jvector.bench.runConfig") != null) throw e;
// Legacy yamlSchemaVersion "0" behavior: no run-config.yml
// - logging disabled
// - console shows compute selection
@@ -121,6 +134,9 @@ public static void main(String[] args) throws IOException {
var loading = BaseVectorLoading.forBenchmark(config);
String datasetName = config.dataset;
+ String pipeline = config.describePipeline();
+ if (pipeline != null) ConsoleSections.pipeline(datasetName, pipeline);
+ ConsoleSections.stage("Dataset");
try (var ds = DataSets.loadDataSet(datasetName, loading).orElseThrow(
() -> new RuntimeException("Could not load dataset:" + datasetName)
).getDataSet()) {
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BuildConfig.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BuildConfig.java
index 44effda63..eae17699a 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BuildConfig.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/BuildConfig.java
@@ -43,6 +43,13 @@ public final class BuildConfig {
public final boolean refineFinalGraph;
/** Compression used for graph construction, or null to build from full-precision vectors. */
public final Function buildCompressor;
+ /** Explicit index settings for schema 3. Legacy callers leave indexType unset. */
+ public final String indexType;
+ public final String writerMode;
+ public final String rerankingType;
+ public final String ashConstructionMode;
+ /** Construction threads for training, encoding, insertion and writing; null for the platform default. */
+ public final Integer workers;
private BuildConfig(Builder b) {
this.M = b.M;
@@ -52,6 +59,11 @@ private BuildConfig(Builder b) {
this.addHierarchy = b.addHierarchy;
this.refineFinalGraph = b.refineFinalGraph;
this.buildCompressor = b.buildCompressor;
+ this.indexType = b.indexType;
+ this.writerMode = b.writerMode;
+ this.rerankingType = b.rerankingType;
+ this.ashConstructionMode = b.ashConstructionMode;
+ this.workers = b.workers;
}
/**
@@ -69,6 +81,11 @@ public Map reportedParameters() {
params.put("neighborOverflow", neighborOverflow);
params.put("addHierarchy", addHierarchy);
params.put("refineFinalGraph", refineFinalGraph);
+ if (indexType != null) params.put("type", indexType);
+ if (writerMode != null) params.put("writerMode", writerMode);
+ if (rerankingType != null) params.put("reranking", rerankingType);
+ if (ashConstructionMode != null) params.put("ashConstruction", ashConstructionMode);
+ if (workers != null) params.put("constructionWorkers", workers);
return Collections.unmodifiableMap(params);
}
@@ -84,6 +101,11 @@ public static final class Builder {
private Boolean addHierarchy;
private Boolean refineFinalGraph;
private Function buildCompressor;
+ private String indexType;
+ private String writerMode;
+ private String rerankingType;
+ private String ashConstructionMode;
+ private Integer workers;
private Builder() {}
@@ -94,6 +116,11 @@ private Builder() {}
public Builder addHierarchy(boolean addHierarchy) { this.addHierarchy = addHierarchy; return this; }
public Builder refineFinalGraph(boolean refineFinalGraph) { this.refineFinalGraph = refineFinalGraph; return this; }
public Builder buildCompressor(Function buildCompressor) { this.buildCompressor = buildCompressor; return this; }
+ public Builder indexType(String indexType) { this.indexType = indexType; return this; }
+ public Builder writerMode(String writerMode) { this.writerMode = writerMode; return this; }
+ public Builder rerankingType(String rerankingType) { this.rerankingType = rerankingType; return this; }
+ public Builder ashConstructionMode(String mode) { this.ashConstructionMode = mode; return this; }
+ public Builder workers(Integer workers) { this.workers = workers; return this; }
/** @throws IllegalStateException if a parameter without a default was not set */
public BuildConfig build() {
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompressedOnDiskBuild.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompressedOnDiskBuild.java
index bdc27d796..4369de0d2 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompressedOnDiskBuild.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompressedOnDiskBuild.java
@@ -21,12 +21,16 @@
import io.github.jbellis.jvector.example.util.OnDiskGraphIndexCache;
import io.github.jbellis.jvector.graph.ImmutableGraphIndex;
import io.github.jbellis.jvector.graph.VectorAccess;
+import io.github.jbellis.jvector.graph.VectorMean;
import io.github.jbellis.jvector.graph.disk.OnDiskGraphIndex;
import io.github.jbellis.jvector.graph.disk.OnDiskParallelGraphIndexWriter;
import io.github.jbellis.jvector.graph.disk.RandomAccessOnDiskGraphIndexWriter;
import io.github.jbellis.jvector.graph.disk.feature.Feature;
import io.github.jbellis.jvector.graph.disk.feature.FeatureId;
import io.github.jbellis.jvector.quantization.BQVectors;
+import io.github.jbellis.jvector.quantization.CompressedVectors;
+import io.github.jbellis.jvector.quantization.ASHVectors;
+import io.github.jbellis.jvector.quantization.AsymmetricHashing;
import io.github.jbellis.jvector.quantization.PQVectors;
import io.github.jbellis.jvector.quantization.ProductQuantization;
import io.github.jbellis.jvector.quantization.VectorCompressor;
@@ -43,30 +47,52 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
-import java.util.function.IntFunction;
+import java.util.function.Supplier;
import java.util.stream.IntStream;
/**
- * Builds the graph from PQ- or BQ-encoded vectors while writing each feature set to disk. Feature sets
+ * Builds the graph from PQ-, BQ- or ASH-encoded vectors while writing each feature set to disk. Feature sets
* already in the index cache are loaded, and only the rest are built.
*/
final class CompressedOnDiskBuild implements IndexBuildStrategy {
private final OnDiskGraphIndexCache cache;
private final Path workDirectory;
- private final VectorCompressor buildCompressor;
+ private VectorCompressor buildCompressor;
+ private final Supplier> buildCodec;
private final ConstructionMetrics constructionMetrics;
+ private final Supplier> fusedCompressor;
+ private final String cacheIdentity;
/** @param constructionMetrics construction metrics sink, or null when not reporting */
CompressedOnDiskBuild(OnDiskGraphIndexCache cache,
Path workDirectory,
VectorCompressor buildCompressor,
ConstructionMetrics constructionMetrics) {
+ this(cache, workDirectory, buildCompressor, constructionMetrics, null, null, null);
+ }
+
+ CompressedOnDiskBuild(OnDiskGraphIndexCache cache, Path workDirectory, Supplier> buildCodec,
+ ConstructionMetrics constructionMetrics, Supplier> fusedCompressor,
+ String cacheIdentity) {
+ this(cache, workDirectory, null, constructionMetrics, buildCodec, fusedCompressor, cacheIdentity);
+ }
+
+ private CompressedOnDiskBuild(OnDiskGraphIndexCache cache, Path workDirectory, VectorCompressor buildCompressor,
+ ConstructionMetrics constructionMetrics, Supplier> buildCodec,
+ Supplier> fusedCompressor, String cacheIdentity) {
+ this.buildCodec = buildCodec;
+ this.fusedCompressor = fusedCompressor;
+ this.cacheIdentity = cacheIdentity;
this.cache = cache;
this.workDirectory = workDirectory;
this.buildCompressor = buildCompressor;
this.constructionMetrics = constructionMetrics;
}
+ VectorCompressor buildCompressor() {
+ return buildCompressor;
+ }
+
@Override
public BuiltIndexes build(List> featureSets, BuildConfig build, DataSet ds) throws IOException {
Map, ImmutableGraphIndex> indexes = new LinkedHashMap<>();
@@ -79,45 +105,51 @@ public BuiltIndexes build(List> featureSets, BuildConfi
// Map feature sets to their cache handles for the build phase
Map, OnDiskGraphIndexCache.WriteHandle> handles = new LinkedHashMap<>();
- for (Set fs : featureSets) {
- var key = cache.key(ds.getName(), fs, build.M, build.efConstruction, build.neighborOverflow, build.alpha,
- build.addHierarchy, build.refineFinalGraph, buildCompressor);
- var cached = cache.tryLoad(key);
+ try {
+ int baseCount = ds.getBaseRavv().size(), dimension = ds.getDimension();
+ for (Set fs : featureSets) {
+ int nvq = fs.contains(FeatureId.NVQ_VECTORS) ? NVQProvider.subvectors(dimension) : 0;
+ var key = cacheIdentity == null
+ ? cache.key(ds.getName(), fs, build.M, build.efConstruction, build.neighborOverflow, build.alpha,
+ build.addHierarchy, build.refineFinalGraph, buildCompressor, baseCount, dimension, nvq)
+ : cache.keyWithCompressorId(ds.getName(), fs, build.M, build.efConstruction, build.neighborOverflow,
+ build.alpha, build.addHierarchy, build.refineFinalGraph, cacheIdentity,
+ baseCount, dimension, nvq);
+ var cached = cache.tryLoad(key);
- if (cached.isPresent()) {
- System.out.printf("%s: Using cached graph index for %s%n", key.datasetName, fs);
- indexes.put(fs, cached.get());
- try { fileSizes.put(fs, Files.size(cache.resolve(key).finalPath)); } catch (IOException ignored) {}
- } else {
- missing.add(fs);
- if (cache.isEnabled()) {
- var handle = cache.beginWrite(key, OnDiskGraphIndexCache.Overwrite.ALLOW);
- // Log cache miss / build start
- System.out.printf("%s: Building graph index (cached enabled) for %s%n", key.datasetName, fs);
- // Prepare the atomic write handle immediately
- handles.put(fs, handle);
+ if (cached.isPresent()) {
+ System.out.printf("%s: Using cached graph index for %s%n", key.datasetName, fs);
+ indexes.put(fs, cached.get());
+ try { fileSizes.put(fs, Files.size(cache.resolve(key).finalPath)); } catch (IOException ignored) {}
} else {
- System.out.printf("%s: Building graph index (cache disabled) for %s%n", key.datasetName, fs);
+ missing.add(fs);
+ if (cache.isEnabled()) {
+ var handle = cache.beginWrite(key, OnDiskGraphIndexCache.Overwrite.ALLOW);
+ // Log cache miss / build start
+ System.out.printf("%s: Building graph index (cached enabled) for %s%n", key.datasetName, fs);
+ // Prepare the atomic write handle immediately
+ handles.put(fs, handle);
+ } else {
+ System.out.printf("%s: Building graph index (cache disabled) for %s%n", key.datasetName, fs);
+ }
}
}
- }
- Double buildTimeS = null;
- if (!missing.isEmpty()) {
- // At least one index needs to be built (b/c not in cache or cache is disabled)
- // We pass the handles map so buildMissing knows exactly where to write
- try {
- buildTimeS = buildMissing(missing, build, ds, outputDir, handles, indexes, fileSizes);
- } catch (IOException | RuntimeException e) {
- // A rejected or failed build must not leave the cache's temp files behind.
- for (var handle : handles.values()) {
- try { handle.close(); }
- catch (IOException suppressed) { e.addSuppressed(suppressed); }
- }
- throw e;
+ Double buildTimeS = missing.isEmpty() ? null
+ : buildMissing(missing, build, ds, outputDir, handles, indexes, fileSizes);
+ return new BuiltIndexes(indexes, fileSizes, buildTimeS);
+ } catch (IOException | RuntimeException e) {
+ // A failed build must release cached readers and discard uncommitted writes.
+ for (var handle : handles.values()) {
+ try { handle.close(); }
+ catch (IOException suppressed) { e.addSuppressed(suppressed); }
+ }
+ for (var index : indexes.values()) {
+ try { index.close(); }
+ catch (IOException suppressed) { e.addSuppressed(suppressed); }
}
+ throw e;
}
- return new BuiltIndexes(indexes, fileSizes, buildTimeS);
}
/** Builds one graph for {@code featureSets}, loads each written index into {@code indexes}, and returns the build time. */
@@ -131,27 +163,46 @@ private double buildMissing(List> featureSets,
{
Files.createDirectories(outputDir);
+ if (build.indexType != null) {
+ buildCompressor = buildCodec == null ? null : GraphBuilderFactory.inConstructionPool(build, buildCodec::get);
+ double seconds = buildAndFlush(featureSets, build, ds, outputDir, handles);
+ long reopenStart = System.nanoTime();
+ reopen(featureSets, outputDir, handles, indexes, fileSizes);
+ if (constructionMetrics != null) constructionMetrics.reopenTimeS = (System.nanoTime() - reopenStart) / 1e9;
+ return seconds;
+ }
+
var floatVectors = ds.getBaseRavv();
- // Record the encoding time if asked for....
- var encoded = (constructionMetrics != null)
- ? constructionMetrics.index(buildCompressor instanceof ProductQuantization ? "PQ" : "BQ")
- .timeEncode(() -> buildCompressor.encodeAll(floatVectors))
- : buildCompressor.encodeAll(floatVectors);
- if (!(encoded instanceof PQVectors) && !(encoded instanceof BQVectors))
- throw new IllegalArgumentException("Disk construction supports PQ or BQ build compression");
+ String quantType = buildCompressor instanceof ProductQuantization ? "PQ"
+ : buildCompressor instanceof AsymmetricHashing ? "ASH" : "BQ";
+ VectorMean mean = featureSets.stream().anyMatch(f -> f.contains(FeatureId.NVQ_VECTORS))
+ ? new VectorMean(floatVectors.size(), floatVectors.dimension()) : null;
+ var encodingVectors = mean == null ? floatVectors : mean.duringScan(floatVectors);
+ var encoded = constructionMetrics == null ? buildCompressor.encodeAll(encodingVectors)
+ : constructionMetrics.index(quantType).timeEncode(() -> buildCompressor.encodeAll(encodingVectors));
+ if (!(encoded instanceof PQVectors) && !(encoded instanceof BQVectors) && !(encoded instanceof ASHVectors))
+ throw new IllegalArgumentException("Disk construction supports PQ, BQ or ASH build compression");
+ var inputMean = mean != null && mean.isComplete() ? mean.mean() : null;
PQVectors pq = encoded instanceof PQVectors ? (PQVectors) encoded : null;
- for (var features : featureSets) {
- for (var id : features) {
- var provider = FeatureProviders.get(id);
- if (provider != null) provider.checkCompressedBuild(pq);
- }
+ ASHVectors ash = encoded instanceof ASHVectors ? (ASHVectors) encoded : null;
+ if (ash != null && Boolean.getBoolean("jvector.ash.debugSymmetric")) {
+ io.github.jbellis.jvector.example.diagnostics.ASHScoreDebug.run(floatVectors, ash,
+ Integer.getInteger("jvector.ash.debugSymmetric.pairs", 2000),
+ Integer.getInteger("jvector.ash.debugSymmetric.subsetN", 50000),
+ Integer.getInteger("jvector.ash.debugSymmetric.topK", 100), new java.util.Random(123));
}
+ for (var features : featureSets) for (var id : features) FeatureProviders.require(id).checkCompressedBuild(encoded);
+ boolean allSymmetric = ash != null && "symmetric-all".equals(
+ System.getProperty("jvector.bench.ashConstruction", "symmetric-all"));
// Compressed scores only guide construction; searches rerank with a feature's stored vectors.
if (featureSets.stream().noneMatch(fs -> FeatureProviders.anyMatch(fs, FeatureProvider::providesRerankVectors))) {
throw new IllegalStateException("Bench looks for either NVQ_VECTORS or INLINE_VECTORS feature set for scoring compressed builds.");
}
- var builder = GraphBuilderFactory.forCompressed(encoded, ds.getSimilarityFunction(), floatVectors.dimension(), build);
+ long providerStart = System.nanoTime();
+ var bsp = GraphBuilderFactory.compressedScores(encoded, ds.getSimilarityFunction(), floatVectors);
+ var builder = GraphBuilderFactory.forCompressed(bsp, floatVectors.dimension(), build);
+ System.out.printf("Build score provider setup: %.6f seconds%n", (System.nanoTime() - providerStart) / 1e9);
Map, RandomAccessOnDiskGraphIndexWriter> writers = new LinkedHashMap<>();
Map, FeatureWriteContext> contexts = new LinkedHashMap<>();
@@ -165,7 +216,7 @@ private double buildMissing(List> featureSets,
} else {
graphPath = outputDir.resolve("graph" + n++);
}
- var ctx = FeatureProviders.configureWriter(features, builder.getGraph(), graphPath, floatVectors, pq, constructionMetrics);
+ var ctx = FeatureProviders.configureWriter(features, builder.getGraph(), graphPath, floatVectors, pq, ash, inputMean, constructionMetrics);
writers.put(features, ctx.buildWriter());
contexts.put(features, ctx);
}
@@ -187,7 +238,8 @@ private double buildMissing(List> featureSets,
if (ds.useDatasetPrefetching()) {
VectorAccess.forEach(floatVectors, PhysicalCoreExecutor.pool(), (node, vector) -> {
if (!deferFeatures) writeInlineStates(writers, contexts, node, vector);
- builder.addGraphNode(node, vector);
+ if (allSymmetric) builder.addGraphNode(node, bsp.searchProviderFor(node));
+ else builder.addGraphNode(node, vector);
});
} else {
var vv = floatVectors.threadLocalSupplier();
@@ -195,11 +247,14 @@ private double buildMissing(List> featureSets,
IntStream.range(0, floatVectors.size()).parallel().forEach(node -> {
var vector = vv.get().getVector(node);
if (!deferFeatures) writeInlineStates(writers, contexts, node, vector);
- builder.addGraphNode(node, vector);
+ if (allSymmetric) builder.addGraphNode(node, bsp.searchProviderFor(node));
+ else builder.addGraphNode(node, vector);
});
}).join();
}
+ long insertionEnd = System.nanoTime();
builder.cleanup();
+ long cleanupEnd = System.nanoTime();
// With deferred features, supply every feature so the parallel writer emits complete records.
writers.entrySet().stream().parallel().forEach(entry -> {
@@ -208,21 +263,72 @@ private double buildMissing(List> featureSets,
try (var vectors = new FinalWriteVectors(floatVectors);
var view = builder.getGraph().getView();
var closeWriter = writer) {
- Map> writeSuppliers = new EnumMap<>(FeatureId.class);
- if (deferFeatures) {
- ctx.vectorSuppliers().forEach((id, supplier) ->
- writeSuppliers.put(id, ordinal -> supplier.apply(vectors.get(ordinal))));
- }
- ctx.afterBuildSuppliers().forEach((id, supplier) -> writeSuppliers.put(id, supplier.apply(view)));
- writer.write(writeSuppliers);
+ writer.write(ctx.completedSuppliers(view, deferFeatures ? vectors::get : null));
} catch (IOException e) {
throw new UncheckedIOException(e);
}
});
builder.close();
double totalTime = (System.nanoTime() - startTime) / 1_000_000_000.0;
+ System.out.printf("Graph phases: insertion=%.6f cleanup=%.6f flush=%.6f seconds%n",
+ (insertionEnd-startTime)/1e9, (cleanupEnd-insertionEnd)/1e9,
+ (System.nanoTime()-cleanupEnd)/1e9);
System.out.format("Build and write %s in %ss%n", featureSets, totalTime);
+ reopen(featureSets, outputDir, handles, indexes, fileSizes);
+ return totalTime;
+ }
+
+ /** Builds the topology before preparing search features and writing the disk index. */
+ private double buildAndFlush(List> featureSets, BuildConfig build, DataSet ds, Path outputDir,
+ Map, OnDiskGraphIndexCache.WriteHandle> handles) throws IOException {
+ var vectors = ds.getBaseRavv();
+ var mean = featureSets.stream().anyMatch(features -> features.contains(FeatureId.NVQ_VECTORS))
+ ? new VectorMean(vectors.size(), vectors.dimension()) : null;
+ var buildVectors = mean == null ? vectors : mean.duringScan(vectors);
+ var encoded = Grid.encodeForConstruction(buildCompressor, buildVectors, constructionMetrics,
+ GraphBuilderFactory.constructionPool(build));
+
+ long start = System.nanoTime();
+ try (var graph = GraphBuilderFactory.buildGraph(ds, build, encoded)) {
+ double buildSeconds = (System.nanoTime() - start) / 1e9;
+ var codec = fusedCompressor == null ? null
+ : GraphBuilderFactory.inConstructionPool(build, fusedCompressor::get);
+ CompressedVectors searchCodes = codec == null ? null : codec == buildCompressor
+ ? encoded
+ : Grid.encodeForConstruction(codec, mean != null && !mean.isComplete() ? buildVectors : vectors,
+ constructionMetrics, GraphBuilderFactory.constructionPool(build));
+ var inputMean = mean != null && mean.isComplete() ? mean.mean() : null;
+ double flushSeconds = 0;
+ int n = 0;
+ for (var features : featureSets) {
+ Path path = handles.containsKey(features) ? handles.get(features).writePath() : outputDir.resolve("graph" + n++);
+ var ctx = FeatureProviders.configureWriter(features, graph, path, vectors,
+ searchCodes instanceof PQVectors ? (PQVectors) searchCodes : null,
+ searchCodes instanceof ASHVectors ? (ASHVectors) searchCodes : null,
+ inputMean, constructionMetrics, build.writerMode, GraphBuilderFactory.constructionThreads(build));
+ start = System.nanoTime();
+ try (var progress = PhaseProgress.timed("Index write " + features);
+ var input = new FinalWriteVectors(vectors);
+ var view = graph.getView();
+ var writer = ctx.buildWriter()) {
+ writer.write(ctx.completedSuppliers(view, input::get));
+ }
+ flushSeconds += (System.nanoTime() - start) / 1e9;
+ }
+ double totalTime = buildSeconds + flushSeconds;
+ if (constructionMetrics != null) {
+ constructionMetrics.graphBuildTimeS = buildSeconds;
+ constructionMetrics.flushTimeS = flushSeconds;
+ }
+ return totalTime;
+ }
+ }
+
+ private static void reopen(List> featureSets, Path outputDir,
+ Map, OnDiskGraphIndexCache.WriteHandle> handles,
+ Map, ImmutableGraphIndex> indexes,
+ Map, Long> fileSizes) throws IOException {
// Commit the index to cache atomically
for (var entry : handles.entrySet()) {
System.out.printf("Committing %s build to graph index cache...%n", entry.getKey());
@@ -230,7 +336,7 @@ private double buildMissing(List> featureSets,
}
// open indexes and capture per-feature-set file sizes
- n = 0;
+ int n = 0;
for (var features : featureSets) {
Path loadPath = handles.containsKey(features)
? handles.get(features).finalPath()
@@ -239,7 +345,6 @@ private double buildMissing(List> featureSets,
indexes.put(features, OnDiskGraphIndex.load(ReaderSupplierFactory.open(loadPath)));
fileSizes.put(features, Files.size(loadPath));
}
- return totalTime;
}
/// Write one node's original-vector features to every writer during insertion. The borrowed
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConsoleReportListener.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConsoleReportListener.java
index 07c605674..1e134109d 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConsoleReportListener.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConsoleReportListener.java
@@ -23,12 +23,15 @@
import java.util.LinkedHashMap;
import java.util.Map;
+import java.util.ArrayList;
+import java.util.List;
/** Prints a results table per pruning setting and topK, showing the console selection from run-config.yml. */
final class ConsoleReportListener implements GridListener {
private final RunArtifacts artifacts;
// Per (index, search compressor), so missing-metric warnings repeat for each configuration.
+ private final List rows = new ArrayList<>();
private SearchSelection selection;
private BenchmarkTablePrinter printer;
private ReportingSelectionResolver.ResolvedSelection resolved;
@@ -50,7 +53,10 @@ public void onPruningStart(SearchContext ctx, boolean usePruning) {
Map indexConfig = new LinkedHashMap<>();
indexConfig.put("featureSetForIndex", ctx.indexFeatures);
indexConfig.putAll(ctx.build.reportedParameters());
- printer.printConfig(indexConfig, Map.of("usePruning", usePruning));
+ var queryConfig = new LinkedHashMap();
+ queryConfig.put("usePruning", usePruning);
+ queryConfig.putAll(ctx.queryParameters);
+ printer.printConfig(indexConfig, queryConfig);
}
@Override
@@ -59,17 +65,19 @@ public void onTopKStart(SearchContext ctx, boolean usePruning, int topK) {
resolved = selection.resolveForTopK(topK);
// Force header to re-print for this topK (columns change)
printer.resetTable();
+ rows.clear();
}
@Override
public void onQueryResult(SearchContext ctx, QueryRow row) {
// Best-effort runtime availability warnings (warn once per key)
selection.warnMissing(row.metrics, resolved);
- printer.printRow(row.overquery, selection.apply(row.metrics, resolved));
+ rows.add(row);
}
@Override
public void onTopKEnd(SearchContext ctx, boolean usePruning, int topK) {
+ for (var row : rows) printer.printRow(row.overquery, selection.apply(row.metrics, resolved));
printer.printFooter();
}
}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConsoleSections.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConsoleSections.java
new file mode 100644
index 000000000..66fb5ba07
--- /dev/null
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConsoleSections.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.github.jbellis.jvector.example;
+
+/** Console structure for BenchYAML runs: a banner per experiment and a header per stage. */
+public final class ConsoleSections {
+ private static final String RULE = "=".repeat(80);
+
+ private ConsoleSections() {}
+
+ /** The experiment's pipeline, framed so it stands out at the top of its output. */
+ public static void pipeline(String dataset, String pipeline) {
+ System.out.println();
+ System.out.println(RULE);
+ System.out.println(" PIPELINE " + dataset + ": " + pipeline);
+ System.out.println(RULE);
+ }
+
+ /** Starts a stage of the run, for example Dataset, Construction or Search. */
+ public static void stage(String title) {
+ System.out.println();
+ System.out.println("== " + title + " ==");
+ }
+}
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConstructionMetrics.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConstructionMetrics.java
index 809ffa3e5..9105a5b75 100644
--- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConstructionMetrics.java
+++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/ConstructionMetrics.java
@@ -36,6 +36,11 @@ final class ConstructionMetrics {
// null means “not applicable / not measured”
public Double indexBuildTimeS;
public Long indexFileSizeBytes;
+ // Schema 3 on-disk phases: graph construction and cleanup, writing the completed graph,
+ // and reopening the written files. Graph build plus flush equals indexBuildTimeS.
+ public Double graphBuildTimeS;
+ public Double flushTimeS;
+ public Double reopenTimeS;
// Index-construction phase (single pass per run)
private final Map indexByType = new HashMap<>();
@@ -43,7 +48,13 @@ final class ConstructionMetrics {
// Search phase (reset per evaluated config)
private final Map searchByType = new HashMap<>();
- ConstructionMetrics() {}
+ private final boolean sumIndexQuantTimes;
+
+ ConstructionMetrics() { this(false); }
+
+ ConstructionMetrics(boolean sumIndexQuantTimes) {
+ this.sumIndexQuantTimes = sumIndexQuantTimes;
+ }
/** Clear search-phase quant timings for the next evaluated (index, compression) configuration. */
void resetSearch() {
@@ -70,17 +81,25 @@ void appendTo(List out) {
"Index file size (MB)", ".2f", indexFileSizeBytes / 1024.0 / 1024.0));
}
+ appendPhase(out, "construction.graph_build_time_s", "Graph build time (s)", graphBuildTimeS);
+ appendPhase(out, "construction.flush_time_s", "Flush time (s)", flushTimeS);
+ appendPhase(out, "construction.reopen_time_s", "Reopen time (s)", reopenTimeS);
+
// Index-construction quant timings
- appendQuant(out, "construction.index_quant_time_s", "Idx", indexByType);
+ appendQuant(out, "construction.index_quant_time_s", "Idx", indexByType, sumIndexQuantTimes);
// Search-phase quant timings (note: search.* namespace matches other query-time metrics)
- appendQuant(out, "search.search_quant_time_s", "Search", searchByType);
+ appendQuant(out, "search.search_quant_time_s", "Search", searchByType, false);
+ }
+
+ private static void appendPhase(List out, String key, String header, Double seconds) {
+ if (seconds != null) out.add(Metric.of(key, header, ".2f", seconds));
}
private static void appendQuant(List out,
String prefix,
String phaseLabel,
- Map byType) {
+ Map byType, boolean sumTimes) {
final String hPrefix = (phaseLabel == null || phaseLabel.isBlank()) ? "" : (phaseLabel + " ");
for (var e : byType.entrySet()) {
@@ -88,16 +107,19 @@ private static void appendQuant(List out,
QuantStats qs = e.getValue();
// Compute
- if (qs.computeTimeS != null) {
+ if (!qs.computeTimesS.isEmpty()) {
+ double computeTime = sumTimes ? qs.computeTimesS.stream().mapToDouble(Double::doubleValue).sum()
+ : qs.computeTimesS.get(qs.computeTimesS.size() - 1);
out.add(Metric.of(prefix + "." + qt + ".compute_time_s",
- hPrefix + qt + " compute (s)", ".2f", qs.computeTimeS));
+ hPrefix + qt + " compute (s)", ".2f", computeTime));
}
// Encoding (usually single; suffix only if multiple)
int n = qs.encodeTimesS.size();
- if (n == 1) {
+ if (n > 0 && (n == 1 || sumTimes)) {
+ double encodeTime = qs.encodeTimesS.stream().mapToDouble(Double::doubleValue).sum();
out.add(Metric.of(prefix + "." + qt + ".encoding_time_s",
- hPrefix + qt + " encoding (s)", ".2f", qs.encodeTimesS.get(0)));
+ hPrefix + qt + " encoding (s)", ".2f", encodeTime));
} else if (n > 1) {
for (int i = 0; i < n; i++) {
out.add(Metric.of(prefix + "." + qt + ".encoding_time_s." + i,
@@ -108,7 +130,7 @@ private static void appendQuant(List