support edge iterator (#521)

Author: QiZhang1997

geaflow/geaflow-core/geaflow-api/src/main/java/com/antgroup/geaflow/api/graph/function/vc/base/VertexCentricFunction.java

@@ -21,6 +21,7 @@
 
 import com.antgroup.geaflow.api.context.RuntimeContext;
 import com.antgroup.geaflow.api.function.Function;
+import com.antgroup.geaflow.common.iterator.CloseableIterator;
 import com.antgroup.geaflow.model.graph.edge.IEdge;
 import com.antgroup.geaflow.model.graph.vertex.IVertex;
 import com.antgroup.geaflow.state.pushdown.filter.IFilter;
@@ -93,16 +94,31 @@ interface EdgeQuery<K, EV> {
          */
         List<IEdge<K, EV>> getEdges();
 
+        /**
+         * Returns the both edges iterator.
+         */
+        CloseableIterator<IEdge<K, EV>> getEdgesIterator();
+
         /**
          * Returns the out edges.
          */
         List<IEdge<K, EV>> getOutEdges();
 
+        /**
+         * Returns the out edges iterator.
+         */
+        CloseableIterator<IEdge<K, EV>> getOutEdgesIterator();
+
         /**
          * Returns the in edges.
          */
         List<IEdge<K, EV>> getInEdges();
 
+        /**
+         * Returns the in edges iterator.
+         */
+        CloseableIterator<IEdge<K, EV>> getInEdgesIterator();
+
         /**
          * Get the edges which satisfies filter condition.
          */

geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/main/java/com/antgroup/geaflow/operator/impl/graph/algo/vc/context/dynamic/DynamicEdgeQueryImpl.java

@@ -21,6 +21,7 @@
 
 
 import com.antgroup.geaflow.api.graph.function.vc.base.VertexCentricFunction.EdgeQuery;
+import com.antgroup.geaflow.common.iterator.CloseableIterator;
 import com.antgroup.geaflow.model.graph.edge.IEdge;
 import com.antgroup.geaflow.state.DynamicEdgeState;
 import com.antgroup.geaflow.state.GraphState;
@@ -57,18 +58,36 @@ public List<IEdge<K, EV>> getEdges() {
         return edgeState.query(versionId, vId).asList();
     }
 
+    @Override
+    public CloseableIterator<IEdge<K, EV>> getEdgesIterator() {
+        DynamicEdgeState<K, VV, EV> edgeState = graphState.dynamicGraph().E();
+        return edgeState.query(versionId, vId).iterator();
+    }
+
     @Override
     public List<IEdge<K, EV>> getOutEdges() {
         DynamicEdgeState<K, VV, EV> edgeState = graphState.dynamicGraph().E();
         return edgeState.query(versionId, vId).by(OutEdgeFilter.instance()).asList();
     }
 
+    @Override
+    public CloseableIterator<IEdge<K, EV>> getOutEdgesIterator() {
+        DynamicEdgeState<K, VV, EV> edgeState = graphState.dynamicGraph().E();
+        return edgeState.query(versionId, vId).by(OutEdgeFilter.instance()).iterator();
+    }
+
     @Override
     public List<IEdge<K, EV>> getInEdges() {
         DynamicEdgeState<K, VV, EV> edgeState = graphState.dynamicGraph().E();
         return edgeState.query(versionId, vId).by(InEdgeFilter.instance()).asList();
     }
 
+    @Override
+    public CloseableIterator<IEdge<K, EV>> getInEdgesIterator() {
+        DynamicEdgeState<K, VV, EV> edgeState = graphState.dynamicGraph().E();
+        return edgeState.query(versionId, vId).by(InEdgeFilter.instance()).iterator();
+    }
+
     @Override
     public List<IEdge<K, EV>> getEdges(IFilter edgeFilter) {
         DynamicEdgeState<K, VV, EV> edgeState = graphState.dynamicGraph().E();

geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/main/java/com/antgroup/geaflow/operator/impl/graph/algo/vc/context/statical/StaticEdgeQueryImpl.java

@@ -20,6 +20,7 @@
 package com.antgroup.geaflow.operator.impl.graph.algo.vc.context.statical;
 
 import com.antgroup.geaflow.api.graph.function.vc.base.VertexCentricFunction.EdgeQuery;
+import com.antgroup.geaflow.common.iterator.CloseableIterator;
 import com.antgroup.geaflow.model.graph.edge.IEdge;
 import com.antgroup.geaflow.state.GraphState;
 import com.antgroup.geaflow.state.pushdown.filter.IFilter;
@@ -50,16 +51,31 @@ public List<IEdge<K, EV>> getEdges() {
         return graphState.staticGraph().E().query(vId).asList();
     }
 
+    @Override
+    public CloseableIterator<IEdge<K, EV>> getEdgesIterator() {
+        return graphState.staticGraph().E().query(vId).iterator();
+    }
+
     @Override
     public List<IEdge<K, EV>> getOutEdges() {
         return graphState.staticGraph().E().query(vId).by(OutEdgeFilter.instance()).asList();
     }
 
+    @Override
+    public CloseableIterator<IEdge<K, EV>> getOutEdgesIterator() {
+        return graphState.staticGraph().E().query(vId).by(OutEdgeFilter.instance()).iterator();
+    }
+
     @Override
     public List<IEdge<K, EV>> getInEdges() {
         return graphState.staticGraph().E().query(vId).by(InEdgeFilter.instance()).asList();
     }
 
+    @Override
+    public CloseableIterator<IEdge<K, EV>> getInEdgesIterator() {
+        return graphState.staticGraph().E().query(vId).by(InEdgeFilter.instance()).iterator();
+    }
+
     @Override
     public List<IEdge<K, EV>> getEdges(IFilter edgeFilter) {
         return graphState.staticGraph().E().query(vId).by(edgeFilter).asList();

geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/com/antgroup/geaflow/dsl/common/algo/AlgorithmRuntimeContext.java

@@ -20,6 +20,7 @@
 package com.antgroup.geaflow.dsl.common.algo;
 
 import com.antgroup.geaflow.common.config.Configuration;
+import com.antgroup.geaflow.common.iterator.CloseableIterator;
 import com.antgroup.geaflow.dsl.common.data.Row;
 import com.antgroup.geaflow.dsl.common.data.RowEdge;
 import com.antgroup.geaflow.dsl.common.types.GraphSchema;
@@ -34,6 +35,8 @@ public interface AlgorithmRuntimeContext<K, M> {
 
     List<RowEdge> loadDynamicEdges(EdgeDirection direction);
 
+    CloseableIterator<RowEdge> loadStaticEdgesIterator(EdgeDirection direction);
+
     void sendMessage(K vertexId, M message);
 
     void updateVertexValue(Row value);

geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/com/antgroup/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java

@@ -169,6 +169,20 @@ public List<RowEdge> loadStaticEdges(EdgeDirection direction) {
         return rowEdges;
     }
 
+    @Override
+    public CloseableIterator<RowEdge> loadStaticEdgesIterator(EdgeDirection direction) {
+        switch (direction) {
+            case OUT:
+                return (CloseableIterator) edgeQuery.getOutEdgesIterator();
+            case IN:
+                return (CloseableIterator) edgeQuery.getInEdgesIterator();
+            case BOTH:
+                return (CloseableIterator) edgeQuery.getEdgesIterator();
+            default:
+                throw new GeaFlowDSLException("Illegal edge direction: " + direction);
+        }
+    }
+
     @Override
     public void sendMessage(Object vertexId, Object message) {
         incVCTraversalCtx.sendMessage(vertexId, message);

geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/com/antgroup/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java

@@ -23,6 +23,7 @@
 import com.antgroup.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.TraversalEdgeQuery;
 import com.antgroup.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.VertexCentricTraversalFuncContext;
 import com.antgroup.geaflow.common.config.Configuration;
+import com.antgroup.geaflow.common.iterator.CloseableIterator;
 import com.antgroup.geaflow.dsl.common.algo.AlgorithmRuntimeContext;
 import com.antgroup.geaflow.dsl.common.data.Row;
 import com.antgroup.geaflow.dsl.common.data.RowEdge;
@@ -93,6 +94,20 @@ public List<RowEdge> loadDynamicEdges(EdgeDirection direction) {
         throw new RuntimeException("GeaFlowAlgorithmRuntimeContext not support loadDynamicEdges");
     }
 
+    @Override
+    public CloseableIterator<RowEdge> loadStaticEdgesIterator(EdgeDirection direction) {
+        switch (direction) {
+            case OUT:
+                return (CloseableIterator) edgeQuery.getOutEdgesIterator();
+            case IN:
+                return (CloseableIterator) edgeQuery.getInEdgesIterator();
+            case BOTH:
+                return (CloseableIterator) edgeQuery.getEdgesIterator();
+            default:
+                throw new GeaFlowDSLException("Illegal edge direction: " + direction);
+        }
+    }
+
     @Override
     public void sendMessage(Object vertexId, Object message) {
         traversalContext.sendMessage(vertexId, message);

geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/com/antgroup/geaflow/dsl/runtime/query/GQLAlgorithmTest.java

@@ -258,6 +258,15 @@ public void testAlgorithmCommonNeighbors() throws Exception {
             .checkSinkResult();
     }
 
+    @Test
+    public void testEdgeIterator() throws Exception {
+        QueryTester
+                .build()
+                .withQueryPath("/query/gql_edge_iterator_test.sql")
+                .execute()
+                .checkSinkResult();
+    }
+
     private void clearGraph() throws IOException {
         File file = new File(TEST_GRAPH_PATH);
         if (file.exists()) {

geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/com/antgroup/geaflow/dsl/runtime/query/udf/TestEdgeIteratorUdf.java

@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.antgroup.geaflow.dsl.runtime.query.udf;
+
+import com.antgroup.geaflow.common.iterator.CloseableIterator;
+import com.antgroup.geaflow.common.type.primitive.LongType;
+import com.antgroup.geaflow.dsl.common.algo.AlgorithmRuntimeContext;
+import com.antgroup.geaflow.dsl.common.algo.AlgorithmUserFunction;
+import com.antgroup.geaflow.dsl.common.data.Row;
+import com.antgroup.geaflow.dsl.common.data.RowEdge;
+import com.antgroup.geaflow.dsl.common.data.RowVertex;
+import com.antgroup.geaflow.dsl.common.data.impl.ObjectRow;
+import com.antgroup.geaflow.dsl.common.function.Description;
+import com.antgroup.geaflow.dsl.common.types.GraphSchema;
+import com.antgroup.geaflow.dsl.common.types.StructType;
+import com.antgroup.geaflow.dsl.common.types.TableField;
+import com.antgroup.geaflow.model.graph.edge.EdgeDirection;
+import java.util.Iterator;
+import java.util.Optional;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+@Description(name = "test_edge_iterator", description = "built-in udga for WeakConnectedComponents")
+public class TestEdgeIteratorUdf implements AlgorithmUserFunction<Object, Long> {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(TestEdgeIteratorUdf.class);
+
+    private AlgorithmRuntimeContext<Object, Long> context;
+
+    private int iteration = 5;
+
+    private int edgeLimit = 100;
+
+    @Override
+    public void init(AlgorithmRuntimeContext<Object, Long> context, Object[] parameters) {
+        this.context = context;
+        if (parameters.length > 0) {
+            iteration = Integer.parseInt(String.valueOf(parameters[0]));
+        }
+        if (parameters.length > 1) {
+            edgeLimit = Integer.parseInt(String.valueOf(parameters[1]));
+        }
+    }
+
+    @Override
+    public void process(RowVertex vertex, Optional<Row> updatedValues, Iterator<Long> messages) {
+        updatedValues.ifPresent(vertex::setValue);
+        CloseableIterator<RowEdge> edgesIterator = context.loadStaticEdgesIterator(EdgeDirection.BOTH);
+        if (context.getCurrentIterationId() < iteration) {
+            int count = 0;
+            while (edgesIterator.hasNext() && count < edgeLimit) {
+                RowEdge next = edgesIterator.next();
+                context.sendMessage(next.getTargetId(), context.getCurrentIterationId());
+                count++;
+            }
+        }
+    }
+
+    @Override
+    public void finish(RowVertex graphVertex, Optional<Row> updatedValues) {
+        updatedValues.ifPresent(graphVertex::setValue);
+        long iteration = (long) graphVertex.getValue().getField(0, LongType.INSTANCE);
+        context.take(ObjectRow.create(graphVertex.getId(), iteration));
+    }
+
+    @Override
+    public StructType getOutputType(GraphSchema graphSchema) {
+        return new StructType(
+            new TableField("id", graphSchema.getIdType(), false),
+            new TableField("iteration", LongType.INSTANCE, false)
+        );
+    }
+}

geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/resources/expect/gql_edge_iterator_test.txt

@@ -0,0 +1,6 @@
+1,1
+5,1
+3,1
+4,1
+2,1
+6,1

geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/resources/query/gql_edge_iterator_test.sql

@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+set geaflow.dsl.window.size = -1;
+set geaflow.dsl.ignore.exception = true;
+create function test_edge_iterator as 'com.antgroup.geaflow.dsl.runtime.query.udf.TestEdgeIteratorUdf';
+
+
+CREATE GRAPH IF NOT EXISTS g4 (
+  Vertex v4 (
+    vid varchar ID,
+    vvalue int
+  ),
+  Edge e4 (
+    srcId varchar SOURCE ID,
+    targetId varchar DESTINATION ID
+  )
+) WITH (
+  storeType='rocksdb',
+  shardCount = 1
+);
+
+CREATE TABLE IF NOT EXISTS v_source (
+    v_id varchar,
+    v_value int,
+    ts varchar,
+    type varchar
+) WITH (
+  type='file',
+  geaflow.dsl.file.path = 'resource:///input/test_vertex'
+);
+
+CREATE TABLE IF NOT EXISTS e_source (
+    src_id varchar,
+    dst_id varchar
+) WITH (
+  type='file',
+  geaflow.dsl.file.path = 'resource:///input/test_edge'
+);
+
+CREATE TABLE IF NOT EXISTS tbl_result (
+  v_id varchar,
+  iteration long
+) WITH (
+  type='file',
+   geaflow.dsl.file.path = '${target}'
+);
+
+USE GRAPH g4;
+
+INSERT INTO g4.v4(vid, vvalue)
+SELECT
+v_id, v_value
+FROM v_source;
+
+INSERT INTO g4.e4(srcId, targetId)
+SELECT
+ src_id, dst_id
+FROM e_source;
+
+INSERT INTO tbl_result(v_id, iteration)
+CALL test_edge_iterator(3, 2) YIELD (vid, iteration)
+RETURN vid, iteration
+;