keep LimitIterator in GraphTransaction to avoid upgrade common versio…

…n from 1.8.1 to 1.8.4 Change-Id: If29ecf55f02d9e8bbb351052d1760b69e575ded8
apache · May 27, 2021 · 11d2c92 · 11d2c92
1 parent 0427427
commit 11d2c92
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/tx/GraphTransaction.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/tx/GraphTransaction.java
@@ -70,7 +70,6 @@
 import com.baidu.hugegraph.iterator.ExtendableIterator;
 import com.baidu.hugegraph.iterator.FilterIterator;
 import com.baidu.hugegraph.iterator.FlatMapperIterator;
-import com.baidu.hugegraph.iterator.LimitIterator;
 import com.baidu.hugegraph.iterator.ListIterator;
 import com.baidu.hugegraph.iterator.MapperIterator;
 import com.baidu.hugegraph.iterator.WrappedIterator;
@@ -1988,4 +1987,47 @@ private <T> void traverseByLabel(SchemaLabel label,
             } while (page != null);
         }
     }
+
+    // TODO: move to common module
+    public static class LimitIterator<T> extends WrappedIterator<T> {
+
+        private final Iterator<T> originIterator;
+        private final Function<T, Boolean> filterCallback;
+
+        public LimitIterator(Iterator<T> origin, Function<T, Boolean> filter) {
+            this.originIterator = origin;
+            this.filterCallback = filter;
+        }
+
+        @Override
+        protected Iterator<T> originIterator() {
+            return this.originIterator;
+        }
+
+        @Override
+        protected final boolean fetch() {
+            while (this.originIterator.hasNext()) {
+                T next = this.originIterator.next();
+                // Do filter
+                boolean reachLimit = this.filterCallback.apply(next);
+                if (reachLimit) {
+                    this.closeOriginIterator();
+                    return false;
+                }
+                if (next != null) {
+                    assert this.current == none();
+                    this.current = next;
+                    return true;
+                }
+            }
+            return false;
+        }
+
+        protected final void closeOriginIterator() {
+            if (this.originIterator == null) {
+                return;
+            }
+            close(this.originIterator);
+        }
+    }
 }
diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/HugeTraverser.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/traversal/algorithm/HugeTraverser.java
@@ -44,11 +44,11 @@
 import com.baidu.hugegraph.backend.query.Query;
 import com.baidu.hugegraph.backend.query.QueryResults;
 import com.baidu.hugegraph.backend.tx.GraphTransaction;
+import com.baidu.hugegraph.backend.tx.GraphTransaction.LimitIterator;
 import com.baidu.hugegraph.config.CoreOptions;
 import com.baidu.hugegraph.exception.NotFoundException;
 import com.baidu.hugegraph.iterator.ExtendableIterator;
 import com.baidu.hugegraph.iterator.FilterIterator;
-import com.baidu.hugegraph.iterator.LimitIterator;
 import com.baidu.hugegraph.iterator.MapperIterator;
 import com.baidu.hugegraph.schema.SchemaLabel;
 import com.baidu.hugegraph.structure.HugeEdge;