1 files changed, 20 insertions, 2 deletions
diff --git a/tensorflow/core/lib/io/record_reader.h b/tensorflow/core/lib/io/record_reader.h
index f96d24aa31..b4c56451be 100644
--- a/tensorflow/core/lib/io/record_reader.h
+++ b/tensorflow/core/lib/io/record_reader.h
@@ -18,6 +18,9 @@ limitations under the License.
 
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/io/inputbuffer.h"
+#include "tensorflow/core/lib/io/zlib_compression_options.h"
+#include "tensorflow/core/lib/io/zlib_inputbuffer.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -27,13 +30,23 @@ class RandomAccessFile;
 
 namespace io {
 
+class RecordReaderOptions {
+ public:
+  enum CompressionType { NONE = 0, ZLIB_COMPRESSION = 1 };
+  CompressionType compression_type = NONE;
+
+  // Options specific to zlib compression.
+  ZlibCompressionOptions zlib_options;
+};
+
 class RecordReader {
  public:
   // Create a reader that will return log records from "*file".
   // "*file" must remain live while this Reader is in use.
-  explicit RecordReader(RandomAccessFile* file);
+  RecordReader(RandomAccessFile* file,
+               const RecordReaderOptions& options = RecordReaderOptions());
 
-  ~RecordReader();
+  virtual ~RecordReader();
 
   // Read the record at "*offset" into *record and update *offset to
   // point to the offset of the next record.  Returns OK on success,
@@ -41,7 +54,12 @@ class RecordReader {
   Status ReadRecord(uint64* offset, string* record);
 
  private:
+  Status ReadChecksummed(uint64 offset, size_t n, StringPiece* result,
+                         string* storage);
+
   RandomAccessFile* src_;
+  RecordReaderOptions options_;
+  std::unique_ptr<ZlibInputBuffer> zlib_input_buffer_;
 
   TF_DISALLOW_COPY_AND_ASSIGN(RecordReader);
 };