Copilot commented on code in PR #18091:
URL: https://github.com/apache/pinot/pull/18091#discussion_r3035212818
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSimilarityFilterOperator.java:
##########
@@ -136,12 +137,27 @@ public List<Operator> getChildOperators() {
@Override
public String toExplainString() {
- return EXPLAIN_NAME + "(indexLookUp:vector_index"
- + ", operator:" + _predicate.getType()
- + ", vector identifier:" + _predicate.getLhs().getIdentifier()
- + ", vector literal:" + Arrays.toString(_predicate.getValue())
- + ", topK to search:" + _predicate.getTopK()
- + ')';
+ Map<String, String> debugInfo = _vectorIndexReader.getIndexDebugInfo();
+ StringBuilder sb = new StringBuilder();
+ sb.append(EXPLAIN_NAME);
+ sb.append("(indexLookUp:vector_index");
+ sb.append(", operator:").append(_predicate.getType());
+ sb.append(", vector
identifier:").append(_predicate.getLhs().getIdentifier());
+ sb.append(", vector
literal:").append(Arrays.toString(_predicate.getValue()));
+ sb.append(", topK to search:").append(_predicate.getTopK());
+ if (!debugInfo.isEmpty()) {
+ sb.append(", backend:").append(debugInfo.getOrDefault("backend",
"UNKNOWN"));
+ String nprobe = debugInfo.get("nprobe");
+ if (nprobe != null) {
+ sb.append(", nprobe:").append(nprobe);
+ }
Review Comment:
`toExplainString()` only appends backend/nprobe when `getIndexDebugInfo()`
is non-empty. Since `IvfFlatVectorIndexReader` does not implement
`getIndexDebugInfo()`, EXPLAIN output for IVF_FLAT will omit backend/nprobe
even though the PR description suggests these fields are present. Consider
adding a fallback (e.g., infer backend from reader type) or implementing
`getIndexDebugInfo()` for IVF_FLAT.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/VectorIndexHandler.java:
##########
@@ -62,12 +63,26 @@ public boolean needUpdateIndices(SegmentDirectory.Reader
segmentReader) {
String segmentName = _segmentDirectory.getSegmentMetadata().getName();
Set<String> columnsToAddIdx = new HashSet<>(_vectorConfigs.keySet());
Set<String> existingColumns =
segmentReader.toSegmentDirectory().getColumnsWithIndex(StandardIndexes.vector());
+ File indexDir = _segmentDirectory.getSegmentMetadata().getIndexDir();
// Check if any existing index need to be removed.
for (String column : existingColumns) {
if (!columnsToAddIdx.remove(column)) {
LOGGER.info("Need to remove existing Vector index from segment: {},
column: {}", segmentName, column);
return true;
}
+ // Check if the backend type changed (e.g., HNSW -> IVF_PQ or vice versa)
+ VectorIndexConfig desiredConfig = _vectorConfigs.get(column);
+ if (desiredConfig != null) {
+ // Resolve desired backend: null/empty defaults to HNSW
+ String rawDesired = desiredConfig.getVectorIndexType();
+ String desiredBackend = (rawDesired == null || rawDesired.isEmpty()) ?
"HNSW" : rawDesired.toUpperCase();
Review Comment:
Backend-change detection compares `desiredConfig.getVectorIndexType()` (may
be null when config relies on the default HNSW backend) to the on-disk backend.
If `vectorIndexType` is omitted but an IVF_* index exists on disk, this logic
will not trigger a rebuild to HNSW. Consider comparing resolved backend types
(e.g., `desiredConfig.resolveBackendType().name()`) instead of the raw string
field.
```suggestion
// Compare against the resolved backend so omitted config still uses
the effective default.
String desiredBackend = desiredConfig.resolveBackendType().name();
```
##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/VectorIndexConfigValidator.java:
##########
@@ -174,6 +182,51 @@ private static void validateIvfFlatProperties(Map<String,
String> properties) {
}
}
+ /**
+ * Validates IVF_PQ-specific property values.
+ */
+ private static void validateIvfPqProperties(VectorIndexConfig config,
Map<String, String> properties) {
+ validatePositiveIntProperty(properties, "nlist", "IVF_PQ nlist");
+ validatePositiveIntProperty(properties, "nprobe", "IVF_PQ nprobe");
+ validatePositiveIntProperty(properties, "trainSampleSize", "IVF_PQ
trainSampleSize");
Review Comment:
`validateIvfPqProperties()` does not validate `trainSampleSize` (or other
numeric IVF_PQ properties beyond nlist/nprobe/pqM/pqNbits). A non-positive
`trainSampleSize` can later cause failures during index building (e.g.,
negative array size / invalid sampling). Consider validating `trainSampleSize`
as a positive integer (similar to IVF_FLAT) and optionally validating
`trainingSeed` is parseable.
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSimilarityFilterOperator.java:
##########
@@ -201,11 +225,11 @@ private ImmutableRoaringBitmap executeSearch() {
* Configures backend-specific search parameters on the reader if it
supports them.
*/
private void configureBackendParams(String column) {
- // Set nprobe on IVF_FLAT readers
+ // Set nprobe on IVF readers (IVF_FLAT and IVF_PQ)
if (_vectorIndexReader instanceof NprobeAware) {
int nprobe = _searchParams.getNprobe();
((NprobeAware) _vectorIndexReader).setNprobe(nprobe);
- LOGGER.debug("Set nprobe={} on IVF_FLAT reader for column: {}", nprobe,
column);
+ LOGGER.debug("Set nprobe={} on IVF reader for column: {}", nprobe,
column);
}
Review Comment:
`configureBackendParams()` always calls
`setNprobe(_searchParams.getNprobe())` for any `NprobeAware` reader. Since
`VectorSearchParams` defaults nprobe to 4 (documented as IVF_FLAT default),
this will override `IvfPqVectorIndexReader`’s default nprobe (8) even when the
user did not set `vectorNprobe`, changing IVF_PQ default behavior. Consider
only applying nprobe when the query option is explicitly provided, or make the
default nprobe backend-specific.
##########
pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/creator/impl/vector/pq/benchmark/VectorIndexBenchmark.java:
##########
@@ -0,0 +1,289 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.benchmark;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+import org.apache.commons.io.FileUtils;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.IvfPqVectorIndexCreator;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.VectorDistanceUtil;
+import
org.apache.pinot.segment.local.segment.index.readers.vector.IvfPqVectorIndexReader;
+import org.apache.pinot.segment.spi.index.creator.VectorIndexConfig;
+import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
+import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
+import org.testng.annotations.Test;
+
+
+/**
+ * Reproducible benchmark suite for vector index backends.
+ *
+ * <p>Compares: exact scan, IVF_PQ (various configs).
+ * HNSW benchmark requires Lucene integration which is tested separately.
+ *
+ * <p>This benchmark is designed to run as a unit test but produces a formatted
+ * results table on stdout. It is NOT run as part of normal CI — invoke
explicitly with:
+ * {@code mvn -pl pinot-segment-local -Dtest=VectorIndexBenchmark test}
+ */
+public class VectorIndexBenchmark {
+ private static final int DIMENSION = 64;
+ private static final int NUM_VECTORS = 10_000;
+ private static final int NUM_QUERIES = 100;
+ private static final int[] TOP_K_VALUES = {10, 100};
+
+ @Test(enabled = false, description = "Manual benchmark — run with
-Dtest=VectorIndexBenchmark")
+ public void runBenchmark()
+ throws IOException {
+ System.out.println("=== Vector Index Benchmark ===");
+ System.out.println("Vectors: " + NUM_VECTORS + ", Dimension: " + DIMENSION
+ ", Queries: " + NUM_QUERIES);
+ System.out.println();
+
+ // Generate L2-oriented dataset (uniform random)
+ Random random = new Random(42);
+ float[][] l2Vectors = generateUniformVectors(NUM_VECTORS, DIMENSION,
random);
+ float[][] l2Queries = generateUniformVectors(NUM_QUERIES, DIMENSION, new
Random(999));
+
+ // Generate normalized cosine dataset
+ float[][] cosineVectors = generateNormalizedVectors(NUM_VECTORS,
DIMENSION, random);
+ float[][] cosineQueries = generateNormalizedVectors(NUM_QUERIES,
DIMENSION, new Random(999));
+
+ System.out.println("## L2/Euclidean Dataset");
+ benchmarkDataset("L2", l2Vectors, l2Queries, "EUCLIDEAN",
+ VectorIndexConfig.VectorDistanceFunction.EUCLIDEAN);
+
+ System.out.println();
+ System.out.println("## Cosine Dataset (normalized vectors)");
+ benchmarkDataset("Cosine", cosineVectors, cosineQueries, "COSINE",
+ VectorIndexConfig.VectorDistanceFunction.COSINE);
+ }
+
+ private void benchmarkDataset(String label, float[][] vectors, float[][]
queries, String distFuncStr,
+ VectorIndexConfig.VectorDistanceFunction distFunc)
+ throws IOException {
+ int distFuncCode = distFunc ==
VectorIndexConfig.VectorDistanceFunction.EUCLIDEAN ? 1 : 0;
+
+ // Precompute exact results for recall computation
+ Map<Integer, int[][]> exactResults = new HashMap<>();
+ for (int topK : TOP_K_VALUES) {
+ int[][] exact = new int[queries.length][];
+ for (int q = 0; q < queries.length; q++) {
+ exact[q] = findExactTopK(queries[q], vectors, topK, distFuncCode);
+ }
+ exactResults.put(topK, exact);
+ }
+
+ // Exact scan benchmark
+ benchmarkExactScan(label, vectors, queries, exactResults, distFuncCode);
+
+ // IVF_PQ configs to benchmark
+ int[][] configs = {
+ // {nlist, pqM, pqNbits, nprobe}
+ {64, 8, 8, 4},
+ {64, 8, 8, 8},
+ {64, 8, 8, 16},
+ {128, 16, 8, 8},
+ {128, 16, 8, 16},
+ {256, 16, 8, 16},
+ {256, 16, 8, 32},
+ };
+
+ printHeader();
+ for (int[] cfg : configs) {
+ benchmarkIvfPq(label, vectors, queries, exactResults, distFuncStr,
distFunc, cfg[0], cfg[1], cfg[2], cfg[3]);
+ }
+
+ // Rerank on/off comparison
+ System.out.println();
+ System.out.println("### Rerank Effect (nlist=128, pqM=16, pqNbits=8,
nprobe=16)");
+ printHeader();
+ benchmarkIvfPqRerank(label, vectors, queries, exactResults, distFuncStr,
distFunc, 128, 16, 8, 16, true);
+ benchmarkIvfPqRerank(label, vectors, queries, exactResults, distFuncStr,
distFunc, 128, 16, 8, 16, false);
+ }
+
+ private void benchmarkExactScan(String label, float[][] vectors, float[][]
queries,
+ Map<Integer, int[][]> exactResults, int distFuncCode) {
+ System.out.println("### Exact Scan");
+ System.out.printf("| %-35s | %10s | %10s | %10s | %10s | %10s | %10s |%n",
+ "Config", "BuildMs", "SizeMB", "Recall@10", "Recall@100", "p50us",
"p95us");
+ System.out.println("|" + "-".repeat(37) + "|" + ("-".repeat(12) +
"|").repeat(6));
+
+ long[] latencies = new long[queries.length];
+ for (int q = 0; q < queries.length; q++) {
+ long start = System.nanoTime();
+ findExactTopK(queries[q], vectors, 100, distFuncCode);
+ latencies[q] = (System.nanoTime() - start) / 1000; // microseconds
+ }
+ Arrays.sort(latencies);
+ System.out.printf("| %-35s | %10d | %10.2f | %10.4f | %10.4f | %10d | %10d
|%n",
+ "ExactScan", 0, 0.0, 1.0, 1.0,
+ latencies[(int) (queries.length * 0.5)], latencies[(int)
(queries.length * 0.95)]);
+ System.out.println();
+ }
+
+ private void printHeader() {
+ System.out.printf("| %-35s | %10s | %10s | %10s | %10s | %10s | %10s |%n",
+ "Config", "BuildMs", "SizeMB", "Recall@10", "Recall@100", "p50us",
"p95us");
+ System.out.println("|" + "-".repeat(37) + "|" + ("-".repeat(12) +
"|").repeat(6));
+ }
+
+ private void benchmarkIvfPq(String label, float[][] vectors, float[][]
queries,
+ Map<Integer, int[][]> exactResults, String distFuncStr,
VectorIndexConfig.VectorDistanceFunction distFunc,
+ int nlist, int pqM, int pqNbits, int nprobe)
+ throws IOException {
+ benchmarkIvfPqRerank(label, vectors, queries, exactResults, distFuncStr,
distFunc,
+ nlist, pqM, pqNbits, nprobe, true);
+ }
+
+ private void benchmarkIvfPqRerank(String label, float[][] vectors, float[][]
queries,
+ Map<Integer, int[][]> exactResults, String distFuncStr,
VectorIndexConfig.VectorDistanceFunction distFunc,
+ int nlist, int pqM, int pqNbits, int nprobe, boolean rerank)
+ throws IOException {
+ File tempDir = new File(FileUtils.getTempDirectory(), "bench_" +
System.nanoTime());
+ tempDir.mkdirs();
+
+ try {
+ VectorIndexConfig config = createConfig(distFuncStr, distFunc, nlist,
pqM, pqNbits);
+
+ // Build
+ long buildStart = System.nanoTime();
+ try (IvfPqVectorIndexCreator creator = new
IvfPqVectorIndexCreator("emb", tempDir, config)) {
+ for (float[] v : vectors) {
+ creator.add(v);
+ }
+ creator.seal();
+ }
+ long buildMs = (System.nanoTime() - buildStart) / 1_000_000;
+
+ File indexFile = new File(tempDir, "emb" +
IvfPqVectorIndexCreator.FILE_EXTENSION);
+ double sizeMB = indexFile.length() / (1024.0 * 1024.0);
+
+ // Search
+ Map<String, String> opts = new HashMap<>();
+ opts.put(QueryOptionKey.VECTOR_NPROBE, String.valueOf(nprobe));
+ opts.put(QueryOptionKey.VECTOR_EXACT_RERANK, String.valueOf(rerank));
+
+ long[] latencies = new long[queries.length];
+ double recall10 = 0;
+ double recall100 = 0;
+
+ try (IvfPqVectorIndexReader reader = new IvfPqVectorIndexReader("emb",
tempDir, vectors.length, nprobe)) {
+ for (int q = 0; q < queries.length; q++) {
+ long start = System.nanoTime();
+ ImmutableRoaringBitmap result10 = reader.getDocIds(queries[q], 10,
opts);
+ latencies[q] = (System.nanoTime() - start) / 1000;
+ recall10 += computeRecall(result10, exactResults.get(10)[q]);
+
+ ImmutableRoaringBitmap result100 = reader.getDocIds(queries[q], 100,
opts);
Review Comment:
This benchmark passes query options (including `VECTOR_EXACT_RERANK`) into
`IvfPqVectorIndexReader#getDocIds(..., searchParams)`, but
`IvfPqVectorIndexReader` does not override the 3-arg method, so the default
interface implementation ignores `opts`. As a result, the “rerank effect”
comparison here won’t actually enable/disable rerank and the reported numbers
will be misleading.
##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/creator/VectorIndexConfigValidator.java:
##########
@@ -121,6 +125,10 @@ private static void
validateBackendSpecificProperties(VectorIndexConfig config,
validateNoForeignProperties(properties, HNSW_PROPERTIES, "IVF_FLAT",
"HNSW");
validateIvfFlatProperties(properties);
break;
+ case IVF_PQ:
+ validateNoForeignProperties(properties, HNSW_PROPERTIES, "IVF_PQ",
"HNSW");
Review Comment:
Backend-specific property mixing between IVF_FLAT and IVF_PQ is not
currently rejected: the IVF_PQ branch only checks for HNSW foreign keys, so
IVF_FLAT-only keys like `minRowsForIndex` would be accepted; similarly IVF_FLAT
only rejects HNSW keys so IVF_PQ-only keys (`pqM`, `pqNbits`, `nprobe`) would
be accepted. If the intended rule is “no cross-backend properties”, add
corresponding `validateNoForeignProperties` checks between IVF_FLAT_PROPERTIES
and IVF_PQ_PROPERTIES.
```suggestion
validateNoForeignProperties(properties, HNSW_PROPERTIES, "IVF_FLAT",
"HNSW");
validateNoForeignProperties(properties, IVF_PQ_PROPERTIES,
"IVF_FLAT", "IVF_PQ");
validateIvfFlatProperties(properties);
break;
case IVF_PQ:
validateNoForeignProperties(properties, HNSW_PROPERTIES, "IVF_PQ",
"HNSW");
validateNoForeignProperties(properties, IVF_FLAT_PROPERTIES,
"IVF_PQ", "IVF_FLAT");
```
##########
pinot-core/src/main/java/org/apache/pinot/core/operator/filter/VectorSimilarityFilterOperator.java:
##########
@@ -160,6 +176,14 @@ protected void explainAttributes(ExplainAttributeBuilder
attributeBuilder) {
if (_searchParams.isExactRerank()) {
attributeBuilder.putString("exactRerank", "true");
}
Review Comment:
The operator advertises/uses multiple distance functions (COSINE /
INNER_PRODUCT), but exact rerank currently computes distances using L2 only.
This makes `vectorExactRerank=true` return incorrectly ordered results for
non-EUCLIDEAN indexes; rerank should compute exact distance using the same
distance function as the index/predicate.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/IvfPqVectorIndexReader.java:
##########
@@ -0,0 +1,292 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pinot.segment.local.segment.index.readers.vector;
+
+import java.io.Closeable;
+import java.io.File;
+import java.io.IOException;
+import java.io.RandomAccessFile;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.PriorityQueue;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.IvfPqIndexFormat;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.IvfPqVectorIndexCreator;
+import org.apache.pinot.segment.local.segment.creator.impl.vector.pq.KMeans;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.ProductQuantizer;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.VectorDistanceUtil;
+import org.apache.pinot.segment.spi.index.reader.NprobeAware;
+import org.apache.pinot.segment.spi.index.reader.VectorIndexReader;
+import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Reader for IVF_PQ vector index. Loads the compact index (no original
vectors)
+ * and performs approximate nearest-neighbor search using IVF + Product
Quantization.
+ *
+ * <p>Exact rerank is handled by the operator layer via the forward index
reader,
+ * not by this reader. This keeps the index file compact (only PQ codes +
docIds).</p>
+ *
+ * <p>Search flow:
+ * <ol>
+ * <li>Find the nprobe closest coarse centroids to the query vector</li>
+ * <li>Score all candidates in probed lists using approximate PQ distances
(ADC)</li>
+ * <li>Return top-K approximate candidates as doc IDs</li>
+ * </ol>
+ */
+public class IvfPqVectorIndexReader implements VectorIndexReader, NprobeAware {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(IvfPqVectorIndexReader.class);
+ private static final int DEFAULT_NPROBE = 8;
+
+ private final String _column;
+ private final int _dimension;
+ private final int _nlist;
+ private final int _pqM;
+ private final int _pqNbits;
+ private final int _numVectors;
+ private final int _distanceFunctionCode;
+
+ private final float[][] _coarseCentroids;
+ private final ProductQuantizer _pq;
+
+ // Inverted lists: docIds and PQ codes per list
+ private final int[][] _listDocIds;
+ private final byte[][][] _listPqCodes;
+
+ private volatile int _nprobe;
+ private final Closeable _closeHandle;
+
+ public IvfPqVectorIndexReader(String column, File segmentDir, int numDocs,
int nprobe) {
+ _column = column;
+ _nprobe = nprobe > 0 ? nprobe : DEFAULT_NPROBE;
+
+ File indexFile = findIndexFile(segmentDir);
+ if (indexFile == null) {
+ throw new IllegalStateException("Failed to find IVF_PQ index file for
column: " + column);
+ }
+
+ try {
+ RandomAccessFile raf = new RandomAccessFile(indexFile, "r");
+ MappedByteBuffer buffer =
raf.getChannel().map(FileChannel.MapMode.READ_ONLY, 0, raf.length());
+ buffer.order(ByteOrder.LITTLE_ENDIAN);
+ _closeHandle = raf;
+
+ // Read header
+ int magic = buffer.getInt();
+ if (magic != IvfPqIndexFormat.MAGIC) {
+ throw new IOException("Invalid IVF_PQ index magic: " +
Integer.toHexString(magic));
+ }
+ int version = buffer.getInt();
+ if (version != IvfPqIndexFormat.VERSION) {
+ throw new IOException("Unsupported IVF_PQ index version: " + version);
+ }
+ _dimension = buffer.getInt();
+ _nlist = buffer.getInt();
+ _pqM = buffer.getInt();
+ _pqNbits = buffer.getInt();
+ _distanceFunctionCode = buffer.getInt();
+ _numVectors = buffer.getInt();
+
+ int ksub = 1 << _pqNbits;
+ int dsub = _pqM > 0 ? _dimension / _pqM : 1;
+
+ if (_numVectors == 0 || _nlist == 0) {
+ _coarseCentroids = new float[0][];
+ int safeDim = Math.max(_dimension, 1);
+ int safeM = Math.max(_pqM, 1);
+ if (safeDim % safeM != 0) {
+ safeM = 1;
+ }
+ _pq = new ProductQuantizer(safeDim, safeM, Math.max(_pqNbits, 1));
+ _listDocIds = new int[0][];
+ _listPqCodes = new byte[0][][];
+ return;
+ }
+
+ // Read coarse centroids
+ _coarseCentroids = new float[_nlist][_dimension];
+ for (int c = 0; c < _nlist; c++) {
+ for (int d = 0; d < _dimension; d++) {
+ _coarseCentroids[c][d] = buffer.getFloat();
+ }
+ }
+
+ // Read PQ codebooks
+ float[][][] codebooks = new float[_pqM][ksub][dsub];
+ for (int sub = 0; sub < _pqM; sub++) {
+ for (int code = 0; code < ksub; code++) {
+ for (int d = 0; d < dsub; d++) {
+ codebooks[sub][code][d] = buffer.getFloat();
+ }
+ }
+ }
+ _pq = new ProductQuantizer(_dimension, _pqM, _pqNbits);
+ _pq.setCodebooks(codebooks);
+
+ // Read inverted lists (docIds + PQ codes only — no original vectors)
+ _listDocIds = new int[_nlist][];
+ _listPqCodes = new byte[_nlist][][];
+ for (int list = 0; list < _nlist; list++) {
+ int listSize = buffer.getInt();
+ _listDocIds[list] = new int[listSize];
+ for (int i = 0; i < listSize; i++) {
+ _listDocIds[list][i] = buffer.getInt();
+ }
+ _listPqCodes[list] = new byte[listSize][_pqM];
+ for (int i = 0; i < listSize; i++) {
+ buffer.get(_listPqCodes[list][i]);
+ }
+ }
+
+ LOGGER.info("Loaded IVF_PQ index for column: {} with {} lists, {}
vectors, dimension={}, pqM={}, pqNbits={}",
+ column, _nlist, _numVectors, _dimension, _pqM, _pqNbits);
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to load IVF_PQ index for column: " +
column, e);
+ }
+ }
+
+ @Override
+ public MutableRoaringBitmap getDocIds(float[] queryVector, int topK) {
+ if (_numVectors == 0 || _nlist == 0) {
+ return new MutableRoaringBitmap();
+ }
+
+ // For COSINE and INNER_PRODUCT, normalize query to match the normalized
index vectors.
+ // L2 on unit-normalized vectors preserves both cosine and inner-product
ranking.
+ float[] searchVector = (_distanceFunctionCode ==
IvfPqIndexFormat.DIST_COSINE
Review Comment:
`getDocIds()` lacks basic argument validation that other vector readers
enforce (e.g., IVF_FLAT): it doesn’t check `topK > 0` or that
`queryVector.length == _dimension`. With a mismatched query length this will
throw (or compute incorrect residuals) and `topK <= 0` will create an invalid
heap size. Add explicit precondition checks to fail fast with a clear error.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/IvfPqVectorIndexReader.java:
##########
@@ -0,0 +1,292 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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 org.apache.pinot.segment.local.segment.index.readers.vector;
+
+import java.io.Closeable;
+import java.io.File;
+import java.io.IOException;
+import java.io.RandomAccessFile;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.PriorityQueue;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.IvfPqIndexFormat;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.IvfPqVectorIndexCreator;
+import org.apache.pinot.segment.local.segment.creator.impl.vector.pq.KMeans;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.ProductQuantizer;
+import
org.apache.pinot.segment.local.segment.creator.impl.vector.pq.VectorDistanceUtil;
+import org.apache.pinot.segment.spi.index.reader.NprobeAware;
+import org.apache.pinot.segment.spi.index.reader.VectorIndexReader;
+import org.apache.pinot.segment.spi.store.SegmentDirectoryPaths;
+import org.roaringbitmap.buffer.MutableRoaringBitmap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * Reader for IVF_PQ vector index. Loads the compact index (no original
vectors)
+ * and performs approximate nearest-neighbor search using IVF + Product
Quantization.
+ *
+ * <p>Exact rerank is handled by the operator layer via the forward index
reader,
+ * not by this reader. This keeps the index file compact (only PQ codes +
docIds).</p>
+ *
+ * <p>Search flow:
+ * <ol>
+ * <li>Find the nprobe closest coarse centroids to the query vector</li>
+ * <li>Score all candidates in probed lists using approximate PQ distances
(ADC)</li>
+ * <li>Return top-K approximate candidates as doc IDs</li>
+ * </ol>
+ */
+public class IvfPqVectorIndexReader implements VectorIndexReader, NprobeAware {
+ private static final Logger LOGGER =
LoggerFactory.getLogger(IvfPqVectorIndexReader.class);
+ private static final int DEFAULT_NPROBE = 8;
+
+ private final String _column;
+ private final int _dimension;
+ private final int _nlist;
+ private final int _pqM;
+ private final int _pqNbits;
+ private final int _numVectors;
+ private final int _distanceFunctionCode;
+
+ private final float[][] _coarseCentroids;
+ private final ProductQuantizer _pq;
+
+ // Inverted lists: docIds and PQ codes per list
+ private final int[][] _listDocIds;
+ private final byte[][][] _listPqCodes;
+
+ private volatile int _nprobe;
+ private final Closeable _closeHandle;
+
+ public IvfPqVectorIndexReader(String column, File segmentDir, int numDocs,
int nprobe) {
+ _column = column;
+ _nprobe = nprobe > 0 ? nprobe : DEFAULT_NPROBE;
+
+ File indexFile = findIndexFile(segmentDir);
+ if (indexFile == null) {
+ throw new IllegalStateException("Failed to find IVF_PQ index file for
column: " + column);
+ }
+
+ try {
+ RandomAccessFile raf = new RandomAccessFile(indexFile, "r");
+ MappedByteBuffer buffer =
raf.getChannel().map(FileChannel.MapMode.READ_ONLY, 0, raf.length());
+ buffer.order(ByteOrder.LITTLE_ENDIAN);
+ _closeHandle = raf;
+
+ // Read header
+ int magic = buffer.getInt();
+ if (magic != IvfPqIndexFormat.MAGIC) {
+ throw new IOException("Invalid IVF_PQ index magic: " +
Integer.toHexString(magic));
+ }
+ int version = buffer.getInt();
+ if (version != IvfPqIndexFormat.VERSION) {
+ throw new IOException("Unsupported IVF_PQ index version: " + version);
+ }
+ _dimension = buffer.getInt();
+ _nlist = buffer.getInt();
+ _pqM = buffer.getInt();
+ _pqNbits = buffer.getInt();
+ _distanceFunctionCode = buffer.getInt();
+ _numVectors = buffer.getInt();
+
+ int ksub = 1 << _pqNbits;
+ int dsub = _pqM > 0 ? _dimension / _pqM : 1;
+
+ if (_numVectors == 0 || _nlist == 0) {
+ _coarseCentroids = new float[0][];
+ int safeDim = Math.max(_dimension, 1);
+ int safeM = Math.max(_pqM, 1);
+ if (safeDim % safeM != 0) {
+ safeM = 1;
+ }
+ _pq = new ProductQuantizer(safeDim, safeM, Math.max(_pqNbits, 1));
+ _listDocIds = new int[0][];
+ _listPqCodes = new byte[0][][];
+ return;
+ }
+
+ // Read coarse centroids
+ _coarseCentroids = new float[_nlist][_dimension];
+ for (int c = 0; c < _nlist; c++) {
+ for (int d = 0; d < _dimension; d++) {
+ _coarseCentroids[c][d] = buffer.getFloat();
+ }
+ }
+
+ // Read PQ codebooks
+ float[][][] codebooks = new float[_pqM][ksub][dsub];
+ for (int sub = 0; sub < _pqM; sub++) {
+ for (int code = 0; code < ksub; code++) {
+ for (int d = 0; d < dsub; d++) {
+ codebooks[sub][code][d] = buffer.getFloat();
+ }
+ }
+ }
+ _pq = new ProductQuantizer(_dimension, _pqM, _pqNbits);
+ _pq.setCodebooks(codebooks);
+
+ // Read inverted lists (docIds + PQ codes only — no original vectors)
+ _listDocIds = new int[_nlist][];
+ _listPqCodes = new byte[_nlist][][];
+ for (int list = 0; list < _nlist; list++) {
+ int listSize = buffer.getInt();
+ _listDocIds[list] = new int[listSize];
+ for (int i = 0; i < listSize; i++) {
+ _listDocIds[list][i] = buffer.getInt();
+ }
+ _listPqCodes[list] = new byte[listSize][_pqM];
+ for (int i = 0; i < listSize; i++) {
+ buffer.get(_listPqCodes[list][i]);
+ }
+ }
+
+ LOGGER.info("Loaded IVF_PQ index for column: {} with {} lists, {}
vectors, dimension={}, pqM={}, pqNbits={}",
+ column, _nlist, _numVectors, _dimension, _pqM, _pqNbits);
+ } catch (IOException e) {
+ throw new RuntimeException("Failed to load IVF_PQ index for column: " +
column, e);
+ }
+ }
+
+ @Override
+ public MutableRoaringBitmap getDocIds(float[] queryVector, int topK) {
+ if (_numVectors == 0 || _nlist == 0) {
+ return new MutableRoaringBitmap();
+ }
+
+ // For COSINE and INNER_PRODUCT, normalize query to match the normalized
index vectors.
+ // L2 on unit-normalized vectors preserves both cosine and inner-product
ranking.
+ float[] searchVector = (_distanceFunctionCode ==
IvfPqIndexFormat.DIST_COSINE
+ || _distanceFunctionCode == IvfPqIndexFormat.DIST_INNER_PRODUCT)
+ ? VectorDistanceUtil.normalize(queryVector) : queryVector;
+
+ int effectiveNprobe = Math.min(_nprobe, _nlist);
+
+ // Step 1: Find closest coarse centroids
+ int[] probedLists = KMeans.findKNearest(searchVector, _coarseCentroids,
_dimension, effectiveNprobe);
+
+ // Step 2: Score candidates using PQ ADC distances
+ PriorityQueue<ScoredDoc> topKHeap = new PriorityQueue<>(topK + 1,
Review Comment:
`topK` is used directly to size the heap (`new PriorityQueue<>(topK + 1,
...)`) without clamping to the number of indexed vectors. For very large `topK`
this can allocate an excessively large heap even though at most `_numVectors`
results exist. Consider using `effectiveTopK = Math.min(topK, _numVectors)` (as
IVF_FLAT does) for sizing and loop conditions.
##########
pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/VectorIndexHandler.java:
##########
@@ -84,14 +99,27 @@ public boolean needUpdateIndices(SegmentDirectory.Reader
segmentReader) {
public void updateIndices(SegmentDirectory.Writer segmentWriter)
throws Exception {
Set<String> columnsToAddIdx = new HashSet<>(_vectorConfigs.keySet());
- // Remove indices not set in table config any more
+ // Remove indices not set in table config any more, or where backend
changed
String segmentName = _segmentDirectory.getSegmentMetadata().getName();
+ File indexDir = _segmentDirectory.getSegmentMetadata().getIndexDir();
Set<String> existingColumns =
segmentWriter.toSegmentDirectory().getColumnsWithIndex(StandardIndexes.vector());
for (String column : existingColumns) {
if (!columnsToAddIdx.remove(column)) {
LOGGER.info("Removing existing Vector index from segment: {}, column:
{}", segmentName, column);
segmentWriter.removeIndex(column, StandardIndexes.vector());
LOGGER.info("Removed existing Vector index from segment: {}, column:
{}", segmentName, column);
+ } else {
+ // Check if backend type changed — if so, remove old index and rebuild
+ VectorIndexConfig desiredConfig = _vectorConfigs.get(column);
+ String rawDesired = desiredConfig != null ?
desiredConfig.getVectorIndexType() : null;
+ String desiredBackend = (rawDesired == null || rawDesired.isEmpty()) ?
"HNSW" : rawDesired.toUpperCase();
+ String existingBackend =
VectorIndexUtils.detectVectorIndexBackend(indexDir, column);
+ if (existingBackend != null &&
!desiredBackend.equalsIgnoreCase(existingBackend)) {
+ LOGGER.info("Rebuilding Vector index for segment: {}, column: {}
(backend changed from {} to {})",
Review Comment:
Same issue in rebuild path: `desiredBackend` is derived from
`getVectorIndexType()` and can be null for default-HNSW configs, which prevents
detecting and rebuilding when an IVF_* backend exists on disk. Use the resolved
backend type (defaults included) for the comparison to avoid stale indexes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
