fix(sqlengine): balance deep OR expression trees

src/db/sqlengine/parser/zvec_sql_parser.cc

@@ -14,7 +14,9 @@
 
 #include "zvec_sql_parser.h"
 #include <exception>
+#include <functional>
 #include <memory>
+#include <vector>
 #include <zvec/ailego/logger/logger.h>
 #include "atn/ParserATNSimulator.h"
 #include "db/sqlengine/antlr/gen/SQLLexer.h"
@@ -211,29 +213,152 @@ SelectInfo::Ptr ZVecSQLParser::select_info(VoidPtr node) {
 }
 
 Node::Ptr ZVecSQLParser::handle_logic_expr_node(VoidPtr node) {
-  SQLParser::Logic_exprContext *logicExprNode =
-      reinterpret_cast<SQLParser::Logic_exprContext *>(node);
-  const std::vector<SQLParser::Logic_exprContext *> &logicExprChildNodes =
-      logicExprNode->logic_expr();
-
-  if (logicExprNode->OR() != nullptr) {
-    Node::Ptr orExpr = std::make_shared<Node>(NodeOp::T_OR);
-    orExpr->set_left(handle_logic_expr_node(logicExprChildNodes[0]));
-    orExpr->set_right(handle_logic_expr_node(logicExprChildNodes[1]));
-    return orExpr;
-  } else if (logicExprNode->AND() != nullptr) {
-    Node::Ptr andExpr = std::make_shared<Node>(NodeOp::T_AND);
-    andExpr->set_left(handle_logic_expr_node(logicExprChildNodes[0]));
-    andExpr->set_right(handle_logic_expr_node(logicExprChildNodes[1]));
-    return andExpr;
-  } else if (logicExprNode->enclosed_expr() != nullptr) {
-    // enclosed_expr is represented by sub-tree structure
-    return handle_logic_expr_node(logicExprNode->enclosed_expr()->logic_expr());
-  } else if (logicExprNode->relation_expr() != nullptr) {
-    return handle_rel_expr_node(logicExprNode->relation_expr());
+  // ANTLR represents a flat OR chain as a deeply skewed binary parse tree.
+  // Consume that tree iteratively, flatten each maximal OR run, and rebuild it
+  // as a balanced binary Node tree. This preserves left-to-right evaluation
+  // order while bounding downstream recursive traversals to logarithmic depth.
+  enum class BindType { ROOT, LEFT, RIGHT };
+  struct BindTarget {
+    Node *parent;
+    BindType bind_type;
+  };
+  struct Frame {
+    SQLParser::Logic_exprContext *node;
+    Node *parent;
+    BindType bind_type;
+  };
+  // Preserve the historical shape of ordinary small expressions. Large runs
+  // are normalized so downstream recursive traversals remain stack-safe.
+  constexpr size_t kOrBalanceThreshold = 64;
+
+  Node::Ptr root_result;
+  auto attach_node = [&root_result](Node *parent, Node::Ptr child,
+                                    BindType bind_type) {
+    switch (bind_type) {
+      case BindType::ROOT:
+        root_result = std::move(child);
+        break;
+      case BindType::LEFT:
+        parent->set_left(std::move(child));
+        break;
+      case BindType::RIGHT:
+        parent->set_right(std::move(child));
+        break;
+    }
+  };
+
+  std::vector<Frame> stack;
+  stack.push_back({reinterpret_cast<SQLParser::Logic_exprContext *>(node),
+                   nullptr, BindType::ROOT});
+
+  auto unwrap_enclosed = [](SQLParser::Logic_exprContext *context) {
+    while (context != nullptr && context->enclosed_expr() != nullptr) {
+      context = context->enclosed_expr()->logic_expr();
+    }
+    return context;
+  };
+
+  while (!stack.empty()) {
+    Frame frame = stack.back();
+    stack.pop_back();
+
+    SQLParser::Logic_exprContext *logicExprNode = unwrap_enclosed(frame.node);
+    if (logicExprNode == nullptr) {
+      attach_node(frame.parent, nullptr, frame.bind_type);
+      continue;
+    }
+
+    if (logicExprNode->OR() != nullptr) {
+      // An AND subtree remains one operand, preserving precedence and grouping.
+      // Parentheses around OR are safe to flatten because OR is associative.
+      std::vector<SQLParser::Logic_exprContext *> operands;
+      std::vector<SQLParser::Logic_exprContext *> pending{logicExprNode};
+      while (!pending.empty()) {
+        SQLParser::Logic_exprContext *current = unwrap_enclosed(pending.back());
+        pending.pop_back();
+
+        if (current == nullptr || current->OR() == nullptr) {
+          operands.push_back(current);
+          continue;
+        }
+
+        const auto &children = current->logic_expr();
+        if (children.size() != 2U) {
+          err_msg_ = "Parse failed. Invalid OR expression.";
+          operands.clear();
+          break;
+        }
+        // Push right first so operands retain their original left-to-right
+        // order when consumed from the LIFO stack.
+        pending.push_back(children[1]);
+        pending.push_back(children[0]);
+      }
+
+      if (operands.empty()) {
+        attach_node(frame.parent, nullptr, frame.bind_type);
+        continue;
+      }
+
+      if (operands.size() <= kOrBalanceThreshold) {
+        const auto &children = logicExprNode->logic_expr();
+        Node::Ptr expr = std::make_shared<Node>(NodeOp::T_OR);
+        Node *expr_raw = expr.get();
+        attach_node(frame.parent, std::move(expr), frame.bind_type);
+        stack.push_back({children[1], expr_raw, BindType::RIGHT});
+        stack.push_back({children[0], expr_raw, BindType::LEFT});
+        continue;
+      }
+
+      std::vector<BindTarget> operand_targets(
+          operands.size(), BindTarget{nullptr, BindType::ROOT});
+      std::function<void(size_t, size_t, Node *, BindType)>
+          build_balanced_skeleton;
+      build_balanced_skeleton = [&](size_t begin, size_t end, Node *parent,
+                                    BindType bind_type) {
+        if (end - begin == 1U) {
+          operand_targets[begin] = {parent, bind_type};
+          return;
+        }
+
+        const size_t middle = begin + (end - begin) / 2U;
+        Node::Ptr expr = std::make_shared<Node>(NodeOp::T_OR);
+        Node *expr_raw = expr.get();
+        attach_node(parent, std::move(expr), bind_type);
+        build_balanced_skeleton(begin, middle, expr_raw, BindType::LEFT);
+        build_balanced_skeleton(middle, end, expr_raw, BindType::RIGHT);
+      };
+      build_balanced_skeleton(0, operands.size(), frame.parent,
+                              frame.bind_type);
+
+      for (size_t i = operands.size(); i-- > 0;) {
+        stack.push_back({operands[i], operand_targets[i].parent,
+                         operand_targets[i].bind_type});
+      }
+    } else if (logicExprNode->AND() != nullptr) {
+      const auto &children = logicExprNode->logic_expr();
+      if (children.size() != 2U) {
+        err_msg_ = "Parse failed. Invalid AND expression.";
+        attach_node(frame.parent, nullptr, frame.bind_type);
+        continue;
+      }
+
+      Node::Ptr expr = std::make_shared<Node>(NodeOp::T_AND);
+      Node *expr_raw = expr.get();
+      attach_node(frame.parent, std::move(expr), frame.bind_type);
+      // Preserve the parser's exact AND shape because analyzer/optimizer
+      // subroot selection is currently shape-dependent.
+      stack.push_back({children[1], expr_raw, BindType::RIGHT});
+      stack.push_back({children[0], expr_raw, BindType::LEFT});
+    } else if (logicExprNode->relation_expr() != nullptr) {
+      attach_node(frame.parent,
+                  handle_rel_expr_node(logicExprNode->relation_expr()),
+                  frame.bind_type);
+    } else {
+      attach_node(frame.parent, nullptr, frame.bind_type);
+    }
   }
 
-  return nullptr;
+  return root_result;
 }
 
 Node::Ptr ZVecSQLParser::handle_rel_expr_left_node(VoidPtr node) {

tests/db/sqlengine/query_info_test.cc

@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <algorithm>
 #include <memory>
+#include <vector>
 #include <gmock/gmock-matchers.h>
 #include <gtest/gtest.h>
 #include "db/sqlengine/sqlengine_impl.h"
@@ -441,6 +443,97 @@ TEST_F(QueryInfoTest, QueryRequestWithInFilterNum1024) {
 }
 
 
+namespace {
+
+// QueryAnalyzer accepts at most 4096 filter relations. Q_GT is deliberately
+// used because equality terms may be folded into a Q_IN expression.
+constexpr int kDeepOrCount = 4096;
+constexpr size_t kMaxBalancedLogicDepth = 12;
+
+struct DeepOrTreeStats {
+  bool valid = true;
+  size_t logic_count = 0;
+  size_t relation_count = 0;
+  size_t max_logic_depth = 0;
+};
+
+DeepOrTreeStats CheckDeepOrTree(const QueryNode::Ptr &root) {
+  DeepOrTreeStats stats;
+  if (!root) {
+    stats.valid = false;
+    return stats;
+  }
+
+  struct PendingNode {
+    QueryNode::Ptr node;
+    size_t logic_depth;
+  };
+
+  std::vector<PendingNode> stack{{root, 0}};
+  while (!stack.empty()) {
+    PendingNode pending = std::move(stack.back());
+    stack.pop_back();
+
+    if (pending.node->type() == QueryNode::QueryNodeType::LOGIC_EXPR) {
+      if (pending.node->op() != QueryNodeOp::Q_OR || !pending.node->left() ||
+          !pending.node->right()) {
+        stats.valid = false;
+        return stats;
+      }
+      const size_t logic_depth = pending.logic_depth + 1;
+      stats.max_logic_depth = std::max(stats.max_logic_depth, logic_depth);
+      ++stats.logic_count;
+      stack.push_back({pending.node->right(), logic_depth});
+      stack.push_back({pending.node->left(), logic_depth});
+      continue;
+    }
+
+    if (pending.node->type() != QueryNode::QueryNodeType::REL_EXPR ||
+        pending.node->op() != QueryNodeOp::Q_GT) {
+      stats.valid = false;
+      return stats;
+    }
+    ++stats.relation_count;
+  }
+
+  return stats;
+}
+
+}  // namespace
+
+TEST_F(QueryInfoTest, QueryRequestWithNonFoldableBalancedOr4096) {
+  SearchQuery query;
+  query.output_fields_ = {"*"};
+  query.topk_ = 10;
+  query.target_.set_vector("[0.1, 0.2, 0.3, 0.4]");
+  query.target_.field_name_ = "face_feature";
+  query.target_.query_params_ = std::make_shared<QueryParams>(IndexType::FLAT);
+
+  std::string filter;
+  filter.reserve(kDeepOrCount * 18);
+  for (int i = 0; i < kDeepOrCount; ++i) {
+    if (i != 0) {
+      filter += " or ";
+    }
+    filter += "name>" + std::to_string(i);
+  }
+  query.filter_ = std::move(filter);
+
+  auto engine = std::make_shared<SQLEngineImpl>(std::make_shared<Profiler>());
+  auto ret = engine->build_query_info(schema, query, nullptr);
+  ASSERT_TRUE(ret.has_value()) << ret.error().c_str();
+  ASSERT_TRUE(ret.value()->filter_cond());
+
+  auto root =
+      std::dynamic_pointer_cast<QueryNode>(ret.value()->filter_cond()->right());
+  auto stats = CheckDeepOrTree(root);
+  EXPECT_TRUE(stats.valid);
+  EXPECT_EQ(kDeepOrCount - 1, stats.logic_count);
+  EXPECT_EQ(kDeepOrCount, stats.relation_count);
+  EXPECT_LE(stats.max_logic_depth, kMaxBalancedLogicDepth);
+}
+
+
 TEST_F(QueryInfoTest, QueryRequestWithFilter_contain) {
   SearchQuery query;
   query.output_fields_ = {"*"};