[tf.data] Saveable iterator for TFRecordDatasetOp.

PiperOrigin-RevId: 176020932

2 files changed, 32 insertions, 0 deletions

diff --git a/tensorflow/core/lib/io/record_reader.cc b/tensorflow/core/lib/io/record_reader.cc
index c3b87ee5bf..403c82818e 100644
--- a/tensorflow/core/lib/io/record_reader.cc
+++ b/tensorflow/core/lib/io/record_reader.cc
@@ -196,6 +196,19 @@ Status RecordReader::ReadRecord(uint64* offset, string* record) {
   return Status::OK();
 }
 
+Status RecordReader::SkipNBytes(uint64 offset) {
+#if !defined(IS_SLIM_BUILD)
+  if (zlib_input_stream_) {
+    TF_RETURN_IF_ERROR(zlib_input_stream_->SkipNBytes(offset));
+  } else {
+#endif
+    if (options_.buffer_size > 0) {
+      TF_RETURN_IF_ERROR(input_stream_->SkipNBytes(offset));
+    }
+  }
+  return Status::OK();
+}
+
 SequentialRecordReader::SequentialRecordReader(
     RandomAccessFile* file, const RecordReaderOptions& options)
     : underlying_(file, options), offset_(0) {}
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;