Method Next

oneflow/core/embedding/persistent_table.cpp:752–790

  ~SnapshotIteratorImpl() override = default;

  void Next(uint32_t num_keys, uint32_t* return_keys, void* keys, void* values) override {
    *return_keys = 0;
    while (current_chunk_ < indices_names_.size()) {
      if (!chunk_iterator_) {
        const std::string snapshot_base = table_->SnapshotDirPath(snapshot_name_);
        const uint64_t chunk_id = GetChunkId(indices_names_[current_chunk_], kIndexFileNamePrefix);
        PosixFile index_file(PosixFile::JoinPath(snapshot_base, indices_names_[current_chunk_]),
                             O_RDONLY, 0644);
        const size_t index_file_size = index_file.Size();
        CHECK_EQ(index_file_size % sizeof(uint64_t), 0);
        if (index_file_size == 0) {
          current_chunk_ += 1;
          continue;
        }
        const size_t n_entries = index_file_size / sizeof(uint64_t);
        indices_file_.reset(new PosixMappedFile(std::move(index_file), index_file_size, PROT_READ));
        PosixFile key_file(table_->KeyFilePath(chunk_id), O_RDONLY, 0644);
        keys_file_.reset(new PosixMappedFile(std::move(key_file), key_file.Size(), PROT_READ));
        PosixFile value_file(table_->ValueFilePath(chunk_id), O_RDONLY, 0644);
        values_file_.reset(
            new PosixMappedFile(std::move(value_file), value_file.Size(), PROT_READ));
        chunk_iterator_.reset(new ChunkIteratorImpl<Key>(
            value_size_, logical_block_size_, num_values_per_block_, num_values_per_chunk_,
            chunk_id, n_entries, static_cast<const Key*>(keys_file_->ptr()),
            static_cast<const uint64_t*>(indices_file_->ptr()), values_file_->ptr()));
      }
      chunk_iterator_->Next(num_keys, return_keys, keys, values);
      if (*return_keys == 0) {
        chunk_iterator_.reset();
        keys_file_.reset();
        values_file_.reset();
        indices_file_.reset();
        current_chunk_ += 1;
        continue;
      } else {
        return;
      }
    }
  }

  void Reset() override { UNIMPLEMENTED(); }