xiangfu0 commented on code in PR #19303: URL: https://github.com/apache/pinot/pull/19303#discussion_r3910090660
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/vector/BasePinotDocIdBitmapFilterQuery.java: ########## @@ -0,0 +1,148 @@ +/** + * 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.IOException; +import javax.annotation.Nullable; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.ConstantScoreWeight; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.Weight; +import org.roaringbitmap.buffer.ImmutableRoaringBitmap; + + +/// Base class for Lucene [Query] implementations that accept only documents whose Pinot doc id is present +/// in a [ImmutableRoaringBitmap]. Used to implement pre-filter ANN search by restricting HNSW graph +/// traversal to the filtered document set. +/// +/// Because Lucene uses its own internal doc ids (which differ from Pinot doc ids), subclasses supply the +/// per-leaf iterator that maps Lucene doc ids to Pinot doc ids before testing membership in the bitmap +/// (via a doc-id translator, doc values, etc.). This class owns the constant-score weight/scorer +/// scaffolding, identity-based equality, and cache opt-out, so filter-correctness fixes apply to every +/// implementation at once. +/// +/// Instances are single-use per search and must never be cached by Lucene ([Weight#isCacheable] returns +/// false), since the accepted docs depend on the bitmap instance. +/// +/// **Bitmap ownership.** The bitmap is retained by reference, not copied. [ImmutableRoaringBitmap] only +/// promises that *this* type exposes no mutators -- a caller may pass a `MutableRoaringBitmap`, which is a +/// subtype -- so the caller must not modify it once it has been handed over. Mutating it during a search +/// changes which documents are accepted midway through traversal and yields results matching neither the old +/// nor the new set. Callers that cannot promise that must pass a detached copy. +/// +/// Given that, instances are safe to share across the threads of a single search: the bitmap is only read, +/// and each leaf gets its own iterator. +public abstract class BasePinotDocIdBitmapFilterQuery extends Query { Review Comment: Renamed the shared query and its test back to `BaseFilterQuery`. ########## pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/HnswVectorRealtimeTest.java: ########## @@ -0,0 +1,236 @@ +/** + * 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.integration.tests.custom; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.File; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Random; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.commons.lang3.StringUtils; +import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/// Realtime integration coverage for filtered HNSW search on a **consuming** segment, where the vector index is +/// `MutableVectorIndex` rather than the offline reader. +/// +/// This pins a behavior change that reaches plain realtime tables, not only upsert ones: because the consuming +/// segment now advertises `supportsPreFilter()`, a `vectorSimilarity` predicate combined with a metadata filter +/// plans as `FILTER_THEN_ANN` instead of running an unfiltered ANN and intersecting afterwards. Nothing else +/// covers this -- the other HNSW integration tests are offline-only, and [IvfPqVectorRealtimeTest] asserts the +/// IVF_PQ exact-scan fallback. +/// +/// ## Why the fixture is sized the way it is +/// +/// `VectorSearchStrategy.decide` only wires the optional pre-filter when the filter matches at least +/// `EXACT_SCAN_THRESHOLD` (1000) documents *and* selectivity falls below the mid-range cutoff of 0.105; anything +/// more selective is cheaper as an exact scan, anything less selective is left to post-filtering. A single +/// consuming segment of [#getCountStarResult] rows over [#NUM_CATEGORIES] categories yields 1250 matches at +/// selectivity 0.083, which clears both bounds with margin. One Kafka partition and a flush size above the row +/// count keep every row in one consuming segment, so the per-segment counts the planner sees are the ones +/// computed here. The category column carries an inverted index because the wiring additionally requires every +/// non-vector filter to produce a bitmap -- a scan-based predicate silently leaves the vector operator +/// post-filtering. +@Test(suiteName = "CustomClusterIntegrationTest") +public class HnswVectorRealtimeTest extends CustomDataQueryClusterIntegrationTest { + private static final String DEFAULT_TABLE_NAME = "HnswVectorRealtimeTest"; + private static final String VECTOR_COL = "embedding"; + private static final String CATEGORY = "category"; + private static final int VECTOR_DIM_SIZE = 32; + private static final int NUM_CATEGORIES = 12; + private static final int NUM_ROWS = 15000; + private static final String TARGET_CATEGORY = "cat_3"; + + @Override + protected long getCountStarResult() { + return NUM_ROWS; + } + + @Override + public String getTableName() { + return DEFAULT_TABLE_NAME; + } + + @Override + public boolean isRealtimeTable() { + return true; + } + + /// Keep every row in a single consuming segment: the planner reasons about per-segment counts, so a mid-stream + /// commit would shrink them below the pre-filter thresholds this test depends on. + @Override + protected int getRealtimeSegmentFlushSize() { + return NUM_ROWS * 10; + } + + @Override + protected int getNumKafkaPartitions() { + return 1; + } + + @Override + public Schema createSchema() { + return new Schema.SchemaBuilder().setSchemaName(getTableName()) + .addMultiValueDimension(VECTOR_COL, FieldSpec.DataType.FLOAT) Review Comment: Imported `FieldSpec.DataType` and now use the `DataType` constants directly. ########## pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/readers/vector/BasePinotDocIdBitmapFilterQueryTest.java: ########## @@ -0,0 +1,124 @@ +/** + * 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.util.concurrent.atomic.AtomicReference; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.QueryVisitor; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.Weight; +import org.roaringbitmap.buffer.ImmutableRoaringBitmap; +import org.roaringbitmap.buffer.MutableRoaringBitmap; +import org.testng.Assert; +import org.testng.annotations.Test; + + +/// Tests the shared Lucene constant-score scaffolding and identity contract for Pinot doc-id bitmap filters. +public class BasePinotDocIdBitmapFilterQueryTest { + + @Test + public void testConstantScoreScaffoldingAndCacheOptOut() + throws Exception { + TestQuery query = new TestQuery(MutableRoaringBitmap.bitmapOf(2, 4), new Object(), + DocIdSetIterator.range(1, 3)); + Weight weight = query.createWeight(null, ScoreMode.COMPLETE, 2.5F); + + Assert.assertFalse(weight.isCacheable(null)); Review Comment: Converted the new query test to static TestNG assertion imports. -- 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]
