xiangfu0 commented on code in PR #17994: URL: https://github.com/apache/pinot/pull/17994#discussion_r3004786918
########## 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: Acknowledged — L2-only for phase 1. Multi-distance exact scan tracked for phase 2. ########## 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: Fixed — added Preconditions.checkArgument for dimension match and iterate up to _dimension instead of values.length. ########## 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: Already logs a WARN. Phase 1 scope is immutable only. Mutable IVF_FLAT tracked for phase 2. ########## 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: Valid edge case. In practice, table config changes trigger segment reload which creates a fresh reader. Transitioning backends requires segment rebuild (re-ingestion). This is a known limitation. ########## 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: Safe in practice — ReaderFactory dispatches by VectorBackendType before reader construction. HNSW reader finds its own files first in the fallback chain. -- 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]
