refactor(interactive): Support ScanEarlyStopRule for Query Optimization (#4431)

<!--
Thanks for your contribution! please review
https://github.com/alibaba/GraphScope/blob/main/CONTRIBUTING.md before
opening an issue.
-->

## What do these changes do?
Leverage 3 rules for `Scan Early Stop` Optimization:
1. ScanEarlyStopRule
<img width="939" alt="image"
src="https://github.com/user-attachments/assets/f85098b2-c9e9-4c93-ac51-5dec59484bee"
/>


2. ScanExpandFusionRule
<img width="1133" alt="image"
src="https://github.com/user-attachments/assets/0366e6df-d113-4379-afb4-ce0359534192"
/>

3. TopKPushDownRule
<img width="1212" alt="image"
src="https://github.com/user-attachments/assets/5de53d9d-e8ca-4868-9857-d3ed7dc40b3d"
/>

<!-- Please give a short brief about these changes. -->

## Related issue number

<!-- Are there any issues opened that will be resolved by merging this
change? -->

Fixes #4356

---------

Co-authored-by: xiaolei.zl <[email protected]>
Co-authored-by: liulx20 <[email protected]>

Author: 3 people

docs/interactive_engine/gopt.md

@@ -349,4 +349,21 @@ Design of GOpt
 :::
 
 ### Detailed Introduction
-A comprehensive introduction to GOpt will be provided in subsequent sections. Please stay tuned for detailed updates and information.
+
+#### Rules
+
+Rules play a pivotal role in GOpt’s optimization framework, enabling efficient and effective query transformations. Below is an outline of some key rules implemented in GOpt:
+
+**ScanEarlyStopRule**: Pushes the limit operation down to the scan node. During the scan process, the scan stops as soon as the specified limit count is reached.
+   
+**ScanExpandFusionRule**: This rule transforms edge expansion into edge scan wherever possible. For example, consider the following Cypher query:
+```cypher
+Match (a:PERSON)-[b:KNOWS]->(c:PERSON) Return b.name;
+```
+Although the query involves Scan and GetV steps, their results are not directly utilized by subsequent project operations. The only effectively used data is the edge data produced by the Expand operation. In such cases, we can perform a fusion operation, transforming the pattern
+`(a:PERSON)-[b:KNOWS]->(c:PERSON)` into a scan operation on the KNOWS edge. It is important to note that whether fusion is feasible also depends on the label dependencies between nodes and edges. If the edge label is determined strictly by the triplet (src_label, edge_label, dst_label), fusion cannot be performed. For example, consider the following query:
+```cypher
+Match (a:PERSON)-[b:LIKES]->(c:COMMENT) Return b.name;
+```
+
+**TopKPushDownRule**: This rule pushes down topK operations to the project node and is based on Calcite's [SortProjectTransposeRule](https://calcite.apache.org/javadocAggregate/org/apache/calcite/rel/rules/SortProjectTransposeRule.html), leveraging the original rule wherever possible. However, in our more complex distributed scenario, deferring the execution of the project node can disrupt already sorted data. To address this, we modified the matching logic in `SortProjectTransposeRule`. Currently, the PushDown operation is applied only when the sort fields are empty, which means only the limit is pushed down to the project node.

flex/engines/graph_db/runtime/common/operators/retrieve/scan.cc

@@ -19,7 +19,8 @@ namespace gs {
 namespace runtime {
 
 Context Scan::find_vertex_with_oid(const GraphReadInterface& graph,
-                                   label_t label, const Any& oid, int alias) {
+                                   label_t label, const Any& oid,
+                                   int32_t alias) {
   SLVertexColumnBuilder builder(label);
   vid_t vid;
   if (graph.GetVertexIndex(label, oid, vid)) {
@@ -31,7 +32,7 @@ Context Scan::find_vertex_with_oid(const GraphReadInterface& graph,
 }
 
 Context Scan::find_vertex_with_gid(const GraphReadInterface& graph,
-                                   label_t label, int64_t gid, int alias) {
+                                   label_t label, int64_t gid, int32_t alias) {
   SLVertexColumnBuilder builder(label);
   if (GlobalId::get_label_id(gid) == label) {
     builder.push_back_opt(GlobalId::get_vid(gid));
@@ -44,42 +45,6 @@ Context Scan::find_vertex_with_gid(const GraphReadInterface& graph,
   return ctx;
 }
 
-Context Scan::find_vertex_with_id(const GraphReadInterface& graph,
-                                  label_t label, const Any& pk, int alias,
-                                  bool scan_oid) {
-  if (scan_oid) {
-    SLVertexColumnBuilder builder(label);
-    vid_t vid;
-    if (graph.GetVertexIndex(label, pk, vid)) {
-      builder.push_back_opt(vid);
-    }
-    Context ctx;
-    ctx.set(alias, builder.finish());
-    return ctx;
-  } else {
-    SLVertexColumnBuilder builder(label);
-    vid_t vid{};
-    int64_t gid{};
-    if (pk.type == PropertyType::kInt64) {
-      gid = pk.AsInt64();
-    } else if (pk.type == PropertyType::kInt32) {
-      gid = pk.AsInt32();
-    } else {
-      LOG(FATAL) << "Unsupported primary key type";
-    }
-    if (GlobalId::get_label_id(gid) == label) {
-      vid = GlobalId::get_vid(gid);
-    } else {
-      LOG(ERROR) << "Global id " << gid << " does not match label " << label;
-      return Context();
-    }
-    builder.push_back_opt(vid);
-    Context ctx;
-    ctx.set(alias, builder.finish());
-    return ctx;
-  }
-}
-
 template <typename T>
 static Context _scan_vertex_with_special_vertex_predicate(
     const GraphReadInterface& graph, const ScanParams& params,

flex/engines/graph_db/runtime/common/operators/retrieve/scan.h

@@ -26,6 +26,9 @@ namespace runtime {
 struct ScanParams {
   int alias;
   std::vector<label_t> tables;
+  int32_t limit;
+
+  ScanParams() : alias(-1), limit(std::numeric_limits<int32_t>::max()) {}
 };
 class Scan {
  public:
@@ -61,6 +64,50 @@ class Scan {
     return ctx;
   }
 
+  template <typename PRED_T>
+  static Context scan_vertex_with_limit(const GraphReadInterface& graph,
+                                        const ScanParams& params,
+                                        const PRED_T& predicate) {
+    Context ctx;
+    int32_t cur_limit = params.limit;
+    if (params.tables.size() == 1) {
+      label_t label = params.tables[0];
+      SLVertexColumnBuilder builder(label);
+      auto vertices = graph.GetVertexSet(label);
+      for (auto vid : vertices) {
+        if (cur_limit <= 0) {
+          break;
+        }
+        if (predicate(label, vid)) {
+          builder.push_back_opt(vid);
+          cur_limit--;
+        }
+      }
+      ctx.set(params.alias, builder.finish());
+    } else if (params.tables.size() > 1) {
+      MSVertexColumnBuilder builder;
+
+      for (auto label : params.tables) {
+        if (cur_limit <= 0) {
+          break;
+        }
+        auto vertices = graph.GetVertexSet(label);
+        builder.start_label(label);
+        for (auto vid : vertices) {
+          if (cur_limit <= 0) {
+            break;
+          }
+          if (predicate(label, vid)) {
+            builder.push_back_opt(vid);
+            cur_limit--;
+          }
+        }
+      }
+      ctx.set(params.alias, builder.finish());
+    }
+    return ctx;
+  }
+
   static Context scan_vertex_with_special_vertex_predicate(
       const GraphReadInterface& graph, const ScanParams& params,
       const SPVertexPredicate& pred);
@@ -70,24 +117,36 @@ class Scan {
                              const ScanParams& params, const PRED_T& predicate,
                              const std::vector<int64_t>& gids) {
     Context ctx;
+    int32_t cur_limit = params.limit;
     if (params.tables.size() == 1) {
       label_t label = params.tables[0];
       SLVertexColumnBuilder builder(label);
       for (auto gid : gids) {
+        if (cur_limit <= 0) {
+          break;
+        }
         vid_t vid = GlobalId::get_vid(gid);
         if (GlobalId::get_label_id(gid) == label && predicate(label, vid)) {
           builder.push_back_opt(vid);
+          cur_limit--;
         }
       }
       ctx.set(params.alias, builder.finish());
     } else if (params.tables.size() > 1) {
       MLVertexColumnBuilder builder;
 
       for (auto label : params.tables) {
+        if (cur_limit <= 0) {
+          break;
+        }
         for (auto gid : gids) {
+          if (cur_limit <= 0) {
+            break;
+          }
           vid_t vid = GlobalId::get_vid(gid);
           if (GlobalId::get_label_id(gid) == label && predicate(label, vid)) {
             builder.push_back_vertex({label, vid});
+            cur_limit--;
           }
         }
       }
@@ -105,14 +164,19 @@ class Scan {
                              const ScanParams& params, const PRED_T& predicate,
                              const std::vector<Any>& oids) {
     Context ctx;
+    auto limit = params.limit;
     if (params.tables.size() == 1) {
       label_t label = params.tables[0];
       SLVertexColumnBuilder builder(label);
       for (auto oid : oids) {
+        if (limit <= 0) {
+          break;
+        }
         vid_t vid;
         if (graph.GetVertexIndex(label, oid, vid)) {
           if (predicate(label, vid)) {
             builder.push_back_opt(vid);
+            --limit;
           }
         }
       }
@@ -121,11 +185,18 @@ class Scan {
       std::vector<std::pair<label_t, vid_t>> vids;
 
       for (auto label : params.tables) {
+        if (limit <= 0) {
+          break;
+        }
         for (auto oid : oids) {
+          if (limit <= 0) {
+            break;
+          }
           vid_t vid;
           if (graph.GetVertexIndex(label, oid, vid)) {
             if (predicate(label, vid)) {
               vids.emplace_back(label, vid);
+              --limit;
             }
           }
         }
@@ -149,15 +220,12 @@ class Scan {
       const GraphReadInterface& graph, const ScanParams& params,
       const SPVertexPredicate& predicate, const std::vector<Any>& oids);
 
-  static Context find_vertex_with_id(const GraphReadInterface& graph,
-                                     label_t label, const Any& pk, int alias,
-                                     bool scan_oid);
-
   static Context find_vertex_with_oid(const GraphReadInterface& graph,
-                                      label_t label, const Any& pk, int alias);
+                                      label_t label, const Any& pk,
+                                      int32_t alias);
 
   static Context find_vertex_with_gid(const GraphReadInterface& graph,
-                                      label_t label, int64_t pk, int alias);
+                                      label_t label, int64_t pk, int32_t alias);
 };
 
 }  // namespace runtime

flex/engines/graph_db/runtime/execute/ops/retrieve/scan.cc

@@ -495,17 +495,31 @@ class ScanWithGPredOpr : public IReadOperator {
     auto expr =
         parse_expression(graph, tmp, params, pred_, VarType::kVertexVar);
     if (expr->is_optional()) {
-      auto ret = Scan::scan_vertex(
-          graph, scan_params_, [&expr](label_t label, vid_t vid) {
-            return expr->eval_vertex(label, vid, 0, 0).as_bool();
-          });
-      return ret;
+      if (scan_params_.limit == std::numeric_limits<int32_t>::max()) {
+        return Scan::scan_vertex(
+            graph, scan_params_, [&expr](label_t label, vid_t vid) {
+              return expr->eval_vertex(label, vid, 0, 0).as_bool();
+            });
+      } else {
+        return Scan::scan_vertex_with_limit(
+            graph, scan_params_, [&expr](label_t label, vid_t vid) {
+              return expr->eval_vertex(label, vid, 0, 0).as_bool();
+            });
+      }
     } else {
-      auto ret = Scan::scan_vertex(
-          graph, scan_params_, [&expr](label_t label, vid_t vid) {
-            return expr->eval_vertex(label, vid, 0).as_bool();
-          });
-      return ret;
+      if (scan_params_.limit == std::numeric_limits<int32_t>::max()) {
+        auto ret = Scan::scan_vertex(
+            graph, scan_params_, [&expr](label_t label, vid_t vid) {
+              return expr->eval_vertex(label, vid, 0).as_bool();
+            });
+        return ret;
+      } else {
+        auto ret = Scan::scan_vertex_with_limit(
+            graph, scan_params_, [&expr](label_t label, vid_t vid) {
+              return expr->eval_vertex(label, vid, 0).as_bool();
+            });
+        return ret;
+      }
     }
   }
 
@@ -523,8 +537,13 @@ class ScanWithoutPredOpr : public IReadOperator {
                             const std::map<std::string, std::string>& params,
                             gs::runtime::Context&& ctx,
                             gs::runtime::OprTimer& timer) override {
-    return Scan::scan_vertex(graph, scan_params_,
-                             [](label_t, vid_t) { return true; });
+    if (scan_params_.limit == std::numeric_limits<int32_t>::max()) {
+      return Scan::scan_vertex(graph, scan_params_,
+                               [](label_t, vid_t) { return true; });
+    } else {
+      return Scan::scan_vertex_with_limit(graph, scan_params_,
+                                          [](label_t, vid_t) { return true; });
+    }
   }
 
  private:
@@ -566,6 +585,17 @@ std::pair<std::unique_ptr<IReadOperator>, ContextMeta> ScanOprBuilder::Build(
 
   ScanParams scan_params;
   scan_params.alias = scan_opr.has_alias() ? scan_opr.alias().value() : -1;
+  scan_params.limit = std::numeric_limits<int32_t>::max();
+  if (scan_opr.params().has_limit()) {
+    auto& limit_range = scan_opr.params().limit();
+    if (limit_range.lower() != 0) {
+      LOG(FATAL) << "Scan with lower limit expect 0, but got "
+                 << limit_range.lower();
+    }
+    if (limit_range.upper() > 0) {
+      scan_params.limit = limit_range.upper();
+    }
+  }
   for (auto& table : scan_opr.params().tables()) {
     // bug here, exclude invalid vertex label id
     if (schema.vertex_label_num() <= table.id()) {

flex/interactive/sdk/python/gs_interactive/tests/test_robustness.py

@@ -310,6 +310,40 @@ def test_call_proc_in_cypher(interactive_session, neo4j_session, create_modern_g
     assert cnt == 8
 
 
+@pytest.mark.skipif(
+    os.environ.get("RUN_ON_PROTO", None) != "ON",
+    reason="Scan+Limit fuse only works on proto",
+)
+def test_scan_limit_fuse(interactive_session, neo4j_session, create_modern_graph):
+    print("[Test call procedure in cypher]")
+    import_data_to_full_modern_graph(interactive_session, create_modern_graph)
+    start_service_on_graph(interactive_session, create_modern_graph)
+    ensure_compiler_schema_ready(
+        interactive_session, neo4j_session, create_modern_graph
+    )
+    result = neo4j_session.run(
+        'MATCH(p: person) with p.id as oid CALL k_neighbors("person", oid, 1) return label_name, vertex_oid;'
+    )
+    cnt = 0
+    for record in result:
+        cnt += 1
+    assert cnt == 8
+
+    # Q: Why we could use this query to verify whether Scan+Limit fuse works?
+    # A: If Scan+Limit fuse works, the result of this query should be 2, otherwise it should be 6
+    result = neo4j_session.run("MATCH(n) return n.id limit 2")
+    cnt = 0
+    for record in result:
+        cnt += 1
+    assert cnt == 2
+
+    result = neo4j_session.run("MATCH(n) return n.id limit 0")
+    cnt = 0
+    for record in result:
+        cnt += 1
+    assert cnt == 0
+
+
 def test_custom_pk_name(
     interactive_session, neo4j_session, create_graph_with_custom_pk_name
 ):

flex/tests/hqps/interactive_config_test_cbo.yaml

@@ -20,6 +20,9 @@ compiler:
       - NotMatchToAntiJoinRule
       - ExtendIntersectRule
       - ExpandGetVFusionRule
+      - ScanExpandFusionRule
+      - TopKPushDownRule
+      - ScanEarlyStopRule # This rule must be placed after TopKPushDownRule and ScanExpandFusionRule
   meta:
     reader:
       schema: