EXPLAIN (FORMAT TPL) front-end stuff done. Just need the actual TPL. (#…

…1588)

src/include/parser/explain_statement.h

@@ -7,6 +7,9 @@
 
 namespace noisepage::parser {
 
+/** The format of the explain statement's output. */
+enum class ExplainStatementFormat : uint8_t { JSON, TPL, TBC };
+
 /**
  * Represents the SQL "EXPLAIN ..."
  */
@@ -20,11 +23,19 @@ class ExplainStatement : public SQLStatement {
 
   void Accept(common::ManagedPointer<binder::SqlNodeVisitor> v) override { v->Visit(common::ManagedPointer(this)); }
 
+  /** @brief override the default format of this EXPLAIN statement */
+  void SetFormat(const ExplainStatementFormat format) { format_ = format; }
+
   /** @return the SQL statement to be explained */
-  common::ManagedPointer<SQLStatement> GetSQLStatement() { return common::ManagedPointer(real_sql_stmt_); }
+  common::ManagedPointer<SQLStatement> GetSQLStatement() const { return common::ManagedPointer(real_sql_stmt_); }
+
+  /** @return format of the EXPLAIN */
+  ExplainStatementFormat GetFormat() const { return format_; }
 
  private:
   std::unique_ptr<SQLStatement> real_sql_stmt_;
+  ExplainStatementFormat format_ = ExplainStatementFormat::JSON;  // default to JSON since we rely on serializing the
+                                                                  // physical plan via dumping to JSON
 };
 
 }  // namespace noisepage::parser

src/network/postgres/postgres_network_commands.cpp

@@ -46,7 +46,10 @@ static void ExecutePortal(const common::ManagedPointer<network::ConnectionContex
     }
     if (query_type == network::QueryType::QUERY_CREATE_INDEX) {
       result = t_cop->ExecuteCreateStatement(connection_ctx, physical_plan, query_type);
+      NOISEPAGE_ASSERT(result.type_ == trafficcop::ResultType::COMPLETE,
+                       "Got through the binder as a valid index name, so we don't expect this to fail.");
       result = t_cop->CodegenPhysicalPlan(connection_ctx, out, portal);
+      // TODO(Matt): do something with result here in case codegen fails
       result = t_cop->RunExecutableQuery(connection_ctx, out, portal);
     } else {
       result = t_cop->ExecuteCreateStatement(connection_ctx, physical_plan, query_type);

src/parser/postgresparser.cpp

@@ -1874,9 +1874,35 @@ std::vector<common::ManagedPointer<AbstractExpression>> PostgresParser::ParamLis
 }
 
 std::unique_ptr<ExplainStatement> PostgresParser::ExplainTransform(ParseResult *parse_result, ExplainStmt *root) {
+  static constexpr char k_format_tok[] = "format";
   std::unique_ptr<ExplainStatement> result;
   auto query = NodeTransform(parse_result, root->query_);
   result = std::make_unique<ExplainStatement>(std::move(query));
+
+  if (root->options_ != nullptr) {
+    for (ListCell *cell = root->options_->head; cell != nullptr; cell = cell->next) {
+      NOISEPAGE_ASSERT(reinterpret_cast<Node *>(cell->data.ptr_value)->type == T_DefElem, "Expect a DefElem.");
+      const auto *const def_elem UNUSED_ATTRIBUTE = reinterpret_cast<DefElem *>(cell->data.ptr_value);
+
+      if (strncmp(def_elem->defname_, k_format_tok, sizeof(k_format_tok)) == 0) {
+        const auto *const format_cstr = reinterpret_cast<value *>(def_elem->arg_)->val_.str_;
+        // lowercase
+        if (strncmp(format_cstr, "tpl", 3) == 0) {
+          result->SetFormat(ExplainStatementFormat::TPL);
+        } else if (strncmp(format_cstr, "tbc", 3) == 0) {
+          result->SetFormat(ExplainStatementFormat::TBC);
+        } else if (strncmp(format_cstr, "json", 4) == 0) {
+          // this is the default format for us anyway so it's a noop
+          NOISEPAGE_ASSERT(result->GetFormat() == ExplainStatementFormat::JSON,
+                           "We assume this is the default format.");
+        } else {
+          throw ParserException("Unsupported format string for EXPLAIN.", __FILE__, __LINE__, def_elem->location_);
+        }
+      } else {
+        throw ParserException("Unsupported option string for EXPLAIN.", __FILE__, __LINE__, def_elem->location_);
+      }
+    }
+  }
   return result;
 }
 

src/traffic_cop/traffic_cop.cpp

@@ -365,9 +365,15 @@ TrafficCopResult TrafficCop::ExecuteExplainStatement(
   std::vector<planner::OutputSchema::Column> output_columns;
   output_columns.emplace_back("QUERY PLAN", type::TypeId::VARCHAR, nullptr);
 
-  // TODO(WAN): Integrate with Matt's PR.
-#if 0
-  {
+  const auto format =
+      portal->GetStatement()->RootStatement().CastManagedPointerTo<parser::ExplainStatement>()->GetFormat();
+  std::string plan_string;
+  if (format == parser::ExplainStatementFormat::JSON) {
+    plan_string = portal->OptimizeResult()->GetPlanNode()->ToJson().dump(4);
+  } else {
+    NOISEPAGE_ASSERT(format == parser::ExplainStatementFormat::TPL || format == parser::ExplainStatementFormat::TBC,
+                     "We only support JSON, TPL, and TBC formats.");
+
     // Codegen must happen for certain types of EXPLAIN metadata to be collected, e.g., collection of TPL.
     auto codegen = CodegenPhysicalPlan(connection_ctx, out, portal);
     if (codegen.type_ != ResultType::COMPLETE) {
@@ -378,15 +384,15 @@ TrafficCopResult TrafficCop::ExecuteExplainStatement(
     const auto &fragments = portal->GetStatement()->GetExecutableQuery()->GetFragments();
     NOISEPAGE_ASSERT(fragments.size() == 1, "We currently always compile with just one query fragment.");
     const auto &metadata = fragments.at(0)->GetModuleMetadata().GetCompileTimeMetadata();
-    std::string output = metadata.GetTPL() + "\n\n" + metadata.GetTBC();
-    const execution::sql::StringVal plan_string_val = execution::sql::StringVal(output.c_str(), output.length());
-  }
-#endif
 
-  const auto physical_plan = portal->OptimizeResult()->GetPlanNode();
-  const std::string plan_string = physical_plan->ToJson().dump(4);
-  const execution::sql::StringVal plan_string_val =
-      execution::sql::StringVal(plan_string.c_str(), plan_string.length());
+    if (format == parser::ExplainStatementFormat::TPL) {
+      plan_string = metadata.GetTPL();
+    } else {
+      NOISEPAGE_ASSERT(format == parser::ExplainStatementFormat::TBC, "Did you add a new case?");
+      plan_string = metadata.GetTBC();
+    }
+  }
+  const auto plan_string_val = execution::sql::StringVal(plan_string.c_str(), plan_string.length());
   out->WriteDataRow(reinterpret_cast<const byte *const>(&plan_string_val), output_columns,
                     {network::FieldFormat::text});
 
@@ -481,13 +487,15 @@ TrafficCopResult TrafficCop::CodegenPhysicalPlan(
   execution::exec::ExecutionSettings exec_settings{};
   exec_settings.UpdateFromSettingsManager(settings_manager_);
 
-  // TODO(WAN): Integrate with Matt's PR.
   // Set any compilation settings based on the original query type.
   if (portal->GetStatement()->GetQueryType() == network::QueryType::QUERY_EXPLAIN) {
     execution::compiler::CompilerSettings settings;
-    // if (portal->GetStatement()->RootStatement().CastManagedPointerTo<parser::ExplainStatement>().GetFormat()...)
-    settings.SetShouldCaptureTPL(true);
-    settings.SetShouldCaptureTBC(true);
+    auto stmt = portal->GetStatement()->RootStatement().CastManagedPointerTo<parser::ExplainStatement>();
+    if (stmt->GetFormat() == parser::ExplainStatementFormat::TPL) {
+      settings.SetShouldCaptureTPL(true);
+    } else if (stmt->GetFormat() == parser::ExplainStatementFormat::TBC) {
+      settings.SetShouldCaptureTBC(true);
+    }
     exec_settings.SetCompilerSettings(settings);
   }