perf: opt nested agg calls (#4022)

* perf(expr pass): cache aggregate function call in single projection

* test: add two test cases for aggregate calls

1. nested aggregate call
2. duplicated aggregate call

* feat(lambdafy): cache same expression

Initially we only intend to cache udaf calls under same window, this
change seem to cache all expressions under correct context.

This cache mechanism may not work since its not fully tested.

* feat(expr pass): letify expression to avoid repeated agg calls in nested
aggregate expression

A nested agg function call, e.g.

  agg_fn1(col1, agg_fn2(col2))

Will evaluate like a let expression:

  LET fn2 = agg_fn2(col2) IN agg_fn1(col1, fn2)

TODO:
- fix agg call cache

* feat(expr pass): equivalent agg call cache

TOOD:

- there are some cases still need fix

* fix: set output to id node for equivalent agg calls

Author: aceforeverd

cases/query/udaf_query.yaml

@@ -216,3 +216,54 @@ cases:
         - [3, "aa", 1, 8]
         - [4, "aa", 0, 21]
 
+  - id: 6
+    desc: |
+      nested aggregate function call
+    inputs:
+      - columns: [ "id int","ts timestamp","group1 string","val1 string" ]
+        indexs: [ "index1:group1:ts" ]
+        name: t1
+        data: |
+          1, 1612130400000, g1, 1
+          2, 1612130401000, g1, 2
+          3, 1612130402000, g1, 4
+    sql: |
+      select
+      `id`,
+      `val1`,
+      first_value(val1) over w1 as ft,
+      count_where(id, val1 = first_value(val1)) over w1 as agg1,
+      count_where(id, val1 != first_value(val1)) over w1 as agg2,
+      from `t1` WINDOW
+        w1 as (partition by `group1` order by `ts` rows_range between 5s preceding and 0s preceding)
+    expect:
+      columns: ["id int", "val1 string", "ft string", "agg1 int64", "agg2 int64"]
+      order: id
+      rows:
+         - [1, 1, 1, 1, 0]
+         - [2, 2, 1, 1, 1]
+         - [3, 4, 1, 1, 2]
+  - id: 7
+    desc: |
+      duplicated aggregate call
+    inputs:
+      - columns: [ "id int","ts timestamp","group1 string","val1 string" ]
+        indexs: [ "index1:group1:ts" ]
+        name: t1
+        data: |
+          1, 1612130400000, g1, 1
+          2, 1612130401000, g1, 2
+    sql: |
+      select
+      `id`,
+      `val1`,
+      first_value(val1) over w1 as agg1,
+      first_value(val1) over w1 as agg2,
+      from `t1` WINDOW
+        w1 as (partition by `group1` order by `ts` rows_range between 5s preceding and 0s preceding)
+    expect:
+      columns: ["id int", "val1 string", "agg1 string", "agg2 string"]
+      order: id
+      rows:
+         - [1, 1, "1", "1"]
+         - [2, 2, "1", "1"]

hybridse/include/node/node_enum.h

@@ -149,6 +149,7 @@ enum ExprType {
     kExprArray,
     kExprArrayElement,      // extract value from a array or map, with `[]` operator
     kExprStructCtorParens,  // (expr1, expr2, ...)
+    kExprLet,               // LET [bindings] in <output expr>
     kExprFake,              // not a real one
     kExprLast = kExprFake,
 };

hybridse/include/node/sql_node.h

@@ -17,17 +17,21 @@
 #ifndef HYBRIDSE_INCLUDE_NODE_SQL_NODE_H_
 #define HYBRIDSE_INCLUDE_NODE_SQL_NODE_H_
 
+#include <absl/status/status.h>
 #include <iostream>
 #include <map>
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "base/fe_status.h"
 #include "boost/algorithm/string.hpp"
 #include "boost/algorithm/string/predicate.hpp"
 #include "boost/filesystem/operations.hpp"
@@ -618,6 +622,7 @@ class ArrayElementExpr : public ExprNode {
     Status InferAttr(ExprAnalysisContext *ctx) override;
 };
 
+
 class FnNode : public SqlNode {
  public:
     FnNode() : SqlNode(kFn, 0, 0), indent(0) {}
@@ -1869,6 +1874,59 @@ class EscapedExpr : public ExprNode {
     }
 };
 
+class LetExpr : public ExprNode {
+ public:
+    class LetCtxEntry {
+     public:
+        LetCtxEntry(ExprIdNode *id_node, ExprNode *expr, const FrameNode *frame)
+            : id_node(id_node), expr(expr), frame(frame) {}
+
+        ExprIdNode *id_node;
+        ExprNode *expr;  // referred udaf call
+        const FrameNode *frame;
+    };
+
+    class LetContext {
+     public:
+        base::Status Append(ExprIdNode *k, ExprNode *v, const FrameNode *frame) {
+            auto it = cache.find(k);
+            if (it == cache.end()) {
+                bindings.emplace_back(k, v, frame);
+                cache.emplace(k, v);
+            } else {
+                CHECK_TRUE(v == it->second, common::kPlanError,
+                           "let context: try mapping id node to two different resolved nodes");
+            }
+
+            return base::Status::OK();
+        }
+
+        bool empty() const noexcept { return bindings.empty(); }
+
+        // necessary let bindings, not all cached entry will appear in bindings
+        std::vector<LetCtxEntry> bindings;
+        absl::flat_hash_map<ExprIdNode *, ExprNode *> cache;
+    };
+
+    LetExpr(ExprNode *expr, const LetContext &ctx) : ExprNode(kExprLet), ctx_(ctx), expr_(expr) {
+        AddChild(expr);
+    }
+    ~LetExpr() override {}
+
+
+    const LetContext& ctx() const {return ctx_; }
+    ExprNode *expr() const { return expr_; }
+
+    void Print(std::ostream &output, const std::string &org_tab) const override;
+    const std::string GetExprString() const override;
+    LetExpr *ShadowCopy(NodeManager *nm) const override;
+    Status InferAttr(ExprAnalysisContext *ctx) override;
+
+ private:
+    LetContext ctx_;
+    ExprNode* expr_;
+};
+
 class ResTarget : public SqlNode {
  public:
     ResTarget() : SqlNode(kResTarget, 0, 0), name_(""), val_(nullptr) {}

hybridse/include/vm/physical_op.h

@@ -330,6 +330,8 @@ class ColumnProjects : public FnComponent {
         primary_frame_ = frame;
     }
 
+    std::vector<const node::FrameNode *> frames() const { return frames_; }
+
     const node::FrameNode *GetPrimaryFrame() const { return primary_frame_; }
 
     base::Status ReplaceExpr(const passes::ExprReplacer &replacer,

hybridse/src/codegen/fn_let_ir_builder.cc

@@ -109,10 +109,31 @@ Status RowFnLetIRBuilder::Build(
     std::map<std::string, AggregateIRBuilder> window_agg_builder;
     uint32_t agg_builder_id = 0;
 
-    auto expr_list = compile_func->body();
+    auto fn_body = compile_func->body();
+    auto expr_list = fn_body;
+    auto maybe_let_expr = fn_body->GetAsOrNull<node::LetExpr>();
+    bool is_let_fn_body = maybe_let_expr != nullptr;
+    if (is_let_fn_body) {
+        expr_list = maybe_let_expr->expr();
+    }
     CHECK_TRUE(project_frames.size() == expr_list->GetChildNum(), kCodegenError,
                "Frame num should match expr num");
 
+    if (is_let_fn_body) {
+        for (auto& entry : maybe_let_expr->ctx().bindings) {
+            auto key = entry.id_node;
+            auto exp = entry.expr;
+            auto frame = entry.frame;
+            CHECK_STATUS(BindProjectFrame(&expr_ir_builder, frame, compile_func, ctx_->GetCurrentBlock(), sv));
+            NativeValue exp_out;
+            CHECK_STATUS(expr_ir_builder.Build(exp, &exp_out));
+
+            base::Status s;
+            variable_ir_builder.StoreValue(key->GetExprString(), exp_out, s);
+            CHECK_STATUS(s);
+        }
+    }
+
     for (size_t i = 0; i < expr_list->GetChildNum(); ++i) {
         const ::hybridse::node::ExprNode* expr = expr_list->GetChild(i);
         auto frame = project_frames[i];

hybridse/src/node/expr_node.cc

@@ -17,14 +17,15 @@
 #include "node/expr_node.h"
 
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
 #include "absl/strings/substitute.h"
 #include "codec/fe_row_codec.h"
+#include "codegen/ir_base_builder.h"
 #include "node/node_manager.h"
 #include "node/sql_node.h"
 #include "passes/expression/expr_pass.h"
 #include "passes/resolve_fn_and_attrs.h"
 #include "vm/schemas_context.h"
-#include "codegen/ir_base_builder.h"
 
 using ::hybridse::common::kTypeError;
 
@@ -1231,6 +1232,30 @@ Status ArrayElementExpr::InferAttr(ExprAnalysisContext* ctx) {
 ExprNode *ArrayElementExpr::array() const { return GetChild(0); }
 ExprNode *ArrayElementExpr::position() const { return GetChild(1); }
 
+
+void LetExpr::Print(std::ostream &output, const std::string &org_tab) const {}
+const std::string LetExpr::GetExprString() const {
+    return absl::Substitute("LET $0 IN $1",
+                            absl::StrJoin(ctx_.bindings, ",",
+                                          [](std::string* out, const decltype(LetContext::bindings)::value_type& kv) {
+                                              absl::StrAppend(out, kv.id_node->GetExprString(), "=",
+                                                              kv.expr->GetExprString());
+                                          }),
+                            expr_->GetExprString());
+}
+LetExpr* LetExpr::ShadowCopy(NodeManager* nm) const { return nm->MakeNode<LetExpr>(expr_, ctx_); }
+
+Status LetExpr::InferAttr(ExprAnalysisContext* ctx) {
+    for (auto& [k, v, _] : ctx_.bindings) {
+        CHECK_TRUE(k->GetOutputType() != nullptr, common::kTypeError, "expr id node not resolved");
+        CHECK_TRUE(k->GetOutputType() == v->GetOutputType(), common::kTypeError,
+                   "expr id node return type does not match binding expr");
+    }
+    SetOutputType(expr_->GetOutputType());
+    SetNullable(expr_->nullable());
+    return {};
+}
+
 StructCtorWithParens* StructCtorWithParens::ShadowCopy(NodeManager* nm) const {
     return nm->MakeNode<StructCtorWithParens>(fields());
 }

hybridse/src/node/sql_node.cc

@@ -117,38 +117,39 @@ static const absl::flat_hash_map<DataType, absl::string_view>& GetDataTypeNamesM
 }
 
 static absl::flat_hash_map<ExprType, absl::string_view> CreateExprTypeNamesMap() {
-  absl::flat_hash_map<ExprType, absl::string_view> map = {
-      {kExprPrimary, "primary"},
-      {kExprParameter, "parameter"},
-      {kExprId, "id"},
-      {kExprBinary, "binary"},
-      {kExprUnary, "unary"},
-      {kExprCall, "function"},
-      {kExprCase, "case"},
-      {kExprWhen, "when"},
-      {kExprBetween, "between"},
-      {kExprColumnRef, "column ref"},
-      {kExprColumnId, "column id"},
-      {kExprCast, "cast"},
-      {kExprAll, "all"},
-      {kExprStruct, "struct"},
-      {kExprQuery, "query"},
-      {kExprOrder, "order"},
-      {kExprGetField, "get field"},
-      {kExprCond, "cond"},
-      {kExprUnknow, "unknow"},
-      {kExprIn, "in"},
-      {kExprList, "expr_list"},
-      {kExprForIn, "for_in"},
-      {kExprRange, "range"},
-      {kExprOrderExpression, "order"},
-      {kExprEscaped, "escape"},
-      {kExprArray, "array"},
-      {kExprArrayElement, "array element"},
-      {kExprStructCtorParens, "struct with parens"},
-  };
-  for (auto kind = 0; kind < ExprType::kExprLast; ++kind) {
-      DCHECK(map.find(static_cast<ExprType>(kind)) != map.end());
+    absl::flat_hash_map<ExprType, absl::string_view> map = {
+        {kExprPrimary, "primary"},
+        {kExprParameter, "parameter"},
+        {kExprId, "id"},
+        {kExprBinary, "binary"},
+        {kExprUnary, "unary"},
+        {kExprCall, "function"},
+        {kExprCase, "case"},
+        {kExprWhen, "when"},
+        {kExprBetween, "between"},
+        {kExprColumnRef, "column ref"},
+        {kExprColumnId, "column id"},
+        {kExprCast, "cast"},
+        {kExprAll, "all"},
+        {kExprStruct, "struct"},
+        {kExprQuery, "query"},
+        {kExprOrder, "order"},
+        {kExprGetField, "get field"},
+        {kExprCond, "cond"},
+        {kExprUnknow, "unknow"},
+        {kExprIn, "in"},
+        {kExprList, "expr_list"},
+        {kExprForIn, "for_in"},
+        {kExprRange, "range"},
+        {kExprOrderExpression, "order"},
+        {kExprEscaped, "escape"},
+        {kExprArray, "array"},
+        {kExprArrayElement, "array element"},
+        {kExprStructCtorParens, "struct with parens"},
+        {kExprLet, "let"},
+    };
+    for (auto kind = 0; kind < ExprType::kExprLast; ++kind) {
+        DCHECK(map.find(static_cast<ExprType>(kind)) != map.end());
   }
   return map;
 }

hybridse/src/passes/expression/cache_expressions.cc

@@ -0,0 +1,70 @@
+/**
+ * Copyright (c) 2025 Ace <[email protected]>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "passes/expression/cache_expressions.h"
+
+namespace hybridse {
+namespace passes {
+
+static std::string CallExprKey(const node::CallExprNode* call);
+
+std::string CallExprKey(const node::CallExprNode* call) {
+    std::string str = call->GetFnDef()->GetName();
+    str.append("(");
+    for (size_t i = 0; i < call->children_.size(); ++i) {
+        str.append(call->children_[i]->GetExprString());
+        if (i < call->children_.size() - 1) {
+            str.append(", ");
+        }
+    }
+    str.append(")");
+
+    return str;
+}
+
+base::Status CacheExpressions::Apply(node::ExprAnalysisContext* ctx, node::ExprNode* expr, node::ExprNode** out) {
+    *out = expr;
+    for (int i = 0; i < expr->GetChildNum(); ++i) {
+        node::ExprNode* co = nullptr;
+        CHECK_STATUS(Apply(ctx, expr->GetChild(i), &co));
+        if (co != nullptr && co != expr->GetChild(i)) {
+            expr->SetChild(i, co);
+        }
+    }
+
+    if (expr ->GetExprType() != node::kExprCall) {
+        return {};
+    }
+
+    auto call = expr->GetAsOrNull<node::CallExprNode>();
+    if (call == nullptr) {
+        return {};
+    }
+
+    auto key = CallExprKey(call);
+
+    auto it = expr_cache_.find(key);
+    if (it != expr_cache_.end()) {
+        *out = it->second;
+    } else {
+        expr_cache_.emplace(key, call);
+    }
+
+    return {};
+}
+
+}  // namespace passes
+}  // namespace hybridse