Skip to content

Commit

Permalink
Add clone method to KQIR Node (#2234)
Browse files Browse the repository at this point in the history
  • Loading branch information
@PragmaTwice
PragmaTwice committed Apr 10, 2024
1 parent 70581d0 commit fd5c376
Show file tree
Hide file tree
Showing 2 changed files with 131 additions and 9 deletions.
82 changes: 76 additions & 6 deletions src/search/ir.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ struct Node {
virtual NodeIterator ChildBegin() { return {}; };
virtual NodeIterator ChildEnd() { return {}; };

virtual std::unique_ptr<Node> Clone() const = 0;

virtual ~Node() = default;

template <typename T, typename U = Node, typename... Args>
Expand All @@ -70,27 +72,35 @@ struct Node {
}
};

struct FieldRef : Node {
struct Ref : Node {};

struct FieldRef : Ref {
std::string name;

explicit FieldRef(std::string name) : name(std::move(name)) {}

std::string_view Name() const override { return "FieldRef"; }
std::string Dump() const override { return name; }
std::string Content() const override { return Dump(); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<FieldRef>(*this); }
};

struct StringLiteral : Node {
struct Literal : virtual Node {};

struct StringLiteral : Literal {
std::string val;

explicit StringLiteral(std::string val) : val(std::move(val)) {}

std::string_view Name() const override { return "StringLiteral"; }
std::string Dump() const override { return fmt::format("\"{}\"", util::EscapeString(val)); }
std::string Content() const override { return Dump(); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<StringLiteral>(*this); }
};

struct QueryExpr : Node {};
struct QueryExpr : virtual Node {};

struct BoolAtomExpr : QueryExpr {};

Expand All @@ -106,16 +116,23 @@ struct TagContainExpr : BoolAtomExpr {

NodeIterator ChildBegin() override { return {field.get(), tag.get()}; };
NodeIterator ChildEnd() override { return {}; };

std::unique_ptr<Node> Clone() const override {
return std::make_unique<TagContainExpr>(Node::MustAs<FieldRef>(field->Clone()),
Node::MustAs<StringLiteral>(tag->Clone()));
}
};

struct NumericLiteral : Node {
struct NumericLiteral : Literal {
double val;

explicit NumericLiteral(double val) : val(val) {}

std::string_view Name() const override { return "NumericLiteral"; }
std::string Dump() const override { return fmt::format("{}", val); }
std::string Content() const override { return Dump(); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<NumericLiteral>(*this); }
};

// NOLINTNEXTLINE
Expand Down Expand Up @@ -189,16 +206,23 @@ struct NumericCompareExpr : BoolAtomExpr {

NodeIterator ChildBegin() override { return {field.get(), num.get()}; };
NodeIterator ChildEnd() override { return {}; };

std::unique_ptr<Node> Clone() const override {
return std::make_unique<NumericCompareExpr>(op, Node::MustAs<FieldRef>(field->Clone()),
Node::MustAs<NumericLiteral>(num->Clone()));
}
};

struct BoolLiteral : BoolAtomExpr {
struct BoolLiteral : BoolAtomExpr, Literal {
bool val;

explicit BoolLiteral(bool val) : val(val) {}

std::string_view Name() const override { return "BoolLiteral"; }
std::string Dump() const override { return val ? "true" : "false"; }
std::string Content() const override { return Dump(); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<BoolLiteral>(*this); }
};

struct QueryExpr;
Expand All @@ -213,6 +237,10 @@ struct NotExpr : QueryExpr {

NodeIterator ChildBegin() override { return NodeIterator{inner.get()}; };
NodeIterator ChildEnd() override { return {}; };

std::unique_ptr<Node> Clone() const override {
return std::make_unique<NotExpr>(Node::MustAs<QueryExpr>(inner->Clone()));
}
};

struct AndExpr : QueryExpr {
Expand All @@ -227,6 +255,15 @@ struct AndExpr : QueryExpr {

NodeIterator ChildBegin() override { return NodeIterator(inners.begin()); };
NodeIterator ChildEnd() override { return NodeIterator(inners.end()); };

std::unique_ptr<Node> Clone() const override {
std::vector<std::unique_ptr<QueryExpr>> res;
res.reserve(inners.size());
for (const auto &n : inners) {
res.push_back(Node::MustAs<QueryExpr>(n->Clone()));
}
return std::make_unique<AndExpr>(std::move(res));
}
};

struct OrExpr : QueryExpr {
Expand All @@ -241,6 +278,15 @@ struct OrExpr : QueryExpr {

NodeIterator ChildBegin() override { return NodeIterator(inners.begin()); };
NodeIterator ChildEnd() override { return NodeIterator(inners.end()); };

std::unique_ptr<Node> Clone() const override {
std::vector<std::unique_ptr<QueryExpr>> res;
res.reserve(inners.size());
for (const auto &n : inners) {
res.push_back(Node::MustAs<QueryExpr>(n->Clone()));
}
return std::make_unique<OrExpr>(std::move(res));
}
};

struct LimitClause : Node {
Expand All @@ -252,6 +298,8 @@ struct LimitClause : Node {
std::string_view Name() const override { return "LimitClause"; }
std::string Dump() const override { return fmt::format("limit {}, {}", offset, count); }
std::string Content() const override { return fmt::format("{}, {}", offset, count); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<LimitClause>(*this); }
};

struct SortByClause : Node {
Expand All @@ -268,6 +316,10 @@ struct SortByClause : Node {

NodeIterator ChildBegin() override { return NodeIterator(field.get()); };
NodeIterator ChildEnd() override { return {}; };

std::unique_ptr<Node> Clone() const override {
return std::make_unique<SortByClause>(order, Node::MustAs<FieldRef>(field->Clone()));
}
};

struct SelectExpr : Node {
Expand All @@ -283,16 +335,27 @@ struct SelectExpr : Node {

NodeIterator ChildBegin() override { return NodeIterator(fields.begin()); };
NodeIterator ChildEnd() override { return NodeIterator(fields.end()); };

std::unique_ptr<Node> Clone() const override {
std::vector<std::unique_ptr<FieldRef>> res;
res.reserve(fields.size());
for (const auto &f : fields) {
res.push_back(Node::MustAs<FieldRef>(f->Clone()));
}
return std::make_unique<SelectExpr>(std::move(res));
}
};

struct IndexRef : Node {
struct IndexRef : Ref {
std::string name;

explicit IndexRef(std::string name) : name(std::move(name)) {}

std::string_view Name() const override { return "IndexRef"; }
std::string Dump() const override { return name; }
std::string Content() const override { return Dump(); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<IndexRef>(*this); }
};

struct SearchStmt : Node {
Expand Down Expand Up @@ -328,6 +391,13 @@ struct SearchStmt : Node {

NodeIterator ChildBegin() override { return NodeIterator(this, ChildMap.begin()); };
NodeIterator ChildEnd() override { return NodeIterator(this, ChildMap.end()); };

std::unique_ptr<Node> Clone() const override {
return std::make_unique<SearchStmt>(
Node::MustAs<IndexRef>(index->Clone()), Node::MustAs<QueryExpr>(query_expr->Clone()),
Node::MustAs<LimitClause>(limit->Clone()), Node::MustAs<SortByClause>(sort_by->Clone()),
Node::MustAs<SelectExpr>(select_expr->Clone()));
}
};

} // namespace kqir
58 changes: 55 additions & 3 deletions src/search/ir_plan.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,80 +34,132 @@ struct PlanOperator : Node {};
struct FullIndexScan : PlanOperator {
IndexInfo *index;

explicit FullIndexScan(IndexInfo *index) : index(index) {}

std::string_view Name() const override { return "FullIndexScan"; };
std::string Content() const override { return index->name; };
std::string Dump() const override { return fmt::format("full-scan {}", Content()); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<FullIndexScan>(*this); }
};

struct FieldScan : PlanOperator {
FieldInfo *field;

explicit FieldScan(FieldInfo *field) : field(field) {}
};

struct Interval {
double l, r; // [l, r)

explicit Interval(double l, double r) : l(l), r(r) {}

std::string ToString() const { return fmt::format("[{}, {})", l, r); }
};

struct NumericFieldScan : FieldScan {
Interval range;

NumericFieldScan(FieldInfo *field, Interval range) : FieldScan(field), range(range) {}

std::string_view Name() const override { return "NumericFieldScan"; };
std::string Content() const override { return fmt::format("{}, {}", field->name, range.ToString()); };
std::string Dump() const override { return fmt::format("numeric-scan {}", Content()); }

std::unique_ptr<Node> Clone() const override { return std::make_unique<NumericFieldScan>(*this); }
};

struct TagFieldScan : FieldScan {
std::string tag;

TagFieldScan(FieldInfo *field, std::string tag) : FieldScan(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::unique_ptr<Node> Clone() const override { return std::make_unique<TagFieldScan>(*this); }
};

struct Filter : PlanOperator {
std::unique_ptr<PlanOperator> source;
std::unique_ptr<QueryExpr> filter_expr;

Filter(std::unique_ptr<PlanOperator> &&source, std::unique_ptr<QueryExpr> &&filter_expr)
: source(std::move(source)), filter_expr(std::move(filter_expr)) {}

std::string_view Name() const override { return "Filter"; };
std::string Dump() const override { return fmt::format("(filter {}, {})", source->Dump(), Content()); }

NodeIterator ChildBegin() override { return {source.get(), filter_expr.get()}; }
NodeIterator ChildEnd() override { return {}; }

std::unique_ptr<Node> Clone() const override {
return std::make_unique<Filter>(Node::MustAs<PlanOperator>(source->Clone()),
Node::MustAs<QueryExpr>(filter_expr->Clone()));
}
};

struct Merge : PlanOperator {
std::vector<std::unique_ptr<PlanOperator>> ops;

explicit Merge(std::vector<std::unique_ptr<PlanOperator>> &&ops) : ops(std::move(ops)) {}

std::string_view Name() const override { return "Merge"; };
std::string Dump() const override {
return fmt::format("(merge {})", util::StringJoin(ops, [](const auto &v) { return v->Dump(); }));
}

NodeIterator ChildBegin() override { return NodeIterator(ops.begin()); }
NodeIterator ChildEnd() override { return NodeIterator(ops.end()); }

std::unique_ptr<Node> Clone() const override {
std::vector<std::unique_ptr<PlanOperator>> res;
res.reserve(ops.size());
for (const auto &op : ops) {
res.push_back(Node::MustAs<PlanOperator>(op->Clone()));
}
return std::make_unique<Merge>(std::move(res));
}
};

struct Limit : PlanOperator {
std::unique_ptr<PlanOperator> op;
size_t offset = 0, count = std::numeric_limits<size_t>::max();
std::unique_ptr<LimitClause> limit;

Limit(std::unique_ptr<PlanOperator> &&op, std::unique_ptr<LimitClause> &&limit)
: op(std::move(op)), limit(std::move(limit)) {}

std::string_view Name() const override { return "Limit"; };
std::string Dump() const override { return fmt::format("(limit {}, {}: {})", offset, count, op->Dump()); }
std::string Dump() const override {
return fmt::format("(limit {}, {}: {})", limit->offset, limit->count, op->Dump());
}

NodeIterator ChildBegin() override { return NodeIterator{op.get()}; }
NodeIterator ChildBegin() override { return NodeIterator{op.get(), limit.get()}; }
NodeIterator ChildEnd() override { return {}; }

std::unique_ptr<Node> Clone() const override {
return std::make_unique<Limit>(Node::MustAs<PlanOperator>(op->Clone()), Node::MustAs<LimitClause>(limit->Clone()));
}
};

struct Projection : PlanOperator {
std::unique_ptr<PlanOperator> source;
std::unique_ptr<SelectExpr> select;

Projection(std::unique_ptr<PlanOperator> &&source, std::unique_ptr<SelectExpr> &&select)
: source(std::move(source)), select(std::move(select)) {}

std::string_view Name() const override { return "Projection"; };
std::string Dump() const override { return fmt::format("(project {}: {})", select, source); }

NodeIterator ChildBegin() override { return {source.get(), select.get()}; }
NodeIterator ChildEnd() override { return {}; }

std::unique_ptr<Node> Clone() const override {
return std::make_unique<Projection>(Node::MustAs<PlanOperator>(source->Clone()),
Node::MustAs<SelectExpr>(select->Clone()));
}
};

} // namespace kqir

0 comments on commit fd5c376

Please sign in to comment.