diff options
author | A. Unique TensorFlower <gardener@tensorflow.org> | 2016-09-22 10:00:16 -0800 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2016-09-22 11:02:46 -0700 |
commit | b6c8def190bc3d885dc7ae9c3c570bdf94031378 (patch) | |
tree | 20c72c0d06f204e5ca2e2a23b75aef65a180f0e1 /tensorflow/python/lib/io/py_record_reader.cc | |
parent | e4a63b578f97c9dca26fd4d3a364f90a94cb45b5 (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.cc | 11 |
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; } |