alessandrobenedetti commented on code in PR #4919: URL: https://github.com/apache/solr/pull/4919#discussion_r4153626313
########## solr/core/src/test/org/apache/solr/search/vector/KnnQParserOversampleRerankTest.java: ########## @@ -0,0 +1,164 @@ +/* + * 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.solr.search.vector; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.solr.SolrTestCaseJ4; +import org.apache.solr.common.SolrInputDocument; +import org.apache.solr.common.params.CommonParams; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Tests the {@code rerankOversample} local param of the {@code knn} query parser against a + * quantized dense vector field, where the HNSW search ranks candidates using lossy quantized + * vectors and the re-ranking phase re-scores them against the raw full precision vectors. + */ +public class KnnQParserOversampleRerankTest extends SolrTestCaseJ4 { + + private static final String IDField = "id"; + + /** Not quantized: the knn search already scores against the raw vectors. */ + private static final String exactField = "vector"; + + /** 4 bit scalar quantized: the knn search scores against lossy quantized vectors. */ + private static final String quantizedField = "v_scalar_half_byte"; + + /** + * The 5 nearest neighbours of {@code [1.0, 2.0, 3.0, 4.0]} by exact cosine similarity. The top + * four are separated by less than 0.003, so they are easily reordered by quantization. + */ + private static final String[] EXPECTED_EXACT_TOP_5 = + new String[] { + "//result[@numFound='5']", + "//result/doc[1]/str[@name='id'][.='1']", + "//result/doc[2]/str[@name='id'][.='4']", + "//result/doc[3]/str[@name='id'][.='2']", + "//result/doc[4]/str[@name='id'][.='10']", + "//result/doc[5]/str[@name='id'][.='3']" + }; + + @Before + public void prepareIndex() throws Exception { + initCore("solrconfig_codec.xml", "schema-densevector-quantized.xml"); + + for (SolrInputDocument doc : prepareDocs()) { + assertU(adoc(doc)); + } + assertU(commit()); + } + + /** Indexes the same 10 vectors into both an exact and a quantized field. */ + private List<SolrInputDocument> prepareDocs() { + List<List<Float>> vectors = + List.of( + Arrays.asList(1f, 2f, 3f, 4f), // id 1, cosine = 1.0 + Arrays.asList(1.5f, 2.5f, 3.5f, 4.5f), // id 2, cosine = 0.998 + Arrays.asList(7.5f, 15.5f, 17.5f, 22.5f), // id 3, cosine = 0.992 + Arrays.asList(1.4f, 2.4f, 3.4f, 4.4f), // id 4, cosine = 0.999 + Arrays.asList(30f, 22f, 35f, 20f), // id 5, cosine = 0.862 + Arrays.asList(40f, 1f, 1f, 200f), // id 6, cosine = 0.756 + Arrays.asList(5f, 10f, 20f, 40f), // id 7, cosine = 0.970 + Arrays.asList(120f, 60f, 30f, 15f), // id 8, cosine = 0.515 + Arrays.asList(200f, 50f, 100f, 25f), // id 9, cosine = 0.554 + Arrays.asList(1.8f, 2.5f, 3.7f, 4.9f)); //id 10, cosine = 0.997 + + List<SolrInputDocument> docs = new ArrayList<>(vectors.size()); + for (int i = 0; i < vectors.size(); i++) { + SolrInputDocument doc = new SolrInputDocument(); + doc.addField(IDField, i + 1); + doc.addField(exactField, vectors.get(i)); + doc.addField(quantizedField, vectors.get(i)); + docs.add(doc); + } + return docs; + } + + @After + public void cleanUp() { + clearIndex(); + deleteCore(); + } + + @Test + public void exactField_isTheReferenceRanking() { + // sanity check: the un-quantized field produces the exact cosine ranking + assertQ( + req(CommonParams.Q, "{!knn f=" + exactField + " topK=5}[1.0, 2.0, 3.0, 4.0]", "fl", "id"), + EXPECTED_EXACT_TOP_5); + } + + @Test + public void rerankOversampledQuantizedSearch_shouldRestoreExactRanking() { + // topK * rerankOversample = 25 candidates covers all 10 documents, so every document is + // re-scored against its raw vector: the ranking must match the exact one, whatever the + // quantized vectors ranked them as + assertQ( + req( + CommonParams.Q, + "{!knn f=" + quantizedField + " topK=5 rerankOversample=5}[1.0, 2.0, 3.0, 4.0]", + "fl", + "id"), + EXPECTED_EXACT_TOP_5); Review Comment: should we add here an assertion with the same query but no oversampling? showing the difference? ########## solr/core/src/java/org/apache/solr/search/vector/KnnQParser.java: ########## @@ -161,19 +187,36 @@ public Query parse() throws SyntaxError { req, subQuery(allParentsQuery, null).getQuery()); final BooleanQuery acceptedParents = getParentsFilter(parentsFilterQueries); + denseVectorType.checkRerankOversampleSupported(vectorField, rerankOversample); + Query acceptedChildren = getChildrenFilter(getFilterQuery(), acceptedParents, allParentsBitSet); switch (vectorEncoding) { case FLOAT32: - return new DiversifyingChildrenFloatKnnVectorQuery( + // The diversifying query returns the best matching child per parent, so collecting + // candidateTopK of them and re-ranking down to topK only ever narrows an already + // diversified set: at most one child per parent is preserved. Note that which child + // represents a parent is still picked using the (possibly quantized) approximate score, + // re-ranking only reorders the representatives that were chosen. + final float[] target = vectorBuilder.getFloatVector(); + final Query diversified = + new DiversifyingChildrenFloatKnnVectorQuery( + vectorField, target, acceptedChildren, candidateTopK, allParentsBitSet); + if (rerankOversample <= 1) { + return diversified; + } + return new SolrRescoreTopNQuery( + diversified, + new FullPrecisionFloatVectorSimilarityValuesSource( + target, vectorField, denseVectorType.getSimilarityFunction()), + topK); + case BYTE: + return new DiversifyingChildrenByteKnnVectorQuery( vectorField, - vectorBuilder.getFloatVector(), + vectorBuilder.getByteVector(), acceptedChildren, - topK, + candidateTopK, Review Comment: here candidateTopK could be misleading as BYTE encoding won't be supported in this contribution? ########## solr/solr-ref-guide/modules/query-guide/pages/dense-vector-search.adoc: ########## @@ -481,6 +483,20 @@ Here's an example of a `knn` search using the early termination with input param [source,text] ?q={!knn f=vector topK=10 earlyTermination=true saturationThreshold=0.989 patience=10 efSearchScaleFactor=3.0}[1.0, 2.0, 3.0, 4.0] +`rerankOversample`:: ++ +[%autowidth,frame=none] +|=== +|Optional |Default: 1 +|=== ++ +If provided, the query will retrieve a candidate pool of documents equal to topK multiplied by the rerankOversample value provided. The candidate set of documents will then be rescored with their raw vector values, and reranked down to a topK result set. Review Comment: I would specify at the moment that is compatible only with scalar quantised or binary quantised vector fields ########## solr/core/src/test/org/apache/solr/search/vector/KnnQParserTest.java: ########## @@ -197,6 +197,270 @@ public void efSearchScaleFactorSet_shouldWorkCorrectly() { "//result/doc[5]/str[@name='id'][.='3']"); } + @Test + public void incorrectOversample_shouldThrowException() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQEx( + "String rerankOversample should throw Exception", + "For input string: \"string\"", + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=string}" + vectorToSearch, + "fl", + "id"), + SolrException.ErrorCode.BAD_REQUEST); + + assertQEx( + "Double rerankOversample should throw Exception", + "For input string: \"2.5\"", + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=2.5}" + vectorToSearch, + "fl", + "id"), + SolrException.ErrorCode.BAD_REQUEST); + } + + @Test + public void rerankOversampleLessThanOne_shouldThrowException() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQEx( + "rerankOversample = 0 should throw Exception", + "rerankOversample (0) must be >= 1", + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=0}" + vectorToSearch, + "fl", + "id"), + SolrException.ErrorCode.BAD_REQUEST); + + assertQEx( + "Negative rerankOversample should throw Exception", + "rerankOversample (-1) must be >= 1", + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=-1}" + vectorToSearch, + "fl", + "id"), + SolrException.ErrorCode.BAD_REQUEST); + } + + @Test + public void rerankOversampleOverflowingTopK_shouldThrowException() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQEx( + "topK * rerankOversample overflowing an integer should throw Exception", + "topK (2000000000) * rerankOversample (3) overflows an integer", + req( + CommonParams.Q, + "{!knn f=vector topK=2000000000 rerankOversample=3}" + vectorToSearch, + "fl", + "id"), + SolrException.ErrorCode.BAD_REQUEST); + } + + @Test + public void rerankOversampleOnByteEncodedField_shouldThrowException() { + String vectorToSearch = "[1, 2, 3, 4]"; + + assertQEx( + "rerankOversample on a BYTE encoded field should throw Exception", + "rerankOversample is only supported for FLOAT32 vector encoding; field 'vector_byte_encoding' uses BYTE", + req( + CommonParams.Q, + "{!knn f=vector_byte_encoding topK=3 rerankOversample=2}" + vectorToSearch, + "fl", + "id"), + SolrException.ErrorCode.BAD_REQUEST); + } + + @Test + public void rerankOversampleOneOnByteEncodedField_shouldNotThrow() { + String vectorToSearch = "[1, 2, 3, 4]"; + + assertQ( + req( + CommonParams.Q, + "{!knn f=vector_byte_encoding topK=3 rerankOversample=1}" + vectorToSearch, + "fl", + "id"), + "//result[@numFound='3']"); + } + + @Test + public void rerankOversampleSet_shouldReturnTopKResults() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + // rerankOversample widens the candidate pool, but exactly topK results are returned + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=3}" + vectorToSearch, + "fl", + "id"), + "//result[@numFound='5']", + "//result/doc[1]/str[@name='id'][.='1']", + "//result/doc[2]/str[@name='id'][.='4']", + "//result/doc[3]/str[@name='id'][.='2']", + "//result/doc[4]/str[@name='id'][.='10']", + "//result/doc[5]/str[@name='id'][.='3']"); + } + + @Test + public void rerankOversampleOnNonQuantizedField_shouldNotChangeRanking() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + // 'vector' is not quantized, so the knn search already scores against the raw vectors and + // re-ranking them cannot reorder anything: oversampling must be a no-op on the final ranking + String[] expected = + new String[] { + "//result[@numFound='5']", + "//result/doc[1]/str[@name='id'][.='1']", + "//result/doc[2]/str[@name='id'][.='4']", + "//result/doc[3]/str[@name='id'][.='2']", + "//result/doc[4]/str[@name='id'][.='10']", + "//result/doc[5]/str[@name='id'][.='3']" + }; + + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=1}" + vectorToSearch, + "fl", + "id"), + expected); + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=4}" + vectorToSearch, + "fl", + "id"), + expected); + } + + @Test + public void rerankOversampleWithPreFilter_shouldReturnTopKFilteredResults() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=4 rerankOversample=3 preFilter='id:(1 4 7 8 9 10)'}" + + vectorToSearch, + "fl", + "id"), + "//result[@numFound='4']", + "//result/doc[1]/str[@name='id'][.='1']", + "//result/doc[2]/str[@name='id'][.='4']", + "//result/doc[3]/str[@name='id'][.='10']", + "//result/doc[4]/str[@name='id'][.='7']"); + } + + @Test + public void rerankOversampleWithSeedQuery_shouldReturnTopKResults() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=4 rerankOversample=3 seedQuery='id:(1 4 7 8 9)'}" + vectorToSearch, + "fl", + "id"), + "//result[@numFound='4']", + "//result/doc[1]/str[@name='id'][.='1']", + "//result/doc[2]/str[@name='id'][.='4']", + "//result/doc[3]/str[@name='id'][.='2']", + "//result/doc[4]/str[@name='id'][.='10']"); + } + + @Test + public void rerankOversampleWithEarlyTermination_shouldReturnTopKResults() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=3 earlyTermination=true saturationThreshold=0.989 patience=10}" + + vectorToSearch, + "fl", + "id"), + "//result[@numFound='5']"); + } + + @Test + public void rerankOversampleWithFilterMatchingNoDocs_shouldReturnNoResults() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + // A plain fq is folded into the knn query's pre-filter, and a filter matching nothing makes + // the knn query rewrite to MatchNoDocsQuery. The re-ranking wrapper has to cope with an empty + // candidate set rather than failing. + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=3}" + vectorToSearch, + "fq", + "id:nonexistent", + "fl", + "id"), + "//result[@numFound='0']"); + } + + @Test + public void rerankOversampleWithDebugQuery_matchingNoDocs_shouldNotThrow() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + // debugQuery stringifies the parsed query. The re-ranking query resolves its similarity + // function lazily during the search, which never happens when nothing matches, so the + // similarity function has to be set up front or toString blows up here. + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=3 preFilter='id:nonexistent'}" + vectorToSearch, + "fl", + "id", + CommonParams.DEBUG_QUERY, + "true"), + "//result[@numFound='0']", + "//str[@name='parsedquery_toString']"); + } + + @Test + public void rerankOversampleWithDebugQuery_shouldNotThrow() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=3}" + vectorToSearch, + "fl", + "id", + CommonParams.DEBUG_QUERY, + "true"), + "//result[@numFound='5']", + "//str[@name='parsedquery_toString']"); + } + + @Test + public void rerankOversampleWithEfSearchScaleFactor_shouldReturnTopKResults() { + String vectorToSearch = "[1.0, 2.0, 3.0, 4.0]"; + + assertQ( + req( + CommonParams.Q, + "{!knn f=vector topK=5 rerankOversample=2 efSearchScaleFactor=2.0}" + vectorToSearch, + "fl", + "id"), + "//result[@numFound='5']", + "//result/doc[1]/str[@name='id'][.='1']", + "//result/doc[2]/str[@name='id'][.='4']", + "//result/doc[3]/str[@name='id'][.='2']", + "//result/doc[4]/str[@name='id'][.='10']", + "//result/doc[5]/str[@name='id'][.='3']"); + } + Review Comment: What's the differentiator between some of these tests and the dedicated oversampling KNN tests? I get the throw/not throw exception, but the others feel similar to the ones in the other class. ########## solr/core/src/java/org/apache/solr/search/vector/SolrRescoreTopNQuery.java: ########## @@ -0,0 +1,63 @@ +/* + * 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.solr.search.vector; + +import java.io.IOException; +import org.apache.lucene.search.DoubleValuesSource; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchNoDocsQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.RescoreTopNQuery; + +/** + * A {@link RescoreTopNQuery} that tolerates an inner query matching no documents. + * + * <p>{@link RescoreTopNQuery#rewrite} unconditionally hands its collected hits to {@code + * DocAndScoreQuery#createDocAndScoreQuery}, which requires at least one hit: it asserts as much, + * and without assertions enabled it reads element zero of an empty array. {@code + * AbstractKnnVectorQuery#rewrite} guards the very same call by returning {@link MatchNoDocsQuery}, + * but the re-ranking query has no equivalent guard, so an oversampled knn query that matches + * nothing (a restrictive {@code preFilter}, say) fails instead of returning no results. + * + * <p>TODO: remove this class once the guard is added upstream in Lucene and Solr picks up a release + * containing it. + */ +public class SolrRescoreTopNQuery extends RescoreTopNQuery { + + private final Query innerQuery; + private final DoubleValuesSource valuesSource; + private final int n; + + public SolrRescoreTopNQuery(Query query, DoubleValuesSource valuesSource, int n) { + super(query, valuesSource, n); + this.innerQuery = query; + this.valuesSource = valuesSource; + this.n = n; + } + + @Override + public Query rewrite(IndexSearcher indexSearcher) throws IOException { + final Query rewrittenInner = indexSearcher.rewrite(innerQuery); + if (rewrittenInner instanceof MatchNoDocsQuery) { + return rewrittenInner; + } + // Delegate using the already rewritten inner query rather than calling super.rewrite(), which Review Comment: I suspect this is a 'if... else...' maybe better to make it explicit? ########## solr/core/src/java/org/apache/solr/schema/DenseVectorField.java: ########## @@ -524,25 +570,29 @@ public Query getKnnVectorQuery( ? new SolrKnnFloatVectorQuery( fieldName, vectorBuilder.getFloatVector(), - topK, + candidateTopK, efSearch, filterQuery, searchStrategy) : new SolrKnnFloatVectorQuery( - fieldName, vectorBuilder.getFloatVector(), topK, efSearch, filterQuery); + fieldName, + vectorBuilder.getFloatVector(), + candidateTopK, + efSearch, + filterQuery); break; case BYTE: baseQuery = searchStrategy != null ? new SolrKnnByteVectorQuery( fieldName, vectorBuilder.getByteVector(), - topK, + candidateTopK, Review Comment: same comment as below, isn't this a bit confusing given the fact in this PR BYTE encoding won't be supported? -- 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]
