@nxtedition/rocksdb 13.1.4 → 13.2.0

package/deps/rocksdb/rocksdb/utilities/secondary_index/faiss_ivf_index_test.cc

@@ -0,0 +1,124 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include "utilities/secondary_index/faiss_ivf_index.h"
+
+#include <charconv>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "faiss/IndexFlat.h"
+#include "faiss/IndexIVFFlat.h"
+#include "faiss/utils/random.h"
+#include "rocksdb/utilities/transaction_db.h"
+#include "test_util/testharness.h"
+#include "util/coding.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+TEST(FaissIVFIndexTest, Basic) {
+  constexpr size_t dim = 128;
+  auto quantizer = std::make_unique<faiss::IndexFlatL2>(dim);
+
+  constexpr size_t num_lists = 16;
+  auto index =
+      std::make_unique<faiss::IndexIVFFlat>(quantizer.get(), dim, num_lists);
+
+  constexpr faiss::idx_t num_vectors = 1024;
+  std::vector<float> embeddings(dim * num_vectors);
+  faiss::float_rand(embeddings.data(), dim * num_vectors, 42);
+
+  index->train(num_vectors, embeddings.data());
+
+  index->nprobe = 2;
+
+  const std::string db_name = test::PerThreadDBPath("faiss_ivf_index_test");
+  EXPECT_OK(DestroyDB(db_name, Options()));
+
+  Options options;
+  options.create_if_missing = true;
+
+  TransactionDBOptions txn_db_options;
+  const std::string primary_column_name = "embedding";
+  txn_db_options.secondary_indices.emplace_back(
+      std::make_shared<FaissIVFIndex>(std::move(index), primary_column_name));
+
+  TransactionDB* db = nullptr;
+  ASSERT_OK(TransactionDB::Open(options, txn_db_options, db_name, &db));
+
+  std::unique_ptr<TransactionDB> db_guard(db);
+
+  ColumnFamilyOptions cf1_opts;
+  ColumnFamilyHandle* cfh1 = nullptr;
+  ASSERT_OK(db->CreateColumnFamily(cf1_opts, "cf1", &cfh1));
+  std::unique_ptr<ColumnFamilyHandle> cfh1_guard(cfh1);
+
+  ColumnFamilyOptions cf2_opts;
+  ColumnFamilyHandle* cfh2 = nullptr;
+  ASSERT_OK(db->CreateColumnFamily(cf2_opts, "cf2", &cfh2));
+  std::unique_ptr<ColumnFamilyHandle> cfh2_guard(cfh2);
+
+  const auto& secondary_index = txn_db_options.secondary_indices.back();
+  secondary_index->SetPrimaryColumnFamily(cfh1);
+  secondary_index->SetSecondaryColumnFamily(cfh2);
+
+  {
+    std::unique_ptr<Transaction> txn(db->BeginTransaction(WriteOptions()));
+
+    for (faiss::idx_t i = 0; i < num_vectors; ++i) {
+      const std::string primary_key = std::to_string(i);
+
+      ASSERT_OK(txn->PutEntity(
+          cfh1, primary_key,
+          WideColumns{
+              {primary_column_name,
+               Slice(reinterpret_cast<const char*>(embeddings.data() + i * dim),
+                     dim * sizeof(float))}}));
+    }
+
+    ASSERT_OK(txn->Commit());
+  }
+
+  {
+    size_t num_found = 0;
+
+    std::unique_ptr<Iterator> it(db->NewIterator(ReadOptions(), cfh2));
+
+    for (it->SeekToFirst(); it->Valid(); it->Next()) {
+      Slice key = it->key();
+      faiss::idx_t label = -1;
+      ASSERT_TRUE(GetVarsignedint64(&key, &label));
+      ASSERT_GE(label, 0);
+      ASSERT_LT(label, num_lists);
+
+      faiss::idx_t id = -1;
+      ASSERT_EQ(std::from_chars(key.data(), key.data() + key.size(), id).ec,
+                std::errc());
+      ASSERT_GE(id, 0);
+      ASSERT_LT(id, num_vectors);
+
+      // Since we use IndexIVFFlat, there is no fine quantization, so the code
+      // is actually just the original embedding
+      ASSERT_EQ(
+          it->value(),
+          Slice(reinterpret_cast<const char*>(embeddings.data() + id * dim),
+                dim * sizeof(float)));
+
+      ++num_found;
+    }
+
+    ASSERT_OK(it->status());
+    ASSERT_EQ(num_found, num_vectors);
+  }
+}
+
+}  // namespace ROCKSDB_NAMESPACE
+
+int main(int argc, char** argv) {
+  ROCKSDB_NAMESPACE::port::InstallStackTraceHandler();
+  ::testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}

package/deps/rocksdb/rocksdb/utilities/secondary_index/secondary_index_mixin.h

@@ -0,0 +1,441 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#include <cassert>
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "db/wide/wide_columns_helper.h"
+#include "rocksdb/options.h"
+#include "rocksdb/utilities/secondary_index.h"
+#include "rocksdb/wide_columns.h"
+#include "util/autovector.h"
+#include "util/overload.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+template <typename Txn>
+class SecondaryIndexMixin : public Txn {
+ public:
+  template <typename... Args>
+  explicit SecondaryIndexMixin(
+      const std::vector<std::shared_ptr<SecondaryIndex>>* secondary_indices,
+      Args&&... args)
+      : Txn(std::forward<Args>(args)...),
+        secondary_indices_(secondary_indices) {
+    assert(secondary_indices_);
+    assert(!secondary_indices_->empty());
+  }
+
+  using Txn::Put;
+  Status Put(ColumnFamilyHandle* /* column_family */, const Slice& /* key */,
+             const Slice& /* value */,
+             const bool /* assume_tracked */ = false) override {
+    return Status::NotSupported("Put with secondary indices not yet supported");
+  }
+  Status Put(ColumnFamilyHandle* /* column_family */,
+             const SliceParts& /* key */, const SliceParts& /* value */,
+             const bool /* assume_tracked */ = false) override {
+    return Status::NotSupported("Put with secondary indices not yet supported");
+  }
+
+  Status PutEntity(ColumnFamilyHandle* column_family, const Slice& key,
+                   const WideColumns& columns,
+                   bool assume_tracked = false) override {
+    return PerformWithSavePoint([&]() {
+      const bool do_validate = !assume_tracked;
+      return PutEntityWithSecondaryIndices(column_family, key, columns,
+                                           do_validate);
+    });
+  }
+
+  using Txn::Merge;
+  Status Merge(ColumnFamilyHandle* /* column_family */, const Slice& /* key */,
+               const Slice& /* value */,
+               const bool /* assume_tracked */ = false) override {
+    return Status::NotSupported(
+        "Merge with secondary indices not yet supported");
+  }
+
+  using Txn::Delete;
+  Status Delete(ColumnFamilyHandle* /* column_family */, const Slice& /* key */,
+                const bool /* assume_tracked */ = false) override {
+    return Status::NotSupported(
+        "Delete with secondary indices not yet supported");
+  }
+  Status Delete(ColumnFamilyHandle* /* column_family */,
+                const SliceParts& /* key */,
+                const bool /* assume_tracked */ = false) override {
+    return Status::NotSupported(
+        "Delete with secondary indices not yet supported");
+  }
+
+  using Txn::SingleDelete;
+  Status SingleDelete(ColumnFamilyHandle* /* column_family */,
+                      const Slice& /* key */,
+                      const bool /* assume_tracked */ = false) override {
+    return Status::NotSupported(
+        "SingleDelete with secondary indices not yet supported");
+  }
+  Status SingleDelete(ColumnFamilyHandle* /* column_family */,
+                      const SliceParts& /* key */,
+                      const bool /* assume_tracked */ = false) override {
+    return Status::NotSupported(
+        "SingleDelete with secondary indices not yet supported");
+  }
+
+  using Txn::PutUntracked;
+  Status PutUntracked(ColumnFamilyHandle* /* column_family */,
+                      const Slice& /* key */,
+                      const Slice& /* value */) override {
+    return Status::NotSupported(
+        "PutUntracked with secondary indices not yet supported");
+  }
+  Status PutUntracked(ColumnFamilyHandle* /* column_family */,
+                      const SliceParts& /* key */,
+                      const SliceParts& /* value */) override {
+    return Status::NotSupported(
+        "PutUntracked with secondary indices not yet supported");
+  }
+
+  Status PutEntityUntracked(ColumnFamilyHandle* column_family, const Slice& key,
+                            const WideColumns& columns) override {
+    return PerformWithSavePoint([&]() {
+      constexpr bool do_validate = false;
+      return PutEntityWithSecondaryIndices(column_family, key, columns,
+                                           do_validate);
+    });
+  }
+
+  using Txn::MergeUntracked;
+  Status MergeUntracked(ColumnFamilyHandle* /* column_family */,
+                        const Slice& /* key */,
+                        const Slice& /* value */) override {
+    return Status::NotSupported(
+        "MergeUntracked with secondary indices not yet supported");
+  }
+
+  using Txn::DeleteUntracked;
+  Status DeleteUntracked(ColumnFamilyHandle* /* column_family */,
+                         const Slice& /* key */) override {
+    return Status::NotSupported(
+        "DeleteUntracked with secondary indices not yet supported");
+  }
+  Status DeleteUntracked(ColumnFamilyHandle* /* column_family */,
+                         const SliceParts& /* key */) override {
+    return Status::NotSupported(
+        "DeleteUntracked with secondary indices not yet supported");
+  }
+
+  using Txn::SingleDeleteUntracked;
+  Status SingleDeleteUntracked(ColumnFamilyHandle* /* column_family */,
+                               const Slice& /* key */) override {
+    return Status::NotSupported(
+        "SingleDeleteUntracked with secondary indices not yet supported");
+  }
+
+ private:
+  class IndexData {
+   public:
+    IndexData(const SecondaryIndex* index, const Slice& previous_column_value)
+        : index_(index), previous_column_value_(previous_column_value) {
+      assert(index_);
+    }
+
+    const SecondaryIndex* index() const { return index_; }
+    const Slice& previous_column_value() const {
+      return previous_column_value_;
+    }
+    std::optional<std::variant<Slice, std::string>>& updated_column_value() {
+      return updated_column_value_;
+    }
+    Slice primary_column_value() const {
+      return updated_column_value_.has_value() ? AsSlice(*updated_column_value_)
+                                               : previous_column_value_;
+    }
+
+   private:
+    const SecondaryIndex* index_;
+    Slice previous_column_value_;
+    std::optional<std::variant<Slice, std::string>> updated_column_value_;
+  };
+
+  static Slice AsSlice(const std::variant<Slice, std::string>& var) {
+    return std::visit([](const auto& value) -> Slice { return value; }, var);
+  }
+
+  static std::string AsString(const std::variant<Slice, std::string>& var) {
+    return std::visit(
+        overload{
+            [](const Slice& value) -> std::string { return value.ToString(); },
+            [](const std::string& value) -> std::string { return value; }},
+        var);
+  }
+
+  template <typename Iterator>
+  static Iterator FindPrimaryColumn(const SecondaryIndex* secondary_index,
+                                    ColumnFamilyHandle* column_family,
+                                    Iterator begin, Iterator end) {
+    assert(secondary_index);
+    assert(column_family);
+
+    if (column_family != secondary_index->GetPrimaryColumnFamily()) {
+      return end;
+    }
+
+    return WideColumnsHelper::Find(begin, end,
+                                   secondary_index->GetPrimaryColumnName());
+  }
+
+  template <typename Operation>
+  Status PerformWithSavePoint(Operation&& operation) {
+    Txn::SetSavePoint();
+
+    const Status s = operation();
+
+    if (!s.ok()) {
+      [[maybe_unused]] const Status st = Txn::RollbackToSavePoint();
+      assert(st.ok());
+
+      return s;
+    }
+
+    [[maybe_unused]] const Status st = Txn::PopSavePoint();
+    assert(st.ok());
+
+    return Status::OK();
+  }
+
+  Status GetPrimaryEntryForUpdate(ColumnFamilyHandle* column_family,
+                                  const Slice& primary_key,
+                                  PinnableWideColumns* existing_primary_columns,
+                                  bool do_validate) {
+    assert(column_family);
+    assert(existing_primary_columns);
+
+    constexpr bool exclusive = true;
+
+    return Txn::GetEntityForUpdate(ReadOptions(), column_family, primary_key,
+                                   existing_primary_columns, exclusive,
+                                   do_validate);
+  }
+
+  Status RemoveSecondaryEntry(const SecondaryIndex* secondary_index,
+                              const Slice& primary_key,
+                              const Slice& existing_primary_column_value) {
+    assert(secondary_index);
+
+    std::variant<Slice, std::string> secondary_key_prefix;
+
+    const Status s = secondary_index->GetSecondaryKeyPrefix(
+        primary_key, existing_primary_column_value, &secondary_key_prefix);
+    if (!s.ok()) {
+      return s;
+    }
+
+    const std::string secondary_key =
+        AsString(secondary_key_prefix) + primary_key.ToString();
+
+    return Txn::SingleDelete(secondary_index->GetSecondaryColumnFamily(),
+                             secondary_key);
+  }
+
+  Status AddPrimaryEntry(ColumnFamilyHandle* column_family,
+                         const Slice& primary_key,
+                         const WideColumns& primary_columns) {
+    assert(column_family);
+
+    constexpr bool assume_tracked = true;
+
+    return Txn::PutEntity(column_family, primary_key, primary_columns,
+                          assume_tracked);
+  }
+
+  Status AddSecondaryEntry(const SecondaryIndex* secondary_index,
+                           const Slice& primary_key,
+                           const Slice& primary_column_value,
+                           const Slice& previous_column_value) {
+    assert(secondary_index);
+
+    std::variant<Slice, std::string> secondary_key_prefix;
+
+    {
+      const Status s = secondary_index->GetSecondaryKeyPrefix(
+          primary_key, primary_column_value, &secondary_key_prefix);
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    std::optional<std::variant<Slice, std::string>> secondary_value;
+
+    {
+      const Status s = secondary_index->GetSecondaryValue(
+          primary_key, primary_column_value, previous_column_value,
+          &secondary_value);
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    {
+      const std::string secondary_key =
+          AsString(secondary_key_prefix) + primary_key.ToString();
+
+      const Status s = Txn::Put(
+          secondary_index->GetSecondaryColumnFamily(), secondary_key,
+          secondary_value.has_value() ? AsSlice(*secondary_value) : Slice());
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    return Status::OK();
+  }
+
+  Status RemoveSecondaryEntries(ColumnFamilyHandle* column_family,
+                                const Slice& primary_key, bool do_validate) {
+    assert(column_family);
+
+    PinnableWideColumns existing_primary_columns;
+
+    const Status s = GetPrimaryEntryForUpdate(
+        column_family, primary_key, &existing_primary_columns, do_validate);
+    if (s.IsNotFound()) {
+      return Status::OK();
+    }
+    if (!s.ok()) {
+      return s;
+    }
+
+    const auto& existing_columns = existing_primary_columns.columns();
+
+    for (const auto& secondary_index : *secondary_indices_) {
+      const auto it =
+          FindPrimaryColumn(secondary_index.get(), column_family,
+                            existing_columns.cbegin(), existing_columns.cend());
+      if (it == existing_columns.cend()) {
+        continue;
+      }
+
+      const Status st =
+          RemoveSecondaryEntry(secondary_index.get(), primary_key, it->value());
+      if (!st.ok()) {
+        return st;
+      }
+    }
+
+    return Status::OK();
+  }
+
+  Status UpdatePrimaryColumnValues(ColumnFamilyHandle* column_family,
+                                   const Slice& primary_key,
+                                   WideColumns& primary_columns,
+                                   autovector<IndexData>& applicable_indices) {
+    assert(applicable_indices.empty());
+
+    applicable_indices.reserve(secondary_indices_->size());
+
+    for (const auto& secondary_index : *secondary_indices_) {
+      const auto it =
+          FindPrimaryColumn(secondary_index.get(), column_family,
+                            primary_columns.begin(), primary_columns.end());
+      if (it == primary_columns.end()) {
+        continue;
+      }
+
+      applicable_indices.emplace_back(
+          IndexData(secondary_index.get(), it->value()));
+
+      auto& index_data = applicable_indices.back();
+
+      const Status s = secondary_index->UpdatePrimaryColumnValue(
+          primary_key, index_data.previous_column_value(),
+          &index_data.updated_column_value());
+      if (!s.ok()) {
+        return s;
+      }
+
+      it->value() = index_data.primary_column_value();
+    }
+
+    return Status::OK();
+  }
+
+  Status AddSecondaryEntries(const Slice& primary_key,
+                             const autovector<IndexData>& applicable_indices) {
+    for (const auto& index_data : applicable_indices) {
+      const Status s = AddSecondaryEntry(index_data.index(), primary_key,
+                                         index_data.primary_column_value(),
+                                         index_data.previous_column_value());
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    return Status::OK();
+  }
+
+  Status PutEntityWithSecondaryIndices(ColumnFamilyHandle* column_family,
+                                       const Slice& key,
+                                       const WideColumns& columns,
+                                       bool do_validate) {
+    // TODO: we could avoid removing and recreating secondary entries for
+    // which neither the secondary key prefix nor the value has changed
+
+    if (!column_family) {
+      column_family = Txn::DefaultColumnFamily();
+    }
+
+    const Slice& primary_key = key;
+
+    {
+      const Status s =
+          RemoveSecondaryEntries(column_family, primary_key, do_validate);
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    autovector<IndexData> applicable_indices;
+
+    WideColumns primary_columns(columns);
+    WideColumnsHelper::SortColumns(primary_columns);
+
+    {
+      const Status s = UpdatePrimaryColumnValues(
+          column_family, primary_key, primary_columns, applicable_indices);
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    {
+      const Status s =
+          AddPrimaryEntry(column_family, primary_key, primary_columns);
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    {
+      const Status s = AddSecondaryEntries(primary_key, applicable_indices);
+      if (!s.ok()) {
+        return s;
+      }
+    }
+
+    return Status::OK();
+  }
+
+  const std::vector<std::shared_ptr<SecondaryIndex>>* secondary_indices_;
+};
+
+}  // namespace ROCKSDB_NAMESPACE

package/deps/rocksdb/rocksdb/utilities/table_properties_collectors/compact_for_tiering_collector.cc

@@ -21,14 +21,29 @@ namespace ROCKSDB_NAMESPACE {
 const std::string
     CompactForTieringCollector::kNumEligibleLastLevelEntriesPropertyName =
         "rocksdb.eligible.last.level.entries";
+const std::string
+    CompactForTieringCollector::kAverageDataUnixWriteTimePropertyName =
+        "rocksdb.data.unix.write.time.average";
+const std::string
+    CompactForTieringCollector::kMaxDataUnixWriteTimePropertyName =
+        "rocksdb.data.unix.write.time.max";
+const std::string
+    CompactForTieringCollector::kMinDataUnixWriteTimePropertyName =
+        "rocksdb.data.unix.write.time.min";
+const std::string
+    CompactForTieringCollector::kNumInfinitelyOldEntriesPropertyName =
+        "rocksdb.num.infinitely.old.entries";
 
 CompactForTieringCollector::CompactForTieringCollector(
     SequenceNumber last_level_inclusive_max_seqno_threshold,
-    double compaction_trigger_ratio)
+    double compaction_trigger_ratio, bool collect_data_age_stats)
     : last_level_inclusive_max_seqno_threshold_(
           last_level_inclusive_max_seqno_threshold),
-      compaction_trigger_ratio_(compaction_trigger_ratio) {
+      compaction_trigger_ratio_(compaction_trigger_ratio),
+      collect_data_age_stats_(collect_data_age_stats) {
   assert(last_level_inclusive_max_seqno_threshold_ != kMaxSequenceNumber);
+  // TODO(yuzhangyu): implement collect the data age stats.
+  (void)collect_data_age_stats_;
 }
 
 Status CompactForTieringCollector::AddUserKey(const Slice& /*key*/,
@@ -93,9 +108,11 @@ CompactForTieringCollectorFactory::CreateTablePropertiesCollector(
       context.last_level_inclusive_max_seqno_threshold == kMaxSequenceNumber) {
     return nullptr;
   }
+  // TODO(yuzhangyu): pass actual value.
   return new CompactForTieringCollector(
       context.last_level_inclusive_max_seqno_threshold,
-      compaction_trigger_ratio);
+      compaction_trigger_ratio,
+      /*collect_data_age_stats*/ false);
 }
 
 static std::unordered_map<std::string, OptionTypeInfo>
@@ -141,4 +158,18 @@ NewCompactForTieringCollectorFactory(double compaction_trigger_ratio) {
       compaction_trigger_ratio);
 }
 
+Status GetDataCollectionUnixWriteTimeInfoForFile(
+    const std::shared_ptr<const TableProperties>& /* table_properties */,
+    std::unique_ptr<DataCollectionUnixWriteTimeInfo>* /* file_info */) {
+  return Status::NotSupported();
+}
+
+Status GetDataCollectionUnixWriteTimeInfoForLevels(
+    const std::vector<std::unique_ptr<
+        TablePropertiesCollection>>& /* levels_table_properties */,
+    std::vector<
+        std::unique_ptr<DataCollectionUnixWriteTimeInfo>>* /* levels_info */) {
+  return Status::NotSupported();
+}
+
 }  // namespace ROCKSDB_NAMESPACE

package/deps/rocksdb/rocksdb/utilities/table_properties_collectors/compact_for_tiering_collector.h

@@ -16,10 +16,14 @@ namespace ROCKSDB_NAMESPACE {
 class CompactForTieringCollector : public TablePropertiesCollector {
  public:
   static const std::string kNumEligibleLastLevelEntriesPropertyName;
+  static const std::string kAverageDataUnixWriteTimePropertyName;
+  static const std::string kMaxDataUnixWriteTimePropertyName;
+  static const std::string kMinDataUnixWriteTimePropertyName;
+  static const std::string kNumInfinitelyOldEntriesPropertyName;
 
   CompactForTieringCollector(
-      SequenceNumber last_level_inclusive_max_seqno_threshold_,
-      double compaction_trigger_ratio);
+      SequenceNumber last_level_inclusive_max_seqno_threshold,
+      double compaction_trigger_ratio, bool collect_data_age_stats);
 
   Status AddUserKey(const Slice& key, const Slice& value, EntryType type,
                     SequenceNumber seq, uint64_t file_size) override;
@@ -41,5 +45,6 @@ class CompactForTieringCollector : public TablePropertiesCollector {
   size_t total_entries_counter_ = 0;
   bool finish_called_ = false;
   bool need_compaction_ = false;
+  bool collect_data_age_stats_ = false;
 };
 }  // namespace ROCKSDB_NAMESPACE