Copilot commented on code in PR #17994: URL: https://github.com/apache/pinot/pull/17994#discussion_r3004689619
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/vector/IvfFlatVectorIndexCreator.java: ########## @@ -0,0 +1,561 @@ +/** + * 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.vector; + +import com.google.common.base.Preconditions; +import java.io.BufferedOutputStream; +import java.io.DataOutputStream; +import java.io.File; +import java.io.FileOutputStream; +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Random; +import javax.annotation.Nullable; +import org.apache.pinot.common.function.scalar.VectorFunctions; +import org.apache.pinot.segment.spi.V1Constants; +import org.apache.pinot.segment.spi.index.creator.VectorIndexConfig; +import org.apache.pinot.segment.spi.index.creator.VectorIndexCreator; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Creates an IVF_FLAT (Inverted File with flat vectors) index for immutable segments. + * + * <p>The creator buffers all vectors in memory during {@link #add(float[])} calls, then + * trains k-means centroids, assigns vectors to their nearest centroids, and serializes + * the complete index to a single {@code .ivfflat.index} file during {@link #seal()}.</p> + * + * <h3>Thread safety</h3> + * <p>This class is NOT thread-safe. It is designed for single-threaded segment creation.</p> + * + * <h3>File format (version 1)</h3> + * <pre> + * [Header] + * magic: 4 bytes (0x49564646 = "IVFF") + * version: 4 bytes (1) + * dimension: 4 bytes + * numVectors: 4 bytes + * nlist: 4 bytes + * distanceFunctionOrd: 4 bytes + * + * [Centroids Section] + * nlist x dimension x 4 bytes (float32) + * + * [Inverted Lists Section] + * For each centroid i (0..nlist-1): + * listSize_i: 4 bytes + * docIds_i: listSize_i x 4 bytes (int32) + * vectors_i: listSize_i x dimension x 4 bytes (float32) + * + * [Inverted List Offsets] + * nlist x 8 bytes (long offset to start of each inverted list) + * + * [Footer] + * offsetToOffsets: 8 bytes (position of the offsets section) + * </pre> + * + * <p>All multi-byte values are written in big-endian order (Java {@link DataOutputStream} default).</p> + */ +public class IvfFlatVectorIndexCreator implements VectorIndexCreator { + private static final Logger LOGGER = LoggerFactory.getLogger(IvfFlatVectorIndexCreator.class); + + /** Magic bytes identifying an IVF_FLAT index file: ASCII "IVFF". */ + public static final int MAGIC = 0x49564646; + + /** Current file format version. */ + public static final int FORMAT_VERSION = 1; + + /** Default number of Voronoi cells (centroids). */ + public static final int DEFAULT_NLIST = 128; + + /** Maximum number of k-means iterations. */ + static final int MAX_KMEANS_ITERATIONS = 50; + + /** Convergence threshold: stop when centroid movement is below this fraction. */ + static final float CONVERGENCE_THRESHOLD = 1e-5f; + + /** Default training sample size multiplier relative to nlist. */ + static final int DEFAULT_TRAIN_SAMPLE_MULTIPLIER = 40; + + /** Minimum training sample size. */ + static final int DEFAULT_MIN_TRAIN_SAMPLE_SIZE = 10000; + + private final String _column; + private final File _indexDir; + private final int _dimension; + private final int _nlist; + private final int _trainSampleSize; + private final long _trainingSeed; + private final VectorIndexConfig.VectorDistanceFunction _distanceFunction; + + /** All vectors collected during add(), indexed by docId (ordinal). */ + private final List<float[]> _vectors = new ArrayList<>(); + + private boolean _sealed = false; + + /** + * Creates a new IVF_FLAT index creator. + * + * @param column the column name + * @param indexDir the segment index directory + * @param config the vector index configuration + */ + public IvfFlatVectorIndexCreator(String column, File indexDir, VectorIndexConfig config) { + _column = column; + _indexDir = indexDir; + _dimension = config.getVectorDimension(); + _distanceFunction = config.getVectorDistanceFunction(); + + Map<String, String> properties = config.getProperties(); + _nlist = properties != null && properties.containsKey("nlist") + ? Integer.parseInt(properties.get("nlist")) + : DEFAULT_NLIST; + _trainSampleSize = properties != null && properties.containsKey("trainSampleSize") + ? Integer.parseInt(properties.get("trainSampleSize")) + : Math.max(_nlist * DEFAULT_TRAIN_SAMPLE_MULTIPLIER, DEFAULT_MIN_TRAIN_SAMPLE_SIZE); + _trainingSeed = properties != null && properties.containsKey("trainingSeed") + ? Long.parseLong(properties.get("trainingSeed")) + : System.nanoTime(); + + Preconditions.checkArgument(_dimension > 0, "Vector dimension must be positive, got: %s", _dimension); + Preconditions.checkArgument(_nlist > 0, "nlist must be positive, got: %s", _nlist); + + LOGGER.info("Creating IVF_FLAT index for column: {} in dir: {}, dimension={}, nlist={}, distance={}", + column, indexDir.getAbsolutePath(), _dimension, _nlist, _distanceFunction); + } + + @Override + public void add(Object[] values, @Nullable int[] dictIds) { + // The segment builder calls this overload for multi-value columns. + // Convert Object[] (boxed Floats) to float[] and delegate to add(float[]). + float[] floatValues = new float[_dimension]; + for (int i = 0; i < values.length; i++) { + floatValues[i] = (Float) values[i]; + } + add(floatValues); Review Comment: add(Object[] values, ...) copies values into a float[] of size _dimension but iterates up to values.length. If values.length > _dimension this will throw ArrayIndexOutOfBoundsException; if values.length < _dimension it silently pads with zeros. Add an explicit length check (values.length == _dimension) and fail with a clear IllegalArgumentException to avoid corrupt index data or unexpected runtime errors. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/vector/VectorIndexType.java: ########## @@ -141,12 +173,20 @@ public VectorIndexReader createIndexReader(SegmentDirectory.Reader segmentReader throws IndexReaderConstraintException { if (metadata.getDataType() != FieldSpec.DataType.FLOAT || metadata.getFieldSpec().isSingleValueField()) { throw new IndexReaderConstraintException(metadata.getColumnName(), StandardIndexes.vector(), - "HNSW Vector index is currently only supported on float array type columns"); + "Vector index is currently only supported on float array type columns"); } File segmentDir = segmentReader.toSegmentDirectory().getPath().toFile(); - VectorIndexConfig indexConfig = fieldIndexConfigs.getConfig(StandardIndexes.vector()); - return new HnswVectorIndexReader(metadata.getColumnName(), segmentDir, metadata.getTotalDocs(), indexConfig); + VectorBackendType backendType = indexConfig.resolveBackendType(); + + switch (backendType) { + case HNSW: + return new HnswVectorIndexReader(metadata.getColumnName(), segmentDir, metadata.getTotalDocs(), indexConfig); + case IVF_FLAT: + return new IvfFlatVectorIndexReader(metadata.getColumnName(), segmentDir, indexConfig); + default: + throw new IllegalStateException("Unsupported vector backend type: " + backendType); Review Comment: ReaderFactory chooses the reader implementation solely based on the current VectorIndexConfig (table config). If a table transitions from HNSW to IVF_FLAT (or vice versa), older segments will still have the previous on-disk format but segmentReader.hasIndexFor(...) will return true (because VectorIndexType now advertises both extensions). In that scenario, this factory will instantiate the wrong reader and segment load/query will fail. Consider selecting the backend by checking which on-disk index file(s) actually exist for the segment/column (and only then validating that the backend is supported), so mixed-backend segments can coexist during config transitions. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/vector/VectorIndexType.java: ########## @@ -161,10 +201,19 @@ public MutableIndex createMutableIndex(MutableIndexContext context, VectorIndexC return null; } - return new MutableVectorIndex(context.getSegmentName(), context.getFieldSpec().getName(), config); - } - - public enum IndexType { - HNSW + VectorBackendType backendType = config.resolveBackendType(); + switch (backendType) { + case HNSW: + return new MutableVectorIndex(context.getSegmentName(), context.getFieldSpec().getName(), config); + case IVF_FLAT: + // IVF_FLAT does not support mutable indexes in phase 1. + LOGGER.warn("IVF_FLAT vector index does not support mutable/realtime segments. " + + "No vector index will be built for column: {} in segment: {}. " + + "Queries will fall back to exact scan.", + context.getFieldSpec().getName(), context.getSegmentName()); + return null; Review Comment: For realtime/mutable segments, IVF_FLAT currently logs a warning and returns null (no index built), which can silently route queries to the expensive exact-scan fallback. Since IVF_FLAT is explicitly unsupported for mutable segments in phase 1, consider failing fast during validation when the table type is REALTIME (or when createMutableIndex is invoked) so misconfiguration is caught early and doesn’t degrade query latency unexpectedly. ```suggestion // IVF_FLAT does not support mutable indexes in phase 1; fail fast to surface misconfiguration. throw new IllegalStateException( "IVF_FLAT vector index is not supported for mutable/realtime segments. " + "Cannot build vector index for column: " + context.getFieldSpec().getName() + " in segment: " + context.getSegmentName()); ``` ########## pinot-core/src/main/java/org/apache/pinot/core/operator/filter/ExactVectorScanFilterOperator.java: ########## @@ -0,0 +1,223 @@ +/** + * 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.core.operator.filter; + +import com.google.common.base.CaseFormat; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.PriorityQueue; +import org.apache.pinot.common.function.scalar.VectorFunctions; +import org.apache.pinot.common.request.context.predicate.VectorSimilarityPredicate; +import org.apache.pinot.core.common.BlockDocIdSet; +import org.apache.pinot.core.common.Operator; +import org.apache.pinot.core.operator.ExplainAttributeBuilder; +import org.apache.pinot.core.operator.docidsets.BitmapDocIdSet; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; +import org.apache.pinot.segment.spi.index.reader.ForwardIndexReaderContext; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.trace.FilterType; +import org.apache.pinot.spi.trace.InvocationRecording; +import org.apache.pinot.spi.trace.Tracing; +import org.roaringbitmap.buffer.ImmutableRoaringBitmap; +import org.roaringbitmap.buffer.MutableRoaringBitmap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Fallback operator that performs exact brute-force vector similarity search by scanning the forward index. + * + * <p>This operator is used when no ANN vector index exists on a segment for the target column + * (e.g., the segment was built before the vector index was added, or the index type is not + * supported). It reads all vectors from the forward index, computes exact distances to the + * query vector, and returns the top-K closest document IDs.</p> + * + * <p>The distance computation uses L2 (Euclidean) squared distance. For COSINE similarity, + * vectors should be pre-normalized. This matches the behavior of Lucene's HNSW implementation.</p> + * + * <p>This operator is intentionally simple and correct rather than fast -- it is a safety net. + * A warning is logged when this operator is used because it scans all documents in the segment.</p> + * + * <p>This class is thread-safe for single-threaded execution per query (same as other filter operators).</p> + */ +public class ExactVectorScanFilterOperator extends BaseFilterOperator { + private static final Logger LOGGER = LoggerFactory.getLogger(ExactVectorScanFilterOperator.class); + private static final String EXPLAIN_NAME = "VECTOR_SIMILARITY_EXACT_SCAN"; + + private final ForwardIndexReader<?> _forwardIndexReader; + private final VectorSimilarityPredicate _predicate; + private final String _column; + private ImmutableRoaringBitmap _matches; + + /** + * Creates an exact scan operator. + * + * @param forwardIndexReader the forward index reader for the vector column + * @param predicate the vector similarity predicate containing query vector and top-K + * @param column the column name (for logging and explain) + * @param numDocs the total number of documents in the segment + */ + public ExactVectorScanFilterOperator(ForwardIndexReader<?> forwardIndexReader, + VectorSimilarityPredicate predicate, String column, int numDocs) { + super(numDocs, false); + _forwardIndexReader = forwardIndexReader; + _predicate = predicate; + _column = column; + } + + @Override + protected BlockDocIdSet getTrues() { + if (_matches == null) { + _matches = computeExactTopK(); + } + return new BitmapDocIdSet(_matches, _numDocs); + } + + @Override + public int getNumMatchingDocs() { + if (_matches == null) { + _matches = computeExactTopK(); + } + return _matches.getCardinality(); + } + + @Override + public boolean canProduceBitmaps() { + return true; + } + + @Override + public BitmapCollection getBitmaps() { + if (_matches == null) { + _matches = computeExactTopK(); + } + record(_matches); + return new BitmapCollection(_numDocs, false, _matches); + } + + @Override + public List<Operator> getChildOperators() { + return Collections.emptyList(); + } + + @Override + public String toExplainString() { + return EXPLAIN_NAME + "(indexLookUp:exact_scan" + + ", operator:" + _predicate.getType() + + ", vector identifier:" + _column + + ", vector literal:" + Arrays.toString(_predicate.getValue()) + + ", topK to search:" + _predicate.getTopK() + + ')'; + } + + @Override + protected String getExplainName() { + return CaseFormat.UPPER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, EXPLAIN_NAME); + } + + @Override + protected void explainAttributes(ExplainAttributeBuilder attributeBuilder) { + super.explainAttributes(attributeBuilder); + attributeBuilder.putString("indexLookUp", "exact_scan"); + attributeBuilder.putString("operator", _predicate.getType().name()); + attributeBuilder.putString("vectorIdentifier", _column); + attributeBuilder.putString("vectorLiteral", Arrays.toString(_predicate.getValue())); + attributeBuilder.putLongIdempotent("topKtoSearch", _predicate.getTopK()); + } + + /** + * Performs brute-force exact search over all documents in the segment. + * Uses a max-heap to maintain the top-K closest vectors. + */ + @SuppressWarnings("unchecked") + private ImmutableRoaringBitmap computeExactTopK() { + LOGGER.warn("Performing exact vector scan fallback on column: {} for segment with {} docs. " + + "This is expensive -- consider adding a vector index.", _column, _numDocs); + + float[] queryVector = _predicate.getValue(); + int topK = _predicate.getTopK(); + + // Max-heap: entry with largest distance is at the top so we can efficiently evict it + PriorityQueue<DocDistance> maxHeap = new PriorityQueue<>(topK + 1, + (a, b) -> Float.compare(b._distance, a._distance)); + + ForwardIndexReader rawReader = _forwardIndexReader; + try (ForwardIndexReaderContext context = rawReader.createContext()) { + for (int docId = 0; docId < _numDocs; docId++) { + float[] docVector = rawReader.getFloatMV(docId, context); + if (docVector == null || docVector.length == 0) { + continue; + } + float distance = computeL2SquaredDistance(queryVector, docVector); + if (maxHeap.size() < topK) { + maxHeap.add(new DocDistance(docId, distance)); + } else if (distance < maxHeap.peek()._distance) { + maxHeap.poll(); + maxHeap.add(new DocDistance(docId, distance)); + } Review Comment: ExactVectorScanFilterOperator always ranks candidates using squared L2 distance. This makes the fallback results incorrect for tables configured with COSINE, DOT_PRODUCT, or INNER_PRODUCT vector distance functions (the same query can yield different docIds depending on whether a segment has an ANN index). Consider passing the configured VectorDistanceFunction into this operator (constructed from the vector index config/table config) and computing the corresponding exact distance here. ########## pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/store/SegmentDirectoryPaths.java: ########## @@ -156,6 +156,12 @@ public static File findVectorIndexIndexFile(File segmentIndexDir, String column) vectorIndexDirectory = column + V1Constants.Indexes.VECTOR_HNSW_INDEX_FILE_EXTENSION; formatFile = findFormatFile(segmentIndexDir, vectorIndexDirectory); } + + // check for IVF_FLAT index, if null + if (formatFile == null) { + String ivfFlatFile = column + V1Constants.Indexes.VECTOR_IVF_FLAT_INDEX_FILE_EXTENSION; + formatFile = findFormatFile(segmentIndexDir, ivfFlatFile); + } Review Comment: SegmentDirectoryPaths.findVectorIndexIndexFile now falls back to returning the IVF_FLAT index file when no HNSW index exists. This method is also used by HnswVectorIndexReader, which expects a Lucene directory; if a segment has only an IVF_FLAT file (e.g., after changing table config across time, or mixed segments), HNSW reader creation will fail. Consider splitting this into backend-specific helpers (e.g., findHnswVectorIndexFile / findIvfFlatVectorIndexFile) and/or making VectorIndexType.ReaderFactory detect which on-disk backend exists and instantiate the matching reader, rather than returning an arbitrary vector index file. -- 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]
