Skip to content

Commit bbe949c

Browse files
prudhvigodithiprudhvigodithi
committed
Use point tree clone and improve the coverage
Signed-off-by: Prudhvi Godithi <[email protected]>
1 parent 757aebf commit bbe949c

File tree

2 files changed

+178
-5
lines changed

2 files changed

+178
-5
lines changed

‎server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -298,16 +298,15 @@ public void intersectRight(PointValues.IntersectVisitor visitor, PointValues.Poi
298298
case CELL_INSIDE_QUERY:
299299
case CELL_CROSSES_QUERY:
300300
if (pointTree.moveToChild() && docCount[0] < size) {
301+
PointValues.PointTree leftChild = pointTree.clone();
301302
// BKD is binary today, so one moveToSibling() is enough to land on the right child.
302303
// If PointTree ever becomes n-ary, update the traversal below to visit all siblings or re-enable a full loop.
303304
pointTree.moveToSibling();
304305
assert pointTree.moveToSibling() == false;
305306
intersectRight(visitor, pointTree, docCount);
306307
pointTree.moveToParent();
307308
if (docCount[0] < size) {
308-
pointTree.moveToChild();
309-
intersectRight(visitor, pointTree, docCount);
310-
pointTree.moveToParent();
309+
intersectRight(visitor, leftChild, docCount);
311310
}
312311
} else {
313312
if (docCount[0] < size) {
@@ -324,7 +323,6 @@ public void intersectRight(PointValues.IntersectVisitor visitor, PointValues.Poi
324323
default:
325324
throw new IllegalArgumentException("Unreachable code");
326325
}
327-
328326
}
329327

330328
@Override

‎server/src/test/java/org/opensearch/search/approximate/ApproximatePointRangeQueryTests.java

Lines changed: 176 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -392,7 +392,6 @@ public void testApproximateRangeShortCircuitDescSort() throws IOException {
392392
try (Directory directory = newDirectory()) {
393393
try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory, new WhitespaceAnalyzer())) {
394394
int dims = 1;
395-
396395
long[] scratch = new long[dims];
397396
int numPoints = 1000;
398397
for (int i = 0; i < numPoints; i++) {
@@ -436,4 +435,180 @@ public void testApproximateRangeShortCircuitDescSort() throws IOException {
436435
}
437436
}
438437
}
438+
439+
// Test to cover the left child traversal in intersectRight with CELL_INSIDE_QUERY
440+
public void testIntersectRightLeftChildTraversal() throws IOException {
441+
try (Directory directory = newDirectory()) {
442+
try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory, new WhitespaceAnalyzer())) {
443+
int dims = 1;
444+
long[] scratch = new long[dims];
445+
// Create a dataset that will create a multi-level BKD tree
446+
// We need enough documents to create internal nodes (not just leaves)
447+
int numPoints = 2000;
448+
for (int i = 0; i < numPoints; i++) {
449+
Document doc = new Document();
450+
scratch[0] = i;
451+
doc.add(new LongPoint("point", scratch[0]));
452+
iw.addDocument(asList(new LongPoint("point", scratch[0]), new NumericDocValuesField("point", scratch[0])));
453+
if (i % 100 == 0) {
454+
iw.flush();
455+
}
456+
}
457+
iw.flush();
458+
iw.forceMerge(1);
459+
460+
try (IndexReader reader = iw.getReader()) {
461+
// Query that will match many documents and require tree traversal
462+
long lower = 1000;
463+
long upper = 1999;
464+
ApproximatePointRangeQuery query = new ApproximatePointRangeQuery(
465+
"point",
466+
pack(lower).bytes,
467+
pack(upper).bytes,
468+
dims,
469+
50, // Small size to ensure we hit the left child traversal condition
470+
SortOrder.DESC,
471+
ApproximatePointRangeQuery.LONG_FORMAT
472+
);
473+
IndexSearcher searcher = new IndexSearcher(reader);
474+
Sort sort = new Sort(new SortField("point", SortField.Type.LONG, true)); // DESC
475+
TopDocs topDocs = searcher.search(query, 50, sort);
476+
assertEquals("Should return exactly 50 documents", 50, topDocs.scoreDocs.length);
477+
}
478+
}
479+
}
480+
}
481+
482+
// Test to cover pointTree.visitDocIDs(visitor) in CELL_INSIDE_QUERY leaf case for intersectRight
483+
public void testIntersectRightCellInsideQueryLeaf() throws IOException {
484+
try (Directory directory = newDirectory()) {
485+
try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory, new WhitespaceAnalyzer())) {
486+
int dims = 1;
487+
long[] scratch = new long[dims];
488+
// Create a smaller dataset that will result in leaf nodes that are completely inside the query range
489+
for (int i = 900; i <= 999; i++) {
490+
Document doc = new Document();
491+
scratch[0] = i;
492+
iw.addDocument(asList(new LongPoint("point", scratch[0]), new NumericDocValuesField("point", scratch[0])));
493+
}
494+
iw.flush();
495+
iw.forceMerge(1);
496+
497+
try (IndexReader reader = iw.getReader()) {
498+
// Query that completely contains all documents (CELL_INSIDE_QUERY)
499+
long lower = 800;
500+
long upper = 1100;
501+
502+
ApproximatePointRangeQuery query = new ApproximatePointRangeQuery(
503+
"point",
504+
pack(lower).bytes,
505+
pack(upper).bytes,
506+
dims,
507+
200,
508+
SortOrder.DESC,
509+
ApproximatePointRangeQuery.LONG_FORMAT
510+
);
511+
512+
IndexSearcher searcher = new IndexSearcher(reader);
513+
Sort sort = new Sort(new SortField("point", SortField.Type.LONG, true)); // DESC
514+
TopDocs topDocs = searcher.search(query, 200, sort);
515+
516+
assertEquals("Should find all documents", 100, topDocs.totalHits.value());
517+
// Should return all the indexed point values from 900 to 999 which tests CELL_INSIDE_QUERY
518+
assertEquals("Should return exactly 100 documents", 100, topDocs.scoreDocs.length);
519+
}
520+
}
521+
}
522+
}
523+
524+
// Test to cover CELL_OUTSIDE_QUERY break case for intersectRight
525+
public void testIntersectRightCellOutsideQuery() throws IOException {
526+
try (Directory directory = newDirectory()) {
527+
try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory, new WhitespaceAnalyzer())) {
528+
int dims = 1;
529+
long[] scratch = new long[dims];
530+
// Create documents in two separate ranges to ensure some cells are outside query
531+
// Range 1: 0-99
532+
for (int i = 0; i < 100; i++) {
533+
Document doc = new Document();
534+
scratch[0] = i;
535+
iw.addDocument(asList(new LongPoint("point", scratch[0]), new NumericDocValuesField("point", scratch[0])));
536+
}
537+
// Range 2: 500-599 (gap ensures some tree nodes will be completely outside query)
538+
for (int i = 500; i < 600; i++) {
539+
Document doc = new Document();
540+
scratch[0] = i;
541+
iw.addDocument(asList(new LongPoint("point", scratch[0]), new NumericDocValuesField("point", scratch[0])));
542+
}
543+
iw.flush();
544+
iw.forceMerge(1);
545+
546+
try (IndexReader reader = iw.getReader()) {
547+
// Query only the middle range - this should create CELL_OUTSIDE_QUERY for some nodes
548+
long lower = 200;
549+
long upper = 400;
550+
551+
ApproximatePointRangeQuery query = new ApproximatePointRangeQuery(
552+
"point",
553+
pack(lower).bytes,
554+
pack(upper).bytes,
555+
dims,
556+
50,
557+
SortOrder.DESC,
558+
ApproximatePointRangeQuery.LONG_FORMAT
559+
);
560+
561+
IndexSearcher searcher = new IndexSearcher(reader);
562+
Sort sort = new Sort(new SortField("point", SortField.Type.LONG, true)); // DESC
563+
TopDocs topDocs = searcher.search(query, 50, sort);
564+
565+
// Should find no documents since our query range (200-400) has no documents
566+
assertEquals("Should find no documents in the gap range", 0, topDocs.totalHits.value());
567+
}
568+
}
569+
}
570+
}
571+
572+
// Test to cover intersectRight with CELL_CROSSES_QUERY case and ensure comprehensive coverage for intersectRight
573+
public void testIntersectRightCellCrossesQuery() throws IOException {
574+
try (Directory directory = newDirectory()) {
575+
try (RandomIndexWriter iw = new RandomIndexWriter(random(), directory, new WhitespaceAnalyzer())) {
576+
int dims = 1;
577+
long[] scratch = new long[dims];
578+
// Create documents that will result in cells that cross the query boundary
579+
for (int i = 0; i < 1000; i++) {
580+
Document doc = new Document();
581+
scratch[0] = i;
582+
iw.addDocument(asList(new LongPoint("point", scratch[0]), new NumericDocValuesField("point", scratch[0])));
583+
}
584+
iw.flush();
585+
iw.forceMerge(1);
586+
587+
try (IndexReader reader = iw.getReader()) {
588+
// Query that will partially overlap with tree nodes (CELL_CROSSES_QUERY)
589+
// This range will intersect with some tree nodes but not completely contain them
590+
long lower = 250;
591+
long upper = 750;
592+
593+
ApproximatePointRangeQuery query = new ApproximatePointRangeQuery(
594+
"point",
595+
pack(lower).bytes,
596+
pack(upper).bytes,
597+
dims,
598+
100,
599+
SortOrder.DESC,
600+
ApproximatePointRangeQuery.LONG_FORMAT
601+
);
602+
603+
IndexSearcher searcher = new IndexSearcher(reader);
604+
Sort sort = new Sort(new SortField("point", SortField.Type.LONG, true)); // DESC
605+
TopDocs topDocs = searcher.search(query, 100, sort);
606+
607+
assertEquals("Should return exactly 100 documents", 100, topDocs.scoreDocs.length);
608+
// For Desc sort the ApproximatePointRangeQuery will slightly over collect to retain the highest matched docs
609+
assertTrue("Should collect at least requested number of documents", topDocs.totalHits.value() >= 100);
610+
}
611+
}
612+
}
613+
}
439614
}

0 commit comments

Comments
 (0)