msokolov commented on code in PR #12434:
URL: https://github.com/apache/lucene/pull/12434#discussion_r1278615147
##########
lucene/core/src/test/org/apache/lucene/util/hnsw/TestNeighborQueue.java:
##########
@@ -114,6 +114,38 @@ public void testUnboundedQueue() {
assertEquals(maxNode, nn.topNode());
}
+ public void testCollectAndProvideResultsSameIds() {
Review Comment:
This is a strange test case? We don't usually expect to store the same node
id multiple times do we?
##########
lucene/core/src/java/org/apache/lucene/search/KnnCollector.java:
##########
@@ -0,0 +1,88 @@
+/*
+ * 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;
+
+/**
+ * KnnCollector is a knn collector used for gathering kNN results and
providing topDocs from the
+ * gathered neighbors
+ *
+ * @lucene.experimental
+ */
+public interface KnnCollector {
Review Comment:
Aside from supporting this change, the new interface seems to allow for a
lot of consolidation, +1
##########
lucene/core/src/java/org/apache/lucene/util/hnsw/RandomAccessVectorValues.java:
##########
@@ -46,4 +46,14 @@ public interface RandomAccessVectorValues<T> {
* {@link RandomAccessVectorValues#vectorValue}.
*/
RandomAccessVectorValues<T> copy() throws IOException;
+
+ /**
+ * Translates vector ordinal to the correct document ID. By default, this is
an identity function.
+ *
+ * @param ord the vector ordinal
+ * @return the document Id for that vector ordinal
+ */
+ default int ordToDoc(int ord) {
Review Comment:
I think there might be some useless overrides after this?
##########
lucene/join/src/java/org/apache/lucene/search/join/ToParentJoinKnnCollector.java:
##########
@@ -0,0 +1,294 @@
+/*
+ * 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.util.HashMap;
+import java.util.Map;
+import org.apache.lucene.search.AbstractKnnCollector;
+import org.apache.lucene.search.ScoreDoc;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.TotalHits;
+import org.apache.lucene.util.ArrayUtil;
+import org.apache.lucene.util.BitSet;
+
+/** parent joining knn collector, vector docIds are deduplicated according to
the parent bit set. */
+class ToParentJoinKnnCollector extends AbstractKnnCollector {
+
+ private final BitSet parentBitSet;
+ private final NodeIdCachingHeap heap;
+
+ /**
+ * Create a new object for joining nearest child kNN documents with a parent
bitset
+ *
+ * @param k The number of joined parent documents to collect
+ * @param visitLimit how many child vectors can be visited
+ * @param parentBitSet The leaf parent bitset
+ */
+ public ToParentJoinKnnCollector(int k, int visitLimit, BitSet parentBitSet) {
+ super(k, visitLimit);
+ this.parentBitSet = parentBitSet;
+ this.heap = new NodeIdCachingHeap(k);
+ }
+
+ /**
+ * If the heap is not full (size is less than the initialSize provided to
the constructor), adds a
+ * new node-and-score element. If the heap is full, compares the score
against the current top
+ * score, and replaces the top element if newScore is better than (greater
than unless the heap is
+ * reversed), the current top score.
+ *
+ * <p>If docId's parent node has previously been collected and the provided
nodeScore is less than
+ * the stored score it will not be collected.
+ *
+ * @param docId the neighbor docId
+ * @param nodeScore the score of the neighbor, relative to some other node
+ */
+ @Override
+ public boolean collect(int docId, float nodeScore) {
+ assert !parentBitSet.get(docId);
+ int nodeId = parentBitSet.nextSetBit(docId);
+ return heap.insertWithOverflow(nodeId, nodeScore);
+ }
+
+ @Override
+ public float minCompetitiveSimilarity() {
+ return heap.size >= k() ? heap.topScore() : Float.NEGATIVE_INFINITY;
+ }
+
+ @Override
+ public String toString() {
+ return "ToParentJoinKnnCollector[k=" + k() + ", size=" + heap.size() + "]";
+ }
+
+ @Override
+ public TopDocs topDocs() {
+ assert heap.size() <= k() : "Tried to collect more results than the
maximum number allowed";
+ while (heap.size() > k()) {
+ heap.popToDrain();
+ }
+ ScoreDoc[] scoreDocs = new ScoreDoc[heap.size()];
+ for (int i = 1; i <= scoreDocs.length; i++) {
+ scoreDocs[scoreDocs.length - i] = new ScoreDoc(heap.topNode(),
heap.topScore());
+ heap.popToDrain();
+ }
+
+ TotalHits.Relation relation =
+ earlyTerminated()
+ ? TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO
+ : TotalHits.Relation.EQUAL_TO;
+ return new TopDocs(new TotalHits(visitedCount(), relation), scoreDocs);
+ }
+
+ /**
+ * This is a minimum binary heap, inspired by {@link
org.apache.lucene.util.LongHeap}. But instead
+ * of encoding and using `long` values. Node ids and scores are kept
separate. Additionally, this
+ * prevents duplicate nodes from being added.
+ *
+ * <p>So, for every node added, we will update its score if the newly
provided score is better.
+ * Every time we update a node's stored score, we ensure the heap's order.
+ */
+ private static class NodeIdCachingHeap {
+ private final int maxSize;
+ private int[] heapNodes;
+ private float[] heapScores;
+ private int size = 0;
+
+ // Used to keep track of nodeId -> positionInHeap. This way when new
scores are added for a
Review Comment:
did you consider having a heap of (node, score, position) tuples? I don't
think there's any theoretical complexity difference (it's all O(1) insert into
heap vs. updating a record) but it might be a little neater, requiring fewer
parallel data structures.
##########
lucene/core/src/java/org/apache/lucene/codecs/lucene95/OffHeapFloatVectorValues.java:
##########
@@ -62,8 +62,6 @@ public float[] vectorValue(int targetOrd) throws IOException {
return value;
}
- public abstract int ordToDoc(int ord);
Review Comment:
can we also delete the overrides in the Dense* subclasses?
##########
lucene/core/src/java/org/apache/lucene/search/KnnCollector.java:
##########
@@ -0,0 +1,88 @@
+/*
+ * 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;
+
+/**
+ * KnnCollector is a knn collector used for gathering kNN results and
providing topDocs from the
+ * gathered neighbors
+ *
+ * @lucene.experimental
+ */
+public interface KnnCollector {
+
+ /**
+ * If search visits too many documents, the results collector will terminate
early. Usually, this
+ * is due to some restricted filter on the document set.
+ *
+ * <p>When collection is earlyTerminated, the results are not a correct
representation of k
+ * nearest neighbors.
+ *
+ * @return is the current result set marked as incomplete?
+ */
+ boolean earlyTerminated();
+
+ /**
+ * @param count increments the visited vector count, must be greater than 0.
+ */
+ void incVisitedCount(int count);
+
+ /**
+ * @return the current visited vector count
+ */
+ long visitedCount();
+
+ /**
+ * @return the visited vector limit
+ */
+ long visitLimit();
+
+ /**
+ * @return the expected number of collected results
+ */
+ int k();
+
+ /**
+ * Collect the provided docId and include in the result set.
Review Comment:
I think we are collecting vector ordinals, not docids. The javadoc should be
clear about that so future implementors don't get confused!
--
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: issues-unsubscr...@lucene.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscr...@lucene.apache.org
For additional commands, e-mail: issues-h...@lucene.apache.org
