Improve the speed of text file position restore

Author: cbalioglu

native/python/src/fairseq2n/bindings/data/text/text_reader.cc

@@ -53,14 +53,14 @@ def_text_reader(py::module_ &text_module)
             return read_text(std::move(path), std::move(key), std::move(opts));
         },
         py::arg("path"),
-        py::arg("key")         = std::nullopt,
-        py::arg("encoding")    = std::nullopt,
-        py::arg("line_ending") = line_ending::infer,
-        py::arg("ltrim")       = false,
-        py::arg("rtrim")       = false,
-        py::arg("skip_empty")  = false,
-        py::arg("memory_map")  = false,
-        py::arg("block_size")  = std::nullopt);
+        py::arg("key")          = std::nullopt,
+        py::arg("encoding")     = std::nullopt,
+        py::arg("line_ending")  = line_ending::infer,
+        py::arg("ltrim")        = false,
+        py::arg("rtrim")        = false,
+        py::arg("skip_empty")   = false,
+        py::arg("memory_map")   = false,
+        py::arg("block_size")   = std::nullopt);
 }
 
 }  // namespace fairseq2n

native/src/fairseq2n/data/byte_stream.h

@@ -10,6 +10,7 @@
 
 #include "fairseq2n/api.h"
 #include "fairseq2n/memory.h"
+#include "fairseq2n/data/tape.h"
 
 namespace fairseq2n {
 
@@ -29,8 +30,23 @@ class FAIRSEQ2_API byte_stream {
     virtual memory_block
     read_chunk() = 0;
 
+    virtual void
+    seek(std::size_t offset) = 0;
+
+    virtual std::size_t
+    position() const = 0;
+
     virtual void
     reset() = 0;
+
+    virtual void
+    record_position(tape &t) const = 0;
+
+    virtual void
+    reload_position(tape &t) = 0;
+
+    virtual bool
+    supports_seek() const noexcept = 0;
 };
 
 class FAIRSEQ2_API byte_stream_error : public std::runtime_error {

native/src/fairseq2n/data/file.cc

@@ -83,9 +83,10 @@ open_file(const std::filesystem::path &path, const file_options &opts)
     } else
         stream = std::make_unique<file_stream>(std::move(fd), path, chunk_size);
 
-    if (opts.mode() == file_mode::text)
-        stream = std::make_unique<utf8_stream>(
-            std::move(stream), opts.maybe_text_encoding(), chunk_size);
+// TODO: fix!
+//    if (opts.mode() == file_mode::text)
+//        stream = std::make_unique<utf8_stream>(
+//            std::move(stream), opts.maybe_text_encoding(), chunk_size);
 
     return stream;
 }

native/src/fairseq2n/data/file.h

@@ -83,6 +83,7 @@ class file_options {
         return maybe_text_encoding_;
     }
 
+
 private:
     file_mode mode_ = file_mode::binary;
     bool memory_map_ = false;

native/src/fairseq2n/data/file_stream.cc

@@ -61,23 +61,48 @@ file_stream::read_chunk()
     return chunk.share_first(chunk.size() - remaining_space.size());
 }
 
+void
+file_stream::seek(std::size_t offset)
+{
+    seek(offset, /*from_current=*/false);
+}
+
+std::size_t
+file_stream::position() const
+{
+    return seek(0, /*from_current=*/true);
+}
+
 void
 file_stream::reset()
 {
-    ::off_t offset = ::lseek(fd_.get(), 0, SEEK_SET);
-    if (offset == -1) {
-        std::error_code err = last_error();
+    seek(0);
 
-        if (err == std::errc::invalid_seek)
-            throw_<byte_stream_error>("'{}' is not seekable and cannot be reset.", path_.string());
+    is_eod_ = false;
+}
 
-        throw_system_error(err,
-            "'{}' cannot be reset", path_.string());
-    }
+void
+file_stream::record_position(tape &t) const
+{
+    std::size_t offset = position();
+
+    t.record(offset);
+}
+
+void
+file_stream::reload_position(tape &t)
+{
+    seek(t.read<std::size_t>());
 
     is_eod_ = false;
 }
 
+bool
+file_stream::supports_seek() const noexcept
+{
+    return true;
+}
+
 std::size_t
 file_stream::fill_chunk(writable_memory_span chunk)
 {
@@ -89,4 +114,22 @@ file_stream::fill_chunk(writable_memory_span chunk)
     return static_cast<std::size_t>(num_bytes_read);
 }
 
+std::size_t
+file_stream::seek(std::size_t offset, bool from_current) const
+{
+    ::off_t off = ::lseek(
+        fd_.get(), static_cast<::off_t>(offset), from_current ? SEEK_CUR : SEEK_SET);
+
+    if (off != -1)
+        return static_cast<std::size_t>(off);
+
+    std::error_code err = last_error();
+
+    if (err == std::errc::invalid_seek)
+        throw_<byte_stream_error>("'{}' is not seekable.", path_.string());
+
+    throw_system_error(err,
+        "'{}' cannot be read", path_.string());
+}
+
 }  // namespace fairseq2n::detail

native/src/fairseq2n/data/file_stream.h

@@ -29,13 +29,31 @@ class file_stream final : public byte_stream {
     memory_block
     read_chunk() override;
 
+    void
+    seek(std::size_t offset) override;
+
+    std::size_t
+    position() const override;
+
     void
     reset() override;
 
+    void
+    record_position(tape &t) const override;
+
+    void
+    reload_position(tape &t) override;
+
+    bool
+    supports_seek() const noexcept override;
+
 private:
     std::size_t
     fill_chunk(writable_memory_span chunk);
 
+    std::size_t
+    seek(std::size_t offset, bool from_current) const;
+
 private:
     file_desc fd_;
     std::filesystem::path path_;

native/src/fairseq2n/data/memory_stream.cc

@@ -14,10 +14,47 @@ memory_stream::read_chunk()
     return std::exchange(block_, {});
 }
 
+void
+memory_stream::seek(std::size_t offset)
+{
+    if (offset >= block_.size())
+        block_ = {};
+    else
+        block_ = block_.share_slice(offset);
+}
+
+std::size_t
+memory_stream::position() const
+{
+    return original_block_.size() - block_.size();
+}
+
 void
 memory_stream::reset()
 {
     block_ = original_block_;
 }
 
+void
+memory_stream::record_position(tape &t) const
+{
+    std::size_t offset = position();
+
+    t.record(offset);
+}
+
+void
+memory_stream::reload_position(tape &t)
+{
+    block_ = original_block_;
+
+    seek(t.read<std::size_t>());
+}
+
+bool
+memory_stream::supports_seek() const noexcept
+{
+    return true;
+}
+
 }  // namespace fairseq2n::detail

native/src/fairseq2n/data/memory_stream.h

@@ -25,9 +25,24 @@ class memory_stream final : public byte_stream {
     memory_block
     read_chunk() override;
 
+    void
+    seek(std::size_t offset) override;
+
+    std::size_t
+    position() const override;
+
     void
     reset() override;
 
+    void
+    record_position(tape &t) const override;
+
+    void
+    reload_position(tape &t) override;
+
+    bool
+    supports_seek() const noexcept override;
+
 private:
     memory_block block_;
     memory_block original_block_;

native/src/fairseq2n/data/record_reader.cc

@@ -7,6 +7,7 @@
 #include "fairseq2n/data/record_reader.h"
 
 #include <algorithm>
+#include <stdexcept>
 
 #include <fmt/format.h>
 
@@ -34,11 +35,41 @@ record_reader::next()
 void
 record_reader::reset()
 {
+    stream_->reset();
+
     current_chunk_ = {};
 
-    previous_chunks_.clear();
+    stream_offset_ = 0;
 
-    stream_->reset();
+    chunk_offset_ = 0;
+}
+
+void
+record_reader::record_position(tape &t) const
+{
+    t.record(stream_offset_);
+
+    t.record(chunk_offset_);
+}
+
+void
+record_reader::reload_position(tape &t)
+{
+    stream_offset_ = t.read<std::size_t>();
+
+    stream_->seek(stream_offset_);
+
+    chunk_offset_ = t.read<std::size_t>();
+    if (chunk_offset_ > 0) {
+        current_chunk_ = stream_->read_chunk();
+
+        if (chunk_offset_ >= current_chunk_.size())
+            throw_<std::invalid_argument>(
+                "The tape is corrupt. The state of the record reader cannot be restored.");
+
+        current_chunk_ = current_chunk_.share_slice(chunk_offset_);
+    } else
+        current_chunk_ = {};
 }
 
 bool
@@ -52,6 +83,8 @@ record_reader::load_next_record()
 
     // Load and store memory chunks until we find the end of the next record.
     while (!(maybe_record_end_offset = maybe_find_record_end(current_chunk_, first_chunk))) {
+        stream_offset_ = stream_->position();
+
         memory_block next_chunk = stream_->read_chunk();
         if (next_chunk.empty()) {
             // If `next_chunk` is empty and we don't have any partial record
@@ -71,6 +104,8 @@ record_reader::load_next_record()
 
         current_chunk_ = std::move(next_chunk);
 
+        chunk_offset_ = 0;
+
         first_chunk = false;
     }
 
@@ -117,6 +152,8 @@ record_reader::move_to_next_record()
 {
     current_chunk_ = current_chunk_.share_slice(record_end_offset_);
 
+    chunk_offset_ += record_end_offset_;
+
     previous_chunks_.clear();
 }
 

native/src/fairseq2n/data/record_reader.h

@@ -42,6 +42,12 @@ class FAIRSEQ2_API record_reader {
     void
     reset();
 
+    void
+    record_position(tape &t) const;
+
+    void
+    reload_position(tape &);
+
 private:
     bool
     load_next_record();
@@ -62,6 +68,8 @@ class FAIRSEQ2_API record_reader {
     std::unique_ptr<byte_stream> stream_;
     memory_block current_chunk_{};
     std::vector<memory_block> previous_chunks_{};
+    std::size_t stream_offset_ = 0;
+    std::size_t chunk_offset_ = 0;
     std::size_t record_len_ = 0;
     std::size_t record_end_offset_ = 0;
 };