Function InsertTable

tensorflow/core/kernels/save_restore_tensor.cc:488–567

}

Status InsertTable(
    BundleReader* reader, HashTable* table,
    const std::vector<TensibleVariable*>& tensibles,
    const string& table_name, const std::vector<string>& tensibles_name,
    const std::vector<TensorSliceProto>& table_slices,
    const std::vector<std::vector<TensorSliceProto>>& tensible_slices,
    const RestoreHashTableSlice& slice, mutex* mu,
    int64 slice_beg, int64 slice_length, int64 slice_size,
    const std::function<void(int64*,size_t)>& revise_fn = nullptr) {
  std::unique_ptr<SegmentBundleReader> table_bundle_reader;
  std::vector<std::unique_ptr<SegmentBundleReader>> tensible_bundle_readers;
  {
    mutex_lock lock(*mu);
    string xname = checkpoint::EncodeTensorNameSlice(
        table_name, TensorSlice(table_slices[slice.slice_id]));
    table_bundle_reader.reset(new SegmentBundleReader(
            reader, xname, slice.beg, slice.len));
    TF_RETURN_IF_ERROR(table_bundle_reader->Begin());

    for (size_t j = 0; j < tensibles.size(); j++) {
      if (tensible_slices[j].empty()) {
        tensible_bundle_readers.emplace_back(nullptr);
        continue;
      }
      string xname = checkpoint::EncodeTensorNameSlice(
          tensibles_name[j], TensorSlice(tensible_slices[j][slice.slice_id]));
      tensible_bundle_readers.emplace_back(new SegmentBundleReader(
              reader, xname, slice.beg, slice.len));
      TF_RETURN_IF_ERROR(tensible_bundle_readers.back()->Begin());
    }
  }
  int64 id_size = slice.len;
  constexpr int64 kSlice = 1 << 15;
  int64 slice_end = slice_beg + slice_length;
  while (id_size > 0) {
    int64 xid = std::min(id_size, kSlice);
    id_size -= xid;
    int64 key[kSlice];
    TF_RETURN_IF_ERROR(table_bundle_reader->Read(key, xid * sizeof(int64)));
    if (revise_fn) {
      revise_fn(key, xid);
    }
    std::vector<int64> real_key, real_id, real_offset;
    for (int j = 0; j < xid; j++) {
      int64 slice_idx = (uint64_t)(key[j]) % (uint64_t)(slice_size);
      if (slice_beg <= slice_idx && slice_idx < slice_end) {
        real_key.push_back(key[j]);
        real_offset.push_back(j);
      }
    }
    real_id.resize(real_key.size());
    int64 size = table->GetIdsWithoutResize(
        &real_key[0], &real_id[0], real_key.size());
    for (size_t j = 0; j < tensibles.size(); j++) {
      tensibles[j]->ZeroCostResize(size);
    }
    real_offset.push_back(-1);
    int64 idx = 0;