(pinot) 01/01: Fix stale realtime Lucene searchable document counts

[xi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

xiangfu0 pushed a commit to branch codex/fix-realtime-lucene-doc-count
in repository https://gitbox.apache.org/repos/asf/pinot.git

commit 09569cbe55d9a951eafc46fe58bd66014d1673ca
Author: Xiang Fu <xi...@apache.org>
AuthorDate: Mon Oct 19 12:37:37 2026 -0700

    Fix stale realtime Lucene searchable document counts
---
 .../invertedindex/RealtimeLuceneTextIndex.java     | 21 +++++++++++-
 .../invertedindex/LuceneMutableTextIndexTest.java  | 40 +++++++++++++++++++++-
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java
index d2ea109f74b..af961697636 100644
--- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java
+++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java
@@ -19,6 +19,7 @@
 package org.apache.pinot.segment.local.realtime.impl.invertedindex;
 
 import java.io.File;
+import java.io.IOException;
 import java.lang.reflect.Constructor;
 import java.util.Map;
 import java.util.concurrent.Future;
@@ -339,8 +340,26 @@ public class RealtimeLuceneTextIndex implements MutableTextIndex {
   public SearcherManager getSearcherManager() {
     return _searcherManager;
   }
+
   @Override
   public int getSearchableDocCount() {
-    return _refreshListener.getLastRefreshNumDocs();
+    // The refresh listener records a lower bound that can miss documents added during a refresh.
+    SearcherManager searcherManager = _searcherManager;
+    IndexSearcher indexSearcher = null;
+    try {
+      indexSearcher = searcherManager.acquire();
+      return indexSearcher.getIndexReader().numDocs();
+    } catch (IOException e) {
+      throw new RuntimeException("Failed to get searchable document count for realtime text index for segment "
+          + _segmentName + " for column " + _column, e);
+    } finally {
+      if (indexSearcher != null) {
+        try {
+          searcherManager.release(indexSearcher);
+        } catch (IOException e) {
+          LOGGER.error("Failed while releasing searcher manager for column {}", _column, e);
+        }
+      }
+    }
   }
 }
diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/LuceneMutableTextIndexTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/LuceneMutableTextIndexTest.java
index 8d71f6a21f0..247590e627e 100644
--- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/LuceneMutableTextIndexTest.java
+++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/LuceneMutableTextIndexTest.java
@@ -27,11 +27,13 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.commons.io.FileUtils;
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.analysis.core.KeywordTokenizer;
 import org.apache.lucene.queryparser.classic.QueryParser;
+import org.apache.lucene.search.ReferenceManager.RefreshListener;
 import org.apache.pinot.common.metrics.ServerMetrics;
 import org.apache.pinot.segment.local.segment.index.text.TextIndexConfigBuilder;
 import org.apache.pinot.segment.spi.index.TextIndexConfig;
@@ -295,7 +297,7 @@ public class LuceneMutableTextIndexTest {
       index.add(new String[]{"foo bar"});
       index.add(new String[]{"baz qux"});
 
-      // Block until the refresh attempt completes so the listener has recorded the current doc count
+      // Block until the documents are visible to the current searcher.
       index.getSearcherManager().maybeRefreshBlocking();
 
       assertEquals(index.getSearchableDocCount(), 3);
@@ -304,6 +306,42 @@ public class LuceneMutableTextIndexTest {
     }
   }
 
+  @Test
+  public void testGetSearchableDocCountWhenDocumentIsAddedDuringRefresh()
+      throws IOException {
+    TextIndexConfig config = new TextIndexConfigBuilder().withUseANDForMultiTermQueries(false).build();
+    RealtimeLuceneTextIndex index = new RealtimeLuceneTextIndex(TEXT_COLUMN_NAME, INDEX_DIR,
+        "table__0__" + SEGMENT_NAME_SUFFIX_COUNTER.getAndIncrement() + "__20240601T1818Z", config);
+    try {
+      index.add("hello world");
+      index.add("foo bar");
+      AtomicBoolean documentAdded = new AtomicBoolean();
+      index.getSearcherManager().addListener(new RefreshListener() {
+        @Override
+        public void beforeRefresh() {
+          // The production listener has captured two docs, but the refreshed reader will include this third doc.
+          if (documentAdded.compareAndSet(false, true)) {
+            index.add("baz qux");
+          }
+        }
+
+        @Override
+        public void afterRefresh(boolean didRefresh) {
+        }
+      });
+      index.getSearcherManager().maybeRefreshBlocking();
+
+      assertEquals(index.getDocIds("baz"), ImmutableRoaringBitmap.bitmapOf(2));
+      assertEquals(index.getSearchableDocCount(), 3);
+
+      // A refresh with no further writes must retain the full searchable document count.
+      index.getSearcherManager().maybeRefreshBlocking();
+      assertEquals(index.getSearchableDocCount(), 3);
+    } finally {
+      index.close();
+    }
+  }
+
   @Test
   public void testQueries() {
     configureIndex(null, null, null, null);


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org