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 // BlockBuilder generates blocks where keys are prefix-compressed:
17 //
18 // When we store a key, we drop the prefix shared with the previous
19 // string. This helps reduce the space requirement significantly.
20 // Furthermore, once every K keys, we do not apply the prefix
21 // compression and store the entire key. We call this a "restart
22 // point". The tail end of the block stores the offsets of all of the
23 // restart points, and can be used to do a binary search when looking
24 // for a particular key. Values are stored as-is (without compression)
25 // immediately following the corresponding key.
26 //
27 // An entry for a particular key-value pair has the form:
28 // shared_bytes: varint32
29 // unshared_bytes: varint32
30 // value_length: varint32
31 // key_delta: char[unshared_bytes]
32 // value: char[value_length]
33 // shared_bytes == 0 for restart points.
34 //
35 // The trailer of the block has the form:
36 // restarts: uint32[num_restarts]
37 // num_restarts: uint32
38 // restarts[i] contains the offset within the block of the ith restart point.
39
40 #include "tensorflow/core/lib/io/block_builder.h"
41
42 #include <assert.h>
43 #include <algorithm>
44 #include "tensorflow/core/lib/core/coding.h"
45 #include "tensorflow/core/lib/io/table_builder.h"
46
47 namespace tensorflow {
48 namespace table {
49
BlockBuilder(const Options * options)50 BlockBuilder::BlockBuilder(const Options* options)
51 : options_(options), restarts_(), counter_(0), finished_(false) {
52 assert(options->block_restart_interval >= 1);
53 restarts_.push_back(0); // First restart point is at offset 0
54 }
55
Reset()56 void BlockBuilder::Reset() {
57 buffer_.clear();
58 restarts_.clear();
59 restarts_.push_back(0); // First restart point is at offset 0
60 counter_ = 0;
61 finished_ = false;
62 last_key_.clear();
63 }
64
CurrentSizeEstimate() const65 size_t BlockBuilder::CurrentSizeEstimate() const {
66 return (buffer_.size() + // Raw data buffer
67 restarts_.size() * sizeof(uint32) + // Restart array
68 sizeof(uint32)); // Restart array length
69 }
70
Finish()71 StringPiece BlockBuilder::Finish() {
72 // Append restart array
73 CHECK_LE(restarts_.size(), std::numeric_limits<uint32_t>::max());
74 for (const auto r : restarts_) {
75 core::PutFixed32(&buffer_, r);
76 }
77 // Downcast safe because of the CHECK.
78 core::PutFixed32(&buffer_, static_cast<uint32_t>(restarts_.size()));
79 finished_ = true;
80 return StringPiece(buffer_);
81 }
82
Add(const StringPiece & key,const StringPiece & value)83 void BlockBuilder::Add(const StringPiece& key, const StringPiece& value) {
84 StringPiece last_key_piece(last_key_);
85 assert(!finished_);
86 assert(counter_ <= options_->block_restart_interval);
87 assert(buffer_.empty() // No values yet?
88 || key.compare(last_key_piece) > 0);
89 size_t shared = 0;
90 if (counter_ < options_->block_restart_interval) {
91 // See how much sharing to do with previous string
92 const size_t min_length = std::min(last_key_piece.size(), key.size());
93 while ((shared < min_length) && (last_key_piece[shared] == key[shared])) {
94 shared++;
95 }
96 } else {
97 // Restart compression
98 CHECK_LE(buffer_.size(), std::numeric_limits<uint32_t>::max());
99 restarts_.push_back(static_cast<uint32_t>(buffer_.size()));
100 counter_ = 0;
101 }
102 const size_t non_shared = key.size() - shared;
103
104 CHECK_LE(shared, std::numeric_limits<uint32_t>::max());
105 CHECK_LE(non_shared, std::numeric_limits<uint32_t>::max());
106 CHECK_LE(value.size(), std::numeric_limits<uint32_t>::max());
107
108 // Add "<shared><non_shared><value_size>" to buffer_
109 core::PutVarint32(&buffer_, static_cast<uint32_t>(shared));
110 core::PutVarint32(&buffer_, static_cast<uint32_t>(non_shared));
111 core::PutVarint32(&buffer_, static_cast<uint32_t>(value.size()));
112
113 // Add string delta to buffer_ followed by value
114 buffer_.append(key.data() + shared, non_shared);
115 buffer_.append(value.data(), static_cast<uint32_t>(value.size()));
116
117 // Update state
118 last_key_.resize(shared);
119 last_key_.append(key.data() + shared, non_shared);
120 assert(StringPiece(last_key_) == key);
121 counter_++;
122 }
123
124 } // namespace table
125 } // namespace tensorflow
126