1 files changed, 19 insertions, 0 deletions
diff --git a/tensorflow/core/lib/io/record_reader.h b/tensorflow/core/lib/io/record_reader.h
index e4f6a5b492..62dd2efb79 100644
--- a/tensorflow/core/lib/io/record_reader.h
+++ b/tensorflow/core/lib/io/record_reader.h
@@ -74,6 +74,10 @@ class RecordReader {
   // sequential.
   Status ReadRecord(uint64* offset, string* record);
 
+  // Skip the records till "offset". Returns OK on success,
+  // OUT_OF_RANGE for end of file, or something else for an error.
+  Status SkipNBytes(uint64 offset);
+
  private:
   Status ReadChecksummed(uint64 offset, size_t n, StringPiece* result,
                          string* storage);
@@ -107,6 +111,21 @@ class SequentialRecordReader {
     return underlying_.ReadRecord(&offset_, record);
   }
 
+  // Returns the current offset in the file.
+  uint64 TellOffset() { return offset_; }
+
+  // Seek to this offset within the file and set this offset as the current
+  // offset. Trying to seek backward will throw error.
+  Status SeekOffset(uint64 offset) {
+    if (offset < offset_)
+      return errors::InvalidArgument(
+          "Trying to seek offset: ", offset,
+          " which is less than the current offset: ", offset_);
+    TF_RETURN_IF_ERROR(underlying_.SkipNBytes(offset - offset_));
+    offset_ = offset;
+    return Status::OK();
+  }
+
  private:
   RecordReader underlying_;
   uint64 offset_ = 0;