fix(reasoner): fix PER_NODE_LIMIT (#368)

Co-authored-by: peilong <[email protected]>
OpenSPG · Oct 21, 2024 · 8364df9 · 8364df9
1 parent 5dc03cb
commit 8364df9
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 8 deletions.
diff --git a/...src/test/java/com/antgroup/openspg/reasoner/runner/local/main/KgReasonerTopKFilmTest.java b/...src/test/java/com/antgroup/openspg/reasoner/runner/local/main/KgReasonerTopKFilmTest.java
@@ -457,4 +457,17 @@ private void doTest14() {
     Assert.assertEquals("root", result.get(0)[0]);
     Assert.assertEquals("L1_1_star", result.get(0)[1]);
   }
+
+  @Test
+  public void test15() {
+    FileMutex.runTestWithMutex(this::doTest15);
+  }
+
+  private void doTest15() {
+    String dsl =
+        "match (s:Film)-[p:starOfFilm|directOfFilm  PER_NODE_LIMIT 1]->(o:FilmStar|FilmDirector) where s.id = 'root' return s.id, o.id";
+    List<String[]> result = runTestResult(dsl);
+    Assert.assertEquals(2, result.size());
+    Assert.assertEquals(2, result.get(0).length);
+  }
 }
diff --git a/...ner/runner-common/src/main/java/com/antgroup/openspg/reasoner/pattern/PatternMatcher.java b/...ner/runner-common/src/main/java/com/antgroup/openspg/reasoner/pattern/PatternMatcher.java
@@ -256,9 +256,6 @@ public boolean test(IEdge<IVertexId, IProperty> e) {
                   + JSON.toJSONString(dstVertexRuleList));
         }
       }
-      if (patternConnection.limit() != null && patternConnection.limit() > 0) {
-        limit = new Long(patternConnection.limit());
-      }
       List<IEdge<IVertexId, IProperty>> validEdges =
           matchEdges(
               vertexContext, willMatchEdgeList, patternConnection, pattern, edgeRuleMap, limit);
@@ -299,10 +296,15 @@ private List<IEdge<IVertexId, IProperty>> matchEdges(
       Connection patternConnection,
       Pattern pattern,
       Map<String, List<String>> edgeRuleMap,
-      Long limit) {
+      Long totalLimit) {
     ArrayList<IEdge<IVertexId, IProperty>> result = new ArrayList<>();
-    long oneTypeEdgeCount = 0;
+    Map<String, Long> edgeTypeCountMap = new HashMap<>();
+    Long totalCount = 0L;
     for (IEdge<IVertexId, IProperty> edge : edgeList) {
+      String edgeType = edge.getType();
+      if (!edgeTypeCountMap.containsKey(edgeType)) {
+        edgeTypeCountMap.put(edgeType, 0L);
+      }
       if (!isEdgeMatch(
           vertexContext,
           edge,
@@ -311,11 +313,17 @@ private List<IEdge<IVertexId, IProperty>> matchEdges(
           edgeRuleMap.get(patternConnection.alias()))) {
         continue;
       }
-      oneTypeEdgeCount++;
-      if (null != limit && oneTypeEdgeCount > limit) {
-        // reach max path limit
+      totalCount = totalCount + 1;
+      if (null != totalLimit && totalCount > totalLimit) {
         break;
       }
+      long currentEdgeTypeCount = edgeTypeCountMap.get(edgeType) + 1;
+      edgeTypeCountMap.put(edgeType, currentEdgeTypeCount);
+      if (null != patternConnection.limit()
+          && patternConnection.limit() > 0
+          && currentEdgeTypeCount > patternConnection.limit()) {
+        continue;
+      }
       result.add(edge);
     }
     result.trimToSize();