You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by jp...@apache.org on 2024/09/13 05:45:17 UTC

(lucene) branch main updated: Add prefetching support to term vectors. (#13758)

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

jpountz pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git


The following commit(s) were added to refs/heads/main by this push:
     new e8e179cd455 Add prefetching support to term vectors. (#13758)
e8e179cd455 is described below

commit e8e179cd4553b503dc529179beb7c6290202acf0
Author: Adrien Grand <jp...@gmail.com>
AuthorDate: Fri Sep 13 07:45:12 2024 +0200

    Add prefetching support to term vectors. (#13758)
    
    This follows the same approach that we used for stored fields.
---
 .../Lucene90CompressingTermVectorsReader.java      | 31 +++++++
 .../apache/lucene/index/BaseCompositeReader.java   |  9 ++
 .../java/org/apache/lucene/index/CheckIndex.java   |  3 +
 .../apache/lucene/index/ParallelLeafReader.java    | 23 ++++-
 .../lucene/index/SlowCodecReaderWrapper.java       |  5 ++
 .../index/SlowCompositeCodecReaderWrapper.java     |  9 ++
 .../apache/lucene/index/SortingCodecReader.java    |  5 ++
 .../java/org/apache/lucene/index/TermVectors.java  | 13 +++
 .../lucene90/TestLucene90TermVectorsFormat.java    | 97 ++++++++++++++++++++++
 .../lucene/tests/index/AssertingLeafReader.java    |  6 ++
 10 files changed, 199 insertions(+), 2 deletions(-)

diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingTermVectorsReader.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingTermVectorsReader.java
index e84ca079cb4..06af857c58c 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingTermVectorsReader.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingTermVectorsReader.java
@@ -31,6 +31,7 @@ import static org.apache.lucene.codecs.lucene90.compressing.Lucene90CompressingT
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
@@ -77,6 +78,9 @@ import org.apache.lucene.util.packed.PackedInts;
  */
 public final class Lucene90CompressingTermVectorsReader extends TermVectorsReader {
 
+  private static final int PREFETCH_CACHE_SIZE = 1 << 4;
+  private static final int PREFETCH_CACHE_MASK = PREFETCH_CACHE_SIZE - 1;
+
   private final FieldInfos fieldInfos;
   final FieldsIndex indexReader;
   final IndexInput vectorsStream;
@@ -93,6 +97,11 @@ public final class Lucene90CompressingTermVectorsReader extends TermVectorsReade
   private final long numDirtyDocs; // cumulative number of docs in incomplete chunks
   private final long maxPointer; // end of the data section
   private BlockState blockState = new BlockState(-1, -1, 0);
+  // Cache of recently prefetched block IDs. This helps reduce chances of prefetching the same block
+  // multiple times, which is otherwise likely due to index sorting or recursive graph bisection
+  // clustering similar documents together. NOTE: this cache must be small since it's fully scanned.
+  private final long[] prefetchedBlockIDCache;
+  private int prefetchedBlockIDCacheIndex;
 
   // used by clone
   private Lucene90CompressingTermVectorsReader(Lucene90CompressingTermVectorsReader reader) {
@@ -111,6 +120,8 @@ public final class Lucene90CompressingTermVectorsReader extends TermVectorsReade
     this.numDirtyChunks = reader.numDirtyChunks;
     this.numDirtyDocs = reader.numDirtyDocs;
     this.maxPointer = reader.maxPointer;
+    this.prefetchedBlockIDCache = new long[PREFETCH_CACHE_SIZE];
+    Arrays.fill(prefetchedBlockIDCache, -1);
     this.closed = false;
   }
 
@@ -212,6 +223,9 @@ public final class Lucene90CompressingTermVectorsReader extends TermVectorsReade
       CodecUtil.checkFooter(metaIn, null);
       metaIn.close();
 
+      this.prefetchedBlockIDCache = new long[PREFETCH_CACHE_SIZE];
+      Arrays.fill(prefetchedBlockIDCache, -1);
+
       success = true;
     } catch (Throwable t) {
       if (metaIn != null) {
@@ -327,6 +341,23 @@ public final class Lucene90CompressingTermVectorsReader extends TermVectorsReade
 
   private record BlockState(long startPointer, int docBase, int chunkDocs) {}
 
+  @Override
+  public void prefetch(int docID) throws IOException {
+    final long blockID = indexReader.getBlockID(docID);
+
+    for (long prefetchedBlockID : prefetchedBlockIDCache) {
+      if (prefetchedBlockID == blockID) {
+        return;
+      }
+    }
+
+    final long blockStartPointer = indexReader.getBlockStartPointer(blockID);
+    final long blockLength = indexReader.getBlockLength(blockID);
+    vectorsStream.prefetch(blockStartPointer, blockLength);
+
+    prefetchedBlockIDCache[prefetchedBlockIDCacheIndex++ & PREFETCH_CACHE_MASK] = blockID;
+  }
+
   @Override
   public Fields get(int doc) throws IOException {
     ensureOpen();
diff --git a/lucene/core/src/java/org/apache/lucene/index/BaseCompositeReader.java b/lucene/core/src/java/org/apache/lucene/index/BaseCompositeReader.java
index 68fc36120a6..dcda91a4c5d 100644
--- a/lucene/core/src/java/org/apache/lucene/index/BaseCompositeReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/BaseCompositeReader.java
@@ -117,6 +117,15 @@ public abstract class BaseCompositeReader<R extends IndexReader> extends Composi
     ensureOpen();
     TermVectors[] subVectors = new TermVectors[subReaders.length];
     return new TermVectors() {
+      @Override
+      public void prefetch(int docID) throws IOException {
+        final int i = readerIndex(docID); // find subreader num
+        if (subVectors[i] == null) {
+          subVectors[i] = subReaders[i].termVectors();
+        }
+        subVectors[i].prefetch(docID - starts[i]);
+      }
+
       @Override
       public Fields get(int docID) throws IOException {
         final int i = readerIndex(docID); // find subreader num
diff --git a/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java b/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java
index d768aaf1ab2..e61501a1244 100644
--- a/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java
+++ b/lucene/core/src/java/org/apache/lucene/index/CheckIndex.java
@@ -3810,6 +3810,9 @@ public final class CheckIndex implements Closeable {
       if (vectorsReader != null) {
         vectorsReader = vectorsReader.getMergeInstance();
         for (int j = 0; j < reader.maxDoc(); ++j) {
+          if ((j & 0x03) == 0) {
+            vectorsReader.prefetch(j);
+          }
           // Intentionally pull/visit (but don't count in
           // stats) deleted documents to make sure they too
           // are not corrupt:
diff --git a/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java b/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java
index 68bc8b1f807..e7b13e821a2 100644
--- a/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/ParallelLeafReader.java
@@ -48,6 +48,7 @@ import org.apache.lucene.util.Version;
  * behavior</em>.
  */
 public class ParallelLeafReader extends LeafReader {
+
   private final FieldInfos fieldInfos;
   private final LeafReader[] parallelReaders, storedFieldsReaders;
   private final Set<LeafReader> completeReaderSet =
@@ -325,14 +326,32 @@ public class ParallelLeafReader extends LeafReader {
   @Override
   public TermVectors termVectors() throws IOException {
     ensureOpen();
-    // TODO: optimize
+
+    Map<LeafReader, TermVectors> readerToTermVectors = new IdentityHashMap<>();
+    for (LeafReader reader : parallelReaders) {
+      if (reader.getFieldInfos().hasVectors()) {
+        TermVectors termVectors = reader.termVectors();
+        readerToTermVectors.put(reader, termVectors);
+      }
+    }
+
     return new TermVectors() {
+      @Override
+      public void prefetch(int docID) throws IOException {
+        // Prefetch all vectors. Note that this may be wasteful if the consumer doesn't need to read
+        // all the fields but we have no way to know what fields the consumer needs.
+        for (TermVectors termVectors : readerToTermVectors.values()) {
+          termVectors.prefetch(docID);
+        }
+      }
+
       @Override
       public Fields get(int docID) throws IOException {
         ParallelFields fields = null;
         for (Map.Entry<String, LeafReader> ent : tvFieldToReader.entrySet()) {
           String fieldName = ent.getKey();
-          Terms vector = ent.getValue().termVectors().get(docID, fieldName);
+          TermVectors termVectors = readerToTermVectors.get(ent.getValue());
+          Terms vector = termVectors.get(docID, fieldName);
           if (vector != null) {
             if (fields == null) {
               fields = new ParallelFields();
diff --git a/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java b/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
index 4d05d241e69..57836ae482f 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SlowCodecReaderWrapper.java
@@ -296,6 +296,11 @@ public final class SlowCodecReaderWrapper {
       throw new UncheckedIOException(e);
     }
     return new TermVectorsReader() {
+      @Override
+      public void prefetch(int docID) throws IOException {
+        termVectors.prefetch(docID);
+      }
+
       @Override
       public Fields get(int docID) throws IOException {
         return termVectors.get(docID);
diff --git a/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java b/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java
index 6d8342908ef..148ead9cb2e 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SlowCompositeCodecReaderWrapper.java
@@ -244,6 +244,15 @@ final class SlowCompositeCodecReaderWrapper extends CodecReader {
       }
     }
 
+    @Override
+    public void prefetch(int doc) throws IOException {
+      int readerId = docIdToReaderId(doc);
+      TermVectorsReader reader = readers[readerId];
+      if (reader != null) {
+        reader.prefetch(doc - docStarts[readerId]);
+      }
+    }
+
     @Override
     public Fields get(int doc) throws IOException {
       int readerId = docIdToReaderId(doc);
diff --git a/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java b/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java
index 191111c0ad3..fee0fc2f730 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SortingCodecReader.java
@@ -655,6 +655,11 @@ public final class SortingCodecReader extends FilterCodecReader {
 
   private TermVectorsReader newTermVectorsReader(TermVectorsReader delegate) {
     return new TermVectorsReader() {
+      @Override
+      public void prefetch(int doc) throws IOException {
+        delegate.prefetch(docMap.newToOld(doc));
+      }
+
       @Override
       public Fields get(int doc) throws IOException {
         return delegate.get(docMap.newToOld(doc));
diff --git a/lucene/core/src/java/org/apache/lucene/index/TermVectors.java b/lucene/core/src/java/org/apache/lucene/index/TermVectors.java
index fe2e8b56681..fde57b48c93 100644
--- a/lucene/core/src/java/org/apache/lucene/index/TermVectors.java
+++ b/lucene/core/src/java/org/apache/lucene/index/TermVectors.java
@@ -18,6 +18,7 @@ package org.apache.lucene.index;
 
 import java.io.IOException;
 import org.apache.lucene.analysis.tokenattributes.OffsetAttribute; // javadocs
+import org.apache.lucene.store.IndexInput;
 
 /**
  * API for reading term vectors.
@@ -30,6 +31,18 @@ public abstract class TermVectors {
   /** Sole constructor. (For invocation by subclass constructors, typically implicit.) */
   protected TermVectors() {}
 
+  /**
+   * Optional method: Give a hint to this {@link TermVectors} instance that the given document will
+   * be read in the near future. This typically delegates to {@link IndexInput#prefetch} and is
+   * useful to parallelize I/O across multiple documents.
+   *
+   * <p>NOTE: This API is expected to be called on a small enough set of doc IDs that they could all
+   * fit in the page cache. If you plan on retrieving a very large number of documents, it may be a
+   * good idea to perform calls to {@link #prefetch} and {@link #get} in batches instead of
+   * prefetching all documents up-front.
+   */
+  public void prefetch(int docID) throws IOException {}
+
   /**
    * Returns term vectors for this document, or null if term vectors were not indexed.
    *
diff --git a/lucene/core/src/test/org/apache/lucene/codecs/lucene90/TestLucene90TermVectorsFormat.java b/lucene/core/src/test/org/apache/lucene/codecs/lucene90/TestLucene90TermVectorsFormat.java
index 6660b9b6429..8cc6dda1c1d 100644
--- a/lucene/core/src/test/org/apache/lucene/codecs/lucene90/TestLucene90TermVectorsFormat.java
+++ b/lucene/core/src/test/org/apache/lucene/codecs/lucene90/TestLucene90TermVectorsFormat.java
@@ -16,7 +16,24 @@
  */
 package org.apache.lucene.codecs.lucene90;
 
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.FieldType;
+import org.apache.lucene.document.TextField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.IndexWriter;
+import org.apache.lucene.index.IndexWriterConfig;
+import org.apache.lucene.index.TermVectors;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.store.FilterDirectory;
+import org.apache.lucene.store.FilterIndexInput;
+import org.apache.lucene.store.IOContext;
+import org.apache.lucene.store.IndexInput;
+import org.apache.lucene.tests.codecs.compressing.dummy.DummyCompressingCodec;
 import org.apache.lucene.tests.index.BaseTermVectorsFormatTestCase;
 import org.apache.lucene.tests.util.TestUtil;
 
@@ -25,4 +42,84 @@ public class TestLucene90TermVectorsFormat extends BaseTermVectorsFormatTestCase
   protected Codec getCodec() {
     return TestUtil.getDefaultCodec();
   }
+
+  private static class CountingPrefetchDirectory extends FilterDirectory {
+
+    private final AtomicInteger counter;
+
+    CountingPrefetchDirectory(Directory in, AtomicInteger counter) {
+      super(in);
+      this.counter = counter;
+    }
+
+    @Override
+    public IndexInput openInput(String name, IOContext context) throws IOException {
+      return new CountingPrefetchIndexInput(super.openInput(name, context), counter);
+    }
+  }
+
+  private static class CountingPrefetchIndexInput extends FilterIndexInput {
+
+    private final AtomicInteger counter;
+
+    public CountingPrefetchIndexInput(IndexInput input, AtomicInteger counter) {
+      super(input.toString(), input);
+      this.counter = counter;
+    }
+
+    @Override
+    public void prefetch(long offset, long length) throws IOException {
+      in.prefetch(offset, length);
+      counter.incrementAndGet();
+    }
+
+    @Override
+    public IndexInput clone() {
+      return new CountingPrefetchIndexInput(in.clone(), counter);
+    }
+
+    @Override
+    public IndexInput slice(String sliceDescription, long offset, long length) throws IOException {
+      return new CountingPrefetchIndexInput(in.slice(sliceDescription, offset, length), counter);
+    }
+  }
+
+  public void testSkipRedundantPrefetches() throws IOException {
+    // Use the "dummy" codec, which has the same base class as Lucene90StoredFieldsFormat but allows
+    // configuring the number of docs per chunk.
+    Codec codec = new DummyCompressingCodec(1 << 10, 2, false, 16);
+    try (Directory origDir = newDirectory()) {
+      AtomicInteger counter = new AtomicInteger();
+      Directory dir = new CountingPrefetchDirectory(origDir, counter);
+      try (IndexWriter w = new IndexWriter(dir, new IndexWriterConfig().setCodec(codec))) {
+        FieldType ft = new FieldType(TextField.TYPE_NOT_STORED);
+        ft.setStoreTermVectors(true);
+        for (int i = 0; i < 100; ++i) {
+          Document doc = new Document();
+          doc.add(new Field("content", Integer.toString(i), ft));
+          w.addDocument(doc);
+        }
+        w.forceMerge(1);
+      }
+
+      try (IndexReader reader = DirectoryReader.open(dir)) {
+        TermVectors termVectors = reader.termVectors();
+        counter.set(0);
+        assertEquals(0, counter.get());
+        termVectors.prefetch(0);
+        assertEquals(1, counter.get());
+        termVectors.prefetch(1);
+        // This format has 2 docs per block, so the second prefetch is skipped
+        assertEquals(1, counter.get());
+        termVectors.prefetch(15);
+        assertEquals(2, counter.get());
+        termVectors.prefetch(14);
+        // 14 is in the same block as 15, so the prefetch was skipped
+        assertEquals(2, counter.get());
+        // Already prefetched in the past, so skipped again
+        termVectors.prefetch(1);
+        assertEquals(2, counter.get());
+      }
+    }
+  }
 }
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 3151754a075..02b5c098afd 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
@@ -153,6 +153,12 @@ public class AssertingLeafReader extends FilterLeafReader {
       this.in = in;
     }
 
+    @Override
+    public void prefetch(int docID) throws IOException {
+      assertThread("TermVectors", creationThread);
+      in.prefetch(docID);
+    }
+
     @Override
     public Fields get(int doc) throws IOException {
       assertThread("TermVectors", creationThread);