From 5c7e73d7aedcca1bd7eb14c52bee8c4264c8a442 Mon Sep 17 00:00:00 2001 From: zhouhui Date: Sat, 11 May 2024 15:25:24 +0800 Subject: [PATCH] Implemented visitWithSortedDim for pointCount in PointRangeQuery. --- .../apache/lucene/search/PointRangeQuery.java | 12 +++++ .../lucene/search/TestPointQueries.java | 43 +++++++++++++++ .../tests/index/AssertingLeafReader.java | 54 +++++++++++++++++++ 3 files changed, 109 insertions(+) diff --git a/lucene/core/src/java/org/apache/lucene/search/PointRangeQuery.java b/lucene/core/src/java/org/apache/lucene/search/PointRangeQuery.java index c57aff9ad72..92c85de0914 100644 --- a/lucene/core/src/java/org/apache/lucene/search/PointRangeQuery.java +++ b/lucene/core/src/java/org/apache/lucene/search/PointRangeQuery.java @@ -519,6 +519,18 @@ public abstract class PointRangeQuery extends Query { } } + @Override + public boolean visitWithSortedDim(int docID, byte[] packedValue, int sortedDim) + throws IOException { + int matchState = matchesWithState(packedValue, sortedDim); + if (matchState == PointValues.MatchState.MATCH) { + matchingNodeCount[0]++; + } else if (matchState == PointValues.MatchState.HIGH_IN_SORTED_DIM) { + return false; + } + return true; + } + @Override public Relation compare(byte[] minPackedValue, byte[] maxPackedValue) { return nodeComparator.apply(minPackedValue, maxPackedValue); diff --git a/lucene/core/src/test/org/apache/lucene/search/TestPointQueries.java b/lucene/core/src/test/org/apache/lucene/search/TestPointQueries.java index 0ca2d177b09..7ab7300c49d 100644 --- a/lucene/core/src/test/org/apache/lucene/search/TestPointQueries.java +++ b/lucene/core/src/test/org/apache/lucene/search/TestPointQueries.java @@ -170,6 +170,14 @@ public class TestPointQueries extends LuceneTestCase { doc.add(new LongPoint("point", 3)); w.addDocument(doc); + doc = new Document(); + doc.add(new LongPoint("point", 4)); + w.addDocument(doc); + + doc = new Document(); + doc.add(new LongPoint("point", 5)); + w.addDocument(doc); + DirectoryReader r = DirectoryReader.open(w); IndexSearcher s = new IndexSearcher(r); assertEquals(2, s.count(LongPoint.newRangeQuery("point", -8L, 1L))); @@ -181,6 +189,41 @@ public class TestPointQueries extends LuceneTestCase { dir.close(); } + public void testVisitWithSoredDim() throws Exception { + Directory dir = newDirectory(); + IndexWriter w = new IndexWriter(dir, new IndexWriterConfig(new MockAnalyzer(random()))); + + Document doc = new Document(); + doc.add(new LongPoint("point", -7)); + w.addDocument(doc); + + doc = new Document(); + doc.add(new LongPoint("point", 0)); + w.addDocument(doc); + + doc = new Document(); + doc.add(new LongPoint("point", 3)); + w.addDocument(doc); + + doc = new Document(); + doc.add(new LongPoint("point", 4)); + w.addDocument(doc); + + doc = new Document(); + doc.add(new LongPoint("point", 5)); + w.addDocument(doc); + + DirectoryReader r = DirectoryReader.open(w); + IndexSearcher s = new IndexSearcher(r); + assertEquals(2, s.search(LongPoint.newRangeQuery("point", -8L, 1L), 10).scoreDocs.length); + assertEquals(3, s.search(LongPoint.newRangeQuery("point", -7L, 3L), 10).scoreDocs.length); + assertEquals(1, s.search(LongPoint.newExactQuery("point", -7L), 10).scoreDocs.length); + assertEquals(0, s.search(LongPoint.newExactQuery("point", -6L), 10).scoreDocs.length); + w.close(); + r.close(); + dir.close(); + } + public void testBasicDoubles() throws Exception { Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, new IndexWriterConfig(new MockAnalyzer(random()))); diff --git a/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java b/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java index c61c88ff143..57dcde34d23 100644 --- a/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java +++ b/lucene/test-framework/src/java/org/apache/lucene/tests/index/AssertingLeafReader.java @@ -1382,6 +1382,60 @@ public class AssertingLeafReader extends FilterLeafReader { in.visit(docID, packedValue); } + @Override + public boolean visitWithSortedDim(int docID, byte[] packedValue, int sortedDim) + throws IOException { + assert --docBudget >= 0 : "called add() more times than the last call to grow() reserved"; + + // This method, to filter each doc's value, should only be invoked when the cell crosses the + // query shape: + assert lastCompareResult == null + || lastCompareResult == PointValues.Relation.CELL_CROSSES_QUERY; + + if (lastCompareResult != null) { + // This doc's packed value should be contained in the last cell passed to compare: + for (int dim = 0; dim < numIndexDims; dim++) { + assert Arrays.compareUnsigned( + lastMinPackedValue, + dim * bytesPerDim, + dim * bytesPerDim + bytesPerDim, + packedValue, + dim * bytesPerDim, + dim * bytesPerDim + bytesPerDim) + <= 0 + : "dim=" + dim + " of " + numDataDims + " value=" + new BytesRef(packedValue); + assert Arrays.compareUnsigned( + lastMaxPackedValue, + dim * bytesPerDim, + dim * bytesPerDim + bytesPerDim, + packedValue, + dim * bytesPerDim, + dim * bytesPerDim + bytesPerDim) + >= 0 + : "dim=" + dim + " of " + numDataDims + " value=" + new BytesRef(packedValue); + } + lastCompareResult = null; + } + + // TODO: we should assert that this "matches" whatever relation the last call to compare had + // returned + assert packedValue.length == numDataDims * bytesPerDim; + if (numDataDims == 1) { + int cmp = Arrays.compareUnsigned(lastDocValue, 0, bytesPerDim, packedValue, 0, bytesPerDim); + if (cmp < 0) { + // ok + } else if (cmp == 0) { + assert lastDocID <= docID : "doc ids are out of order when point values are the same!"; + } else { + // out of order! + assert false : "point values are out of order"; + } + System.arraycopy(packedValue, 0, lastDocValue, 0, bytesPerDim); + lastDocID = docID; + } + return in.visitWithSortedDim(docID, packedValue, sortedDim); + } + @Override public void grow(int count) { in.grow(count);