Add verbatim stream copying support (#360)

Summary:
Pull Request resolved: #360

- Add verbatim_copy() function and VRSWriter methods (add_verbatim_copy_streams/copy_verbatim_records) for efficient C++ verbatim VRS stream copying
- Add getRecordFileReader() accessor to OssVRSReader
- Add tests and regenerate pyi stubs

Reviewed By: kongchen1992

Differential Revision: D96868589

Author: danielyan86129

csrc/reader/VRSReader.h

@@ -161,6 +161,11 @@ class OssVRSReader : public VRSReaderBase {
   /// Initialize the module by calling initVrsBindings()
   void init();
 
+  /// Get a reference to the underlying RecordFileReader.
+  RecordFileReader& getRecordFileReader() {
+    return reader_;
+  }
+
   void open(const string& path);
   void open(const PyFileSpec& spec);
 

csrc/writer/VRSWriter.cpp

@@ -138,4 +138,37 @@ int VRSWriter::close() {
   return writer_.waitForFileClosed();
 }
 
+int VRSWriter::addVerbatimCopyStreams(
+    RecordFileReader& reader,
+    const std::vector<std::string>& streamIds) {
+  verbatimReader_ = &reader;
+  verbatimCopyOptions_ = std::make_unique<vrs::utils::CopyOptions>(false);
+  for (const auto& sidStr : streamIds) {
+    auto sid = StreamId::fromNumericName(sidStr);
+    if (!sid.isValid()) {
+      throw py::value_error("Invalid stream ID: " + sidStr);
+    }
+    verbatimCopiers_.push_back(
+        std::make_unique<vrs::utils::Copier>(reader, writer_, sid, *verbatimCopyOptions_));
+    verbatimStreamIds_.insert(sid);
+  }
+  return 0;
+}
+
+int VRSWriter::copyVerbatimRecords() {
+  if (!verbatimReader_ || verbatimStreamIds_.empty()) {
+    return 0;
+  }
+  const auto& index = verbatimReader_->getIndex();
+  for (const auto& record : index) {
+    if (verbatimStreamIds_.count(record.streamId) > 0) {
+      int status = verbatimReader_->readRecord(record);
+      if (status != 0) {
+        return status;
+      }
+    }
+  }
+  return 0;
+}
+
 } // namespace pyvrs

csrc/writer/VRSWriter.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <memory>
+#include <set>
 #include <string>
 #include <vector>
 
@@ -29,8 +30,10 @@
 #endif
 
 #include <vrs/DataLayout.h>
+#include <vrs/RecordFileReader.h>
 #include <vrs/RecordFileWriter.h>
 #include <vrs/Recordable.h>
+#include <vrs/utils/FilterCopyHelpers.h>
 
 #include "PyDataPiece.h"
 
@@ -76,6 +79,18 @@ class VRSWriter {
 
   void addRecordable(Recordable* recordable);
 
+  /// Register streams from a RecordFileReader for verbatim copying.
+  /// Copiers are created that will copy records for the specified streams.
+  /// Must be called before create() / first writeRecords().
+  /// WARNING: The reader must remain alive until after copyVerbatimRecords() is called.
+  int addVerbatimCopyStreams(RecordFileReader& reader, const std::vector<std::string>& streamIds);
+
+  /// Copy all registered verbatim stream records into the writer.
+  /// Reads records from the reader for the registered streams and copies them verbatim.
+  /// Call after all processed records have been written.
+  /// NOTE: Assumes the reader passed to addVerbatimCopyStreams() is still valid.
+  int copyVerbatimRecords();
+
   int writeRecords(double maxTimestamp);
 
   uint64_t getBackgroundThreadQueueByteSize();
@@ -89,6 +104,11 @@ class VRSWriter {
  private:
   RecordFileWriter writer_;
   std::vector<std::unique_ptr<PyStream>> streams_;
+  // Non-owning pointer; caller must keep reader alive until after copyVerbatimRecords() is called.
+  RecordFileReader* verbatimReader_ = nullptr;
+  std::vector<std::unique_ptr<vrs::utils::Copier>> verbatimCopiers_;
+  std::set<StreamId> verbatimStreamIds_;
+  std::unique_ptr<vrs::utils::CopyOptions> verbatimCopyOptions_;
 };
 
 } // namespace pyvrs

csrc/writer/Writer.cpp

@@ -33,7 +33,10 @@
 #include <pybind11/stl.h>
 
 #include <vrs/os/Platform.h>
+#include <vrs/utils/FilterCopy.h>
 
+#include "../VrsBindings.h"
+#include "../reader/VRSReader.h"
 #include "PyRecordable.h"
 #include "VRSWriter.h"
 
@@ -73,6 +76,38 @@ using namespace vrs;
       }));
 
 void pybind_writer(py::module& m) {
+  m.def(
+      "verbatim_copy",
+      [](const std::string& inputPath,
+         const std::string& outputPath,
+         const std::vector<std::string>& streamIds) {
+        initVrsBindings();
+        vrs::utils::FilteredFileReader filteredReader(inputPath);
+        int status = filteredReader.openFile();
+        if (status != 0) {
+          throw py::value_error("Failed to open input file: " + inputPath);
+        }
+        if (!streamIds.empty()) {
+          filteredReader.filter.streams.clear();
+          for (const auto& sid : streamIds) {
+            auto id = StreamId::fromNumericName(sid);
+            if (!id.isValid()) {
+              throw py::value_error("Invalid stream ID: " + sid);
+            }
+            filteredReader.filter.streams.insert(id);
+          }
+        }
+        vrs::utils::CopyOptions copyOptions(false);
+        status = vrs::utils::filterCopy(filteredReader, outputPath, copyOptions);
+        if (status != 0) {
+          throw std::runtime_error("verbatim_copy failed with error code " + to_string(status));
+        }
+        return status;
+      },
+      py::arg("input_path"),
+      py::arg("output_path"),
+      py::arg("stream_ids") = std::vector<std::string>{});
+
   py::class_<pyvrs::PyRecordFormat, std::unique_ptr<pyvrs::PyRecordFormat, py::nodelete>>(
       m, "RecordFormat")
       .def("getMembers", &pyvrs::PyRecordFormat::getMembers)
@@ -112,6 +147,16 @@ void pybind_writer(py::module& m) {
       .def("writeRecords", &pyvrs::VRSWriter::writeRecords)
       .def("getBackgroundThreadQueueByteSize", &pyvrs::VRSWriter::getBackgroundThreadQueueByteSize)
       .def("close", &pyvrs::VRSWriter::close)
+      .def(
+          "addVerbatimCopyStreams",
+          [](pyvrs::VRSWriter& self,
+             PyVRSReader& reader,
+             const std::vector<std::string>& streamIds) {
+            return self.addVerbatimCopyStreams(reader.getRecordFileReader(), streamIds);
+          },
+          py::arg("reader"),
+          py::arg("stream_ids"))
+      .def("copyVerbatimRecords", &pyvrs::VRSWriter::copyVerbatimRecords)
 #if IS_VRS_FB_INTERNAL()
 #include "Writer_fb.hpp"
 #endif

pyvrs/__init__.py

@@ -42,6 +42,7 @@
     StreamNotFoundError,
     TimestampNotFoundError,
     verbatim_checksum,
+    verbatim_copy,
     VRSRecord,
     Writer,
 )
@@ -79,6 +80,7 @@
     "StreamNotFoundError",
     "TimestampNotFoundError",
     "verbatim_checksum",
+    "verbatim_copy",
     "VRSRecord",
     "Writer",
 ]

pyvrs/writer.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import time
-from typing import Dict, List, Union
+from typing import Dict, List, Sequence, Union
 
 import numpy as np
 
@@ -101,6 +101,30 @@ def _create(self):
     def close(self) -> None:
         return self._writer.close()
 
+    def add_verbatim_copy_streams(self, reader, stream_ids: Sequence[str]) -> int:
+        """Register streams from a reader for verbatim copying into this writer.
+
+        Must be called before the first flush_records() call.
+        The reader must remain open until after copy_verbatim_records() is called.
+
+        Args:
+            reader: A pyvrs Reader (C++ binding object, e.g. SyncVRSReader._reader).
+            stream_ids: Stream IDs to copy verbatim (e.g. ["281-1", "282-1"]).
+        """
+        if self.file_created:
+            raise Exception(
+                "add_verbatim_copy_streams must be called before flush_records"
+            )
+        return self._writer.addVerbatimCopyStreams(reader, list(stream_ids))
+
+    def copy_verbatim_records(self) -> int:
+        """Copy all registered verbatim stream records into the output file.
+
+        Call after all processed records have been written, before close().
+        """
+        self._create()
+        return self._writer.copyVerbatimRecords()
+
     # Recordable instance ids are automatically assigned when Recordable objects are created.
     # This guarantees that each Recordable gets a unique ID.
     # WARNING! If your code relies on specific instance IDs, your design is weak, and you are