diff options
author | daniilvasilev <daniilvasilev@yandex-team.com> | 2024-09-21 11:12:07 +0300 |
---|---|---|
committer | daniilvasilev <daniilvasilev@yandex-team.com> | 2024-09-21 11:22:06 +0300 |
commit | 7deed162d9d0d42ef21eea39547befcd9e48a784 (patch) | |
tree | 6f67ad063371e413cbf39715a4b9fba71606aefd /yt | |
parent | 9180a474dc5980e6f950c239117a7c897c13cac0 (diff) | |
download | ydb-7deed162d9d0d42ef21eea39547befcd9e48a784.tar.gz |
YT: Add with totals expression to query builder
commit_hash:6952a02dfac593408c83f2400cad5e81594dc0f8
Diffstat (limited to 'yt')
-rw-r--r-- | yt/yt/client/query_client/query_builder.cpp | 13 | ||||
-rw-r--r-- | yt/yt/client/query_client/query_builder.h | 9 | ||||
-rw-r--r-- | yt/yt/client/unittests/query_builder_ut.cpp | 20 |
3 files changed, 42 insertions, 0 deletions
diff --git a/yt/yt/client/query_client/query_builder.cpp b/yt/yt/client/query_client/query_builder.cpp index 0c2100e899..4cdce71596 100644 --- a/yt/yt/client/query_client/query_builder.cpp +++ b/yt/yt/client/query_client/query_builder.cpp @@ -67,6 +67,11 @@ void TQueryBuilder::AddGroupByExpression(TString expression, TString alias) }); } +void TQueryBuilder::SetWithTotals(EWithTotalsMode withTotalsMode) +{ + WithTotalsMode_ = withTotalsMode; +} + void TQueryBuilder::AddHavingConjunct(TString expression) { HavingConjuncts_.push_back(std::move(expression)); @@ -151,6 +156,10 @@ TString TQueryBuilder::Build() parts.push_back(JoinSeq(", ", GroupByEntries_)); } + if (WithTotalsMode_ == EWithTotalsMode::BeforeHaving) { + parts.push_back("WITH TOTALS"); + } + if (!HavingConjuncts_.empty()) { if (GroupByEntries_.empty()) { THROW_ERROR_EXCEPTION("Having without group by is not valid"); @@ -159,6 +168,10 @@ TString TQueryBuilder::Build() parts.push_back(JoinSeq(" AND ", Parenthesize(HavingConjuncts_))); } + if (WithTotalsMode_ == EWithTotalsMode::AfterHaving) { + parts.push_back("WITH TOTALS"); + } + if (!OrderByEntries_.empty()) { parts.push_back("ORDER BY"); parts.push_back(JoinSeq(", ", OrderByEntries_)); diff --git a/yt/yt/client/query_client/query_builder.h b/yt/yt/client/query_client/query_builder.h index ab509d970c..56ec58458a 100644 --- a/yt/yt/client/query_client/query_builder.h +++ b/yt/yt/client/query_client/query_builder.h @@ -18,6 +18,12 @@ DEFINE_ENUM(ETableJoinType, (Left) ); +DEFINE_ENUM(EWithTotalsMode, + (None) + (BeforeHaving) + (AfterHaving) +); + //////////////////////////////////////////////////////////////////////////////// class TQueryBuilder @@ -34,6 +40,8 @@ public: void AddGroupByExpression(TString expression); void AddGroupByExpression(TString expression, TString alias); + void SetWithTotals(EWithTotalsMode withTotalsMode); + void AddHavingConjunct(TString expression); void AddOrderByExpression(TString expression); @@ -76,6 +84,7 @@ private: std::vector<TString> WhereConjuncts_; std::vector<TOrderByEntry> OrderByEntries_; std::vector<TEntryWithAlias> GroupByEntries_; + EWithTotalsMode WithTotalsMode_ = EWithTotalsMode::None; std::vector<TString> HavingConjuncts_; std::vector<TJoinEntry> JoinEntries_; std::optional<i64> Limit_; diff --git a/yt/yt/client/unittests/query_builder_ut.cpp b/yt/yt/client/unittests/query_builder_ut.cpp index 2a6313c495..687f95d5c0 100644 --- a/yt/yt/client/unittests/query_builder_ut.cpp +++ b/yt/yt/client/unittests/query_builder_ut.cpp @@ -27,6 +27,8 @@ TEST(TQueryBuilderTest, Simple) b.AddGroupByExpression("x + y * z", "group_expr"); b.AddGroupByExpression("x - 1"); + b.SetWithTotals(EWithTotalsMode::BeforeHaving); + b.AddHavingConjunct("group_expr > 42"); b.AddHavingConjunct("group_expr < 420"); @@ -46,6 +48,7 @@ TEST(TQueryBuilderTest, Simple) "LEFT JOIN [table2] AS [lookup2] ON (idx) = (lookup2.idx) " "WHERE (x > y_alias) AND (y = 177 OR y % 2 = 0) " "GROUP BY (x + y * z) AS group_expr, (x - 1) " + "WITH TOTALS " "HAVING (group_expr > 42) AND (group_expr < 420) " "ORDER BY (z) ASC, (x) DESC, (x + y) DESC, (z - y_alias) " "LIMIT 43"); @@ -62,6 +65,23 @@ TEST(TQueryBuilderTest, SourceAlias) "FROM [//t] AS t_alias"); } +TEST(TQueryBuilderTest, TotalsAfterHaving) +{ + TQueryBuilder b; + b.AddSelectExpression("x"); + b.SetSource("//t"); + b.AddGroupByExpression("x - 1"); + b.SetWithTotals(EWithTotalsMode::AfterHaving); + b.AddHavingConjunct("group_expr > 42"); + + EXPECT_EQ(b.Build(), + "(x) " + "FROM [//t] " + "GROUP BY (x - 1) " + "HAVING (group_expr > 42) " + "WITH TOTALS"); +} + //////////////////////////////////////////////////////////////////////////////// } // namespace |