aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/python/lib/io/py_record_reader.cc
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2016-09-22 10:00:16 -0800
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2016-09-22 11:02:46 -0700
commitb6c8def190bc3d885dc7ae9c3c570bdf94031378 (patch)
tree20c72c0d06f204e5ca2e2a23b75aef65a180f0e1 /tensorflow/python/lib/io/py_record_reader.cc
parente4a63b578f97c9dca26fd4d3a364f90a94cb45b5 (diff)
Add support for GZIP (and other future compression types) to TFRecordReader in python.
Also clean up Record{Reader,Writer}Options creation to reuse logic and log errors. Change: 133973785
Diffstat (limited to 'tensorflow/python/lib/io/py_record_reader.cc')
-rw-r--r--tensorflow/python/lib/io/py_record_reader.cc11
1 files changed, 3 insertions, 8 deletions
diff --git a/tensorflow/python/lib/io/py_record_reader.cc b/tensorflow/python/lib/io/py_record_reader.cc
index 47c0878932..d3f557506e 100644
--- a/tensorflow/python/lib/io/py_record_reader.cc
+++ b/tensorflow/python/lib/io/py_record_reader.cc
@@ -43,14 +43,9 @@ PyRecordReader* PyRecordReader::New(const string& filename, uint64 start_offset,
reader->offset_ = start_offset;
reader->file_ = file.release();
- RecordReaderOptions options;
- if (compression_type_string == "ZLIB") {
- options.compression_type = RecordReaderOptions::ZLIB_COMPRESSION;
- options.zlib_options = ZlibCompressionOptions::DEFAULT();
- } else if (compression_type_string == "GZIP") {
- options.compression_type = RecordReaderOptions::ZLIB_COMPRESSION;
- options.zlib_options = ZlibCompressionOptions::GZIP();
- }
+ RecordReaderOptions options =
+ RecordReaderOptions::CreateRecordReaderOptions(compression_type_string);
+
reader->reader_ = new RecordReader(reader->file_, options);
return reader;
}