You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by dn...@apache.org on 2024/09/16 18:43:29 UTC

(lucene) 01/01: Allow copy stored documents with index sorts

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

dnhatn pushed a commit to branch copy-docs
in repository https://gitbox.apache.org/repos/asf/lucene.git

commit b93a2a1f0935a4e5af8e4d07ada2673c2344520a
Author: Nhat Nguyen <nh...@elastic.co>
AuthorDate: Sun Sep 15 22:39:49 2024 -0700

    Allow copy stored documents with index sorts
---
 .../apache/lucene/codecs/StoredFieldsWriter.java   | 71 ++++++++++++++++++++++
 .../Lucene90CompressingStoredFieldsWriter.java     | 16 +++++
 .../lucene/index/SortingStoredFieldsConsumer.java  | 62 +------------------
 3 files changed, 90 insertions(+), 59 deletions(-)

diff --git a/lucene/core/src/java/org/apache/lucene/codecs/StoredFieldsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/StoredFieldsWriter.java
index 6824b227f68..1a94768b3cf 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/StoredFieldsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/StoredFieldsWriter.java
@@ -28,6 +28,7 @@ import org.apache.lucene.index.DocIDMerger;
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.FieldInfos;
 import org.apache.lucene.index.MergeState;
+import org.apache.lucene.index.Sorter;
 import org.apache.lucene.index.StoredFieldVisitor;
 import org.apache.lucene.store.DataInput;
 import org.apache.lucene.util.Accountable;
@@ -157,6 +158,76 @@ public abstract class StoredFieldsWriter implements Closeable, Accountable {
     return docCount;
   }
 
+  /**
+   * Copies stored fields from the documents of the given reader. The documents are reordered based
+   * on the docMap, if provided. Subclasses can override this method to optimize the copy process.
+   *
+   * @param reader the input reader containing the documents
+   * @param docMap the mapping to reorder document IDs, or null
+   * @param numDocs the number of documents to copy
+   * @throws IOException if an I/O error occurs during reading or writing
+   */
+  public void copyDocs(StoredFieldsReader reader, Sorter.DocMap docMap, int numDocs)
+      throws IOException {
+    CopyVisitor visitor = new CopyVisitor(this);
+    for (int docID = 0; docID < numDocs; docID++) {
+      startDocument();
+      reader.document(docMap == null ? docID : docMap.newToOld(docID), visitor);
+      finishDocument();
+    }
+  }
+
+  /** A visitor that copies every field it sees in the provided {@link StoredFieldsWriter}. */
+  private static class CopyVisitor extends StoredFieldVisitor {
+    final StoredFieldsWriter writer;
+
+    CopyVisitor(StoredFieldsWriter writer) {
+      this.writer = writer;
+    }
+
+    @Override
+    public void binaryField(FieldInfo fieldInfo, DataInput value, int length) throws IOException {
+      writer.writeField(fieldInfo, value, length);
+    }
+
+    @Override
+    public void binaryField(FieldInfo fieldInfo, byte[] value) throws IOException {
+      // TODO: can we avoid new BR here?
+      writer.writeField(fieldInfo, new BytesRef(value));
+    }
+
+    @Override
+    public void stringField(FieldInfo fieldInfo, String value) throws IOException {
+      writer.writeField(
+          fieldInfo, Objects.requireNonNull(value, "String value should not be null"));
+    }
+
+    @Override
+    public void intField(FieldInfo fieldInfo, int value) throws IOException {
+      writer.writeField(fieldInfo, value);
+    }
+
+    @Override
+    public void longField(FieldInfo fieldInfo, long value) throws IOException {
+      writer.writeField(fieldInfo, value);
+    }
+
+    @Override
+    public void floatField(FieldInfo fieldInfo, float value) throws IOException {
+      writer.writeField(fieldInfo, value);
+    }
+
+    @Override
+    public void doubleField(FieldInfo fieldInfo, double value) throws IOException {
+      writer.writeField(fieldInfo, value);
+    }
+
+    @Override
+    public Status needsField(FieldInfo fieldInfo) throws IOException {
+      return Status.YES;
+    }
+  }
+
   /**
    * A visitor that adds every field it sees.
    *
diff --git a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingStoredFieldsWriter.java b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingStoredFieldsWriter.java
index 0356a1fa448..954ec6ccd68 100644
--- a/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingStoredFieldsWriter.java
+++ b/lucene/core/src/java/org/apache/lucene/codecs/lucene90/compressing/Lucene90CompressingStoredFieldsWriter.java
@@ -34,6 +34,7 @@ import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.MergeState;
 import org.apache.lucene.index.SegmentInfo;
+import org.apache.lucene.index.Sorter;
 import org.apache.lucene.store.ByteBuffersDataInput;
 import org.apache.lucene.store.ByteBuffersDataOutput;
 import org.apache.lucene.store.DataInput;
@@ -598,6 +599,21 @@ public final class Lucene90CompressingStoredFieldsWriter extends StoredFieldsWri
     }
   }
 
+  @Override
+  public void copyDocs(StoredFieldsReader storedFieldsReader, Sorter.DocMap docMap, int numDocs)
+      throws IOException {
+    if (storedFieldsReader instanceof Lucene90CompressingStoredFieldsReader
+        && ((Lucene90CompressingStoredFieldsReader) storedFieldsReader).getVersion()
+            == VERSION_CURRENT) {
+      var reader = (Lucene90CompressingStoredFieldsReader) storedFieldsReader;
+      for (int i = 0; i < numDocs; i++) {
+        copyOneDoc(reader, docMap != null ? docMap.newToOld(i) : i);
+      }
+    } else {
+      super.copyDocs(storedFieldsReader, docMap, numDocs);
+    }
+  }
+
   @Override
   public int merge(MergeState mergeState) throws IOException {
     final MatchingReaders matchingReaders = new MatchingReaders(mergeState);
diff --git a/lucene/core/src/java/org/apache/lucene/index/SortingStoredFieldsConsumer.java b/lucene/core/src/java/org/apache/lucene/index/SortingStoredFieldsConsumer.java
index c01de610772..5c3f9e13f2c 100644
--- a/lucene/core/src/java/org/apache/lucene/index/SortingStoredFieldsConsumer.java
+++ b/lucene/core/src/java/org/apache/lucene/index/SortingStoredFieldsConsumer.java
@@ -18,7 +18,6 @@
 package org.apache.lucene.index;
 
 import java.io.IOException;
-import java.util.Objects;
 import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.codecs.StoredFieldsFormat;
 import org.apache.lucene.codecs.StoredFieldsReader;
@@ -104,13 +103,9 @@ final class SortingStoredFieldsConsumer extends StoredFieldsConsumer {
         codec.storedFieldsFormat().fieldsWriter(state.directory, state.segmentInfo, state.context);
     try {
       reader.checkIntegrity();
-      CopyVisitor visitor = new CopyVisitor(sortWriter);
-      for (int docID = 0; docID < state.segmentInfo.maxDoc(); docID++) {
-        sortWriter.startDocument();
-        reader.document(sortMap == null ? docID : sortMap.newToOld(docID), visitor);
-        sortWriter.finishDocument();
-      }
-      sortWriter.finish(state.segmentInfo.maxDoc());
+      int maxDocs = state.segmentInfo.maxDoc();
+      sortWriter.copyDocs(reader, sortMap, maxDocs);
+      sortWriter.finish(maxDocs);
     } finally {
       IOUtils.close(reader, sortWriter);
       IOUtils.deleteFiles(tmpDirectory, tmpDirectory.getTemporaryFiles().values());
@@ -128,55 +123,4 @@ final class SortingStoredFieldsConsumer extends StoredFieldsConsumer {
       }
     }
   }
-
-  /** A visitor that copies every field it sees in the provided {@link StoredFieldsWriter}. */
-  private static class CopyVisitor extends StoredFieldVisitor {
-    final StoredFieldsWriter writer;
-
-    CopyVisitor(StoredFieldsWriter writer) {
-      this.writer = writer;
-    }
-
-    @Override
-    public void binaryField(FieldInfo fieldInfo, DataInput value, int length) throws IOException {
-      writer.writeField(fieldInfo, value, length);
-    }
-
-    @Override
-    public void binaryField(FieldInfo fieldInfo, byte[] value) throws IOException {
-      // TODO: can we avoid new BR here?
-      writer.writeField(fieldInfo, new BytesRef(value));
-    }
-
-    @Override
-    public void stringField(FieldInfo fieldInfo, String value) throws IOException {
-      writer.writeField(
-          fieldInfo, Objects.requireNonNull(value, "String value should not be null"));
-    }
-
-    @Override
-    public void intField(FieldInfo fieldInfo, int value) throws IOException {
-      writer.writeField(fieldInfo, value);
-    }
-
-    @Override
-    public void longField(FieldInfo fieldInfo, long value) throws IOException {
-      writer.writeField(fieldInfo, value);
-    }
-
-    @Override
-    public void floatField(FieldInfo fieldInfo, float value) throws IOException {
-      writer.writeField(fieldInfo, value);
-    }
-
-    @Override
-    public void doubleField(FieldInfo fieldInfo, double value) throws IOException {
-      writer.writeField(fieldInfo, value);
-    }
-
-    @Override
-    public Status needsField(FieldInfo fieldInfo) throws IOException {
-      return Status.YES;
-    }
-  }
 }