feat(search): add a flag to dump DOT graphs in FT.EXPLAINSQL (apache#…

…2366)
Beihao-Zhou · Jun 15, 2024 · 717e52f · 717e52f
1 parent 8f1d2ad
commit 717e52f
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 13 deletions.
diff --git a/src/commands/cmd_search.cc b/src/commands/cmd_search.cc
@@ -19,18 +19,21 @@
  */
 
 #include <memory>
+#include <sstream>
 #include <variant>
 
 #include "commander.h"
 #include "commands/command_parser.h"
 #include "search/index_info.h"
 #include "search/ir.h"
+#include "search/ir_dot_dumper.h"
 #include "search/plan_executor.h"
 #include "search/redis_query_transformer.h"
 #include "search/search_encoding.h"
 #include "search/sql_transformer.h"
 #include "server/redis_reply.h"
 #include "server/server.h"
+#include "string_util.h"
 #include "tao/pegtl/string_input.hpp"
 
 namespace redis {
@@ -153,17 +156,45 @@ static void DumpQueryResult(const std::vector<kqir::ExecutorContext::RowType> &r
 }
 
 class CommandFTExplainSQL : public Commander {
+  Status Parse(const std::vector<std::string> &args) override {
+    if (args.size() == 3) {
+      if (util::EqualICase(args[2], "simple")) {
+        format_ = SIMPLE;
+      } else if (util::EqualICase(args[2], "dot")) {
+        format_ = DOT_GRAPH;
+      } else {
+        return {Status::NotOK, "output format should be SIMPLE or DOT"};
+      }
+    }
+
+    if (args.size() > 3) {
+      return {Status::NotOK, "more arguments than expected"};
+    }
+
+    return Status::OK();
+  }
+
   Status Execute(Server *srv, Connection *conn, std::string *output) override {
     const auto &sql = args_[1];
 
     auto ir = GET_OR_RET(kqir::sql::ParseToIR(kqir::peg::string_input(sql, "ft.explainsql")));
 
-    auto result = GET_OR_RET(srv->index_mgr.Explain(std::move(ir), conn->GetNamespace()));
+    auto plan = GET_OR_RET(srv->index_mgr.GeneratePlan(std::move(ir), conn->GetNamespace()));
 
-    output->append(BulkString(result));
+    if (format_ == SIMPLE) {
+      output->append(BulkString(plan->Dump()));
+    } else if (format_ == DOT_GRAPH) {
+      std::ostringstream ss;
+      kqir::DotDumper dumper(ss);
+
+      dumper.Dump(plan.get());
+      output->append(BulkString(ss.str()));
+    }
 
     return Status::OK();
   };
+
+  enum OutputFormat { SIMPLE, DOT_GRAPH } format_ = SIMPLE;
 };
 
 class CommandFTSearchSQL : public Commander {
@@ -233,9 +264,9 @@ class CommandFTExplain : public Commander {
 
   Status Execute(Server *srv, Connection *conn, std::string *output) override {
     CHECK(ir_);
-    auto result = GET_OR_RET(srv->index_mgr.Explain(std::move(ir_), conn->GetNamespace()));
+    auto plan = GET_OR_RET(srv->index_mgr.GeneratePlan(std::move(ir_), conn->GetNamespace()));
 
-    output->append(redis::SimpleString(result));
+    output->append(redis::BulkString(plan->Dump()));
 
     return Status::OK();
   };
@@ -330,7 +361,7 @@ class CommandFTDrop : public Commander {
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandFTCreate>("ft.create", -2, "write exclusive no-multi no-script", 0, 0, 0),
                         MakeCmdAttr<CommandFTSearchSQL>("ft.searchsql", 2, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandFTSearch>("ft.search", -3, "read-only", 0, 0, 0),
-                        MakeCmdAttr<CommandFTExplainSQL>("ft.explainsql", 2, "read-only", 0, 0, 0),
+                        MakeCmdAttr<CommandFTExplainSQL>("ft.explainsql", -2, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandFTExplain>("ft.explain", -3, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandFTInfo>("ft.info", 2, "read-only", 0, 0, 0),
                         MakeCmdAttr<CommandFTList>("ft._list", 1, "read-only", 0, 0, 0),

diff --git a/src/search/index_manager.h b/src/search/index_manager.h
@@ -194,12 +194,6 @@ struct IndexManager {
     return plan_op;
   }
 
-  StatusOr<std::string> Explain(std::unique_ptr<kqir::Node> ir, const std::string &ns) const {
-    auto plan_op = GET_OR_RET(GeneratePlan(std::move(ir), ns));
-
-    return plan_op->Dump();
-  }
-
   StatusOr<std::vector<kqir::ExecutorContext::RowType>> Search(std::unique_ptr<kqir::Node> ir,
                                                                const std::string &ns) const {
     auto plan_op = GET_OR_RET(GeneratePlan(std::move(ir), ns));

diff --git a/src/search/ir.h b/src/search/ir.h
@@ -363,7 +363,7 @@ struct SelectClause : Node {
 
   explicit SelectClause(std::vector<std::unique_ptr<FieldRef>> &&fields) : fields(std::move(fields)) {}
 
-  std::string_view Name() const override { return "SelectExpr"; }
+  std::string_view Name() const override { return "SelectClause"; }
   std::string Dump() const override {
     if (fields.empty()) return "select *";
     return fmt::format("select {}", util::StringJoin(fields, [](const auto &v) { return v->Dump(); }));
@@ -412,7 +412,7 @@ struct SearchExpr : Node {
         limit(std::move(limit)),
         sort_by(std::move(sort_by)) {}
 
-  std::string_view Name() const override { return "SearchStmt"; }
+  std::string_view Name() const override { return "SearchExpr"; }
   std::string Dump() const override {
     std::string opt;
     if (sort_by) opt += " " + sort_by->Dump();