YT: Add with totals expression to query builder

commit_hash:6952a02dfac593408c83f2400cad5e81594dc0f8

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