Skip to content

Commit

Permalink
feat(search): add a flag to dump DOT graphs in FT.EXPLAINSQL
Browse files Browse the repository at this point in the history
  • Loading branch information
PragmaTwice committed Jun 15, 2024
1 parent bf579dc commit 8156705
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 13 deletions.
41 changes: 36 additions & 5 deletions src/commands/cmd_search.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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();
};
Expand Down Expand Up @@ -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),
Expand Down
6 changes: 0 additions & 6 deletions src/search/index_manager.h
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
4 changes: 2 additions & 2 deletions src/search/ir.h
Original file line number Diff line number Diff line change
Expand Up @@ -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(); }));
Expand Down Expand Up @@ -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();
Expand Down

0 comments on commit 8156705

Please sign in to comment.