Home | History | Annotate | Download | only in io
      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/table.h"
     17 
     18 #include "tensorflow/core/lib/core/coding.h"
     19 #include "tensorflow/core/lib/core/errors.h"
     20 #include "tensorflow/core/lib/io/block.h"
     21 #include "tensorflow/core/lib/io/format.h"
     22 #include "tensorflow/core/lib/io/table_options.h"
     23 #include "tensorflow/core/lib/io/two_level_iterator.h"
     24 #include "tensorflow/core/platform/env.h"
     25 
     26 namespace tensorflow {
     27 namespace table {
     28 
     29 struct Table::Rep {
     30   ~Rep() { delete index_block; }
     31 
     32   Options options;
     33   Status status;
     34   RandomAccessFile* file;
     35   // XXX  uint64 cache_id;
     36 
     37   BlockHandle metaindex_handle;  // Handle to metaindex_block: saved from footer
     38   Block* index_block;
     39 };
     40 
     41 Status Table::Open(const Options& options, RandomAccessFile* file, uint64 size,
     42                    Table** table) {
     43   *table = nullptr;
     44   if (size < Footer::kEncodedLength) {
     45     return errors::DataLoss("file is too short to be an sstable");
     46   }
     47 
     48   char footer_space[Footer::kEncodedLength];
     49   StringPiece footer_input;
     50   Status s = file->Read(size - Footer::kEncodedLength, Footer::kEncodedLength,
     51                         &footer_input, footer_space);
     52   if (!s.ok()) return s;
     53 
     54   Footer footer;
     55   s = footer.DecodeFrom(&footer_input);
     56   if (!s.ok()) return s;
     57 
     58   // Read the index block
     59   BlockContents contents;
     60   Block* index_block = nullptr;
     61   if (s.ok()) {
     62     s = ReadBlock(file, footer.index_handle(), &contents);
     63     if (s.ok()) {
     64       index_block = new Block(contents);
     65     }
     66   }
     67 
     68   if (s.ok()) {
     69     // We've successfully read the footer and the index block: we're
     70     // ready to serve requests.
     71     Rep* rep = new Table::Rep;
     72     rep->options = options;
     73     rep->file = file;
     74     rep->metaindex_handle = footer.metaindex_handle();
     75     rep->index_block = index_block;
     76     // XXX    rep->cache_id = (options.block_cache ?
     77     // options.block_cache->NewId() : 0);
     78     *table = new Table(rep);
     79   } else {
     80     if (index_block) delete index_block;
     81   }
     82 
     83   return s;
     84 }
     85 
     86 Table::~Table() { delete rep_; }
     87 
     88 static void DeleteBlock(void* arg, void* ignored) {
     89   delete reinterpret_cast<Block*>(arg);
     90 }
     91 
     92 // Convert an index iterator value (i.e., an encoded BlockHandle)
     93 // into an iterator over the contents of the corresponding block.
     94 Iterator* Table::BlockReader(void* arg, const StringPiece& index_value) {
     95   Table* table = reinterpret_cast<Table*>(arg);
     96   //  Cache* block_cache = table->rep_->options.block_cache;
     97   Block* block = nullptr;
     98   //  Cache::Handle* cache_handle = NULL;
     99 
    100   BlockHandle handle;
    101   StringPiece input = index_value;
    102   Status s = handle.DecodeFrom(&input);
    103   // We intentionally allow extra stuff in index_value so that we
    104   // can add more features in the future.
    105 
    106   if (s.ok()) {
    107     BlockContents contents;
    108     s = ReadBlock(table->rep_->file, handle, &contents);
    109     if (s.ok()) {
    110       block = new Block(contents);
    111     }
    112   }
    113 
    114   Iterator* iter;
    115   if (block != nullptr) {
    116     iter = block->NewIterator();
    117     iter->RegisterCleanup(&DeleteBlock, block, nullptr);
    118   } else {
    119     iter = NewErrorIterator(s);
    120   }
    121   return iter;
    122 }
    123 
    124 Iterator* Table::NewIterator() const {
    125   return NewTwoLevelIterator(rep_->index_block->NewIterator(),
    126                              &Table::BlockReader, const_cast<Table*>(this));
    127 }
    128 
    129 Status Table::InternalGet(const StringPiece& k, void* arg,
    130                           void (*saver)(void*, const StringPiece&,
    131                                         const StringPiece&)) {
    132   Status s;
    133   Iterator* iiter = rep_->index_block->NewIterator();
    134   iiter->Seek(k);
    135   if (iiter->Valid()) {
    136     BlockHandle handle;
    137     Iterator* block_iter = BlockReader(this, iiter->value());
    138     block_iter->Seek(k);
    139     if (block_iter->Valid()) {
    140       (*saver)(arg, block_iter->key(), block_iter->value());
    141     }
    142     s = block_iter->status();
    143     delete block_iter;
    144   }
    145   if (s.ok()) {
    146     s = iiter->status();
    147   }
    148   delete iiter;
    149   return s;
    150 }
    151 
    152 uint64 Table::ApproximateOffsetOf(const StringPiece& key) const {
    153   Iterator* index_iter = rep_->index_block->NewIterator();
    154   index_iter->Seek(key);
    155   uint64 result;
    156   if (index_iter->Valid()) {
    157     BlockHandle handle;
    158     StringPiece input = index_iter->value();
    159     Status s = handle.DecodeFrom(&input);
    160     if (s.ok()) {
    161       result = handle.offset();
    162     } else {
    163       // Strange: we can't decode the block handle in the index block.
    164       // We'll just return the offset of the metaindex block, which is
    165       // close to the whole file size for this case.
    166       result = rep_->metaindex_handle.offset();
    167     }
    168   } else {
    169     // key is past the last key in the file.  Approximate the offset
    170     // by returning the offset of the metaindex block (which is
    171     // right near the end of the file).
    172     result = rep_->metaindex_handle.offset();
    173   }
    174   delete index_iter;
    175   return result;
    176 }
    177 
    178 }  // namespace table
    179 }  // namespace tensorflow
    180