perf: stream equality and range index scans instead of materializing every id

NitriteIndexer.findByFilter returns a LinkedHashSet of every matching id,
so find(k = v).firstOrNull() built the whole match set before handing back
one row, and a bounded page paid for the entire result. On a non-unique
index over a low-cardinality field that set is a large fraction of the
collection on every lookup.

The composite layout already keeps its rows in key order, so the two plan
shapes that map onto one bounded walk of it, an equality on the indexed
field and a two-sided range on it, are now served by a lazy iterator that
starts at the first key inside the bounds and stops at the first key outside
them. It honours the plan's reverse scan order by visiting the key groups
backwards while reading each group forwards, exactly as the materialized
scan orders them, skips entries removed in an open transaction, and returns
a document indexed under several keys once.

NitriteIndex.findNitriteIdStream and NitriteIndexer.findByFilterStream are
new default methods returning null, so every other index type, plugin
indexer and plan shape keeps the materialized path unchanged. ReadOperations
prefers the stream when one is offered; the covered-count shortcut that lets
size() answer without fetching documents is kept by counting the streamed
ids on demand, so size() still reads the index only.

Tests compare the stream with the materialized scan for equality, range and
reverse order, check the shapes it declines, show with a spied map that only
one key is read for the first row, and exercise counts, paging, descending
order, multi-valued fields and removals through the public API.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: brettwooldridge

nitrite/src/main/java/org/dizitart/no2/collection/operation/ReadOperations.java

@@ -111,35 +111,43 @@ private void prepareLogicalFilter(LogicalFilter logicalFilter) {
     }
 
     private DocumentCursor createCursor(FindPlan findPlan) {
-        // -1 means "not an index scan"; the index branch records the exact id-set size here.
-        long[] indexedIdCount = { -1 };
-        RecordStream<Pair<NitriteId, Document>> recordStream = findSuitableStream(findPlan, indexedIdCount);
+        IndexScan scan = new IndexScan();
+        RecordStream<Pair<NitriteId, Document>> recordStream = findSuitableStream(findPlan, scan);
         DocumentStream cursor = new DocumentStream(recordStream, processorChain);
         cursor.setFindPlan(findPlan);
-        cursor.setCoveredCount(computeCoveredCount(findPlan, indexedIdCount[0]));
+        if (isCountCovered(findPlan)) {
+            if (findPlan.getIndexDescriptor() == null) {
+                // pure full scan over the whole collection
+                cursor.setCoveredCount(nitriteMap.size());
+            } else if (scan.lazyStream != null) {
+                // the ids are read lazily: count them from the index only if size() is asked
+                cursor.setCoveredCountSupplier(scan.lazyStream::countIds);
+            } else if (scan.idCount >= 0) {
+                // the index supplied the exact matching id set
+                cursor.setCoveredCount(scan.idCount);
+            }
+        }
         return cursor;
     }
 
+    /** What an index scan handed back: an exact id count, or the lazy stream, or neither. */
+    private static final class IndexScan {
+        private long idCount = -1;
+        private IndexedStream lazyStream;
+    }
+
     /**
      * Returns the exact match count when the query is fully answered without fetching documents,
      * or {@code null} when the cursor must be drained to count. The count is exact only when
      * nothing downstream drops or changes cardinality (a post-filter, skip, or limit); sort does
      * not change the count, and an OR-union needs de-duplication so its count cannot be derived.
      */
-    private Long computeCoveredCount(FindPlan findPlan, long indexedIdCount) {
-        if (!findPlan.getSubPlans().isEmpty()
-            || findPlan.getCollectionScanFilter() != null
-            || findPlan.getSkip() != null
-            || findPlan.getLimit() != null
-            || findPlan.getByIdFilter() != null) {
-            return null;
-        }
-        if (findPlan.getIndexDescriptor() != null) {
-            // the index supplied the exact matching id set
-            return indexedIdCount >= 0 ? indexedIdCount : null;
-        }
-        // pure full scan over the whole collection
-        return nitriteMap.size();
+    private boolean isCountCovered(FindPlan findPlan) {
+        return findPlan.getSubPlans().isEmpty()
+            && findPlan.getCollectionScanFilter() == null
+            && findPlan.getSkip() == null
+            && findPlan.getLimit() == null
+            && findPlan.getByIdFilter() == null;
     }
 
     /**
@@ -193,7 +201,7 @@ private static Object indexedValue(DBValue dbValue) {
         return dbValue == null || dbValue instanceof DBNull ? null : dbValue.getValue();
     }
 
-    private RecordStream<Pair<NitriteId, Document>> findSuitableStream(FindPlan findPlan, long[] indexedIdCount) {
+    private RecordStream<Pair<NitriteId, Document>> findSuitableStream(FindPlan findPlan, IndexScan scan) {
         RecordStream<Pair<NitriteId, Document>> rawStream;
         RecordStream<Pair<NitriteId, Document>> indexSortedStream = null;
 
@@ -202,7 +210,7 @@ private RecordStream<Pair<NitriteId, Document>> findSuitableStream(FindPlan find
             List<RecordStream<Pair<NitriteId, Document>>> subStreams = new ArrayList<>();
             for (FindPlan subPlan : findPlan.getSubPlans()) {
                 // a sub-plan's own id count cannot answer the union's count (dedup), so discard it
-                RecordStream<Pair<NitriteId, Document>> suitableStream = findSuitableStream(subPlan, new long[]{ -1 });
+                RecordStream<Pair<NitriteId, Document>> suitableStream = findSuitableStream(subPlan, new IndexScan());
                 subStreams.add(suitableStream);
             }
 
@@ -233,14 +241,21 @@ private RecordStream<Pair<NitriteId, Document>> findSuitableStream(FindPlan find
                 if (indexDescriptor != null) {
                     // get optimized filter
                     NitriteIndexer indexer = nitriteConfig.findIndexer(indexDescriptor.getIndexType());
-                    LinkedHashSet<NitriteId> nitriteIds = indexer.findByFilter(findPlan, nitriteConfig);
+                    RecordStream<NitriteId> idStream = indexer.findByFilterStream(findPlan, nitriteConfig);
+                    if (idStream != null) {
+                        // the index walks its matches lazily; a size() is counted from it on demand
+                        scan.lazyStream = new IndexedStream(idStream, nitriteMap);
+                        rawStream = scan.lazyStream;
+                    } else {
+                        LinkedHashSet<NitriteId> nitriteIds = indexer.findByFilter(findPlan, nitriteConfig);
 
-                    // the index supplied the exact matching id set; record its size so a size()
-                    // with no row-dropping step downstream can answer from it without fetching
-                    indexedIdCount[0] = nitriteIds.size();
+                        // the index supplied the exact matching id set; record its size so a size()
+                        // with no row-dropping step downstream can answer from it without fetching
+                        scan.idCount = nitriteIds.size();
 
-                    // create indexed stream from optimized filter
-                    rawStream = new IndexedStream(nitriteIds, nitriteMap);
+                        // create indexed stream from optimized filter
+                        rawStream = new IndexedStream(nitriteIds, nitriteMap);
+                    }
                 } else {
                     indexSortedStream = indexSortedStream(findPlan);
                     rawStream = indexSortedStream != null ? indexSortedStream : nitriteMap.entries();

nitrite/src/main/java/org/dizitart/no2/common/streams/DocumentStream.java

@@ -33,6 +33,7 @@
 
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.function.LongSupplier;
 
 /**
  * @since 4.0
@@ -53,6 +54,13 @@ public class DocumentStream implements DocumentCursor {
     @Setter
     private Long coveredCount;
 
+    /**
+     * Answers {@link #size()} from the index on demand when the match count is known to be
+     * covered but the ids are streamed lazily rather than materialized; evaluated once.
+     */
+    @Setter
+    private LongSupplier coveredCountSupplier;
+
     public DocumentStream(RecordStream<Pair<NitriteId, Document>> recordStream,
                           ProcessorChain processorChain) {
         this.recordStream = recordStream;
@@ -64,6 +72,10 @@ public long size() {
         if (coveredCount != null) {
             return coveredCount;
         }
+        if (coveredCountSupplier != null) {
+            coveredCount = coveredCountSupplier.getAsLong();
+            return coveredCount;
+        }
         return Iterables.size(this);
     }
 

nitrite/src/main/java/org/dizitart/no2/common/streams/IndexedStream.java

@@ -24,17 +24,16 @@
 import org.dizitart.no2.store.NitriteMap;
 
 import java.util.Iterator;
-import java.util.Set;
 
 /**
  * @author Anindya Chatterjee
  * @since 4.0
  */
 public class IndexedStream implements RecordStream<Pair<NitriteId, Document>> {
     private final NitriteMap<NitriteId, Document> nitriteMap;
-    private final Set<NitriteId> nitriteIds;
+    private final Iterable<NitriteId> nitriteIds;
 
-    public IndexedStream(Set<NitriteId> nitriteIds,
+    public IndexedStream(Iterable<NitriteId> nitriteIds,
                   NitriteMap<NitriteId, Document> nitriteMap) {
         this.nitriteIds = nitriteIds;
         this.nitriteMap = nitriteMap;
@@ -45,6 +44,20 @@ public Iterator<Pair<NitriteId, Document>> iterator() {
         return new IndexedStreamIterator(nitriteIds.iterator(), nitriteMap);
     }
 
+    /**
+     * Counts the ids the index supplied, walking the id source only, without fetching a
+     * single document.
+     *
+     * @return the number of ids
+     */
+    public long countIds() {
+        long count = 0;
+        for (NitriteId ignored : nitriteIds) {
+            count++;
+        }
+        return count;
+    }
+
     private static class IndexedStreamIterator implements Iterator<Pair<NitriteId, Document>>,
             SkippableIterator {
         private final Iterator<NitriteId> iterator;

nitrite/src/main/java/org/dizitart/no2/index/ComparableIndexer.java

@@ -21,6 +21,7 @@
 import org.dizitart.no2.collection.NitriteId;
 import org.dizitart.no2.common.DBValue;
 import org.dizitart.no2.common.FieldValues;
+import org.dizitart.no2.common.RecordStream;
 import org.dizitart.no2.common.Fields;
 import org.dizitart.no2.common.tuples.Pair;
 import org.dizitart.no2.exceptions.IndexingException;
@@ -66,6 +67,12 @@ public LinkedHashSet<NitriteId> findByFilter(FindPlan findPlan, NitriteConfig ni
         return nitriteIndex.findNitriteIds(findPlan);
     }
 
+    @Override
+    public RecordStream<NitriteId> findByFilterStream(FindPlan findPlan, NitriteConfig nitriteConfig) {
+        NitriteIndex nitriteIndex = findNitriteIndex(findPlan.getIndexDescriptor(), nitriteConfig);
+        return nitriteIndex.findNitriteIdStream(findPlan);
+    }
+
     @Override
     public List<Pair<DBValue, NitriteId>> readSortKeys(IndexDescriptor indexDescriptor,
                                                        NitriteConfig nitriteConfig,

nitrite/src/main/java/org/dizitart/no2/index/NitriteIndex.java

@@ -21,6 +21,7 @@
 import org.dizitart.no2.collection.NitriteId;
 import org.dizitart.no2.common.DBValue;
 import org.dizitart.no2.common.FieldValues;
+import org.dizitart.no2.common.RecordStream;
 import org.dizitart.no2.common.tuples.Pair;
 import org.dizitart.no2.exceptions.UniqueConstraintException;
 import org.dizitart.no2.exceptions.ValidationException;
@@ -74,6 +75,19 @@ public interface NitriteIndex {
      */
     LinkedHashSet<NitriteId> findNitriteIds(FindPlan findPlan);
 
+    /**
+     * Streams the ids matching the plan lazily, in index order and without duplicates, or
+     * returns {@code null} when this index cannot do so for the given plan, in which case the
+     * caller falls back to {@link #findNitriteIds(FindPlan)}. A stream lets a query that only
+     * needs the first rows, or a bounded page, stop reading the index as soon as it has them.
+     *
+     * @param findPlan the find plan
+     * @return a re-iterable stream of ids, or {@code null}
+     */
+    default RecordStream<NitriteId> findNitriteIdStream(FindPlan findPlan) {
+        return null;
+    }
+
     /**
      * Reads every {@code (indexed value, id)} pair out of the index, so a sorted query can
      * decide its order without deserializing a single document.

nitrite/src/main/java/org/dizitart/no2/index/NitriteIndexer.java

@@ -21,6 +21,7 @@
 import org.dizitart.no2.collection.NitriteId;
 import org.dizitart.no2.common.DBValue;
 import org.dizitart.no2.common.FieldValues;
+import org.dizitart.no2.common.RecordStream;
 import org.dizitart.no2.common.Fields;
 import org.dizitart.no2.common.module.NitritePlugin;
 import org.dizitart.no2.common.tuples.Pair;
@@ -88,6 +89,19 @@ public interface NitriteIndexer extends NitritePlugin {
      */
     LinkedHashSet<NitriteId> findByFilter(FindPlan findPlan, NitriteConfig nitriteConfig);
 
+    /**
+     * Streams the ids matching the plan lazily, or returns {@code null} when the indexer has no
+     * lazy path for it and {@link #findByFilter(FindPlan, NitriteConfig)} must be used. The
+     * default is {@code null}, so existing indexer plugins are unaffected.
+     *
+     * @param findPlan      the find plan
+     * @param nitriteConfig the nitrite config
+     * @return a re-iterable stream of ids, or {@code null}
+     */
+    default RecordStream<NitriteId> findByFilterStream(FindPlan findPlan, NitriteConfig nitriteConfig) {
+        return null;
+    }
+
     /**
      * Reads every {@code (indexed value, id)} pair out of the given index, so a sorted query
      * can decide its order without deserializing a single document.