Fix plan operator child iterator in KQIR (apache#2282)

Beihao-Zhou · Apr 29, 2024 · 1834ce8 · 1834ce8
1 parent 9d9dff2
commit 1834ce8
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 5 deletions.
diff --git a/src/search/ir_plan.h b/src/search/ir_plan.h
@@ -44,9 +44,12 @@ struct FullIndexScan : PlanOperator {
 
   explicit FullIndexScan(std::unique_ptr<IndexRef> index) : index(std::move(index)) {}
 
-  std::string_view Name() const override { return "FullIndexScan"; };
+  std::string_view Name() const override { return "FullIndexScan"; }
   std::string Dump() const override { return fmt::format("full-scan {}", index->name); }
 
+  NodeIterator ChildBegin() override { return NodeIterator{index.get()}; }
+  NodeIterator ChildEnd() override { return {}; }
+
   std::unique_ptr<Node> Clone() const override {
     return std::make_unique<FullIndexScan>(Node::MustAs<IndexRef>(index->Clone()));
   }
@@ -56,6 +59,9 @@ struct FieldScan : PlanOperator {
   std::unique_ptr<FieldRef> field;
 
   explicit FieldScan(std::unique_ptr<FieldRef> field) : field(std::move(field)) {}
+
+  NodeIterator ChildBegin() override { return NodeIterator{field.get()}; }
+  NodeIterator ChildEnd() override { return {}; }
 };
 
 struct NumericFieldScan : FieldScan {
@@ -67,9 +73,9 @@ struct NumericFieldScan : FieldScan {
 
   std::string_view Name() const override { return "NumericFieldScan"; };
   std::string Content() const override {
-    return fmt::format("{}, {}, {}", field->name, range.ToString(), SortByClause::OrderToString(order));
+    return fmt::format("{}, {}", range.ToString(), SortByClause::OrderToString(order));
   };
-  std::string Dump() const override { return fmt::format("numeric-scan {}", Content()); }
+  std::string Dump() const override { return fmt::format("numeric-scan {}, {}", field->name, Content()); }
 
   std::unique_ptr<Node> Clone() const override {
     return std::make_unique<NumericFieldScan>(field->CloneAs<FieldRef>(), range, order);
@@ -82,8 +88,8 @@ struct TagFieldScan : FieldScan {
   TagFieldScan(std::unique_ptr<FieldRef> field, std::string tag) : FieldScan(std::move(field)), tag(std::move(tag)) {}
 
   std::string_view Name() const override { return "TagFieldScan"; };
-  std::string Content() const override { return fmt::format("{}, {}", field->name, tag); };
-  std::string Dump() const override { return fmt::format("tag-scan {}", Content()); }
+  std::string Content() const override { return tag; };
+  std::string Dump() const override { return fmt::format("tag-scan {}, {}", field->name, tag); }
 
   std::unique_ptr<Node> Clone() const override {
     return std::make_unique<TagFieldScan>(field->CloneAs<FieldRef>(), tag);

diff --git a/tests/cppunit/ir_dot_dumper_test.cc b/tests/cppunit/ir_dot_dumper_test.cc
@@ -24,6 +24,9 @@
 #include <sstream>
 
 #include "gtest/gtest.h"
+#include "search/ir_plan.h"
+#include "search/ir_sema_checker.h"
+#include "search/passes/manager.h"
 #include "search/sql_transformer.h"
 
 using namespace kqir;
@@ -53,3 +56,59 @@ TEST(DotDumperTest, Simple) {
   ASSERT_NE(dot.find(fmt::format("{} -> {}", search_stmt, or_expr)), std::string::npos);
   ASSERT_NE(dot.find(fmt::format("{} -> {}", or_expr, and_expr)), std::string::npos);
 }
+
+static auto ParseS(SemaChecker& sc, const std::string& in) {
+  auto ir = *Parse(in);
+  EXPECT_EQ(sc.Check(ir.get()).Msg(), Status::ok_msg);
+  return ir;
+}
+
+static IndexMap MakeIndexMap() {
+  auto f1 = FieldInfo("t1", std::make_unique<redis::SearchTagFieldMetadata>());
+  auto f2 = FieldInfo("t2", std::make_unique<redis::SearchTagFieldMetadata>());
+  f2.metadata->noindex = true;
+  auto f3 = FieldInfo("n1", std::make_unique<redis::SearchNumericFieldMetadata>());
+  auto f4 = FieldInfo("n2", std::make_unique<redis::SearchNumericFieldMetadata>());
+  auto f5 = FieldInfo("n3", std::make_unique<redis::SearchNumericFieldMetadata>());
+  f5.metadata->noindex = true;
+  auto ia = IndexInfo("ia", SearchMetadata());
+  ia.Add(std::move(f1));
+  ia.Add(std::move(f2));
+  ia.Add(std::move(f3));
+  ia.Add(std::move(f4));
+  ia.Add(std::move(f5));
+
+  auto& name = ia.name;
+  IndexMap res;
+  res.emplace(name, std::move(ia));
+  return res;
+}
+
+TEST(DotDumperTest, Plan) {
+  auto index_map = MakeIndexMap();
+  SemaChecker sc(index_map);
+  auto plan = PassManager::Execute(
+      PassManager::Default(),
+      ParseS(
+          sc,
+          "select * from ia where (n1 < 2 or n1 >= 3) and (n1 >= 1 and n1 < 4) and not n3 != 1 and t2 hastag \"a\""));
+
+  std::stringstream ss;
+  DotDumper dd(ss);
+  dd.Dump(plan.get());
+
+  std::string dot = ss.str();
+  std::smatch matches;
+
+  std::regex_search(dot, matches, std::regex(R"((\w+) \[ label = "Filter)"));
+  auto filter = matches[1].str();
+
+  std::regex_search(dot, matches, std::regex(R"((\w+) \[ label = "Merge)"));
+  auto merge = matches[1].str();
+
+  std::regex_search(dot, matches, std::regex(R"((\w+) \[ label = "NumericFieldScan)"));
+  auto scan = matches[1].str();
+
+  ASSERT_NE(dot.find(fmt::format("{} -> {}", filter, merge)), std::string::npos);
+  ASSERT_NE(dot.find(fmt::format("{} -> {}", merge, scan)), std::string::npos);
+}