1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/lib/io/record_writer.h"
17
18 #include "tensorflow/core/lib/core/coding.h"
19 #include "tensorflow/core/lib/hash/crc32c.h"
20 #include "tensorflow/core/lib/io/compression.h"
21 #include "tensorflow/core/platform/env.h"
22
23 namespace tensorflow {
24 namespace io {
25 namespace {
IsZlibCompressed(RecordWriterOptions options)26 bool IsZlibCompressed(RecordWriterOptions options) {
27 return options.compression_type == RecordWriterOptions::ZLIB_COMPRESSION;
28 }
29 } // namespace
30
CreateRecordWriterOptions(const string & compression_type)31 RecordWriterOptions RecordWriterOptions::CreateRecordWriterOptions(
32 const string& compression_type) {
33 RecordWriterOptions options;
34 if (compression_type == "ZLIB") {
35 options.compression_type = io::RecordWriterOptions::ZLIB_COMPRESSION;
36 #if defined(IS_SLIM_BUILD)
37 LOG(ERROR) << "Compression is not supported but compression_type is set."
38 << " No compression will be used.";
39 #else
40 options.zlib_options = io::ZlibCompressionOptions::DEFAULT();
41 #endif // IS_SLIM_BUILD
42 } else if (compression_type == compression::kGzip) {
43 options.compression_type = io::RecordWriterOptions::ZLIB_COMPRESSION;
44 #if defined(IS_SLIM_BUILD)
45 LOG(ERROR) << "Compression is not supported but compression_type is set."
46 << " No compression will be used.";
47 #else
48 options.zlib_options = io::ZlibCompressionOptions::GZIP();
49 #endif // IS_SLIM_BUILD
50 } else if (compression_type != compression::kNone) {
51 LOG(ERROR) << "Unsupported compression_type:" << compression_type
52 << ". No compression will be used.";
53 }
54 return options;
55 }
56
RecordWriter(WritableFile * dest,const RecordWriterOptions & options)57 RecordWriter::RecordWriter(WritableFile* dest,
58 const RecordWriterOptions& options)
59 : dest_(dest), options_(options) {
60 if (IsZlibCompressed(options)) {
61 // We don't have zlib available on all embedded platforms, so fail.
62 #if defined(IS_SLIM_BUILD)
63 LOG(FATAL) << "Zlib compression is unsupported on mobile platforms.";
64 #else // IS_SLIM_BUILD
65 ZlibOutputBuffer* zlib_output_buffer = new ZlibOutputBuffer(
66 dest, options.zlib_options.input_buffer_size,
67 options.zlib_options.output_buffer_size, options.zlib_options);
68 Status s = zlib_output_buffer->Init();
69 if (!s.ok()) {
70 LOG(FATAL) << "Failed to initialize Zlib inputbuffer. Error: "
71 << s.ToString();
72 }
73 dest_ = zlib_output_buffer;
74 #endif // IS_SLIM_BUILD
75 } else if (options.compression_type == RecordWriterOptions::NONE) {
76 // Nothing to do
77 } else {
78 LOG(FATAL) << "Unspecified compression type :" << options.compression_type;
79 }
80 }
81
~RecordWriter()82 RecordWriter::~RecordWriter() {
83 if (dest_ != nullptr) {
84 Status s = Close();
85 if (!s.ok()) {
86 LOG(ERROR) << "Could not finish writing file: " << s;
87 }
88 }
89 }
90
WriteRecord(StringPiece data)91 Status RecordWriter::WriteRecord(StringPiece data) {
92 if (dest_ == nullptr) {
93 return Status(::tensorflow::error::FAILED_PRECONDITION,
94 "Writer not initialized or previously closed");
95 }
96 // Format of a single record:
97 // uint64 length
98 // uint32 masked crc of length
99 // byte data[length]
100 // uint32 masked crc of data
101 char header[kHeaderSize];
102 char footer[kFooterSize];
103 PopulateHeader(header, data.data(), data.size());
104 PopulateFooter(footer, data.data(), data.size());
105 TF_RETURN_IF_ERROR(dest_->Append(StringPiece(header, sizeof(header))));
106 TF_RETURN_IF_ERROR(dest_->Append(data));
107 return dest_->Append(StringPiece(footer, sizeof(footer)));
108 }
109
Close()110 Status RecordWriter::Close() {
111 if (dest_ == nullptr) return Status::OK();
112 #if !defined(IS_SLIM_BUILD)
113 if (IsZlibCompressed(options_)) {
114 Status s = dest_->Close();
115 delete dest_;
116 dest_ = nullptr;
117 return s;
118 }
119 #endif // IS_SLIM_BUILD
120 return Status::OK();
121 }
122
Flush()123 Status RecordWriter::Flush() {
124 if (dest_ == nullptr) {
125 return Status(::tensorflow::error::FAILED_PRECONDITION,
126 "Writer not initialized or previously closed");
127 }
128 if (IsZlibCompressed(options_)) {
129 return dest_->Flush();
130 }
131 return Status::OK();
132 }
133
134 } // namespace io
135 } // namespace tensorflow
136