@@ -508,8 +508,9 @@ public AsyncIterator<ItemSlot> scan(@Nonnull final ReadTransaction readTransacti
508
508
*/
509
509
@ SuppressWarnings ("checkstyle:MethodName" ) // method name introduced by paper
510
510
@ Nonnull
511
- private CompletableFuture <GreedyState > kNearestNeighborsSearch (@ Nonnull final ReadTransaction readTransaction ,
512
- @ Nonnull final Vector <Half > queryVector ) {
511
+ private CompletableFuture <SearchResult <NodeReference >> kNearestNeighborsSearch (@ Nonnull final ReadTransaction readTransaction ,
512
+ final int efSearch ,
513
+ @ Nonnull final Vector <Half > queryVector ) {
513
514
return storageAdapter .fetchEntryNodeKey (readTransaction )
514
515
.thenCompose (entryPointAndLayer -> {
515
516
if (entryPointAndLayer == null ) {
@@ -521,7 +522,7 @@ private CompletableFuture<GreedyState> kNearestNeighborsSearch(@Nonnull final Re
521
522
final var entryState = new GreedyState (entryPointAndLayer .getLayer (),
522
523
entryPointAndLayer .getPrimaryKey (),
523
524
Vector .comparativeDistance (metric , entryPointAndLayer .getVector (), queryVector ));
524
- final AtomicReference <GreedyState > greedyResultReference =
525
+ final AtomicReference <GreedyState > greedyStateReference =
525
526
new AtomicReference <>(entryState );
526
527
527
528
if (entryPointAndLayer .getLayer () == 0 ) {
@@ -530,14 +531,22 @@ private CompletableFuture<GreedyState> kNearestNeighborsSearch(@Nonnull final Re
530
531
}
531
532
532
533
return AsyncUtil .whileTrue (() -> {
533
- final var greedyIn = greedyResultReference .get ();
534
+ final var greedyIn = greedyStateReference .get ();
534
535
return greedySearchLayer (readTransaction , greedyIn .toNodeReferenceWithDistance (),
535
536
greedyIn .getLayer (), queryVector )
536
537
.thenApply (greedyState -> {
537
- greedyResultReference .set (greedyState );
538
+ greedyStateReference .set (greedyState );
538
539
return greedyState .getLayer () > 0 ;
539
540
});
540
- }, executor ).thenApply (ignored -> greedyResultReference .get ());
541
+ }, executor ).thenApply (ignored -> greedyStateReference .get ());
542
+ }).thenCompose (greedyState -> {
543
+ if (greedyState == null ) {
544
+ return CompletableFuture .completedFuture (null );
545
+ }
546
+
547
+ return searchLayer (DataNode ::creator , readTransaction ,
548
+ ImmutableList .of (greedyState .toNodeReferenceWithDistance ()), 0 , efSearch ,
549
+ queryVector );
541
550
});
542
551
}
543
552
0 commit comments