Function TestCompaction

tensorflow/core/kernels/embedding_variable_ops_test.cc:1180–1237

}

void TestCompaction() {
  std::string temp_dir = testing::TmpDir();
  auto feat_desc = new embedding::FeatureDescriptor<float>(
      1, 1, ev_allocator(), embedding::StorageType::DRAM_SSDHASH,
      true, true, {false, 0});
  auto hashmap = new SSDHashKV<int64, float>(
      temp_dir, feat_desc);
  feat_desc->InitSlotInfo(0, 124, {nullptr, 1});
  hashmap->Init();
  ASSERT_EQ(hashmap->Size(), 0);
  std::vector<int64> ids;
  for (int i = 0; i < 262144; i++) {
    ids.emplace_back(i);
  }
  auto t1 = std::thread(SingleCommit, hashmap, ids, 3);
  t1.join();
  ids.clear();
  for (int i = 0; i < 131073; i++) {
    ids.emplace_back(i);
  }
  t1 = std::thread(SingleCommit, hashmap, ids, 1);
  t1.join();
  ids.clear();
  sleep(1);
  void* val = nullptr;
  for (int i = 131073; i < 262144; i++) {
    hashmap->Lookup(i, &val);
    float* v = (float*)val;
    for (int j = 0; j < 124; j++){
      ASSERT_EQ(v[j], i+3);
    }
  }
  for (int i = 131073; i < 262144; i++) {
    ids.emplace_back(i);
  }
  t1 = std::thread(SingleCommit, hashmap, ids, 2);
  t1.join();
  ids.clear();
  ids.emplace_back(262155);
  t1 = std::thread(SingleCommit, hashmap, ids, 0);
  t1.join();
  sleep(1);
  for (int i = 0; i < 131073; i++) {
    hashmap->Lookup(i, &val);
    float* v = (float*)val;
    for (int j = 0; j < 124; j++){
      ASSERT_EQ(v[j], i + 1);
    }
  }
  for (int i = 131073; i < 262144; i++) {
    hashmap->Lookup(i, &val);
    float* v = (float*)val;
    for (int j = 0; j < 124; j++){
      ASSERT_EQ(v[j], i + 2);
    }
  }
  delete hashmap;
}