maedhroz commented on code in PR #4353:
URL: https://github.com/apache/cassandra/pull/4353#discussion_r2674129089
##########
src/java/org/apache/cassandra/index/sai/disk/v1/segment/VectorIndexSegmentSearcher.java:
##########
@@ -98,41 +118,44 @@ public KeyRangeIterator search(Expression exp,
AbstractBounds<PartitionPosition>
throw new
IllegalArgumentException(index.identifier().logMessage("Unsupported expression
during ANN index query: " + exp));
int topK = optimizeFor.topKFor(limit);
- BitsOrPostingList bitsOrPostingList =
bitsOrPostingListForKeyRange(context.vectorContext(), keyRange, topK);
- if (bitsOrPostingList.skipANN())
- return toPrimaryKeyIterator(bitsOrPostingList.postingList(),
context);
float[] queryVector =
index.termType().decomposeVector(exp.lower().value.raw.duplicate());
- VectorPostingList vectorPostings = graph.search(queryVector, topK,
limit, bitsOrPostingList.getBits());
- if (bitsOrPostingList.expectedNodesVisited >= 0)
- updateExpectedNodes(vectorPostings.getVisitedCount(),
bitsOrPostingList.expectedNodesVisited);
- return toPrimaryKeyIterator(vectorPostings, context);
+ CloseableIterator<RowIdWithScore> result = searchInternal(keyRange,
queryVector, limit, topK);
+ return toScoreSortedIterator(result);
}
/**
* Return bit set we need to search the graph; otherwise return posting
list to bypass the graph
*/
- private BitsOrPostingList bitsOrPostingListForKeyRange(VectorQueryContext
context, AbstractBounds<PartitionPosition> keyRange, int limit) throws
IOException
+ private CloseableIterator<RowIdWithScore>
searchInternal(AbstractBounds<PartitionPosition> keyRange, float[] queryVector,
int limit, int topK) throws IOException
{
try (PrimaryKeyMap primaryKeyMap =
primaryKeyMapFactory.newPerSSTablePrimaryKeyMap())
{
// not restricted
if (RangeUtil.coversFullRing(keyRange))
- return new
BitsOrPostingList(context.bitsetForShadowedPrimaryKeys(metadata, primaryKeyMap,
graph));
+ {
+ int expectedNodesVisited = expectedNodesVisited(limit,
graph.size(), graph.size());
+ IntConsumer nodesVisitedConsumer = nodesVisited ->
updateExpectedNodes(nodesVisited, expectedNodesVisited);
+ return graph.search(queryVector, topK, limit, new
Bits.MatchAllBits(graph.size()), nodesVisitedConsumer);
+ }
// it will return the next row id if given key is not found.
long minSSTableRowId =
primaryKeyMap.ceiling(keyRange.left.getToken());
// If we didn't find the first key, we won't find the last primary
key either
if (minSSTableRowId < 0)
- return new BitsOrPostingList(PostingList.EMPTY);
+ return CloseableIterator.empty();
long maxSSTableRowId = getMaxSSTableRowId(primaryKeyMap,
keyRange.right);
if (minSSTableRowId > maxSSTableRowId)
- return new BitsOrPostingList(PostingList.EMPTY);
+ return CloseableIterator.empty();
// if it covers entire segment, skip bit set
if (minSSTableRowId <= metadata.minSSTableRowId && maxSSTableRowId
>= metadata.maxSSTableRowId)
- return new
BitsOrPostingList(context.bitsetForShadowedPrimaryKeys(metadata, primaryKeyMap,
graph));
+ {
+ int expectedNodesVisited = expectedNodesVisited(limit,
graph.size(), graph.size());
+ IntConsumer nodesVisitedConsumer = nodesVisited ->
updateExpectedNodes(nodesVisited, expectedNodesVisited);
+ return graph.search(queryVector, topK, limit, new
Bits.MatchAllBits(graph.size()), nodesVisitedConsumer);
Review Comment:
nit: This looks like the same code block as the one above on 136, but I'll
let you decide whether you think factoring out a method makes sense.
--
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]
