add scan_with_limit

alibaba · Jan 22, 2025 · d98c73a · d98c73a
1 parent 88a48b1
commit d98c73a
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 12 deletions.
diff --git a/flex/engines/graph_db/runtime/common/operators/retrieve/scan.h b/flex/engines/graph_db/runtime/common/operators/retrieve/scan.h
@@ -37,6 +37,38 @@ class Scan {
                              const ScanParams& params,
                              const PRED_T& predicate) {
     Context ctx;
+    if (params.tables.size() == 1) {
+      label_t label = params.tables[0];
+      SLVertexColumnBuilder builder(label);
+      auto vertices = graph.GetVertexSet(label);
+      for (auto vid : vertices) {
+        if (predicate(label, vid)) {
+          builder.push_back_opt(vid);
+        }
+      }
+      ctx.set(params.alias, builder.finish());
+    } else if (params.tables.size() > 1) {
+      MSVertexColumnBuilder builder;
+
+      for (auto label : params.tables) {
+        auto vertices = graph.GetVertexSet(label);
+        builder.start_label(label);
+        for (auto vid : vertices) {
+          if (predicate(label, vid)) {
+            builder.push_back_opt(vid);
+          }
+        }
+      }
+      ctx.set(params.alias, builder.finish());
+    }
+    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];

diff --git a/flex/engines/graph_db/runtime/execute/ops/retrieve/scan.cc b/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: