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 == compression::kZlib) {
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
110 #if defined(PLATFORM_GOOGLE)
WriteRecord(const absl::Cord & data)111 Status RecordWriter::WriteRecord(const absl::Cord& data) {
112 if (dest_ == nullptr) {
113 return Status(::tensorflow::error::FAILED_PRECONDITION,
114 "Writer not initialized or previously closed");
115 }
116 // Format of a single record:
117 // uint64 length
118 // uint32 masked crc of length
119 // byte data[length]
120 // uint32 masked crc of data
121 char header[kHeaderSize];
122 char footer[kFooterSize];
123 PopulateHeader(header, data);
124 PopulateFooter(footer, data);
125 TF_RETURN_IF_ERROR(dest_->Append(StringPiece(header, sizeof(header))));
126 TF_RETURN_IF_ERROR(dest_->Append(data));
127 return dest_->Append(StringPiece(footer, sizeof(footer)));
128 }
129 #endif
130
Close()131 Status RecordWriter::Close() {
132 if (dest_ == nullptr) return Status::OK();
133 #if !defined(IS_SLIM_BUILD)
134 if (IsZlibCompressed(options_)) {
135 Status s = dest_->Close();
136 delete dest_;
137 dest_ = nullptr;
138 return s;
139 }
140 #endif // IS_SLIM_BUILD
141 return Status::OK();
142 }
143
Flush()144 Status RecordWriter::Flush() {
145 if (dest_ == nullptr) {
146 return Status(::tensorflow::error::FAILED_PRECONDITION,
147 "Writer not initialized or previously closed");
148 }
149 return dest_->Flush();
150 }
151
152 } // namespace io
153 } // namespace tensorflow
154