kaivalnp commented on code in PR #16738: URL: https://github.com/apache/lucene/pull/16738#discussion_r4158634235
########## lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenVectorScorer.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.lucene.search.join; + +import java.io.IOException; +import org.apache.lucene.index.QueryTimeout; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.HitQueue; +import org.apache.lucene.search.ScoreDoc; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.TotalHits; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.util.BitSet; + +/** + * Iterates the accepted child documents one parent at a time, tracking the best scoring child of + * each parent. Scoring is delegated to the given {@link VectorScorer}. + */ +class DiversifyingChildrenVectorScorer { Review Comment: Can we add the `@lucene.experimental` tag here? (because the queries are marked the same) ########## lucene/join/src/java/org/apache/lucene/search/join/DiversifyingChildrenFloat16KnnVectorQuery.java: ########## @@ -0,0 +1,176 @@ +/* + * 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.lucene.search.join; + +import static org.apache.lucene.search.knn.KnnSearchStrategy.Hnsw.DEFAULT; + +import java.io.IOException; +import java.util.Arrays; +import java.util.Objects; +import org.apache.lucene.index.Float16VectorValues; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.QueryTimeout; +import org.apache.lucene.search.AcceptDocs; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.search.KnnFloat16VectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.search.TopDocsCollector; +import org.apache.lucene.search.VectorScorer; +import org.apache.lucene.search.knn.KnnCollectorManager; +import org.apache.lucene.search.knn.KnnSearchStrategy; +import org.apache.lucene.util.BitSet; + +/** + * kNN float16 vector query that joins matching children vector documents with their parent doc id. + * The top documents returned are the child document ids and the calculated scores. Here is how to + * use this in conjunction with {@link ToParentBlockJoinQuery}. + * + * <pre><code class="language-java"> + * Query knnQuery = new DiversifyingChildrenFloat16KnnVectorQuery(fieldName, queryVector, ...); + * // Rewrite executes kNN search and collects nearest children docIds and their scores + * Query rewrittenKnnQuery = searcher.rewrite(knnQuery); + * // Join the scored children docs with their parents and score the parents + * Query childrenToParents = new ToParentBlockJoinQuery(rewrittenKnnQuery, parentsFilter, ScoreMode.MAX); + * </code></pre> + */ +public class DiversifyingChildrenFloat16KnnVectorQuery extends KnnFloat16VectorQuery { + private static final TopDocs NO_RESULTS = TopDocsCollector.EMPTY_TOPDOCS; + + private final BitSetProducer parentsFilter; + private final Query childFilter; + private final int k; + private final short[] query; + + /** + * Create a DiversifyingChildrenFloat16KnnVectorQuery. + * + * @param field the query field + * @param query the vector query + * @param childFilter the child filter + * @param k how many parent documents to return given the matching children + * @param parentsFilter Filter identifying the parent documents. + */ + public DiversifyingChildrenFloat16KnnVectorQuery( + String field, short[] query, Query childFilter, int k, BitSetProducer parentsFilter) { + this(field, query, childFilter, k, parentsFilter, DEFAULT); + } + + /** + * Create a DiversifyingChildrenFloat16KnnVectorQuery. + * + * @param field the query field + * @param query the vector query + * @param childFilter the child filter + * @param k how many parent documents to return given the matching children + * @param parentsFilter Filter identifying the parent documents. + * @param searchStrategy the search strategy to use. If null, the default strategy will be used. + * The underlying format may not support all strategies and is free to ignore the requested + * strategy. + * @lucene.experimental Review Comment: Pre-existing, but I think this belongs in the javadocs of the class and not the constructor? This can be fixed in a follow-up PR though. -- 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]
