Skip to content
This repository was archived by the owner on May 9, 2024. It is now read-only.

Commit e760e09

Browse files
ienkovichienkovich
committed
Add window functions support in QueryBuilder.
Signed-off-by: ienkovich <[email protected]>
1 parent 2f1e7e7 commit e760e09

File tree

10 files changed

+1326
-4
lines changed

10 files changed

+1326
-4
lines changed

‎omniscidb/IR/Expr.cpp

+14
Original file line numberDiff line numberDiff line change
@@ -1644,6 +1644,20 @@ std::string WindowFunction::toString() const {
16441644
for (const auto& arg : args_) {
16451645
result += " " + arg->toString();
16461646
}
1647+
if (!partition_keys_.empty()) {
1648+
result += " PARTITION BY";
1649+
for (const auto& part_key : partition_keys_) {
1650+
result += " " + part_key->toString();
1651+
}
1652+
}
1653+
if (!order_keys_.empty()) {
1654+
result += " ORDER BY";
1655+
for (size_t i = 0; i < order_keys_.size(); ++i) {
1656+
result += " " + order_keys_[i]->toString();
1657+
result += collation_[i].is_desc ? " DESC" : " ASC";
1658+
result += collation_[i].nulls_first ? " NULLS FIRST" : " NULLS LAST";
1659+
}
1660+
}
16471661
return result + ") ";
16481662
}
16491663

‎omniscidb/IR/Expr.h

+2
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ class Expr : public std::enable_shared_from_this<Expr> {
6464
virtual std::string toString() const = 0;
6565
virtual void print() const;
6666

67+
bool equal(const Expr* rhs) const { return *this == *rhs; }
68+
6769
/*
6870
* @brief decompress adds cast operator to decompress encoded result
6971
*/

‎omniscidb/QueryBuilder/QueryBuilder.cpp

+276
Original file line numberDiff line numberDiff line change
@@ -509,8 +509,72 @@ std::vector<NodePtr> collectNodes(NodePtr node) {
509509
return nodes;
510510
}
511511

512+
/**
513+
* Check if expr is WindowFunction or AggExpr that can be transformed
514+
* into corresponding WindowFunction.
515+
*/
516+
std::shared_ptr<const WindowFunction> checkOrGetWindowFn(ExprPtr expr) {
517+
if (expr->is<WindowFunction>()) {
518+
return std::dynamic_pointer_cast<const WindowFunction>(expr);
519+
}
520+
521+
if (auto agg = expr->as<AggExpr>()) {
522+
switch (agg->aggType()) {
523+
case AggType::kCount:
524+
if (agg->arg()) {
525+
return std::shared_ptr<const WindowFunction>(new WindowFunction(
526+
agg->type(), WindowFunctionKind::Count, {agg->argShared()}, {}, {}, {}));
527+
} else {
528+
return std::shared_ptr<const WindowFunction>(
529+
new WindowFunction(agg->type(), WindowFunctionKind::Count, {}, {}, {}, {}));
530+
}
531+
case AggType::kAvg:
532+
return std::shared_ptr<const WindowFunction>(new WindowFunction(
533+
agg->type(), WindowFunctionKind::Avg, {agg->argShared()}, {}, {}, {}));
534+
case AggType::kMin:
535+
return std::shared_ptr<const WindowFunction>(new WindowFunction(
536+
agg->type(), WindowFunctionKind::Min, {agg->argShared()}, {}, {}, {}));
537+
case AggType::kMax:
538+
return std::shared_ptr<const WindowFunction>(new WindowFunction(
539+
agg->type(), WindowFunctionKind::Max, {agg->argShared()}, {}, {}, {}));
540+
case AggType::kSum:
541+
return std::shared_ptr<const WindowFunction>(new WindowFunction(
542+
agg->type(), WindowFunctionKind::Sum, {agg->argShared()}, {}, {}, {}));
543+
default:
544+
break;
545+
}
546+
}
547+
548+
return nullptr;
549+
}
550+
512551
} // namespace
513552

553+
BuilderOrderByKey::BuilderOrderByKey()
554+
: expr_(nullptr)
555+
, dir_(SortDirection::Ascending)
556+
, null_pos_(NullSortedPosition::Last) {}
557+
558+
BuilderOrderByKey::BuilderOrderByKey(const BuilderExpr& expr,
559+
SortDirection dir,
560+
NullSortedPosition null_pos)
561+
: expr_(expr.expr()), dir_(dir), null_pos_(null_pos) {}
562+
563+
BuilderOrderByKey::BuilderOrderByKey(const BuilderExpr& expr,
564+
const std::string& dir,
565+
const std::string& null_pos)
566+
: expr_(expr.expr())
567+
, dir_(parseSortDirection(dir))
568+
, null_pos_(parseNullPosition(null_pos)) {}
569+
570+
SortDirection BuilderOrderByKey::parseSortDirection(const std::string& val) {
571+
return BuilderSortField::parseSortDirection(val);
572+
}
573+
574+
NullSortedPosition BuilderOrderByKey::parseNullPosition(const std::string& val) {
575+
return BuilderSortField::parseNullPosition(val);
576+
}
577+
514578
BuilderExpr::BuilderExpr() : builder_(nullptr) {}
515579

516580
BuilderExpr::BuilderExpr(const QueryBuilder* builder,
@@ -666,6 +730,42 @@ BuilderExpr BuilderExpr::corr(const BuilderExpr& arg) const {
666730
return {builder_, agg, name, true};
667731
}
668732

733+
BuilderExpr BuilderExpr::lag(int n) const {
734+
ExprPtr expr{new WindowFunction(expr_->type(),
735+
WindowFunctionKind::Lag,
736+
{expr_, builder_->cst(n).expr()},
737+
{},
738+
{},
739+
{})};
740+
auto name = name_.empty() ? "lag" : name_ + "_lag";
741+
return {builder_, expr, name, true};
742+
}
743+
744+
BuilderExpr BuilderExpr::lead(int n) const {
745+
ExprPtr expr{new WindowFunction(expr_->type(),
746+
WindowFunctionKind::Lead,
747+
{expr_, builder_->cst(n).expr()},
748+
{},
749+
{},
750+
{})};
751+
auto name = name_.empty() ? "lead" : name_ + "_lead";
752+
return {builder_, expr, name, true};
753+
}
754+
755+
BuilderExpr BuilderExpr::firstValue() const {
756+
ExprPtr expr{new WindowFunction(
757+
expr_->type(), WindowFunctionKind::FirstValue, {expr_}, {}, {}, {})};
758+
auto name = name_.empty() ? "first_value" : name_ + "_first_value";
759+
return {builder_, expr, name, true};
760+
}
761+
762+
BuilderExpr BuilderExpr::lastValue() const {
763+
ExprPtr expr{new WindowFunction(
764+
expr_->type(), WindowFunctionKind::LastValue, {expr_}, {}, {}, {})};
765+
auto name = name_.empty() ? "last_value" : name_ + "_last_value";
766+
return {builder_, expr, name, true};
767+
}
768+
669769
BuilderExpr BuilderExpr::agg(const std::string& agg_str, const BuilderExpr& arg) const {
670770
static const std::unordered_map<std::string, AggType> agg_names = {
671771
{"count", AggType::kCount},
@@ -1517,6 +1617,144 @@ BuilderExpr BuilderExpr::at(int64_t idx) const {
15171617
return at(builder_->cst(idx, builder_->ctx_.int64(false)));
15181618
}
15191619

1620+
BuilderExpr BuilderExpr::over() const {
1621+
return over(std::vector<BuilderExpr>());
1622+
}
1623+
1624+
BuilderExpr BuilderExpr::over(const BuilderExpr& key) const {
1625+
return over(std::vector<BuilderExpr>({key}));
1626+
}
1627+
1628+
BuilderExpr BuilderExpr::over(const std::vector<BuilderExpr>& keys) const {
1629+
auto wnd_fn = checkOrGetWindowFn(expr_);
1630+
if (!wnd_fn) {
1631+
throw InvalidQueryError()
1632+
<< "Expected window function or supported aggregate (COUNT, AVG, MIN, MAX, SUM) "
1633+
"for OVER. Provided: "
1634+
<< expr_->toString();
1635+
}
1636+
1637+
if (!keys.empty()) {
1638+
for (auto& key : keys) {
1639+
if (!key.expr()->is<ColumnRef>()) {
1640+
throw InvalidQueryError() << "Currently, only column references can be used as a "
1641+
"partition key. Provided: "
1642+
<< key.expr()->toString();
1643+
}
1644+
}
1645+
1646+
ExprPtrVector new_part_keys;
1647+
new_part_keys.reserve(wnd_fn->partitionKeys().size() + keys.size());
1648+
new_part_keys.insert(new_part_keys.end(),
1649+
wnd_fn->partitionKeys().begin(),
1650+
wnd_fn->partitionKeys().end());
1651+
for (auto& expr : keys) {
1652+
new_part_keys.push_back(expr.expr());
1653+
}
1654+
wnd_fn = makeExpr<WindowFunction>(wnd_fn->type(),
1655+
wnd_fn->kind(),
1656+
wnd_fn->args(),
1657+
new_part_keys,
1658+
wnd_fn->orderKeys(),
1659+
wnd_fn->collation());
1660+
}
1661+
1662+
return {builder_, wnd_fn, name_, auto_name_};
1663+
}
1664+
1665+
BuilderExpr BuilderExpr::orderBy(BuilderExpr key,
1666+
SortDirection dir,
1667+
NullSortedPosition null_pos) const {
1668+
return orderBy(std::vector<BuilderExpr>({key}), dir, null_pos);
1669+
}
1670+
1671+
BuilderExpr BuilderExpr::orderBy(BuilderExpr key,
1672+
const std::string& dir,
1673+
const std::string& null_pos) const {
1674+
return orderBy(std::vector<BuilderExpr>({key}), dir, null_pos);
1675+
}
1676+
1677+
BuilderExpr BuilderExpr::orderBy(std::initializer_list<BuilderExpr> keys,
1678+
SortDirection dir,
1679+
NullSortedPosition null_pos) const {
1680+
return orderBy(std::vector<BuilderExpr>(keys), dir, null_pos);
1681+
}
1682+
1683+
BuilderExpr BuilderExpr::orderBy(std::initializer_list<BuilderExpr> keys,
1684+
const std::string& dir,
1685+
const std::string& null_pos) const {
1686+
return orderBy(std::vector<BuilderExpr>(keys), dir, null_pos);
1687+
}
1688+
1689+
BuilderExpr BuilderExpr::orderBy(const std::vector<BuilderExpr>& keys,
1690+
SortDirection dir,
1691+
NullSortedPosition null_pos) const {
1692+
std::vector<BuilderOrderByKey> order_keys;
1693+
order_keys.reserve(keys.size());
1694+
for (auto& key : keys) {
1695+
order_keys.emplace_back(key, dir, null_pos);
1696+
}
1697+
return orderBy(order_keys);
1698+
}
1699+
1700+
BuilderExpr BuilderExpr::orderBy(const std::vector<BuilderExpr>& keys,
1701+
const std::string& dir,
1702+
const std::string& null_pos) const {
1703+
std::vector<BuilderOrderByKey> order_keys;
1704+
order_keys.reserve(keys.size());
1705+
for (auto& key : keys) {
1706+
order_keys.emplace_back(key, dir, null_pos);
1707+
}
1708+
return orderBy(order_keys);
1709+
}
1710+
1711+
BuilderExpr BuilderExpr::orderBy(const BuilderOrderByKey& key) const {
1712+
return orderBy(std::vector<BuilderOrderByKey>({key}));
1713+
}
1714+
1715+
BuilderExpr BuilderExpr::orderBy(const std::vector<BuilderOrderByKey>& keys) const {
1716+
auto wnd_fn = expr_->as<WindowFunction>();
1717+
if (!wnd_fn) {
1718+
throw InvalidQueryError() << "Expected window function for ORDER BY. Provided: "
1719+
<< expr_->toString();
1720+
}
1721+
1722+
for (auto& key : keys) {
1723+
if (!key.expr()->is<ColumnRef>()) {
1724+
throw InvalidQueryError()
1725+
<< "Currently, only column references can be used in ORDER BY. Provided: "
1726+
<< key.expr()->toString();
1727+
}
1728+
}
1729+
1730+
ExprPtrVector new_order_keys = wnd_fn->orderKeys();
1731+
new_order_keys.reserve(wnd_fn->orderKeys().size() + keys.size());
1732+
new_order_keys.insert(
1733+
new_order_keys.end(), wnd_fn->orderKeys().begin(), wnd_fn->orderKeys().end());
1734+
for (auto& key : keys) {
1735+
new_order_keys.push_back(key.expr());
1736+
}
1737+
1738+
std::vector<OrderEntry> new_collation;
1739+
new_collation.reserve(wnd_fn->collation().size() + keys.size());
1740+
new_collation.insert(
1741+
new_collation.end(), wnd_fn->collation().begin(), wnd_fn->collation().end());
1742+
for (auto& key : keys) {
1743+
new_collation.emplace_back(static_cast<int>(new_collation.size()),
1744+
key.dir() == SortDirection::Descending,
1745+
key.nullsPosition() == NullSortedPosition::First);
1746+
}
1747+
1748+
auto res = makeExpr<WindowFunction>(wnd_fn->type(),
1749+
wnd_fn->kind(),
1750+
wnd_fn->args(),
1751+
wnd_fn->partitionKeys(),
1752+
new_order_keys,
1753+
new_collation);
1754+
1755+
return {builder_, res, name_, auto_name_};
1756+
}
1757+
15201758
BuilderExpr BuilderExpr::rewrite(ExprRewriter& rewriter) const {
15211759
return {builder_, rewriter.visit(expr_.get()), name_, auto_name_};
15221760
}
@@ -2628,6 +2866,44 @@ BuilderExpr QueryBuilder::count() const {
26282866
return {this, agg, "count", true};
26292867
}
26302868

2869+
BuilderExpr QueryBuilder::rowNumber() const {
2870+
ExprPtr expr{new WindowFunction(
2871+
ctx_.int64(false), WindowFunctionKind::RowNumber, {}, {}, {}, {})};
2872+
return {this, expr, "row_number", true};
2873+
}
2874+
2875+
BuilderExpr QueryBuilder::rank() const {
2876+
ExprPtr expr{
2877+
new WindowFunction(ctx_.int64(false), WindowFunctionKind::Rank, {}, {}, {}, {})};
2878+
return {this, expr, "rank", true};
2879+
}
2880+
2881+
BuilderExpr QueryBuilder::denseRank() const {
2882+
ExprPtr expr{new WindowFunction(
2883+
ctx_.int64(false), WindowFunctionKind::DenseRank, {}, {}, {}, {})};
2884+
return {this, expr, "dense_rank", true};
2885+
}
2886+
2887+
BuilderExpr QueryBuilder::percentRank() const {
2888+
ExprPtr expr{new WindowFunction(
2889+
ctx_.fp64(false), WindowFunctionKind::PercentRank, {}, {}, {}, {})};
2890+
return {this, expr, "percent_rank", true};
2891+
}
2892+
2893+
BuilderExpr QueryBuilder::nTile(int tile_count) const {
2894+
if (tile_count <= 0) {
2895+
throw InvalidQueryError()
2896+
<< "Expected positive integer for tile count argument. Provided: " << tile_count;
2897+
}
2898+
ExprPtr expr{new WindowFunction(ctx_.int64(false),
2899+
WindowFunctionKind::NTile,
2900+
{cst(tile_count).expr()},
2901+
{},
2902+
{},
2903+
{})};
2904+
return {this, expr, "ntile", true};
2905+
}
2906+
26312907
BuilderExpr QueryBuilder::cst(int val) const {
26322908
return cst(static_cast<int64_t>(val));
26332909
}

0 commit comments

Comments
 (0)