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 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 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 65 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 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 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
