Add memory constraint support for Pilot (#1584)

src/catalog/catalog_accessor.cpp

@@ -163,6 +163,8 @@ common::ManagedPointer<storage::index::Index> CatalogAccessor::GetIndex(index_oi
   return dbc_->GetIndex(txn_, index);
 }
 
+std::string_view CatalogAccessor::GetIndexName(index_oid_t index) const { return dbc_->GetIndexName(txn_, index); }
+
 language_oid_t CatalogAccessor::CreateLanguage(const std::string &lanname) {
   return dbc_->CreateLanguage(txn_, lanname);
 }

src/catalog/database_catalog.cpp

@@ -219,6 +219,14 @@ common::ManagedPointer<storage::index::Index> DatabaseCatalog::GetIndex(
   return common::ManagedPointer(reinterpret_cast<storage::index::Index *>(ptr_pair.first));
 }
 
+std::string_view DatabaseCatalog::GetIndexName(const common::ManagedPointer<transaction::TransactionContext> txn,
+                                               index_oid_t index) {
+  const auto name_pair = pg_core_.GetClassNameKind(txn, index.UnderlyingValue());
+  NOISEPAGE_ASSERT(name_pair.second == postgres::PgClass::RelKind::INDEX,
+                   "Called GetIndexName with an OID for an object that doesn't have type INDEX");
+  return name_pair.first;
+}
+
 const Schema &DatabaseCatalog::GetSchema(const common::ManagedPointer<transaction::TransactionContext> txn,
                                          const table_oid_t table) {
   const auto ptr_pair = pg_core_.GetClassSchemaPtrKind(txn, table.UnderlyingValue());

src/catalog/postgres/pg_core_impl.cpp

@@ -50,6 +50,9 @@ void PgCoreImpl::BootstrapPRIsPgClass() {
   const std::vector<col_oid_t> set_class_schema_oids{PgClass::REL_SCHEMA.oid_};
   set_class_schema_pri_ = classes_->InitializerForProjectedRow(set_class_schema_oids);
 
+  const std::vector<col_oid_t> get_class_name_kind_oids{PgClass::RELNAME.oid_, PgClass::RELKIND.oid_};
+  get_class_name_kind_pri_ = classes_->InitializerForProjectedRow(get_class_name_kind_oids);
+
   const std::vector<col_oid_t> get_class_pointer_kind_oids{PgClass::REL_PTR.oid_, PgClass::RELKIND.oid_};
   get_class_pointer_kind_pri_ = classes_->InitializerForProjectedRow(get_class_pointer_kind_oids);
 
@@ -1292,6 +1295,44 @@ std::pair<uint32_t, PgClass::RelKind> PgCoreImpl::GetClassOidKind(
   return std::make_pair(oid, kind);
 }
 
+std::pair<std::string_view, PgClass::RelKind> PgCoreImpl::GetClassNameKind(
+    common::ManagedPointer<transaction::TransactionContext> txn, uint32_t oid) {
+  // Buffer is large enough to hold all PRs.
+  byte *const buffer = common::AllocationUtil::AllocateAligned(get_class_name_kind_pri_.ProjectedRowSize());
+
+  const auto &oid_pri = classes_oid_index_->GetProjectedRowInitializer();
+
+  // Scan pg_class_oid_index.
+  std::vector<storage::TupleSlot> index_results;
+  {
+    auto *pr = oid_pri.InitializeRow(buffer);
+    pr->Set<table_oid_t, false>(0, table_oid_t(oid), false);
+    classes_oid_index_->ScanKey(*txn, *pr, &index_results);
+  }
+
+  NOISEPAGE_ASSERT(index_results.size() == 1, "Oid not found in pg_class_oid_index.");
+  NOISEPAGE_ASSERT(get_class_name_kind_pri_.ProjectedRowSize() >= oid_pri.ProjectedRowSize(),
+                   "I want to reuse this buffer because I'm lazy and malloc is slow but it needs to be big enough.");
+
+  // Select out the tuple from pg_class.
+  auto *pr = get_class_name_kind_pri_.InitializeRow(buffer);
+  {
+    const auto result UNUSED_ATTRIBUTE = classes_->Select(txn, index_results[0], pr);
+    NOISEPAGE_ASSERT(result, "Index already verified visibility. This shouldn't fail.");
+  }
+
+  // Get the name and kind.
+  storage::VarlenEntry name_varlen;
+  PgClass::RelKind kind;
+  {
+    name_varlen = *pr->Get<storage::VarlenEntry, false>(0, nullptr);
+    kind = static_cast<PgClass::RelKind>(*pr->Get<char, false>(1, nullptr));
+  }
+
+  delete[] buffer;
+  return std::make_pair(name_varlen.StringView(), kind);
+}
+
 template <typename Column, typename ClassOid, typename ColOid>
 bool PgCoreImpl::CreateColumn(const common::ManagedPointer<transaction::TransactionContext> txn,
                               const ClassOid class_oid, const ColOid col_oid, const Column &col) {

src/include/catalog/catalog.h

@@ -21,6 +21,10 @@ class GarbageCollector;
 class RecoveryManager;
 }  // namespace noisepage::storage
 
+namespace noisepage::selfdriving {
+class PilotUtil;
+}  // namespace noisepage::selfdriving
+
 namespace noisepage::catalog {
 
 class CatalogCache;
@@ -137,6 +141,7 @@ class Catalog {
  private:
   DISALLOW_COPY_AND_MOVE(Catalog);
   friend class storage::RecoveryManager;
+  friend class selfdriving::PilotUtil;
   const common::ManagedPointer<transaction::TransactionManager> txn_manager_;
   const common::ManagedPointer<storage::BlockStore> catalog_block_store_;
   const common::ManagedPointer<storage::GarbageCollector> garbage_collector_;

src/include/catalog/catalog_accessor.h

@@ -282,6 +282,13 @@ class EXPORT CatalogAccessor {
    */
   common::ManagedPointer<storage::index::Index> GetIndex(index_oid_t index) const;
 
+  /**
+   * Obtain the name of the index
+   * @param index to which we want the name. The oid must be valid, otherwise triggers assertion
+   * @return name of the index.
+   */
+  std::string_view GetIndexName(index_oid_t index) const;
+
   /**
    * Adds a language to the catalog (with default parameters for now) if
    * it doesn't exist in pg_language already

src/include/catalog/database_catalog.h

@@ -110,6 +110,9 @@ class DatabaseCatalog {
   common::ManagedPointer<storage::index::Index> GetIndex(common::ManagedPointer<transaction::TransactionContext> txn,
                                                          index_oid_t index);
 
+  /** @brief Get the name of the specified index */
+  std::string_view GetIndexName(common::ManagedPointer<transaction::TransactionContext> txn, index_oid_t index);
+
   /** @brief Get the schema for the specified table. */
   const Schema &GetSchema(common::ManagedPointer<transaction::TransactionContext> txn, table_oid_t table);
   /** @brief Get the index schema for the specified index. */

src/include/catalog/postgres/pg_core_impl.h

@@ -250,6 +250,17 @@ class PgCoreImpl {
    */
   std::pair<void *, PgClass::RelKind> GetClassSchemaPtrKind(common::ManagedPointer<transaction::TransactionContext> txn,
                                                             uint32_t oid);
+
+  /**
+   * @brief Get the name string from pg_class.
+   *
+   * @param txn     The transaction to query in.
+   * @param oid     The OID of the object. Must be valid
+   * @return        The name and the RelKind of the object requested.
+   */
+  std::pair<std::string_view, PgClass::RelKind> GetClassNameKind(
+      common::ManagedPointer<transaction::TransactionContext> txn, uint32_t oid);
+
   /**
    * @brief Get the OID and kind from pg_class.
    *
@@ -354,6 +365,7 @@ class PgCoreImpl {
   storage::ProjectionMap pg_class_all_cols_prm_;
   storage::ProjectedRowInitializer get_class_oid_kind_pri_;
   storage::ProjectedRowInitializer set_class_pointer_pri_;
+  storage::ProjectedRowInitializer get_class_name_kind_pri_;
   storage::ProjectedRowInitializer set_class_schema_pri_;
   storage::ProjectedRowInitializer get_class_pointer_kind_pri_;
   storage::ProjectedRowInitializer get_class_schema_pointer_kind_pri_;

src/include/optimizer/logical_operators.h

@@ -1271,6 +1271,7 @@ class LogicalCreateFunction : public OperatorNodeContents<LogicalCreateFunction>
 class LogicalCreateIndex : public OperatorNodeContents<LogicalCreateIndex> {
  public:
   /**
+   * @param database_oid OID of the database
    * @param namespace_oid OID of the namespace
    * @param table_oid OID of the table
    * @param index_type Type of the index
@@ -1279,8 +1280,9 @@ class LogicalCreateIndex : public OperatorNodeContents<LogicalCreateIndex> {
    * @param index_attrs Attributes of the index
    * @return
    */
-  static Operator Make(catalog::namespace_oid_t namespace_oid, catalog::table_oid_t table_oid,
-                       parser::IndexType index_type, bool unique, std::string index_name,
+  static Operator Make(catalog::db_oid_t database_oid, catalog::namespace_oid_t namespace_oid,
+                       catalog::table_oid_t table_oid, parser::IndexType index_type, bool unique,
+                       std::string index_name,
                        std::vector<common::ManagedPointer<parser::AbstractExpression>> index_attrs);
 
   /**
@@ -1297,6 +1299,11 @@ class LogicalCreateIndex : public OperatorNodeContents<LogicalCreateIndex> {
    */
   const catalog::namespace_oid_t &GetNamespaceOid() const { return namespace_oid_; }
 
+  /**
+   * @return OID of the database
+   */
+  const catalog::db_oid_t &GetDatabaseOid() const { return database_oid_; }
+
   /**
    * @return OID of the table
    */
@@ -1323,6 +1330,11 @@ class LogicalCreateIndex : public OperatorNodeContents<LogicalCreateIndex> {
   const std::vector<common::ManagedPointer<parser::AbstractExpression>> &GetIndexAttr() const { return index_attrs_; }
 
  private:
+  /**
+   * OID of the database
+   */
+  catalog::db_oid_t database_oid_;
+
   /**
    * OID of the namespace
    */

src/include/optimizer/statistics/stats_calculator.h

@@ -80,6 +80,12 @@ class StatsCalculator : public OperatorVisitor {
    */
   void Visit(const LogicalDelete *op) override;
 
+  /**
+   * Visit a LogicalCreateIndex
+   * @param op Operator being visited
+   */
+  void Visit(const LogicalCreateIndex *op) override;
+
  private:
   /**
    * Return estimated cardinality for a filter

src/include/self_driving/forecasting/workload_forecast.h

@@ -2,11 +2,13 @@
 
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <unordered_map>
 #include <utility>
 #include <vector>
 
+#include "catalog/catalog_defs.h"
 #include "parser/expression/constant_value_expression.h"
 #include "self_driving/forecasting/workload_forecast_segment.h"
 
@@ -28,7 +30,7 @@ using WorkloadForecastPrediction = std::unordered_map<uint64_t, std::unordered_m
 class WorkloadMetadata {
  public:
   /** Map from query id to database id */
-  std::unordered_map<execution::query_id_t, uint64_t> query_id_to_dboid_;
+  std::unordered_map<execution::query_id_t, catalog::db_oid_t> query_id_to_dboid_;
 
   /** Map from query id to query text */
   std::unordered_map<execution::query_id_t, std::string> query_id_to_text_;
@@ -74,16 +76,26 @@ class WorkloadForecast {
    * Get number of forecasted segments
    * @return number of forecasted segments
    */
-  uint64_t GetNumberOfSegments() { return num_forecast_segment_; }
+  uint64_t GetNumberOfSegments() const { return num_forecast_segment_; }
+
+  /** @brief Get the set of unique db oids that the forecasted workload has queries with */
+  std::set<catalog::db_oid_t> GetDBOidSet() {
+    std::set<catalog::db_oid_t> db_oid_set;
+    for (auto &[qid, db_oid] : workload_metadata_.query_id_to_dboid_) db_oid_set.insert(db_oid);
+    return db_oid_set;
+  }
 
  private:
   friend class PilotUtil;
-  const WorkloadForecastSegment &GetSegmentByIndex(uint64_t segment_index) {
+
+  uint64_t GetForecastInterval() const { return forecast_interval_; }
+
+  const WorkloadForecastSegment &GetSegmentByIndex(uint64_t segment_index) const {
     NOISEPAGE_ASSERT(segment_index < num_forecast_segment_, "invalid index");
     return forecast_segments_[segment_index];
   }
 
-  std::string GetQuerytextByQid(execution::query_id_t qid) {
+  std::string GetQuerytextByQid(execution::query_id_t qid) const {
     NOISEPAGE_ASSERT(workload_metadata_.query_id_to_text_.find(qid) != workload_metadata_.query_id_to_text_.end(),
                      "invalid qid");
     return workload_metadata_.query_id_to_text_.at(qid);
@@ -102,12 +114,14 @@ class WorkloadForecast {
     return &(workload_metadata_.query_id_to_param_types_.at(qid));
   }
 
-  uint64_t GetDboidByQid(execution::query_id_t qid) {
+  catalog::db_oid_t GetDboidByQid(execution::query_id_t qid) const {
     NOISEPAGE_ASSERT(workload_metadata_.query_id_to_dboid_.find(qid) != workload_metadata_.query_id_to_dboid_.end(),
                      "invalid qid");
     return workload_metadata_.query_id_to_dboid_.at(qid);
   }
 
+  const WorkloadMetadata &GetWorkloadMetadata() const { return workload_metadata_; }
+
   /**
    * Initializes segments from inference results
    * @param inference Inference results

src/include/self_driving/planning/action/abstract_action.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "catalog/catalog_defs.h"
@@ -9,6 +10,8 @@
 
 namespace noisepage::selfdriving::pilot {
 
+class ActionState;
+
 /**
  * The abstract class for self-driving actions
  */
@@ -20,26 +23,40 @@ class AbstractAction {
    * @param db_oid The ID of the database that this action belongs to
    */
   explicit AbstractAction(ActionType family, catalog::db_oid_t db_oid)
-      : action_family_(family), db_oid_(db_oid), id_(action_id_counter++) {}
+      : id_(action_id_counter++), action_type_(family), db_oid_(db_oid) {}
 
   virtual ~AbstractAction() = default;
 
   /**
    * Set the estimated runtime metrics for this action
    * @param estimated_metrics The metrics to set to
    */
-  void SetEstimatedMetrics(const common::ResourceTracker::Metrics &estimated_metrics) {
-    estimated_metrics_ = estimated_metrics;
+  void SetEstimatedMetrics(std::vector<double> &&estimated_metrics) {
+    estimated_metrics_ = std::move(estimated_metrics);
   }
 
   /** @return The estimated runtime metrics for this action */
-  const common::ResourceTracker::Metrics &GetEstimatedMetrics() { return estimated_metrics_; }
+  const std::vector<double> &GetEstimatedMetrics() { return estimated_metrics_; }
+
+  /** @return The estimated elapsed time in us for this action */
+  double GetEstimatedElapsedUs() {
+    if (estimated_metrics_.empty()) return 0;
+    // Assumes the elapsed time is the last element
+    return estimated_metrics_.end()[-1];
+  }
+
+  /** @return The estimated memory consumption in bytes for this action */
+  double GetEstimatedMemoryBytes() {
+    if (estimated_metrics_.size() < 2) return 0;
+    // Assumes the memory consumption is the second last element
+    return estimated_metrics_.end()[-2];
+  }
 
   /** @return This action's ID */
   action_id_t GetActionID() const { return id_; }
 
-  /** @return This action's family */
-  ActionType GetActionFamily() const { return action_family_; }
+  /** @return This action's type */
+  ActionType GetActionType() const { return action_type_; }
 
   /** @return This action's database oid */
   catalog::db_oid_t GetDatabaseOid() const { return db_oid_; }
@@ -95,21 +112,27 @@ class AbstractAction {
    */
   virtual bool IsValid() { return true; }
 
+  /**
+   * Modify the action state given the effect of this action
+   * @param action_state the action state to modify
+   */
+  virtual void ModifyActionState(ActionState *action_state) = 0;
+
  protected:
+  /** ID is unique for an action among on planning process (one MCTS) */
+  action_id_t id_;
+
   std::string sql_command_;  ///< The SQL commaned used to apply the action
 
  private:
   static action_id_t action_id_counter;
 
-  common::ResourceTracker::Metrics estimated_metrics_{};
+  std::vector<double> estimated_metrics_;
 
-  ActionType action_family_;
+  ActionType action_type_;
 
   catalog::db_oid_t db_oid_;
 
-  /** ID is unique for an action among on planning process (one MCTS) */
-  action_id_t id_;
-
   std::vector<action_id_t> invalidated_action_ids_;
   std::vector<action_id_t> enabled_action_ids_;
   std::vector<action_id_t> reverse_action_ids_;

src/include/self_driving/planning/action/change_knob_action.h

@@ -30,6 +30,8 @@ class ChangeKnobAction : public AbstractAction {
 
   const std::string &GetSQLCommand() override;
 
+  void ModifyActionState(ActionState *action_state) override;
+
   bool IsValid() override;
 
  private: