diff options
Diffstat (limited to 'tensorflow/core/lib/io/record_writer.cc')
| -rw-r--r-- | tensorflow/core/lib/io/record_writer.cc | 25 |
1 files changed, 25 insertions, 0 deletions
diff --git a/tensorflow/core/lib/io/record_writer.cc b/tensorflow/core/lib/io/record_writer.cc index 516332d2b7..175bfbd827 100644 --- a/tensorflow/core/lib/io/record_writer.cc +++ b/tensorflow/core/lib/io/record_writer.cc @@ -21,6 +21,31 @@ limitations under the License. namespace tensorflow { namespace io { +RecordWriterOptions RecordWriterOptions::CreateRecordWriterOptions( + const string& compression_type) { + RecordWriterOptions options; + if (compression_type == "ZLIB") { + options.compression_type = io::RecordWriterOptions::ZLIB_COMPRESSION; +#if defined(IS_SLIM_BUILD) + LOG(ERROR) << "Compression is not supported but compression_type is set." + << " No compression will be used."; +#else + options.zlib_options = io::ZlibCompressionOptions::DEFAULT(); +#endif // IS_SLIM_BUILD + } else if (compression_type == "GZIP") { + options.compression_type = io::RecordWriterOptions::ZLIB_COMPRESSION; +#if defined(IS_SLIM_BUILD) + LOG(ERROR) << "Compression is not supported but compression_type is set." + << " No compression will be used."; +#else + options.zlib_options = io::ZlibCompressionOptions::GZIP(); +#endif // IS_SLIM_BUILD + } else if (compression_type != "") { + LOG(ERROR) << "Unsupported compression_type:" << compression_type + << ". No comprression will be used."; + } + return options; +} RecordWriter::RecordWriter(WritableFile* dest, const RecordWriterOptions& options) |