Skip to content

Commit e2ab4bc

Browse files
committed
Remove forbidden API
1 parent 96d2987 commit e2ab4bc

File tree

3 files changed

+32
-8
lines changed

3 files changed

+32
-8
lines changed

lucene/core/src/java/org/apache/lucene/search/AbstractKnnVectorQuery.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ private TopDocs searchLeaf(
116116
return results;
117117
}
118118

119-
protected TopDocs getLeafResults(
119+
private TopDocs getLeafResults(
120120
LeafReaderContext ctx,
121121
Weight filterWeight,
122122
TimeLimitingKnnCollectorManager timeLimitingKnnCollectorManager)

lucene/core/src/java/org/apache/lucene/search/TwoPhaseKnnVectorQuery.java

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
import org.apache.lucene.index.FieldInfo;
2323
import org.apache.lucene.index.FloatVectorValues;
2424
import org.apache.lucene.index.LeafReaderContext;
25+
import org.apache.lucene.search.knn.KnnCollectorManager;
2526
import org.apache.lucene.util.ArrayUtil;
27+
import org.apache.lucene.util.Bits;
2628

2729
public class TwoPhaseKnnVectorQuery extends KnnFloatVectorQuery {
2830

@@ -45,12 +47,14 @@ protected TopDocs mergeLeafResults(TopDocs[] perLeafResults) {
4547
}
4648

4749
@Override
48-
protected TopDocs getLeafResults(
50+
protected TopDocs approximateSearch(
4951
LeafReaderContext context,
50-
Weight filterWeight,
51-
TimeLimitingKnnCollectorManager knnCollectorManager)
52+
Bits acceptDocs,
53+
int visitedLimit,
54+
KnnCollectorManager knnCollectorManager)
5255
throws IOException {
53-
TopDocs results = super.getLeafResults(context, filterWeight, knnCollectorManager);
56+
TopDocs results =
57+
super.approximateSearch(context, acceptDocs, visitedLimit, knnCollectorManager);
5458
if (results.scoreDocs.length <= originalK) {
5559
// short-circuit: no re-ranking needed. we got what we need
5660
return results;

lucene/core/src/test/org/apache/lucene/search/TestTwoPhaseKnnVectorQuery.java

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,19 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
117
package org.apache.lucene.search;
218

319
import java.util.HashMap;
@@ -18,11 +34,12 @@
1834
import org.apache.lucene.index.VectorSimilarityFunction;
1935
import org.apache.lucene.store.ByteBuffersDirectory;
2036
import org.apache.lucene.store.Directory;
37+
import org.apache.lucene.tests.util.LuceneTestCase;
2138
import org.junit.Assert;
2239
import org.junit.Before;
2340
import org.junit.Test;
2441

25-
public class TestTwoPhaseKnnVectorQuery {
42+
public class TestTwoPhaseKnnVectorQuery extends LuceneTestCase {
2643

2744
private static final String FIELD = "vector";
2845
public static final VectorSimilarityFunction VECTOR_SIMILARITY_FUNCTION =
@@ -33,7 +50,9 @@ public class TestTwoPhaseKnnVectorQuery {
3350
private static final int VECTOR_DIMENSION = 128;
3451

3552
@Before
53+
@Override
3654
public void setUp() throws Exception {
55+
super.setUp();
3756
directory = new ByteBuffersDirectory();
3857

3958
// Set up the IndexWriterConfig to use quantized vector storage
@@ -45,9 +64,10 @@ public void setUp() throws Exception {
4564
public void testTwoPhaseKnnVectorQuery() throws Exception {
4665
Map<Integer, float[]> vectors = new HashMap<>();
4766

67+
Random random = random();
68+
4869
// Step 1: Index random vectors in quantized format
4970
try (IndexWriter writer = new IndexWriter(directory, config)) {
50-
Random random = new Random();
5171
for (int i = 0; i < NUM_VECTORS; i++) {
5272
float[] vector = randomFloatVector(VECTOR_DIMENSION, random);
5373
Document doc = new Document();
@@ -61,7 +81,7 @@ public void testTwoPhaseKnnVectorQuery() throws Exception {
6181
// Step 2: Run TwoPhaseKnnVectorQuery with a random target vector
6282
try (IndexReader reader = DirectoryReader.open(directory)) {
6383
IndexSearcher searcher = new IndexSearcher(reader);
64-
float[] targetVector = randomFloatVector(VECTOR_DIMENSION, new Random());
84+
float[] targetVector = randomFloatVector(VECTOR_DIMENSION, random);
6585
int k = 10;
6686
double oversample = 1.0;
6787

0 commit comments

Comments
 (0)