From 3e2d1c11fe186b7afe37b25e8c6f888daab86ba4 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 12 Aug 2026 12:16:06 +0800 Subject: [PATCH 1/9] feat(sampling): add bounded one-hop neighbor sampling --- .../graph/sampling/SubgraphSamplingSpec.java | 87 ++++++ .../sampling/SubgraphSamplingSpecTest.java | 60 ++++ .../geaflow/state/data/OneDegreeGraph.java | 49 +++- .../DeterministicNeighborSampler.java | 275 ++++++++++++++++++ .../state/sampling/LocalNeighborhood.java | 111 +++++++ .../DeterministicNeighborSamplerTest.java | 214 ++++++++++++++ 6 files changed, 795 insertions(+), 1 deletion(-) create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpec.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java create mode 100644 geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java create mode 100644 geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java create mode 100644 geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpec.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpec.java new file mode 100644 index 000000000..8f12dd783 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpec.java @@ -0,0 +1,87 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.Objects; +import org.apache.geaflow.model.graph.edge.EdgeDirection; + +/** + * Configuration for bounded iterative multi-hop sampling. + */ +public class SubgraphSamplingSpec implements Serializable { + + public static final long DEFAULT_MAX_SAMPLED_NODES = 10000L; + public static final long DEFAULT_MAX_SAMPLED_EDGES = 100000L; + public static final long DEFAULT_MAX_RETURNED_EDGES = 100000L; + + private final int hops; + private final int fanout; + private final EdgeDirection direction; + private final long maxReturnedEdges; + private final long seed; + + public SubgraphSamplingSpec(int hops, int fanout, EdgeDirection direction) { + this(hops, fanout, direction, DEFAULT_MAX_RETURNED_EDGES); + } + + public SubgraphSamplingSpec(int hops, int fanout, EdgeDirection direction, + long maxReturnedEdges) { + this(hops, fanout, direction, maxReturnedEdges, 0L); + } + + public SubgraphSamplingSpec(int hops, int fanout, EdgeDirection direction, + long maxReturnedEdges, long seed) { + if (hops < 1) { + throw new IllegalArgumentException("sampling hops must be greater than zero"); + } + if (fanout == 0 || fanout < -1) { + throw new IllegalArgumentException("fanout must be -1 or greater than zero"); + } + if (maxReturnedEdges < 1) { + throw new IllegalArgumentException("maxReturnedEdges must be greater than zero"); + } + this.hops = hops; + this.fanout = fanout; + this.direction = Objects.requireNonNull(direction, "direction"); + this.maxReturnedEdges = maxReturnedEdges; + this.seed = seed; + } + + public int getHops() { + return hops; + } + + public int getFanout() { + return fanout; + } + + public EdgeDirection getDirection() { + return direction; + } + + public long getMaxReturnedEdges() { + return maxReturnedEdges; + } + + public long getSeed() { + return seed; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java new file mode 100644 index 000000000..02b106c30 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java @@ -0,0 +1,60 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class SubgraphSamplingSpecTest { + + @Test + public void testHopsAndFanoutAreScalar() { + SubgraphSamplingSpec spec = new SubgraphSamplingSpec(2, 10, EdgeDirection.OUT); + + Assert.assertEquals(spec.getHops(), 2); + Assert.assertEquals(spec.getFanout(), 10); + Assert.assertEquals(spec.getMaxReturnedEdges(), 100000L); + Assert.assertEquals(spec.getSeed(), 0L); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsNonPositiveHops() { + new SubgraphSamplingSpec(0, 10, EdgeDirection.OUT); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsZeroFanout() { + new SubgraphSamplingSpec(2, 0, EdgeDirection.OUT); + } + + @Test + public void testUnlimitedFanoutKeepsPerVertexEdgeBudget() { + SubgraphSamplingSpec spec = new SubgraphSamplingSpec(2, -1, EdgeDirection.BOTH); + + Assert.assertEquals(spec.getFanout(), -1); + Assert.assertEquals(spec.getMaxReturnedEdges(), 100000L); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsInvalidReturnedEdgeBudget() { + new SubgraphSamplingSpec(2, -1, EdgeDirection.BOTH, 0); + } +} diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java index 43552c9ca..d564aec3b 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java @@ -20,15 +20,26 @@ package org.apache.geaflow.state.data; import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import java.util.Objects; import org.apache.geaflow.common.iterator.CloseableIterator; +import org.apache.geaflow.model.graph.edge.EdgeDirection; import org.apache.geaflow.model.graph.edge.IEdge; import org.apache.geaflow.model.graph.vertex.IVertex; +import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; public class OneDegreeGraph implements Serializable { private IVertex vertex; protected CloseableIterator> edgeIterator; protected K key; + private List> sampledEdges; + private EdgeDirection sampledDirection; + private Integer sampledFanout; + private Long sampledMaxReturnedEdges; + private Long sampledSeed; + private Long sampledVersion; public OneDegreeGraph(K key, IVertex vertex, CloseableIterator> edgeIterator) { this.key = key; @@ -47,5 +58,41 @@ public IVertex getVertex() { public CloseableIterator> getEdgeIterator() { return edgeIterator; } -} + /** Samples this vertex's bounded one-hop neighborhood in the state layer. */ + public synchronized List> sampleNeighbors(EdgeDirection direction, int fanout) { + return sampleNeighbors(direction, fanout, + DeterministicNeighborSampler.DEFAULT_MAX_CANDIDATE_EDGES, 0L, 0L); + } + + /** Samples this one-shot edge iterator for one deterministic sampling round. */ + public synchronized List> sampleNeighbors(EdgeDirection direction, int fanout, + long maxReturnedEdges, long seed, + long samplingVersion) { + if (sampledEdges != null) { + if (sampledDirection != direction || !Objects.equals(sampledFanout, fanout) + || !Objects.equals(sampledMaxReturnedEdges, maxReturnedEdges) + || !Objects.equals(sampledSeed, seed) + || !Objects.equals(sampledVersion, samplingVersion)) { + throw new IllegalStateException( + "one-degree edge iterator was already consumed by a different sampling request"); + } + return sampledEdges; + } + try { + Iterable> iterable = () -> edgeIterator; + sampledEdges = Collections.unmodifiableList( + DeterministicNeighborSampler.sample(key, iterable, direction, fanout, + java.util.Comparator.comparing(String::valueOf), maxReturnedEdges, + seed, samplingVersion)); + sampledDirection = direction; + sampledFanout = fanout; + sampledMaxReturnedEdges = maxReturnedEdges; + sampledSeed = seed; + sampledVersion = samplingVersion; + return sampledEdges; + } finally { + edgeIterator.close(); + } + } +} diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java new file mode 100644 index 000000000..a84632451 --- /dev/null +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java @@ -0,0 +1,275 @@ +/* + * 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 org.apache.geaflow.state.sampling; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.PriorityQueue; +import org.apache.geaflow.model.graph.IGraphElementWithLabelField; +import org.apache.geaflow.model.graph.IGraphElementWithTimeField; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; + +/** + * Storage-independent seeded one-hop neighbor sampling. + */ +public final class DeterministicNeighborSampler { + + public static final long DEFAULT_MAX_CANDIDATE_EDGES = 100000L; + + private DeterministicNeighborSampler() { + } + + public static List> sample(K vertexId, + Iterable> edges, + EdgeDirection direction, + int fanout) { + return sample(vertexId, edges, direction, fanout, + Comparator.comparing(String::valueOf), DEFAULT_MAX_CANDIDATE_EDGES, 0L, 0L); + } + + public static List> sample(K vertexId, + Iterable> edges, + EdgeDirection direction, + int fanout, + Comparator idComparator, + long maxReturnedEdges) { + return sample(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, 0L, 0L); + } + + public static List> sample(K vertexId, + Iterable> edges, + EdgeDirection direction, + int fanout, + Comparator idComparator, + long maxReturnedEdges, + long seed, + long samplingVersion) { + return select(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, + seed, samplingVersion, true); + } + + /** Project an already direction-filtered local neighborhood to a smaller fanout. */ + public static List> project(K vertexId, + Iterable> edges, + EdgeDirection direction, + int fanout) { + return project(vertexId, edges, direction, fanout, + Comparator.comparing(String::valueOf), DEFAULT_MAX_CANDIDATE_EDGES, 0L, 0L); + } + + public static List> project(K vertexId, + Iterable> edges, + EdgeDirection direction, + int fanout, + Comparator idComparator, + long maxReturnedEdges) { + return project(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, 0L, 0L); + } + + public static List> project(K vertexId, + Iterable> edges, + EdgeDirection direction, + int fanout, + Comparator idComparator, + long maxReturnedEdges, + long seed, + long samplingVersion) { + return select(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, + seed, samplingVersion, false); + } + + private static List> select(K vertexId, + Iterable> edges, + EdgeDirection direction, + int fanout, + Comparator idComparator, + long maxReturnedEdges, + long seed, + long samplingVersion, + boolean filterAndNormalize) { + validate(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges); + Comparator> rankComparator = (left, right) -> { + int result = Long.compareUnsigned(left.score, right.score); + return result != 0 ? result : compareIds(left.neighborId, right.neighborId, idComparator); + }; + Map> selected = new HashMap<>(); + PriorityQueue> worstFirst = + fanout < 0 ? null : new PriorityQueue<>(fanout, rankComparator.reversed()); + + for (IEdge sourceEdge : edges) { + Objects.requireNonNull(sourceEdge, "edge"); + if (filterAndNormalize && !matchesDirection(sourceEdge, direction)) { + continue; + } + IEdge edge = filterAndNormalize ? normalize(sourceEdge) : sourceEdge; + K neighborId = Objects.requireNonNull(neighborId(vertexId, edge), "neighborId"); + NeighborGroup group = selected.get(neighborId); + if (group != null) { + group.edges.add(edge); + continue; + } + + group = new NeighborGroup<>(neighborId, + sampleScore(seed, samplingVersion, vertexId, direction, neighborId), edge); + if (fanout < 0 || selected.size() < fanout) { + selected.put(neighborId, group); + if (worstFirst != null) { + worstFirst.add(group); + } + } else if (rankComparator.compare(group, worstFirst.peek()) < 0) { + NeighborGroup removed = worstFirst.remove(); + selected.remove(removed.neighborId); + selected.put(neighborId, group); + worstFirst.add(group); + } + } + + List> groups = new ArrayList<>(selected.values()); + groups.sort(rankComparator); + List> result = new ArrayList<>(); + for (NeighborGroup group : groups) { + group.edges.sort((left, right) -> compareEdges(left, right, idComparator)); + result.addAll(group.edges); + if (result.size() > maxReturnedEdges) { + throw new IllegalStateException(String.format( + "one-hop sampling edge limit exceeded, vertexId=%s, actual=%s, limit=%s", + vertexId, result.size(), maxReturnedEdges)); + } + } + return result; + } + + private static void validate(Object vertexId, Iterable edges, EdgeDirection direction, + int fanout, Comparator idComparator, long maxReturnedEdges) { + Objects.requireNonNull(vertexId, "vertexId"); + Objects.requireNonNull(edges, "edges"); + Objects.requireNonNull(direction, "direction"); + Objects.requireNonNull(idComparator, "idComparator"); + if (fanout == 0 || fanout < -1) { + throw new IllegalArgumentException("fanout must be -1 or greater than zero"); + } + if (maxReturnedEdges < 1) { + throw new IllegalArgumentException("maxReturnedEdges must be greater than zero"); + } + } + + private static boolean matchesDirection(IEdge edge, EdgeDirection direction) { + return direction == EdgeDirection.BOTH || edge.getDirect() == direction; + } + + private static IEdge normalize(IEdge edge) { + if (edge.getDirect() != EdgeDirection.IN) { + return edge; + } + IEdge reversed = edge.reverse(); + // Direction remains the sampling-side marker; endpoints are restored to logical order. + reversed.setDirect(edge.getDirect()); + return reversed; + } + + private static K neighborId(K vertexId, IEdge edge) { + if (Objects.equals(vertexId, edge.getSrcId())) { + return edge.getTargetId(); + } + if (Objects.equals(vertexId, edge.getTargetId())) { + return edge.getSrcId(); + } + return edge.getTargetId(); + } + + private static long sampleScore(long seed, long samplingVersion, Object vertexId, + EdgeDirection direction, Object neighborId) { + long value = mix64(seed) ^ Long.rotateLeft(mix64(samplingVersion), 11); + value ^= Long.rotateLeft(stableHash(vertexId), 23); + value ^= Long.rotateLeft(mix64(direction.ordinal()), 37); + value ^= Long.rotateLeft(stableHash(neighborId), 47); + return mix64(value); + } + + private static long stableHash(Object value) { + String text = value.getClass().getName() + ':' + value; + long hash = 0xcbf29ce484222325L; + for (int i = 0; i < text.length(); i++) { + hash ^= text.charAt(i); + hash *= 0x100000001b3L; + } + return mix64(hash); + } + + private static long mix64(long value) { + value = (value ^ (value >>> 30)) * 0xbf58476d1ce4e5b9L; + value = (value ^ (value >>> 27)) * 0x94d049bb133111ebL; + return value ^ (value >>> 31); + } + + private static int compareIds(K left, K right, Comparator idComparator) { + int result = idComparator.compare(left, right); + return result != 0 ? result : String.valueOf(left).compareTo(String.valueOf(right)); + } + + private static int compareEdges(IEdge left, IEdge right, + Comparator idComparator) { + int result = compareIds(left.getSrcId(), right.getSrcId(), idComparator); + if (result == 0) { + result = compareIds(left.getTargetId(), right.getTargetId(), idComparator); + } + if (result == 0) { + result = left.getDirect().compareTo(right.getDirect()); + } + if (result == 0) { + result = String.valueOf(labelOf(left)).compareTo(String.valueOf(labelOf(right))); + } + if (result == 0) { + result = String.valueOf(timeOf(left)).compareTo(String.valueOf(timeOf(right))); + } + if (result == 0) { + result = String.valueOf(left.getValue()).compareTo(String.valueOf(right.getValue())); + } + return result; + } + + private static String labelOf(IEdge edge) { + return edge instanceof IGraphElementWithLabelField + ? ((IGraphElementWithLabelField) edge).getLabel() : null; + } + + private static Long timeOf(IEdge edge) { + return edge instanceof IGraphElementWithTimeField + ? ((IGraphElementWithTimeField) edge).getTime() : null; + } + + private static final class NeighborGroup { + + private final K neighborId; + private final long score; + private final List> edges = new ArrayList<>(); + + private NeighborGroup(K neighborId, long score, IEdge firstEdge) { + this.neighborId = neighborId; + this.score = score; + this.edges.add(firstEdge); + } + } +} diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java new file mode 100644 index 000000000..5f01c4652 --- /dev/null +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java @@ -0,0 +1,111 @@ +/* + * 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 org.apache.geaflow.state.sampling; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.vertex.IVertex; + +/** + * Versioned one-hop state owned by a single vertex. + */ +public class LocalNeighborhood implements Serializable { + + private final IVertex vertex; + private final List> edges; + private final long snapshotVersion; + private final long samplingVersion; + + public LocalNeighborhood(IVertex vertex, List> edges, + long snapshotVersion) { + this(vertex, edges, snapshotVersion, snapshotVersion); + } + + public LocalNeighborhood(IVertex vertex, List> edges, + long snapshotVersion, long samplingVersion) { + this.vertex = Objects.requireNonNull(vertex, "vertex"); + this.edges = new ArrayList<>(Objects.requireNonNull(edges, "edges")); + for (IEdge edge : this.edges) { + Objects.requireNonNull(edge, "edge"); + } + this.snapshotVersion = snapshotVersion; + this.samplingVersion = samplingVersion; + } + + public IVertex getVertex() { + return vertex; + } + + public List> getEdges() { + return Collections.unmodifiableList(edges); + } + + public long getSnapshotVersion() { + return snapshotVersion; + } + + public long getSamplingVersion() { + return samplingVersion; + } + + public boolean matches(long expectedSnapshotVersion, long expectedSamplingVersion) { + return snapshotVersion == expectedSnapshotVersion + && samplingVersion == expectedSamplingVersion; + } + + public LocalNeighborhood revalidate(IVertex currentVertex, + long currentSnapshotVersion) { + if (currentSnapshotVersion < snapshotVersion) { + throw new IllegalArgumentException("cannot revalidate a neighborhood to an older snapshot"); + } + return new LocalNeighborhood<>(currentVertex, edges, currentSnapshotVersion, samplingVersion); + } + + /** + * Create a bounded view of this already direction-filtered neighborhood. + */ + public LocalNeighborhood project(EdgeDirection direction, int fanout) { + return new LocalNeighborhood<>(vertex, + DeterministicNeighborSampler.project(vertex.getId(), edges, direction, fanout), + snapshotVersion, samplingVersion); + } + + public LocalNeighborhood project(EdgeDirection direction, int fanout, + Comparator idComparator, + long maxReturnedEdges) { + return new LocalNeighborhood<>(vertex, DeterministicNeighborSampler.project(vertex.getId(), + edges, direction, fanout, idComparator, maxReturnedEdges), snapshotVersion, samplingVersion); + } + + public LocalNeighborhood project(EdgeDirection direction, int fanout, + Comparator idComparator, + long maxReturnedEdges, + long seed) { + return new LocalNeighborhood<>(vertex, DeterministicNeighborSampler.project(vertex.getId(), + edges, direction, fanout, idComparator, maxReturnedEdges, seed, samplingVersion), + snapshotVersion, samplingVersion); + } +} diff --git a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java new file mode 100644 index 000000000..6de60ab05 --- /dev/null +++ b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java @@ -0,0 +1,214 @@ +/* + * 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 org.apache.geaflow.state.sampling; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import org.apache.geaflow.common.iterator.CloseableIterator; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.apache.geaflow.model.graph.vertex.impl.ValueVertex; +import org.apache.geaflow.state.data.OneDegreeGraph; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class DeterministicNeighborSamplerTest { + + @Test + public void testSamplingIsBoundedAndIndependentOfInputOrder() { + List> first = Arrays.asList( + edge(1L, 2L), edge(1L, 3L), edge(1L, 4L)); + List> second = Arrays.asList( + edge(1L, 4L), edge(1L, 2L), edge(1L, 3L)); + + List> sampledFirst = + DeterministicNeighborSampler.sample(1L, first, EdgeDirection.OUT, 2); + List> sampledSecond = + DeterministicNeighborSampler.sample(1L, second, EdgeDirection.OUT, 2); + + Assert.assertEquals(sampledFirst.size(), 2); + Assert.assertEquals(targetIds(sampledFirst), targetIds(sampledSecond)); + } + + @Test + public void testPositiveFanoutDoesNotMaterializeAllCandidates() { + List> edges = new ArrayList<>(); + for (long target = 2L; target < 102L; target++) { + edges.add(edge(1L, target)); + } + + List> sampled = DeterministicNeighborSampler.sample( + 1L, edges, EdgeDirection.OUT, 3, Long::compare, 3L, 17L, 9L); + + Assert.assertEquals(sampled.stream().map(IEdge::getTargetId).distinct().count(), 3L); + Assert.assertEquals(sampled.size(), 3); + } + + @Test + public void testSeedAndVersionAreStableAndInputOrderIndependent() { + List> first = new ArrayList<>(); + for (long target = 2L; target < 42L; target++) { + first.add(edge(1L, target)); + } + List> reversed = new ArrayList<>(first); + java.util.Collections.reverse(reversed); + + List firstSample = targetIds(DeterministicNeighborSampler.sample( + 1L, first, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 7L)); + List reorderedSample = targetIds(DeterministicNeighborSampler.sample( + 1L, reversed, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 7L)); + List nextVersionSample = targetIds(DeterministicNeighborSampler.sample( + 1L, first, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 8L)); + + Assert.assertEquals(firstSample, reorderedSample); + Assert.assertNotEquals(firstSample, nextVersionSample); + } + + @Test + public void testDirectionAndUnlimitedFanout() { + IEdge out = edge(1L, 2L); + IEdge in = edge(1L, 3L); + in.setDirect(EdgeDirection.IN); + + List> sampledIn = DeterministicNeighborSampler.sample( + 1L, Arrays.asList(out, in), EdgeDirection.IN, -1); + Assert.assertEquals(sampledIn.size(), 1); + Assert.assertEquals(sampledIn.get(0).getSrcId(), Long.valueOf(3L)); + Assert.assertEquals(sampledIn.get(0).getTargetId(), Long.valueOf(1L)); + Assert.assertEquals(sampledIn.get(0).getDirect(), EdgeDirection.IN); + Assert.assertEquals( + DeterministicNeighborSampler.sample(1L, Arrays.asList(out, in), EdgeDirection.BOTH, -1).size(), + 2); + } + + @Test + public void testFanoutCountsNeighborsAndKeepsSelectedParallelEdges() { + List> edges = Arrays.asList( + edge(1L, 2L), edgeWithValue(1L, 2L, "parallel"), edge(1L, 3L), edge(1L, 4L)); + + List> sampled = DeterministicNeighborSampler.sample( + 1L, edges, EdgeDirection.OUT, 2); + Map allCounts = edges.stream().collect(Collectors.groupingBy( + IEdge::getTargetId, Collectors.counting())); + Map sampledCounts = sampled.stream().collect(Collectors.groupingBy( + IEdge::getTargetId, Collectors.counting())); + + Assert.assertEquals(sampledCounts.size(), 2); + for (Map.Entry entry : sampledCounts.entrySet()) { + Assert.assertEquals(entry.getValue(), allCounts.get(entry.getKey())); + } + } + + @Test + public void testOneDegreeStateExposesOneHopSampling() { + TrackingIterator iterator = new TrackingIterator(Arrays.asList( + edge(1L, 2L), edge(1L, 3L), edge(1L, 4L)).iterator()); + OneDegreeGraph oneDegreeGraph = new OneDegreeGraph<>(1L, + new ValueVertex<>(1L, "vertex"), iterator); + + List> sampled = oneDegreeGraph.sampleNeighbors( + EdgeDirection.OUT, 2); + + Assert.assertEquals(sampled.size(), 2); + Assert.assertTrue(iterator.closed); + Assert.assertEquals(oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2), sampled); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testOneDegreeStateRejectsDifferentRequestAfterIteratorConsumption() { + OneDegreeGraph oneDegreeGraph = new OneDegreeGraph<>(1L, + new ValueVertex<>(1L, "vertex"), new TrackingIterator(Arrays.asList( + edge(1L, 2L), edge(1L, 3L), edge(1L, 4L)).iterator())); + + oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2); + oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 1); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testOneDegreeStateRejectsDifferentSamplingVersion() { + OneDegreeGraph oneDegreeGraph = new OneDegreeGraph<>(1L, + new ValueVertex<>(1L, "vertex"), new TrackingIterator(Arrays.asList( + edge(1L, 2L), edge(1L, 3L), edge(1L, 4L)).iterator())); + + oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2, 10L, 17L, 1L); + oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2, 10L, 17L, 2L); + } + + @Test + public void testNeighborhoodMatchesSnapshotAndSamplingVersion() { + LocalNeighborhood neighborhood = new LocalNeighborhood<>( + new ValueVertex<>(1L, "vertex"), Arrays.asList(edge(1L, 2L)), 7L, 3L); + + Assert.assertTrue(neighborhood.matches(7L, 3L)); + Assert.assertFalse(neighborhood.matches(8L, 3L)); + Assert.assertFalse(neighborhood.matches(7L, 4L)); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testRejectsCandidateEdgeOverflow() { + DeterministicNeighborSampler.sample(1L, + Arrays.asList(edge(1L, 2L), edge(1L, 3L)), EdgeDirection.OUT, -1, + Long::compare, 1); + } + + private static IEdge edge(long source, long target) { + return edgeWithValue(source, target, "value"); + } + + private static IEdge edgeWithValue(long source, long target, String value) { + ValueEdge edge = new ValueEdge<>(source, target, value); + edge.setDirect(EdgeDirection.OUT); + return edge; + } + + private static List targetIds(List> edges) { + return edges.stream().map(IEdge::getTargetId).collect(Collectors.toList()); + } + + private static class TrackingIterator implements CloseableIterator> { + + private final Iterator> delegate; + private boolean closed; + + private TrackingIterator(Iterator> delegate) { + this.delegate = delegate; + } + + @Override + public void close() { + closed = true; + } + + @Override + public boolean hasNext() { + return delegate.hasNext(); + } + + @Override + public IEdge next() { + return delegate.next(); + } + } +} From 10db9f5e5f690b90ec6b57377e149d13e73baac8 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 12 Aug 2026 12:16:21 +0800 Subject: [PATCH 2/9] feat(sampling): assemble layered sampled subgraphs --- .../sampling/LayeredSubgraphAssembler.java | 167 ++++++++++++++ .../api/graph/sampling/LogicalEdgeId.java | 87 ++++++++ .../api/graph/sampling/SampledSubgraph.java | 207 ++++++++++++++++++ .../SubgraphSamplingLimitException.java | 29 +++ .../sampling/SubgraphSamplingRequest.java | 46 ++++ .../sampling/SubgraphSamplingResponse.java | 54 +++++ .../LayeredSubgraphAssemblerTest.java | 180 +++++++++++++++ 7 files changed, 770 insertions(+) create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssembler.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LogicalEdgeId.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingLimitException.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingRequest.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingResponse.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssembler.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssembler.java new file mode 100644 index 000000000..7536737f6 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssembler.java @@ -0,0 +1,167 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import org.apache.geaflow.state.sampling.LocalNeighborhood; + +/** + * Window-local assembler. It is intentionally not serializable or checkpointed; reusable data is + * held by each vertex as {@link LocalNeighborhood} instead. + */ +public class LayeredSubgraphAssembler { + + private final Map> assemblies = new LinkedHashMap<>(); + private final long maxSampledNodes; + private final long maxSampledEdges; + private final Comparator idComparator; + + public LayeredSubgraphAssembler() { + this(SubgraphSamplingSpec.DEFAULT_MAX_SAMPLED_NODES, + SubgraphSamplingSpec.DEFAULT_MAX_SAMPLED_EDGES, null); + } + + public LayeredSubgraphAssembler(long maxSampledNodes, long maxSampledEdges, + Comparator idComparator) { + if (maxSampledNodes < 1 || maxSampledEdges < 1) { + throw new IllegalArgumentException("sampling limits must be greater than zero"); + } + this.maxSampledNodes = maxSampledNodes; + this.maxSampledEdges = maxSampledEdges; + this.idComparator = idComparator; + } + + public void start(K rootId, int maxDepth, LocalNeighborhood rootNeighborhood) { + Objects.requireNonNull(rootId, "rootId"); + Objects.requireNonNull(rootNeighborhood, "rootNeighborhood"); + if (maxDepth < 1) { + throw new IllegalArgumentException("maxDepth must be greater than zero"); + } + if (!Objects.equals(rootId, rootNeighborhood.getVertex().getId())) { + throw new IllegalArgumentException("root neighborhood vertex id does not match rootId"); + } + if (assemblies.containsKey(rootId)) { + throw new IllegalStateException("sampling assembly already exists for rootId=" + rootId); + } + SampledSubgraph subgraph = new SampledSubgraph<>(rootId, + rootNeighborhood.getSnapshotVersion(), maxSampledNodes, maxSampledEdges, idComparator); + subgraph.addNeighborhood(0, rootNeighborhood, maxDepth > 0); + Assembly assembly = new Assembly<>(maxDepth, subgraph); + assembly.minDepthByVertex.put(rootId, 0); + assembly.completedVertices.add(rootId); + assemblies.put(rootId, assembly); + } + + public boolean registerRequest(K rootId, K vertexId, int depth) { + Objects.requireNonNull(vertexId, "vertexId"); + Assembly assembly = requireAssembly(rootId); + validateDepth(depth, assembly.maxDepth); + Integer knownDepth = assembly.minDepthByVertex.get(vertexId); + if (knownDepth != null && knownDepth <= depth) { + return false; + } + if (knownDepth == null && assembly.minDepthByVertex.size() + 1L > maxSampledNodes) { + throw new SubgraphSamplingLimitException(rootId, "nodes", + assembly.minDepthByVertex.size() + 1L, maxSampledNodes); + } + assembly.minDepthByVertex.put(vertexId, depth); + assembly.completedVertices.remove(vertexId); + return true; + } + + public boolean add(SubgraphSamplingResponse response) { + Objects.requireNonNull(response, "response"); + Assembly assembly = assemblies.get(response.getRootId()); + if (assembly == null) { + return false; + } + validateDepth(response.getDepth(), assembly.maxDepth); + K vertexId = response.getNeighborhood().getVertex().getId(); + Integer knownDepth = assembly.minDepthByVertex.get(vertexId); + if (knownDepth == null) { + throw new IllegalStateException(String.format( + "sampling response was not requested, rootId=%s, vertexId=%s", + response.getRootId(), vertexId)); + } + if (response.getDepth() > knownDepth || assembly.completedVertices.contains(vertexId)) { + return false; + } + if (response.getDepth() < knownDepth) { + throw new IllegalStateException(String.format( + "sampling response depth precedes request, rootId=%s, vertexId=%s, depth=%s", + response.getRootId(), vertexId, response.getDepth())); + } + assembly.subgraph.addNeighborhood(response.getDepth(), response.getNeighborhood(), + response.getDepth() < assembly.maxDepth); + assembly.completedVertices.add(vertexId); + return true; + } + + public SampledSubgraph take(K rootId) { + Assembly assembly = assemblies.remove(rootId); + if (assembly == null) { + return null; + } + Set pending = new LinkedHashSet<>(assembly.minDepthByVertex.keySet()); + pending.removeAll(assembly.completedVertices); + if (!pending.isEmpty()) { + throw new IllegalStateException("sampling responses missing for rootId=" + rootId + + ", vertexIds=" + pending); + } + assembly.subgraph.validateComplete(); + return assembly.subgraph; + } + + public void clear() { + assemblies.clear(); + } + + private Assembly requireAssembly(K rootId) { + Assembly assembly = assemblies.get(Objects.requireNonNull(rootId, "rootId")); + if (assembly == null) { + throw new IllegalStateException("sampling assembly does not exist for rootId=" + rootId); + } + return assembly; + } + + private void validateDepth(int depth, int maxDepth) { + if (depth < 1 || depth > maxDepth) { + throw new IllegalArgumentException("sampling depth is outside configured range: " + depth); + } + } + + private static class Assembly { + + private final int maxDepth; + private final SampledSubgraph subgraph; + private final Map minDepthByVertex = new LinkedHashMap<>(); + private final Set completedVertices = new LinkedHashSet<>(); + + private Assembly(int maxDepth, SampledSubgraph subgraph) { + this.maxDepth = maxDepth; + this.subgraph = subgraph; + } + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LogicalEdgeId.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LogicalEdgeId.java new file mode 100644 index 000000000..448c0b524 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/LogicalEdgeId.java @@ -0,0 +1,87 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.Objects; +import org.apache.geaflow.model.graph.IGraphElementWithLabelField; +import org.apache.geaflow.model.graph.IGraphElementWithTimeField; +import org.apache.geaflow.model.graph.edge.IEdge; + +/** Stable identity of a normalized logical edge, independent of storage replicas and values. */ +public final class LogicalEdgeId implements Serializable { + + private final K sourceId; + private final K targetId; + private final String label; + private final Long time; + + public LogicalEdgeId(K sourceId, K targetId, String label, Long time) { + this.sourceId = Objects.requireNonNull(sourceId, "sourceId"); + this.targetId = Objects.requireNonNull(targetId, "targetId"); + this.label = label; + this.time = time; + } + + public static LogicalEdgeId fromNormalized(IEdge edge) { + Objects.requireNonNull(edge, "edge"); + String edgeLabel = edge instanceof IGraphElementWithLabelField + ? ((IGraphElementWithLabelField) edge).getLabel() : null; + Long edgeTime = edge instanceof IGraphElementWithTimeField + ? ((IGraphElementWithTimeField) edge).getTime() : null; + return new LogicalEdgeId<>(edge.getSrcId(), edge.getTargetId(), edgeLabel, edgeTime); + } + + public K getSourceId() { + return sourceId; + } + + public K getTargetId() { + return targetId; + } + + public String getLabel() { + return label; + } + + public Long getTime() { + return time; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof LogicalEdgeId)) { + return false; + } + LogicalEdgeId that = (LogicalEdgeId) other; + return Objects.equals(sourceId, that.sourceId) + && Objects.equals(targetId, that.targetId) + && Objects.equals(label, that.label) + && Objects.equals(time, that.time); + } + + @Override + public int hashCode() { + return Objects.hash(sourceId, targetId, label, time); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java new file mode 100644 index 000000000..9187da687 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java @@ -0,0 +1,207 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import org.apache.geaflow.model.graph.IGraphElementWithLabelField; +import org.apache.geaflow.model.graph.IGraphElementWithTimeField; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.vertex.IVertex; +import org.apache.geaflow.state.sampling.LocalNeighborhood; + +/** + * Short-lived layered view assembled from reusable one-hop vertex state. + */ +public class SampledSubgraph implements Serializable { + + private final K rootId; + private final long snapshotVersion; + private final Map> vertices = new LinkedHashMap<>(); + private final List>> edgeLayers = new ArrayList<>(); + private final Set> edgeIdentities = new LinkedHashSet<>(); + private final long maxSampledNodes; + private final long maxSampledEdges; + private final transient Comparator idComparator; + + public SampledSubgraph(K rootId, long snapshotVersion) { + this(rootId, snapshotVersion, SubgraphSamplingSpec.DEFAULT_MAX_SAMPLED_NODES, + SubgraphSamplingSpec.DEFAULT_MAX_SAMPLED_EDGES, null); + } + + public SampledSubgraph(K rootId, long snapshotVersion, long maxSampledNodes, + long maxSampledEdges, Comparator idComparator) { + this.rootId = Objects.requireNonNull(rootId, "rootId"); + this.snapshotVersion = snapshotVersion; + this.maxSampledNodes = maxSampledNodes; + this.maxSampledEdges = maxSampledEdges; + this.idComparator = idComparator; + } + + public K getRootId() { + return rootId; + } + + public long getSnapshotVersion() { + return snapshotVersion; + } + + public void addVertex(IVertex vertex) { + Objects.requireNonNull(vertex, "vertex"); + K vertexId = Objects.requireNonNull(vertex.getId(), "vertexId"); + if (!vertices.containsKey(vertexId) && vertices.size() + 1L > maxSampledNodes) { + throw new SubgraphSamplingLimitException(rootId, "nodes", + vertices.size() + 1L, maxSampledNodes); + } + vertices.put(vertexId, vertex); + } + + public void addLayer(List> edges) { + Objects.requireNonNull(edges, "edges"); + List> layer = new ArrayList<>(); + for (IEdge edge : edges) { + addEdge(layer, edge); + } + edgeLayers.add(layer); + } + + public void addNeighborhood(int depth, LocalNeighborhood neighborhood, + boolean includeEdges) { + if (depth < 0) { + throw new IllegalArgumentException("sampling neighborhood depth must not be negative"); + } + Objects.requireNonNull(neighborhood, "neighborhood"); + if (neighborhood.getSnapshotVersion() != snapshotVersion) { + throw new IllegalArgumentException("neighborhood snapshot does not match assembly snapshot"); + } + addVertex(neighborhood.getVertex()); + if (!includeEdges) { + return; + } + while (edgeLayers.size() <= depth) { + edgeLayers.add(new ArrayList<>()); + } + List> layer = edgeLayers.get(depth); + for (IEdge edge : neighborhood.getEdges()) { + addEdge(layer, edge); + } + } + + public Map> getVertices() { + if (idComparator == null || vertices.size() < 2) { + return Collections.unmodifiableMap(vertices); + } + List ids = new ArrayList<>(vertices.keySet()); + ids.sort((left, right) -> { + if (Objects.equals(left, rootId)) { + return Objects.equals(right, rootId) ? 0 : -1; + } + if (Objects.equals(right, rootId)) { + return 1; + } + return idComparator.compare(left, right); + }); + Map> ordered = new LinkedHashMap<>(); + for (K id : ids) { + ordered.put(id, vertices.get(id)); + } + return Collections.unmodifiableMap(ordered); + } + + public List>> getEdgeLayers() { + List>> layers = new ArrayList<>(edgeLayers.size()); + for (List> layer : edgeLayers) { + List> ordered = new ArrayList<>(layer); + if (idComparator != null) { + ordered.sort(this::compareEdges); + } + layers.add(Collections.unmodifiableList(ordered)); + } + return Collections.unmodifiableList(layers); + } + + public long getEdgeCount() { + return edgeIdentities.size(); + } + + public void validateComplete() { + for (List> layer : edgeLayers) { + for (IEdge edge : layer) { + if (!vertices.containsKey(edge.getSrcId()) || !vertices.containsKey(edge.getTargetId())) { + throw new IllegalStateException(String.format( + "sampled subgraph contains dangling edge, rootId=%s, srcId=%s, targetId=%s", + rootId, edge.getSrcId(), edge.getTargetId())); + } + } + } + } + + private void addEdge(List> layer, IEdge edge) { + Objects.requireNonNull(edge, "edge"); + Objects.requireNonNull(edge.getSrcId(), "edge.srcId"); + Objects.requireNonNull(edge.getTargetId(), "edge.targetId"); + // Incoming storage copies and outgoing copies represent the same logical edge after normalization. + LogicalEdgeId identity = LogicalEdgeId.fromNormalized(edge); + if (edgeIdentities.add(identity)) { + if (edgeIdentities.size() > maxSampledEdges) { + edgeIdentities.remove(identity); + throw new SubgraphSamplingLimitException(rootId, "edges", + edgeIdentities.size() + 1L, maxSampledEdges); + } + layer.add(edge); + } + } + + private int compareEdges(IEdge left, IEdge right) { + int result = idComparator.compare(left.getSrcId(), right.getSrcId()); + if (result == 0) { + result = idComparator.compare(left.getTargetId(), right.getTargetId()); + } + if (result == 0) { + result = String.valueOf(labelOf(left)).compareTo(String.valueOf(labelOf(right))); + } + if (result == 0) { + result = String.valueOf(timeOf(left)).compareTo(String.valueOf(timeOf(right))); + } + if (result == 0) { + result = String.valueOf(left.getValue()).compareTo(String.valueOf(right.getValue())); + } + return result; + } + + private static String labelOf(IEdge edge) { + return edge instanceof IGraphElementWithLabelField + ? ((IGraphElementWithLabelField) edge).getLabel() : null; + } + + private static Long timeOf(IEdge edge) { + return edge instanceof IGraphElementWithTimeField + ? ((IGraphElementWithTimeField) edge).getTime() : null; + } + +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingLimitException.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingLimitException.java new file mode 100644 index 000000000..9be028436 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingLimitException.java @@ -0,0 +1,29 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +/** Raised when a sampled subgraph would exceed its per-root resource budget. */ +public class SubgraphSamplingLimitException extends IllegalStateException { + + public SubgraphSamplingLimitException(Object rootId, String resource, long actual, long limit) { + super(String.format("subgraph sampling limit exceeded, rootId=%s, resource=%s, " + + "actual=%s, limit=%s", rootId, resource, actual, limit)); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingRequest.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingRequest.java new file mode 100644 index 000000000..ce6b0e074 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingRequest.java @@ -0,0 +1,46 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.Objects; + +/** A request to expand one vertex for one sampling layer. */ +public class SubgraphSamplingRequest implements Serializable { + + private final K rootId; + private final int depth; + + public SubgraphSamplingRequest(K rootId, int depth) { + if (depth < 1) { + throw new IllegalArgumentException("sampling request depth must be greater than zero"); + } + this.rootId = Objects.requireNonNull(rootId, "rootId"); + this.depth = depth; + } + + public K getRootId() { + return rootId; + } + + public int getDepth() { + return depth; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingResponse.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingResponse.java new file mode 100644 index 000000000..b845dc753 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingResponse.java @@ -0,0 +1,54 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.Objects; +import org.apache.geaflow.state.sampling.LocalNeighborhood; + +/** One vertex's local sampling result returned to the requesting root. */ +public class SubgraphSamplingResponse implements Serializable { + + private final K rootId; + private final int depth; + private final LocalNeighborhood neighborhood; + + public SubgraphSamplingResponse(K rootId, int depth, + LocalNeighborhood neighborhood) { + if (depth < 1) { + throw new IllegalArgumentException("sampling response depth must be greater than zero"); + } + this.rootId = Objects.requireNonNull(rootId, "rootId"); + this.depth = depth; + this.neighborhood = Objects.requireNonNull(neighborhood, "neighborhood"); + } + + public K getRootId() { + return rootId; + } + + public int getDepth() { + return depth; + } + + public LocalNeighborhood getNeighborhood() { + return neighborhood; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java new file mode 100644 index 000000000..378f5979e --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java @@ -0,0 +1,180 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueLabelEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueLabelTimeEdge; +import org.apache.geaflow.model.graph.vertex.impl.ValueVertex; +import org.apache.geaflow.state.sampling.LocalNeighborhood; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class LayeredSubgraphAssemblerTest { + + @Test + public void testAssemblesOneHopPerLayerAndReleasesResult() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + LocalNeighborhood root = neighborhood(1L, 2L, 7L); + LocalNeighborhood firstHop = neighborhood(2L, 3L, 7L); + LocalNeighborhood frontier = neighborhood(3L, 4L, 7L); + + assembler.start(1L, 2, root); + Assert.assertTrue(assembler.registerRequest(1L, 2L, 1)); + assembler.add(new SubgraphSamplingResponse<>(1L, 1, firstHop)); + Assert.assertTrue(assembler.registerRequest(1L, 3L, 2)); + assembler.add(new SubgraphSamplingResponse<>(1L, 2, frontier)); + + SampledSubgraph subgraph = assembler.take(1L); + Assert.assertEquals(subgraph.getVertices().size(), 3); + Assert.assertEquals(subgraph.getEdgeLayers().size(), 2); + Assert.assertEquals(subgraph.getEdgeLayers().get(0).get(0).getTargetId(), Long.valueOf(2L)); + Assert.assertEquals(subgraph.getEdgeLayers().get(1).get(0).getTargetId(), Long.valueOf(3L)); + Assert.assertNull(assembler.take(1L)); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsNeighborhoodFromPreviousWindow() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 1, neighborhood(1L, 2L, 7L)); + assembler.registerRequest(1L, 2L, 1); + assembler.add(new SubgraphSamplingResponse<>(1L, 1, neighborhood(2L, 3L, 6L))); + } + + @Test + public void testDoesNotDuplicateLogicalEdgeAcrossLayers() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 2, neighborhood(1L, 2L, 7L)); + assembler.registerRequest(1L, 2L, 1); + assembler.add(new SubgraphSamplingResponse<>(1L, 1, + neighborhoodWithEdge(2L, 1L, 2L, 7L))); + + SampledSubgraph subgraph = assembler.take(1L); + long edgeCount = subgraph.getEdgeLayers().stream().mapToLong(List::size).sum(); + Assert.assertEquals(edgeCount, 1L); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsNeighborhoodFromFutureWindow() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 1, neighborhood(1L, 2L, 7L)); + assembler.registerRequest(1L, 2L, 1); + assembler.add(new SubgraphSamplingResponse<>(1L, 1, neighborhood(2L, 3L, 8L))); + } + + @Test + public void testCycleRegistersEachVertexOnlyOnce() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 3, neighborhood(1L, 2L, 7L)); + + Assert.assertTrue(assembler.registerRequest(1L, 2L, 1)); + Assert.assertFalse(assembler.registerRequest(1L, 1L, 2)); + Assert.assertFalse(assembler.registerRequest(1L, 2L, 2)); + assembler.add(new SubgraphSamplingResponse<>(1L, 1, + neighborhoodWithEdge(2L, 2L, 1L, 7L))); + + Assert.assertEquals(assembler.take(1L).getVertices().size(), 2); + } + + @Test(expectedExceptions = SubgraphSamplingLimitException.class) + public void testRejectsNodeBudgetOverflow() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(1, 10, Long::compare); + assembler.start(1L, 1, neighborhood(1L, 2L, 7L)); + assembler.registerRequest(1L, 2L, 1); + } + + @Test(expectedExceptions = SubgraphSamplingLimitException.class) + public void testRejectsEdgeBudgetOverflow() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(10, 1, Long::compare); + LocalNeighborhood root = new LocalNeighborhood<>( + new ValueVertex<>(1L, 1), java.util.Arrays.asList( + new ValueEdge<>(1L, 2L, 1, EdgeDirection.OUT), + new ValueEdge<>(1L, 3L, 1, EdgeDirection.OUT)), 7L); + assembler.start(1L, 1, root); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testRejectsMissingResponse() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 1, neighborhood(1L, 2L, 7L)); + assembler.registerRequest(1L, 2L, 1); + assembler.take(1L); + } + + @Test + public void testStructuralEdgeIdentityPreservesCollisionsAndLabels() { + SampledSubgraph subgraph = new SampledSubgraph<>("root", 1L); + subgraph.addLayer(java.util.Arrays.asList( + new ValueLabelEdge<>("a->b", "c", 1, "first"), + new ValueLabelEdge<>("a", "b->c", 1, "first"), + new ValueLabelEdge<>("a", "b->c", 1, "second"))); + + Assert.assertEquals(subgraph.getEdgeCount(), 3L); + } + + @Test + public void testLogicalEdgeIdentityIgnoresReplicaDirectionAndValue() { + ValueLabelEdge out = new ValueLabelEdge<>(1L, 2L, "first", "knows"); + out.setDirect(EdgeDirection.OUT); + ValueLabelEdge inReplica = new ValueLabelEdge<>(1L, 2L, "second", "knows"); + inReplica.setDirect(EdgeDirection.IN); + ValueLabelEdge reciprocal = new ValueLabelEdge<>(2L, 1L, "third", "knows"); + + SampledSubgraph subgraph = new SampledSubgraph<>(1L, 1L); + subgraph.addLayer(java.util.Arrays.asList(out, inReplica, reciprocal)); + + Assert.assertEquals(subgraph.getEdgeCount(), 2L); + } + + @Test + public void testLogicalEdgeIdentityPreservesTemporalParallelEdges() { + SampledSubgraph subgraph = new SampledSubgraph<>(1L, 1L); + subgraph.addLayer(java.util.Arrays.asList( + new ValueLabelTimeEdge<>(1L, 2L, "same", "knows", 10L), + new ValueLabelTimeEdge<>(1L, 2L, "same", "knows", 11L))); + + Assert.assertEquals(subgraph.getEdgeCount(), 2L); + } + + private LocalNeighborhood neighborhood(long source, long target, + long version) { + return neighborhoodWithEdge(source, source, target, version); + } + + private LocalNeighborhood neighborhoodWithEdge(long vertexId, + long source, + long target, + long version) { + ValueVertex vertex = new ValueVertex<>(vertexId, (int) vertexId); + ValueEdge edge = new ValueEdge<>(source, target, 1, EdgeDirection.OUT); + return new LocalNeighborhood<>(vertex, Collections.singletonList(edge), version); + } +} From 48cb41022498275566c50dd0837e99fca08fd5db Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 12 Aug 2026 12:16:41 +0800 Subject: [PATCH 3/9] feat(sampling): add iterative multi-hop protocol --- .../graph/sampling/EmptySamplingRequest.java | 43 +++ .../graph/sampling/EmptySamplingResponse.java | 43 +++ .../sampling/IterativeSamplingState.java | 103 ++++++ .../graph/sampling/NeighborStateRequest.java | 51 +++ .../graph/sampling/NeighborStateResponse.java | 55 +++ .../graph/sampling/PendingSamplingRound.java | 99 +++++ .../api/graph/sampling/SamplingClock.java | 146 ++++++++ .../api/graph/sampling/SamplingMessage.java | 28 ++ .../api/graph/sampling/SamplingPhase.java | 28 ++ .../sampling/SamplingResponseCollector.java | 105 ++++++ .../sampling/IterativeSamplingStateTest.java | 98 +++++ .../IterativeSubgraphSamplingE2ETest.java | 342 ++++++++++++++++++ .../sampling/PendingSamplingRoundTest.java | 116 ++++++ .../api/graph/sampling/SamplingClockTest.java | 64 ++++ 14 files changed, 1321 insertions(+) create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java new file mode 100644 index 000000000..76eb3f261 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java @@ -0,0 +1,43 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Objects; + +/** First half of the empty-neighborhood self barrier. */ +public final class EmptySamplingRequest implements SamplingMessage { + + private final SamplingClock clock; + private final K vertexId; + + public EmptySamplingRequest(SamplingClock clock, K vertexId) { + this.clock = NeighborStateRequest.requirePhase(clock, SamplingPhase.REQUEST); + this.vertexId = Objects.requireNonNull(vertexId, "vertexId"); + } + + @Override + public SamplingClock getClock() { + return clock; + } + + public K getVertexId() { + return vertexId; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java new file mode 100644 index 000000000..63af66ed2 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java @@ -0,0 +1,43 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Objects; + +/** Second half of the empty-neighborhood self barrier. */ +public final class EmptySamplingResponse implements SamplingMessage { + + private final SamplingClock clock; + private final K vertexId; + + public EmptySamplingResponse(SamplingClock clock, K vertexId) { + this.clock = NeighborStateRequest.requirePhase(clock, SamplingPhase.RESPOND); + this.vertexId = Objects.requireNonNull(vertexId, "vertexId"); + } + + @Override + public SamplingClock getClock() { + return clock; + } + + public K getVertexId() { + return vertexId; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java new file mode 100644 index 000000000..33007c602 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java @@ -0,0 +1,103 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.Objects; + +/** Per-vertex committed payload and bounded in-flight state for one sampling session. */ +public final class IterativeSamplingState implements Serializable { + + private final long snapshotVersion; + private final long sessionId; + private int completedHop; + private P committedPayload; + private PendingSamplingRound pendingRound; + + public IterativeSamplingState(long snapshotVersion, long sessionId, P initialPayload) { + this.snapshotVersion = snapshotVersion; + this.sessionId = sessionId; + this.committedPayload = initialPayload; + } + + public void startRound(PendingSamplingRound pending) { + Objects.requireNonNull(pending, "pending"); + SamplingClock clock = pending.getRequestClock(); + requireSession(clock); + if (pendingRound != null) { + throw new IllegalStateException("a sampling round is already pending"); + } + if (clock.getHop() != completedHop + 1) { + throw new IllegalArgumentException("sampling request hop does not follow committed state"); + } + this.pendingRound = pending; + } + + public NeighborStateResponse respond(K responderId, NeighborStateRequest request) { + Objects.requireNonNull(request, "request"); + requireSession(request.getClock()); + if (request.getClock().getHop() != completedHop + 1) { + throw new IllegalStateException("requested payload is not committed for the preceding hop"); + } + return new NeighborStateResponse<>(request.getClock().responseClock(), + request.getRequesterId(), responderId, committedPayload); + } + + public void commit(SamplingClock commitClock, P payload) { + Objects.requireNonNull(commitClock, "commitClock"); + requireSession(commitClock); + if (commitClock.getPhase() != SamplingPhase.COMMIT_AND_REQUEST + && commitClock.getPhase() != SamplingPhase.COMPLETE) { + throw new IllegalArgumentException("sampling state can only commit in a commit phase"); + } + if (pendingRound == null || !pendingRound.getRequestClock().isSameRound(commitClock)) { + throw new IllegalStateException("sampling commit does not match the pending round"); + } + this.completedHop = commitClock.getHop(); + this.committedPayload = payload; + this.pendingRound = null; + } + + private void requireSession(SamplingClock clock) { + if (clock.getSnapshotVersion() != snapshotVersion || clock.getSessionId() != sessionId) { + throw new IllegalArgumentException("sampling clock does not match vertex session"); + } + } + + public long getSnapshotVersion() { + return snapshotVersion; + } + + public long getSessionId() { + return sessionId; + } + + public int getCompletedHop() { + return completedHop; + } + + public P getCommittedPayload() { + return committedPayload; + } + + public PendingSamplingRound getPendingRound() { + return pendingRound; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java new file mode 100644 index 000000000..41529d482 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java @@ -0,0 +1,51 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Objects; + +/** Request for a neighbor's state committed at the preceding hop. */ +public final class NeighborStateRequest implements SamplingMessage { + + private final SamplingClock clock; + private final K requesterId; + + public NeighborStateRequest(SamplingClock clock, K requesterId) { + this.clock = requirePhase(clock, SamplingPhase.REQUEST); + this.requesterId = Objects.requireNonNull(requesterId, "requesterId"); + } + + static SamplingClock requirePhase(SamplingClock clock, SamplingPhase phase) { + Objects.requireNonNull(clock, "clock"); + if (clock.getPhase() != phase) { + throw new IllegalArgumentException("sampling message requires phase " + phase); + } + return clock; + } + + @Override + public SamplingClock getClock() { + return clock; + } + + public K getRequesterId() { + return requesterId; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java new file mode 100644 index 000000000..05e41f937 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java @@ -0,0 +1,55 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Objects; + +/** State returned by one selected neighbor. */ +public final class NeighborStateResponse implements SamplingMessage { + + private final SamplingClock clock; + private final K requesterId; + private final K responderId; + private final P payload; + + public NeighborStateResponse(SamplingClock clock, K requesterId, K responderId, P payload) { + this.clock = NeighborStateRequest.requirePhase(clock, SamplingPhase.RESPOND); + this.requesterId = Objects.requireNonNull(requesterId, "requesterId"); + this.responderId = Objects.requireNonNull(responderId, "responderId"); + this.payload = payload; + } + + @Override + public SamplingClock getClock() { + return clock; + } + + public K getRequesterId() { + return requesterId; + } + + public K getResponderId() { + return responderId; + } + + public P getPayload() { + return payload; + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java new file mode 100644 index 000000000..042e8ffb3 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java @@ -0,0 +1,99 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import org.apache.geaflow.model.graph.edge.IEdge; + +/** Bounded requester-side state retained between request and commit supersteps. */ +public final class PendingSamplingRound implements Serializable { + + private final SamplingClock requestClock; + private final K requesterId; + private final Map>> edgesByNeighbor; + + public PendingSamplingRound(SamplingClock requestClock, K requesterId, + Iterable> sampledEdges) { + this.requestClock = NeighborStateRequest.requirePhase(requestClock, SamplingPhase.REQUEST); + this.requesterId = Objects.requireNonNull(requesterId, "requesterId"); + Objects.requireNonNull(sampledEdges, "sampledEdges"); + this.edgesByNeighbor = new LinkedHashMap<>(); + for (IEdge edge : sampledEdges) { + Objects.requireNonNull(edge, "edge"); + K neighborId = neighborId(requesterId, edge); + edgesByNeighbor.computeIfAbsent(neighborId, ignored -> new ArrayList<>()).add(edge); + } + } + + private K neighborId(K vertexId, IEdge edge) { + if (Objects.equals(vertexId, edge.getSrcId())) { + return Objects.requireNonNull(edge.getTargetId(), "neighborId"); + } + if (Objects.equals(vertexId, edge.getTargetId())) { + return Objects.requireNonNull(edge.getSrcId(), "neighborId"); + } + throw new IllegalArgumentException("sampled edge is not incident to requesterId=" + vertexId); + } + + public SamplingClock getRequestClock() { + return requestClock; + } + + public K getRequesterId() { + return requesterId; + } + + public boolean isEmpty() { + return edgesByNeighbor.isEmpty(); + } + + public List getNeighborIds() { + return Collections.unmodifiableList(new ArrayList<>(edgesByNeighbor.keySet())); + } + + public Map>> getEdgesByNeighbor() { + Map>> result = new LinkedHashMap<>(); + for (Map.Entry>> entry : edgesByNeighbor.entrySet()) { + result.put(entry.getKey(), Collections.unmodifiableList(entry.getValue())); + } + return Collections.unmodifiableMap(result); + } + + public Map> createRequests() { + Map> requests = new LinkedHashMap<>(); + for (K neighborId : edgesByNeighbor.keySet()) { + requests.put(neighborId, new NeighborStateRequest<>(requestClock, requesterId)); + } + return Collections.unmodifiableMap(requests); + } + + public EmptySamplingRequest createEmptyRequest() { + if (!isEmpty()) { + throw new IllegalStateException("only an empty sampling round uses an empty request"); + } + return new EmptySamplingRequest<>(requestClock, requesterId); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java new file mode 100644 index 000000000..c34174d17 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java @@ -0,0 +1,146 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; +import java.util.Objects; + +/** Immutable logical clock mapped from the runtime BSP iteration. */ +public final class SamplingClock implements Serializable { + + private final long snapshotVersion; + private final long sessionId; + private final int hop; + private final SamplingPhase phase; + + public SamplingClock(long snapshotVersion, long sessionId, int hop, SamplingPhase phase) { + if (hop < 1) { + throw new IllegalArgumentException("sampling hop must be greater than zero"); + } + this.snapshotVersion = snapshotVersion; + this.sessionId = sessionId; + this.hop = hop; + this.phase = Objects.requireNonNull(phase, "phase"); + } + + public static SamplingClock forIteration(long snapshotVersion, long sessionId, int maxHops, + long startIterationId, long iterationId) { + if (maxHops < 1) { + throw new IllegalArgumentException("maxHops must be greater than zero"); + } + if (iterationId < startIterationId) { + throw new IllegalArgumentException("iteration precedes the sampling session"); + } + long offset = iterationId - startIterationId; + if (offset == 0L) { + return new SamplingClock(snapshotVersion, sessionId, 1, SamplingPhase.REQUEST); + } + if ((offset & 1L) == 1L) { + long responseHop = (offset + 1L) / 2L; + requireHopInRange(responseHop, maxHops, iterationId); + return new SamplingClock(snapshotVersion, sessionId, (int) responseHop, + SamplingPhase.RESPOND); + } + long completedHop = offset / 2L; + requireHopInRange(completedHop, maxHops, iterationId); + SamplingPhase phase = completedHop == maxHops + ? SamplingPhase.COMPLETE : SamplingPhase.COMMIT_AND_REQUEST; + return new SamplingClock(snapshotVersion, sessionId, (int) completedHop, phase); + } + + public static long requiredIterations(int maxHops) { + if (maxHops < 1) { + throw new IllegalArgumentException("maxHops must be greater than zero"); + } + return Math.addExact(Math.multiplyExact((long) maxHops, 2L), 1L); + } + + private static void requireHopInRange(long hop, int maxHops, long iterationId) { + if (hop < 1L || hop > maxHops) { + throw new IllegalArgumentException("iteration is outside the sampling session: " + + iterationId); + } + } + + public SamplingClock responseClock() { + if (phase != SamplingPhase.REQUEST) { + throw new IllegalStateException("only a request clock can create a response clock"); + } + return new SamplingClock(snapshotVersion, sessionId, hop, SamplingPhase.RESPOND); + } + + public SamplingClock nextRequestClock() { + if (phase != SamplingPhase.COMMIT_AND_REQUEST) { + throw new IllegalStateException("current clock does not start another sampling hop"); + } + return new SamplingClock(snapshotVersion, sessionId, Math.addExact(hop, 1), + SamplingPhase.REQUEST); + } + + public boolean isSameRound(SamplingClock other) { + return other != null && snapshotVersion == other.snapshotVersion + && sessionId == other.sessionId && hop == other.hop; + } + + public long getSamplingVersion() { + long value = mix64(snapshotVersion) ^ Long.rotateLeft(mix64(sessionId), 21); + return mix64(value ^ Long.rotateLeft(mix64(hop), 42)); + } + + private static long mix64(long value) { + value = (value ^ (value >>> 30)) * 0xbf58476d1ce4e5b9L; + value = (value ^ (value >>> 27)) * 0x94d049bb133111ebL; + return value ^ (value >>> 31); + } + + public long getSnapshotVersion() { + return snapshotVersion; + } + + public long getSessionId() { + return sessionId; + } + + public int getHop() { + return hop; + } + + public SamplingPhase getPhase() { + return phase; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof SamplingClock)) { + return false; + } + SamplingClock that = (SamplingClock) other; + return snapshotVersion == that.snapshotVersion && sessionId == that.sessionId + && hop == that.hop && phase == that.phase; + } + + @Override + public int hashCode() { + return Objects.hash(snapshotVersion, sessionId, hop, phase); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java new file mode 100644 index 000000000..604c58510 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java @@ -0,0 +1,28 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.Serializable; + +/** Marker for clocked iterative sampling messages. */ +public interface SamplingMessage extends Serializable { + + SamplingClock getClock(); +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java new file mode 100644 index 000000000..602454fcc --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java @@ -0,0 +1,28 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +/** Physical BSP phase for one logical sampling hop. */ +public enum SamplingPhase { + REQUEST, + RESPOND, + COMMIT_AND_REQUEST, + COMPLETE +} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java new file mode 100644 index 000000000..a0d7e9251 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java @@ -0,0 +1,105 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** Validates and orders all responses for one pending sampling round. */ +public final class SamplingResponseCollector { + + private final SamplingClock requestClock; + private final K requesterId; + private final List expectedNeighbors; + private final Map> responses = new LinkedHashMap<>(); + private boolean emptyResponseReceived; + + public SamplingResponseCollector(PendingSamplingRound pending) { + Objects.requireNonNull(pending, "pending"); + this.requestClock = pending.getRequestClock(); + this.requesterId = pending.getRequesterId(); + this.expectedNeighbors = pending.getNeighborIds(); + } + + public void add(NeighborStateResponse response) { + Objects.requireNonNull(response, "response"); + requireResponseRound(response.getClock()); + if (!Objects.equals(requesterId, response.getRequesterId())) { + throw new IllegalArgumentException("sampling response requester does not match pending round"); + } + K responderId = response.getResponderId(); + if (!expectedNeighbors.contains(responderId)) { + throw new IllegalArgumentException("sampling response came from an unrequested neighbor: " + + responderId); + } + if (responses.putIfAbsent(responderId, response) != null) { + throw new IllegalStateException("duplicate sampling response from neighbor: " + responderId); + } + } + + public void addEmpty(EmptySamplingResponse response) { + Objects.requireNonNull(response, "response"); + requireResponseRound(response.getClock()); + if (!expectedNeighbors.isEmpty()) { + throw new IllegalStateException("non-empty sampling round cannot accept an empty response"); + } + if (!Objects.equals(requesterId, response.getVertexId())) { + throw new IllegalArgumentException("empty sampling response vertex does not match requester"); + } + if (emptyResponseReceived) { + throw new IllegalStateException("duplicate empty sampling response"); + } + emptyResponseReceived = true; + } + + private void requireResponseRound(SamplingClock responseClock) { + NeighborStateRequest.requirePhase(responseClock, SamplingPhase.RESPOND); + if (!requestClock.isSameRound(responseClock)) { + throw new IllegalArgumentException("sampling response clock does not match pending round"); + } + } + + public boolean isComplete() { + return expectedNeighbors.isEmpty() ? emptyResponseReceived + : responses.size() == expectedNeighbors.size(); + } + + public void validateComplete() { + if (!isComplete()) { + List missing = new ArrayList<>(expectedNeighbors); + missing.removeAll(responses.keySet()); + throw new IllegalStateException("sampling responses are incomplete, requesterId=" + + requesterId + ", missing=" + missing); + } + } + + public List> getResponses() { + validateComplete(); + List> ordered = new ArrayList<>(expectedNeighbors.size()); + for (K neighborId : expectedNeighbors) { + ordered.add(responses.get(neighborId)); + } + return Collections.unmodifiableList(ordered); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java new file mode 100644 index 000000000..462550057 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java @@ -0,0 +1,98 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class IterativeSamplingStateTest { + + @Test + public void testAllVerticesCommitTwoHopsWithoutAccumulatingSubgraph() { + Map>> adjacency = new LinkedHashMap<>(); + adjacency.put(1L, Arrays.asList(edge(1L, 2L))); + adjacency.put(2L, Arrays.asList(edge(2L, 1L), edge(2L, 3L))); + adjacency.put(3L, Arrays.asList(edge(3L, 2L))); + Map> states = new LinkedHashMap<>(); + states.put(1L, new IterativeSamplingState<>(7L, 11L, 1)); + states.put(2L, new IterativeSamplingState<>(7L, 11L, 2)); + states.put(3L, new IterativeSamplingState<>(7L, 11L, 3)); + + SamplingClock request = SamplingClock.forIteration(7L, 11L, 2, 1L, 1L); + for (int hop = 1; hop <= 2; hop++) { + Map nextPayloads = new LinkedHashMap<>(); + for (Map.Entry> entry + : states.entrySet()) { + Long requesterId = entry.getKey(); + PendingSamplingRound pending = new PendingSamplingRound<>(request, + requesterId, adjacency.get(requesterId)); + entry.getValue().startRound(pending); + Assert.assertTrue(pending.getNeighborIds().size() <= 2); + + SamplingResponseCollector collector = + new SamplingResponseCollector<>(pending); + for (Map.Entry> outbound + : pending.createRequests().entrySet()) { + collector.add(states.get(outbound.getKey()).respond(outbound.getKey(), + outbound.getValue())); + } + int next = collector.getResponses().stream() + .mapToInt(NeighborStateResponse::getPayload).sum(); + nextPayloads.put(requesterId, next); + } + + SamplingClock commit = SamplingClock.forIteration(7L, 11L, 2, 1L, hop * 2L + 1L); + for (Map.Entry payload : nextPayloads.entrySet()) { + states.get(payload.getKey()).commit(commit, payload.getValue()); + } + if (hop < 2) { + request = commit.nextRequestClock(); + } + } + + Assert.assertEquals(states.get(1L).getCommittedPayload(), Integer.valueOf(4)); + Assert.assertEquals(states.get(2L).getCommittedPayload(), Integer.valueOf(4)); + Assert.assertEquals(states.get(3L).getCommittedPayload(), Integer.valueOf(4)); + for (IterativeSamplingState state : states.values()) { + Assert.assertEquals(state.getCompletedHop(), 2); + Assert.assertNull(state.getPendingRound()); + } + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testCannotServeNextHopBeforePreviousHopCommit() { + IterativeSamplingState state = + new IterativeSamplingState<>(7L, 11L, 1); + NeighborStateRequest request = new NeighborStateRequest<>( + new SamplingClock(7L, 11L, 2, SamplingPhase.REQUEST), 2L); + state.respond(1L, request); + } + + private IEdge edge(long source, long target) { + return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java new file mode 100644 index 000000000..ba5bd83ff --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java @@ -0,0 +1,342 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class IterativeSubgraphSamplingE2ETest { + + private static final long SNAPSHOT_VERSION = 7L; + private static final long SESSION_ID = 11L; + private static final long START_ITERATION_ID = 21L; + + @Test + public void testRoutesTwoHopSamplingAcrossAllBspPhases() { + Map>> adjacency = new LinkedHashMap<>(); + adjacency.put(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L))); + adjacency.put(2L, Collections.singletonList(edge(2L, 4L))); + adjacency.put(3L, Collections.singletonList(edge(3L, 5L))); + adjacency.put(4L, Collections.emptyList()); + adjacency.put(5L, Collections.emptyList()); + SubgraphSamplingSpec spec = new SubgraphSamplingSpec( + 2, -1, EdgeDirection.OUT, 100L, 17L); + InMemorySamplingScheduler scheduler = new InMemorySamplingScheduler(adjacency, spec); + + List phases = scheduler.run(); + + Assert.assertEquals(phases, Arrays.asList( + SamplingPhase.REQUEST, + SamplingPhase.RESPOND, + SamplingPhase.COMMIT_AND_REQUEST, + SamplingPhase.RESPOND, + SamplingPhase.COMPLETE)); + Assert.assertEquals(scheduler.getRootPayloads(), Arrays.asList( + vertexIds(1L, 2L, 3L), + vertexIds(1L, 2L, 3L, 4L, 5L))); + Assert.assertEquals(scheduler.getNeighborRequestCount(), 8); + Assert.assertEquals(scheduler.getNeighborResponseCount(), 8); + Assert.assertEquals(scheduler.getEmptyRequestCount(), 4); + Assert.assertEquals(scheduler.getEmptyResponseCount(), 4); + Assert.assertEquals(scheduler.getCommitCount(), 10); + Assert.assertEquals(scheduler.getSamplingCallCount(), 10); + scheduler.assertComplete(); + } + + @Test + public void testDrivesCompleteThreeHopMessageSchedule() { + Map>> adjacency = new LinkedHashMap<>(); + adjacency.put(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L))); + adjacency.put(2L, Collections.singletonList(edge(2L, 4L))); + adjacency.put(3L, Collections.singletonList(edge(3L, 5L))); + adjacency.put(4L, Collections.singletonList(edge(4L, 6L))); + adjacency.put(5L, Collections.singletonList(edge(5L, 7L))); + adjacency.put(6L, Collections.emptyList()); + adjacency.put(7L, Collections.emptyList()); + SubgraphSamplingSpec spec = new SubgraphSamplingSpec( + 3, -1, EdgeDirection.OUT, 100L, 17L); + InMemorySamplingScheduler scheduler = new InMemorySamplingScheduler(adjacency, spec); + + Assert.assertEquals(scheduler.run(), Arrays.asList( + SamplingPhase.REQUEST, + SamplingPhase.RESPOND, + SamplingPhase.COMMIT_AND_REQUEST, + SamplingPhase.RESPOND, + SamplingPhase.COMMIT_AND_REQUEST, + SamplingPhase.RESPOND, + SamplingPhase.COMPLETE)); + Assert.assertEquals(scheduler.getRootPayloads(), Arrays.asList( + vertexIds(1L, 2L, 3L), + vertexIds(1L, 2L, 3L, 4L, 5L), + vertexIds(1L, 2L, 3L, 4L, 5L, 6L, 7L))); + Assert.assertEquals(scheduler.getRootMessageTrace(), Arrays.asList( + "request[1] 1->2", + "request[1] 1->3", + "response[1] 2->1", + "response[1] 3->1", + "request[2] 1->2", + "request[2] 1->3", + "response[2] 2->1", + "response[2] 3->1", + "request[3] 1->2", + "request[3] 1->3", + "response[3] 2->1", + "response[3] 3->1")); + Assert.assertEquals(new LinkedHashSet<>(scheduler.getRootSamplingVersions()).size(), 3, + "each hop must use a distinct sampling version"); + + // Every vertex participates in every hop, including explicit empty rounds for leaves. + Assert.assertEquals(scheduler.getNeighborRequestCount(), 18); + Assert.assertEquals(scheduler.getNeighborResponseCount(), 18); + Assert.assertEquals(scheduler.getEmptyRequestCount(), 6); + Assert.assertEquals(scheduler.getEmptyResponseCount(), 6); + Assert.assertEquals(scheduler.getCommitCount(), 21); + Assert.assertEquals(scheduler.getSamplingCallCount(), 21); + scheduler.assertComplete(); + } + + private static Set vertexIds(Long... ids) { + return new LinkedHashSet<>(Arrays.asList(ids)); + } + + private static IEdge edge(long source, long target) { + return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); + } + + private static final class InMemorySamplingScheduler { + + private final Map>> adjacency; + private final SubgraphSamplingSpec spec; + private final Map>> states = + new LinkedHashMap<>(); + private final List> rootPayloads = new ArrayList<>(); + private final List rootMessageTrace = new ArrayList<>(); + private final List rootSamplingVersions = new ArrayList<>(); + private int neighborRequestCount; + private int neighborResponseCount; + private int emptyRequestCount; + private int emptyResponseCount; + private int commitCount; + private int samplingCallCount; + + private InMemorySamplingScheduler(Map>> adjacency, + SubgraphSamplingSpec spec) { + this.adjacency = adjacency; + this.spec = spec; + for (Long vertexId : adjacency.keySet()) { + states.put(vertexId, new IterativeSamplingState<>( + SNAPSHOT_VERSION, SESSION_ID, vertexIds(vertexId))); + } + } + + private List run() { + List phases = new ArrayList<>(); + Map> inbox = Collections.emptyMap(); + long iterations = SamplingClock.requiredIterations(spec.getHops()); + for (long offset = 0L; offset < iterations; offset++) { + long iterationId = START_ITERATION_ID + offset; + SamplingClock clock = SamplingClock.forIteration(SNAPSHOT_VERSION, SESSION_ID, + spec.getHops(), START_ITERATION_ID, iterationId); + phases.add(clock.getPhase()); + switch (clock.getPhase()) { + case REQUEST: + Assert.assertTrue(inbox.isEmpty()); + inbox = startRounds(clock); + break; + case RESPOND: + inbox = respond(clock, inbox); + break; + case COMMIT_AND_REQUEST: + commit(clock, inbox); + inbox = startRounds(clock.nextRequestClock()); + break; + case COMPLETE: + commit(clock, inbox); + inbox = Collections.emptyMap(); + break; + default: + throw new IllegalStateException("unsupported sampling phase: " + + clock.getPhase()); + } + } + Assert.assertTrue(inbox.isEmpty()); + return phases; + } + + private Map> startRounds(SamplingClock requestClock) { + Map> requests = new LinkedHashMap<>(); + for (Map.Entry>> entry + : states.entrySet()) { + Long requesterId = entry.getKey(); + List> sampled = DeterministicNeighborSampler.sample( + requesterId, adjacency.get(requesterId), spec.getDirection(), spec.getFanout(), + Comparator.naturalOrder(), spec.getMaxReturnedEdges(), spec.getSeed(), + requestClock.getSamplingVersion()); + samplingCallCount++; + PendingSamplingRound pending = new PendingSamplingRound<>( + requestClock, requesterId, sampled); + entry.getValue().startRound(pending); + if (requesterId.equals(1L)) { + rootSamplingVersions.add(requestClock.getSamplingVersion()); + } + if (pending.isEmpty()) { + route(requests, requesterId, pending.createEmptyRequest()); + emptyRequestCount++; + } else { + for (Map.Entry> request + : pending.createRequests().entrySet()) { + route(requests, request.getKey(), request.getValue()); + if (requesterId.equals(1L)) { + rootMessageTrace.add("request[" + requestClock.getHop() + "] " + + requesterId + "->" + request.getKey()); + } + neighborRequestCount++; + } + } + } + return requests; + } + + private Map> respond( + SamplingClock responseClock, Map> requests) { + Map> responses = new LinkedHashMap<>(); + for (Map.Entry> inbox : requests.entrySet()) { + Long responderId = inbox.getKey(); + IterativeSamplingState> responder = states.get(responderId); + Assert.assertNotNull(responder, "message routed to an unknown vertex"); + for (SamplingMessage message : inbox.getValue()) { + Assert.assertTrue(message.getClock().isSameRound(responseClock)); + if (message instanceof NeighborStateRequest) { + NeighborStateRequest request = (NeighborStateRequest) message; + route(responses, request.getRequesterId(), + responder.respond(responderId, request)); + if (request.getRequesterId().equals(1L)) { + rootMessageTrace.add("response[" + responseClock.getHop() + "] " + + responderId + "->" + request.getRequesterId()); + } + neighborResponseCount++; + } else if (message instanceof EmptySamplingRequest) { + EmptySamplingRequest request = (EmptySamplingRequest) message; + Assert.assertEquals(request.getVertexId(), responderId); + route(responses, responderId, new EmptySamplingResponse<>( + request.getClock().responseClock(), responderId)); + emptyResponseCount++; + } else { + throw new IllegalStateException("unexpected sampling request: " + message); + } + } + } + return responses; + } + + private void commit(SamplingClock commitClock, + Map> responses) { + Map> nextPayloads = new LinkedHashMap<>(); + for (Map.Entry>> entry + : states.entrySet()) { + Long requesterId = entry.getKey(); + PendingSamplingRound pending = entry.getValue().getPendingRound(); + SamplingResponseCollector> collector = + new SamplingResponseCollector<>(pending); + for (SamplingMessage message + : responses.getOrDefault(requesterId, Collections.emptyList())) { + if (message instanceof NeighborStateResponse) { + collector.add((NeighborStateResponse>) message); + } else if (message instanceof EmptySamplingResponse) { + collector.addEmpty((EmptySamplingResponse) message); + } else { + throw new IllegalStateException("unexpected sampling response: " + message); + } + } + Set nextPayload = vertexIds(requesterId); + for (NeighborStateResponse> response : collector.getResponses()) { + nextPayload.addAll(response.getPayload()); + } + nextPayloads.put(requesterId, nextPayload); + } + for (Map.Entry> payload : nextPayloads.entrySet()) { + states.get(payload.getKey()).commit(commitClock, payload.getValue()); + commitCount++; + } + rootPayloads.add(new LinkedHashSet<>(states.get(1L).getCommittedPayload())); + } + + private void route(Map> messages, Long destination, + SamplingMessage message) { + messages.computeIfAbsent(destination, ignored -> new ArrayList<>()).add(message); + } + + private void assertComplete() { + for (IterativeSamplingState> state : states.values()) { + Assert.assertEquals(state.getCompletedHop(), spec.getHops()); + Assert.assertNull(state.getPendingRound()); + } + } + + private List> getRootPayloads() { + return rootPayloads; + } + + private List getRootMessageTrace() { + return rootMessageTrace; + } + + private List getRootSamplingVersions() { + return rootSamplingVersions; + } + + private int getNeighborRequestCount() { + return neighborRequestCount; + } + + private int getNeighborResponseCount() { + return neighborResponseCount; + } + + private int getEmptyRequestCount() { + return emptyRequestCount; + } + + private int getEmptyResponseCount() { + return emptyResponseCount; + } + + private int getCommitCount() { + return commitCount; + } + + private int getSamplingCallCount() { + return samplingCallCount; + } + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java new file mode 100644 index 000000000..e28be164a --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java @@ -0,0 +1,116 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.io.ByteArrayOutputStream; +import java.io.ObjectOutputStream; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class PendingSamplingRoundTest { + + @Test + public void testGroupsParallelEdgesAndOrdersResponsesBySampledNeighbor() { + PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, + Arrays.asList(edge(1L, 2L, "first"), edge(1L, 2L, "parallel"), + edge(1L, 3L, "third"))); + + Assert.assertEquals(pending.getNeighborIds(), Arrays.asList(2L, 3L)); + Assert.assertEquals(pending.getEdgesByNeighbor().get(2L).size(), 2); + Assert.assertEquals(pending.createRequests().keySet(), + new java.util.LinkedHashSet<>(Arrays.asList(2L, 3L))); + + SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); + collector.add(new NeighborStateResponse<>(requestClock().responseClock(), 1L, 3L, "three")); + collector.add(new NeighborStateResponse<>(requestClock().responseClock(), 1L, 2L, "two")); + + List> responses = collector.getResponses(); + Assert.assertEquals(responses.get(0).getResponderId(), Long.valueOf(2L)); + Assert.assertEquals(responses.get(1).getResponderId(), Long.valueOf(3L)); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testRejectsMissingResponseAtCommit() { + PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, + Arrays.asList(edge(1L, 2L, "first"), edge(1L, 3L, "second"))); + SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); + collector.add(new NeighborStateResponse<>(requestClock().responseClock(), 1L, 2L, "two")); + collector.validateComplete(); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testRejectsDuplicateResponse() { + PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, + Collections.singletonList(edge(1L, 2L, "first"))); + SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); + NeighborStateResponse response = new NeighborStateResponse<>( + requestClock().responseClock(), 1L, 2L, "two"); + collector.add(response); + collector.add(response); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsResponseFromAnotherSession() { + PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, + Collections.singletonList(edge(1L, 2L, "first"))); + SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); + SamplingClock stale = new SamplingClock(7L, 12L, 1, SamplingPhase.RESPOND); + collector.add(new NeighborStateResponse<>(stale, 1L, 2L, "two")); + } + + @Test + public void testEmptyRoundUsesTwoPhaseSelfBarrier() { + PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, + Collections.emptyList()); + SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); + + EmptySamplingRequest request = pending.createEmptyRequest(); + Assert.assertEquals(request.getVertexId(), Long.valueOf(1L)); + collector.addEmpty(new EmptySamplingResponse<>(request.getClock().responseClock(), 1L)); + Assert.assertTrue(collector.isComplete()); + Assert.assertTrue(collector.getResponses().isEmpty()); + } + + @Test + public void testProtocolStateIsSerializable() throws Exception { + PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, + Collections.singletonList(edge(1L, 2L, "first"))); + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + try (ObjectOutputStream output = new ObjectOutputStream(bytes)) { + output.writeObject(pending); + output.writeObject(pending.createRequests().get(2L)); + } + Assert.assertTrue(bytes.size() > 0); + } + + private SamplingClock requestClock() { + return new SamplingClock(7L, 11L, 1, SamplingPhase.REQUEST); + } + + private IEdge edge(long source, long target, String value) { + return new ValueEdge<>(source, target, value, EdgeDirection.OUT); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java new file mode 100644 index 000000000..24c62bd53 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java @@ -0,0 +1,64 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import org.testng.Assert; +import org.testng.annotations.Test; + +public class SamplingClockTest { + + @Test + public void testMapsTwoHopsToFiveIterations() { + SamplingClock firstRequest = SamplingClock.forIteration(7L, 11L, 2, 1L, 1L); + SamplingClock firstResponse = SamplingClock.forIteration(7L, 11L, 2, 1L, 2L); + SamplingClock firstCommit = SamplingClock.forIteration(7L, 11L, 2, 1L, 3L); + SamplingClock secondResponse = SamplingClock.forIteration(7L, 11L, 2, 1L, 4L); + SamplingClock complete = SamplingClock.forIteration(7L, 11L, 2, 1L, 5L); + + assertClock(firstRequest, 1, SamplingPhase.REQUEST); + assertClock(firstResponse, 1, SamplingPhase.RESPOND); + assertClock(firstCommit, 1, SamplingPhase.COMMIT_AND_REQUEST); + assertClock(firstCommit.nextRequestClock(), 2, SamplingPhase.REQUEST); + assertClock(secondResponse, 2, SamplingPhase.RESPOND); + assertClock(complete, 2, SamplingPhase.COMPLETE); + Assert.assertEquals(SamplingClock.requiredIterations(2), 5L); + } + + @Test + public void testSamplingVersionChangesPerHopButNotPhase() { + SamplingClock request = new SamplingClock(7L, 11L, 1, SamplingPhase.REQUEST); + SamplingClock response = request.responseClock(); + SamplingClock next = new SamplingClock(7L, 11L, 2, SamplingPhase.REQUEST); + + Assert.assertEquals(request.getSamplingVersion(), response.getSamplingVersion()); + Assert.assertNotEquals(request.getSamplingVersion(), next.getSamplingVersion()); + Assert.assertTrue(request.isSameRound(response)); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsIterationAfterSessionCompletion() { + SamplingClock.forIteration(7L, 11L, 2, 1L, 6L); + } + + private void assertClock(SamplingClock clock, int hop, SamplingPhase phase) { + Assert.assertEquals(clock.getHop(), hop); + Assert.assertEquals(clock.getPhase(), phase); + } +} From 7936c3d610ffba6fd06204158bece381d46f6339 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 12 Aug 2026 12:16:52 +0800 Subject: [PATCH 4/9] fix(sampling): track edge targets in dynamic cache --- .../dynamic/cache/TemporaryGraphCache.java | 1 + .../cache/TemporaryGraphCacheTest.java | 37 +++++++++++++++++++ 2 files changed, 38 insertions(+) create mode 100644 geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/test/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCacheTest.java diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/main/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCache.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/main/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCache.java index c6030489c..f19cfade1 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/main/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCache.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/main/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCache.java @@ -55,6 +55,7 @@ public IVertex getVertex(K vId) { public void addEdge(IEdge edge) { this.vertexIds.add(edge.getSrcId()); + this.vertexIds.add(edge.getTargetId()); List> edges = this.vertexEdges.getOrDefault(edge.getSrcId(), new ArrayList<>()); edges.add(edge); diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/test/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCacheTest.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/test/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCacheTest.java new file mode 100644 index 000000000..1f3cd1568 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-operator/src/test/java/org/apache/geaflow/operator/impl/graph/compute/dynamic/cache/TemporaryGraphCacheTest.java @@ -0,0 +1,37 @@ +/* + * 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 org.apache.geaflow.operator.impl.graph.compute.dynamic.cache; + +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class TemporaryGraphCacheTest { + + @Test + public void testEdgeTriggersSourceAndTargetVertices() { + TemporaryGraphCache cache = new TemporaryGraphCache<>(); + cache.addEdge(new ValueEdge<>(1L, 2L, 1, EdgeDirection.OUT)); + + Assert.assertTrue(cache.getAllEvolveVId().contains(1L)); + Assert.assertTrue(cache.getAllEvolveVId().contains(2L)); + } +} From 052cb3c54ced447942686cdfd2a52050951d413f Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 12 Aug 2026 12:17:06 +0800 Subject: [PATCH 5/9] feat(sampling): integrate sampling with DSL runtime --- .../algo/AlgorithmSamplingRuntimeContext.java | 82 ++++++++++++ .../common/algo/AlgorithmUserFunction.java | 4 + .../algo/SubgraphSamplingAlgorithm.java | 24 ++++ .../GeaFlowAlgorithmAggTraversalFunction.java | 14 ++- ...wAlgorithmDynamicAggTraversalFunction.java | 46 ++++++- ...GeaFlowAlgorithmDynamicRuntimeContext.java | 60 ++++++++- .../GeaFlowAlgorithmRuntimeContext.java | 52 +++++++- ...FlowAlgorithmAggTraversalFunctionTest.java | 43 +++++++ ...orithmDynamicAggTraversalFunctionTest.java | 117 ++++++++++++++++++ ...lowAlgorithmDynamicRuntimeContextTest.java | 113 +++++++++++++++++ 10 files changed, 546 insertions(+), 9 deletions(-) create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java create mode 100644 geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java diff --git a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java new file mode 100644 index 000000000..2c90440bb --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java @@ -0,0 +1,82 @@ +/* + * 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 org.apache.geaflow.dsl.common.algo; + +import org.apache.geaflow.api.graph.sampling.SamplingClock; +import org.apache.geaflow.api.graph.sampling.SamplingPhase; +import org.apache.geaflow.api.graph.sampling.SubgraphSamplingSpec; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.RowVertex; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.state.sampling.LocalNeighborhood; + +/** Runtime-facing contract for reusable one-hop sampling. */ +public interface AlgorithmSamplingRuntimeContext extends AlgorithmRuntimeContext { + + LocalNeighborhood sampleOneHop(RowVertex vertex, EdgeDirection direction, + int fanout); + + default LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout, + long maxCandidateEdges) { + LocalNeighborhood neighborhood = sampleOneHop(vertex, direction, fanout); + if (neighborhood.getEdges().size() > maxCandidateEdges) { + throw new IllegalStateException(String.format( + "one-hop sampling edge limit exceeded, vertexId=%s, actual=%s, limit=%s", + vertex.getId(), neighborhood.getEdges().size(), maxCandidateEdges)); + } + return neighborhood; + } + + default LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout, + long maxReturnedEdges, + long seed, + long samplingVersion) { + return sampleOneHop(vertex, direction, fanout, maxReturnedEdges); + } + + default SamplingClock getSamplingClock(SubgraphSamplingSpec spec, long sessionId, + long startIterationId) { + return SamplingClock.forIteration(getSamplingSnapshotVersion(), sessionId, + spec.getHops(), startIterationId, getCurrentIterationId()); + } + + default LocalNeighborhood sampleOneHop(RowVertex vertex, + SubgraphSamplingSpec spec, + SamplingClock requestClock) { + if (requestClock.getPhase() != SamplingPhase.REQUEST) { + throw new IllegalArgumentException("one-hop sampling requires a request clock"); + } + if (requestClock.getSnapshotVersion() != getSamplingSnapshotVersion()) { + throw new IllegalArgumentException("sampling clock does not match runtime snapshot"); + } + return sampleOneHop(vertex, spec.getDirection(), spec.getFanout(), + spec.getMaxReturnedEdges(), spec.getSeed(), requestClock.getSamplingVersion()); + } + + long getSamplingSnapshotVersion(); + + default long getNeighborhoodChangeVersion(Object vertexId) { + return Long.MIN_VALUE; + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java index 4058ff6f6..e125ad8a8 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java @@ -60,6 +60,10 @@ public interface AlgorithmUserFunction extends Serializable { default void finish() { } + /** Called before an iteration starts processing vertices and messages. */ + default void initIteration(long iterationId) { + } + /** * Finish Iteration method called after each iteration finished. */ diff --git a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java new file mode 100644 index 000000000..3c6756ff5 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java @@ -0,0 +1,24 @@ +/* + * 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 org.apache.geaflow.dsl.common.algo; + +/** Marks an algorithm whose first sampling iteration requires a stable window snapshot. */ +public interface SubgraphSamplingAlgorithm { +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java index daca980db..217409bd2 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java @@ -27,6 +27,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; +import org.apache.geaflow.api.function.iterator.RichIteratorFunction; import org.apache.geaflow.api.graph.function.vc.VertexCentricAggTraversalFunction; import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; import org.apache.geaflow.dsl.common.data.Row; @@ -43,7 +44,8 @@ import org.apache.geaflow.utils.keygroup.KeyGroupAssignment; public class GeaFlowAlgorithmAggTraversalFunction implements - VertexCentricAggTraversalFunction { + VertexCentricAggTraversalFunction, + RichIteratorFunction { private static final String STATE_SUFFIX = "UpdatedValueState"; @@ -145,6 +147,16 @@ public void close() { algorithmCtx.close(); } + @Override + public void initIteration(long iterationId) { + userFunction.initIteration(iterationId); + } + + @Override + public void finishIteration(long iterationId) { + userFunction.finishIteration(iterationId); + } + @Override public void initContext(VertexCentricAggContext aggContext) { this.algorithmCtx.setAggContext(Objects.requireNonNull(aggContext)); diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java index 98c475b15..72b2d33b5 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java @@ -33,6 +33,7 @@ import org.apache.geaflow.api.graph.function.vc.IncVertexCentricAggTraversalFunction; import org.apache.geaflow.common.config.keys.FrameworkConfigKeys; import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; +import org.apache.geaflow.dsl.common.algo.SubgraphSamplingAlgorithm; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowVertex; import org.apache.geaflow.dsl.common.types.GraphSchema; @@ -57,6 +58,8 @@ public class GeaFlowAlgorithmDynamicAggTraversalFunction private static final Logger LOGGER = LoggerFactory.getLogger(GeaFlowAlgorithmDynamicAggTraversalFunction.class); private static final String STATE_SUFFIX = "UpdatedValueState"; + private static final String NEIGHBORHOOD_CHANGE_VERSION_STATE_SUFFIX = + "NeighborhoodChangeVersionState"; private final AlgorithmUserFunction userFunction; @@ -64,15 +67,16 @@ public class GeaFlowAlgorithmDynamicAggTraversalFunction private GraphSchema graphSchema; - private IncVertexCentricTraversalFuncContext traversalContext; + private transient IncVertexCentricTraversalFuncContext traversalContext; - private GeaFlowAlgorithmDynamicRuntimeContext algorithmCtx; + private transient GeaFlowAlgorithmDynamicRuntimeContext algorithmCtx; - private MutableGraph mutableGraph; + private transient MutableGraph mutableGraph; private transient Set initVertices; private transient KeyValueState vertexUpdateValues; + private transient KeyValueState neighborhoodChangeVersions; private boolean materializeInFinish; @@ -90,6 +94,11 @@ public void open( IncVertexCentricTraversalFuncContext vertexCentricFuncContext) { this.traversalContext = vertexCentricFuncContext; this.materializeInFinish = traversalContext.getRuntimeContext().getConfiguration().getBoolean(FrameworkConfigKeys.UDF_MATERIALIZE_GRAPH_IN_FINISH); + // Sampling must read a stable window snapshot. Apply the window delta before the first + // sampling iteration, then refresh only vertices triggered by that delta. + if (userFunction instanceof SubgraphSamplingAlgorithm) { + this.materializeInFinish = false; + } this.algorithmCtx = new GeaFlowAlgorithmDynamicRuntimeContext(this, traversalContext, graphSchema); this.initVertices = new HashSet<>(); @@ -109,12 +118,21 @@ public void open( IKeyGroupAssigner keyGroupAssigner = KeyGroupAssignerFactory.createKeyGroupAssigner( keyGroup, taskIndex, maxParallelism); descriptor.withKeyGroupAssigner(keyGroupAssigner); - long recoverWindowId = traversalContext.getRuntimeContext().getWindowId(); + final long recoverWindowId = traversalContext.getRuntimeContext().getWindowId(); this.vertexUpdateValues = StateFactory.buildKeyValueState(descriptor, traversalContext.getRuntimeContext().getConfiguration()); + KeyValueStateDescriptor changeVersionDescriptor = KeyValueStateDescriptor.build( + traversalContext.getTraversalOpName() + "_" + NEIGHBORHOOD_CHANGE_VERSION_STATE_SUFFIX, + traversalContext.getRuntimeContext().getConfiguration().getString(SYSTEM_STATE_BACKEND_TYPE)); + changeVersionDescriptor.withKeyGroup(keyGroup); + changeVersionDescriptor.withKeyGroupAssigner(keyGroupAssigner); + this.neighborhoodChangeVersions = StateFactory.buildKeyValueState(changeVersionDescriptor, + traversalContext.getRuntimeContext().getConfiguration()); if (recoverWindowId > 1) { this.vertexUpdateValues.manage().operate().setCheckpointId(recoverWindowId - 1); this.vertexUpdateValues.manage().operate().recover(); + this.neighborhoodChangeVersions.manage().operate().setCheckpointId(recoverWindowId - 1); + this.neighborhoodChangeVersions.manage().operate().recover(); } } @@ -129,6 +147,9 @@ public void init(ITraversalRequest traversalRequest) { // false when called after the first time to avoid redundant invocation. if (vertexId != null && needInit(vertexId)) { RowVertex vertex = (RowVertex) algorithmCtx.loadVertex(); + if (vertex == null) { + vertex = (RowVertex) algorithmCtx.getIncVCTraversalCtx().getTemporaryGraph().getVertex(); + } if (vertex != null) { algorithmCtx.setVertexId(vertex.getId()); Row newValue = getVertexNewValue(vertex.getId()); @@ -146,8 +167,15 @@ public Row getVertexNewValue(Object vertexId) { return vertexUpdateValues.get(vertexId); } + public long getNeighborhoodChangeVersion(Object vertexId) { + Long version = neighborhoodChangeVersions.get(vertexId); + return version == null ? Long.MIN_VALUE : version; + } + @Override public void evolve(Object vertexId, TemporaryGraph temporaryGraph) { + neighborhoodChangeVersions.put(vertexId, + traversalContext.getRuntimeContext().getWindowId()); if (!materializeInFinish) { IVertex vertex = temporaryGraph.getVertex(); List> edges = temporaryGraph.getEdges(); @@ -182,6 +210,9 @@ public void compute(Object vertexId, Iterator messages) { } } else { vertex = (RowVertex) algorithmCtx.loadVertex(); + if (vertex == null) { + vertex = (RowVertex) algorithmCtx.getIncVCTraversalCtx().getTemporaryGraph().getVertex(); + } } if (vertex != null) { Row newValue = getVertexNewValue(vertex.getId()); @@ -193,6 +224,9 @@ public void compute(Object vertexId, Iterator messages) { public void finish(Object vertexId, MutableGraph mutableGraph) { algorithmCtx.setVertexId(vertexId); RowVertex graphVertex = (RowVertex) algorithmCtx.loadVertex(); + if (graphVertex == null) { + graphVertex = (RowVertex) algorithmCtx.getIncVCTraversalCtx().getTemporaryGraph().getVertex(); + } if (graphVertex != null) { Row newValue = getVertexNewValue(graphVertex.getId()); userFunction.finish(graphVertex, Optional.ofNullable(newValue)); @@ -221,6 +255,9 @@ public void finish() { this.vertexUpdateValues.manage().operate().setCheckpointId(windowId); this.vertexUpdateValues.manage().operate().finish(); this.vertexUpdateValues.manage().operate().archive(); + this.neighborhoodChangeVersions.manage().operate().setCheckpointId(windowId); + this.neighborhoodChangeVersions.manage().operate().finish(); + this.neighborhoodChangeVersions.manage().operate().archive(); } @@ -231,6 +268,7 @@ public void close() { @Override public void initIteration(long iterationId) { + userFunction.initIteration(iterationId); } @Override diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java index d929ae441..191632cd4 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java @@ -20,6 +20,7 @@ package org.apache.geaflow.dsl.runtime.engine; import java.util.ArrayList; +import java.util.Comparator; import java.util.List; import java.util.Objects; import org.apache.geaflow.api.graph.function.aggregate.VertexCentricAggContextFunction.VertexCentricAggContext; @@ -29,9 +30,10 @@ import org.apache.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.TraversalVertexQuery; import org.apache.geaflow.common.config.Configuration; import org.apache.geaflow.common.iterator.CloseableIterator; -import org.apache.geaflow.dsl.common.algo.AlgorithmRuntimeContext; +import org.apache.geaflow.dsl.common.algo.AlgorithmSamplingRuntimeContext; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowEdge; +import org.apache.geaflow.dsl.common.data.RowVertex; import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; import org.apache.geaflow.dsl.common.types.GraphSchema; import org.apache.geaflow.dsl.runtime.traversal.message.ITraversalAgg; @@ -45,8 +47,10 @@ import org.apache.geaflow.state.pushdown.filter.IFilter; import org.apache.geaflow.state.pushdown.filter.InEdgeFilter; import org.apache.geaflow.state.pushdown.filter.OutEdgeFilter; +import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; +import org.apache.geaflow.state.sampling.LocalNeighborhood; -public class GeaFlowAlgorithmDynamicRuntimeContext implements AlgorithmRuntimeContext { +public class GeaFlowAlgorithmDynamicRuntimeContext implements AlgorithmSamplingRuntimeContext { private final IncVertexCentricTraversalFuncContext incVCTraversalCtx; @@ -80,6 +84,10 @@ public void setVertexId(Object vertexId) { this.edgeQuery.withId(vertexId); } + public Object getVertexId() { + return vertexId; + } + public IVertex loadVertex() { return vertexQuery.get(); } @@ -175,6 +183,54 @@ public List loadStaticEdges(EdgeDirection direction) { } } + @Override + public LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout) { + List> sampled = DeterministicNeighborSampler.sample(vertex.getId(), + loadStaticEdges(direction), direction, fanout); + return new LocalNeighborhood<>(vertex, sampled, getSamplingSnapshotVersion()); + } + + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout, + long maxCandidateEdges) { + return sampleOneHop(vertex, direction, fanout, maxCandidateEdges, 0L, + getSamplingSnapshotVersion()); + } + + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout, + long maxReturnedEdges, + long seed, + long samplingVersion) { + try (CloseableIterator iterator = loadStaticEdgesIterator(direction)) { + Iterable iterable = () -> iterator; + Comparator comparator = (left, right) -> + ((org.apache.geaflow.common.type.IType) graphSchema.getIdType()).compare(left, right); + List> sampled = (List) DeterministicNeighborSampler.sample( + vertex.getId(), iterable, direction, fanout, comparator, maxReturnedEdges, + seed, samplingVersion); + return new LocalNeighborhood<>(vertex, sampled, getSamplingSnapshotVersion(), samplingVersion); + } + } + + @Override + public long getSamplingSnapshotVersion() { + return incVCTraversalCtx.getRuntimeContext().getWindowId(); + } + + @Override + public long getNeighborhoodChangeVersion(Object vertexId) { + return traversalFunction.getNeighborhoodChangeVersion(vertexId); + } + @Override public CloseableIterator loadStaticEdgesIterator(EdgeDirection direction) { switch (direction) { diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java index 7696b4f10..5113b85d4 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java @@ -20,6 +20,7 @@ package org.apache.geaflow.dsl.runtime.engine; import java.util.ArrayList; +import java.util.Comparator; import java.util.List; import java.util.Objects; import org.apache.geaflow.api.graph.function.aggregate.VertexCentricAggContextFunction.VertexCentricAggContext; @@ -28,21 +29,25 @@ import org.apache.geaflow.common.config.Configuration; import org.apache.geaflow.common.exception.GeaflowRuntimeException; import org.apache.geaflow.common.iterator.CloseableIterator; -import org.apache.geaflow.dsl.common.algo.AlgorithmRuntimeContext; +import org.apache.geaflow.dsl.common.algo.AlgorithmSamplingRuntimeContext; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowEdge; +import org.apache.geaflow.dsl.common.data.RowVertex; import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; import org.apache.geaflow.dsl.common.types.GraphSchema; import org.apache.geaflow.dsl.runtime.traversal.message.ITraversalAgg; import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; import org.apache.geaflow.model.traversal.ITraversalResponse; import org.apache.geaflow.model.traversal.TraversalType.ResponseType; import org.apache.geaflow.state.pushdown.filter.EmptyFilter; import org.apache.geaflow.state.pushdown.filter.IFilter; import org.apache.geaflow.state.pushdown.filter.InEdgeFilter; import org.apache.geaflow.state.pushdown.filter.OutEdgeFilter; +import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; +import org.apache.geaflow.state.sampling.LocalNeighborhood; -public class GeaFlowAlgorithmRuntimeContext implements AlgorithmRuntimeContext { +public class GeaFlowAlgorithmRuntimeContext implements AlgorithmSamplingRuntimeContext { private final VertexCentricTraversalFuncContext traversalContext; @@ -113,6 +118,49 @@ public List loadStaticEdges(EdgeDirection direction) { return loadEdges(direction); } + @Override + public LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout) { + List> edges = DeterministicNeighborSampler.sample(vertex.getId(), + loadStaticEdges(direction), direction, fanout); + return new LocalNeighborhood<>(vertex, edges, getSamplingSnapshotVersion()); + } + + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout, + long maxCandidateEdges) { + return sampleOneHop(vertex, direction, fanout, maxCandidateEdges, 0L, + getSamplingSnapshotVersion()); + } + + @Override + @SuppressWarnings({"unchecked", "rawtypes"}) + public LocalNeighborhood sampleOneHop(RowVertex vertex, + EdgeDirection direction, + int fanout, + long maxReturnedEdges, + long seed, + long samplingVersion) { + try (CloseableIterator iterator = loadStaticEdgesIterator(direction)) { + Iterable iterable = () -> iterator; + Comparator comparator = (left, right) -> + ((org.apache.geaflow.common.type.IType) graphSchema.getIdType()).compare(left, right); + List> edges = (List) DeterministicNeighborSampler.sample( + vertex.getId(), iterable, direction, fanout, comparator, maxReturnedEdges, + seed, samplingVersion); + return new LocalNeighborhood<>(vertex, edges, getSamplingSnapshotVersion(), samplingVersion); + } + } + + @Override + public long getSamplingSnapshotVersion() { + return traversalContext.getRuntimeContext().getWindowId(); + } + @Override public CloseableIterator loadStaticEdgesIterator(EdgeDirection direction) { switch (direction) { diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java new file mode 100644 index 000000000..7aadb2c90 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java @@ -0,0 +1,43 @@ +/* + * 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 org.apache.geaflow.dsl.runtime.engine; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; +import org.apache.geaflow.dsl.common.types.GraphSchema; +import org.testng.annotations.Test; + +public class GeaFlowAlgorithmAggTraversalFunctionTest { + + @Test + public void testForwardsIterationLifecycle() { + AlgorithmUserFunction userFunction = mock(AlgorithmUserFunction.class); + GeaFlowAlgorithmAggTraversalFunction function = new GeaFlowAlgorithmAggTraversalFunction( + mock(GraphSchema.class), userFunction, new Object[0]); + + function.initIteration(3L); + function.finishIteration(3L); + + verify(userFunction).initIteration(3L); + verify(userFunction).finishIteration(3L); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java new file mode 100644 index 000000000..505201c9c --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java @@ -0,0 +1,117 @@ +/* + * 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 org.apache.geaflow.dsl.runtime.engine; + +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.HashSet; +import org.apache.geaflow.api.graph.function.vc.IncVertexCentricTraversalFunction.IncVertexCentricTraversalFuncContext; +import org.apache.geaflow.api.graph.function.vc.base.IncVertexCentricFunction.TemporaryGraph; +import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.types.GraphSchema; +import org.apache.geaflow.state.KeyValueState; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class GeaFlowAlgorithmDynamicAggTraversalFunctionTest { + + @Test + public void testForwardsIterationLifecycle() { + AlgorithmUserFunction userFunction = mock(AlgorithmUserFunction.class); + GeaFlowAlgorithmDynamicAggTraversalFunction function = + new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), + userFunction, new Object[0]); + + function.initIteration(3L); + function.finishIteration(3L); + + verify(userFunction).initIteration(3L); + verify(userFunction).finishIteration(3L); + } + + @Test + public void testEvolvePersistsNeighborhoodChangeVersion() throws Exception { + GeaFlowAlgorithmDynamicAggTraversalFunction function = + new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), + mock(AlgorithmUserFunction.class), new Object[0]); + IncVertexCentricTraversalFuncContext traversalContext = + mock(IncVertexCentricTraversalFuncContext.class, RETURNS_DEEP_STUBS); + TemporaryGraph temporaryGraph = mock(TemporaryGraph.class); + KeyValueState changeVersions = mock(KeyValueState.class); + when(traversalContext.getRuntimeContext().getWindowId()).thenReturn(7L); + when(temporaryGraph.getEdges()).thenReturn(Collections.emptyList()); + when(changeVersions.get(2L)).thenReturn(5L); + setField(function, "traversalContext", traversalContext); + setField(function, "neighborhoodChangeVersions", changeVersions); + + function.evolve(2L, temporaryGraph); + + verify(changeVersions).put(2L, 7L); + Assert.assertEquals(function.getNeighborhoodChangeVersion(2L), 5L); + } + + @Test + public void testMissingChangeVersionIsStatic() throws Exception { + GeaFlowAlgorithmDynamicAggTraversalFunction function = + new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), + mock(AlgorithmUserFunction.class), new Object[0]); + setField(function, "neighborhoodChangeVersions", mock(KeyValueState.class)); + + Assert.assertEquals(function.getNeighborhoodChangeVersion(1L), Long.MIN_VALUE); + } + + @Test + public void testFinishCheckpointsNeighborhoodChangeVersions() throws Exception { + AlgorithmUserFunction userFunction = mock(AlgorithmUserFunction.class); + GeaFlowAlgorithmDynamicAggTraversalFunction function = + new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), + userFunction, new Object[0]); + IncVertexCentricTraversalFuncContext traversalContext = + mock(IncVertexCentricTraversalFuncContext.class, RETURNS_DEEP_STUBS); + KeyValueState vertexValues = mock(KeyValueState.class, RETURNS_DEEP_STUBS); + KeyValueState changeVersions = mock(KeyValueState.class, RETURNS_DEEP_STUBS); + GeaFlowAlgorithmDynamicRuntimeContext algorithmContext = + mock(GeaFlowAlgorithmDynamicRuntimeContext.class); + when(traversalContext.getRuntimeContext().getWindowId()).thenReturn(7L); + setField(function, "traversalContext", traversalContext); + setField(function, "algorithmCtx", algorithmContext); + setField(function, "initVertices", new HashSet<>()); + setField(function, "vertexUpdateValues", vertexValues); + setField(function, "neighborhoodChangeVersions", changeVersions); + + function.finish(); + + verify(changeVersions.manage().operate()).setCheckpointId(7L); + verify(changeVersions.manage().operate()).finish(); + verify(changeVersions.manage().operate()).archive(); + } + + private void setField(Object target, String name, Object value) throws Exception { + Field field = target.getClass().getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } +} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java new file mode 100644 index 000000000..e1c382126 --- /dev/null +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java @@ -0,0 +1,113 @@ +/* + * 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 org.apache.geaflow.dsl.runtime.engine; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.Arrays; +import java.util.Collections; +import org.apache.geaflow.api.context.RuntimeContext; +import org.apache.geaflow.api.graph.function.vc.IncVertexCentricTraversalFunction.IncVertexCentricTraversalFuncContext; +import org.apache.geaflow.api.graph.function.vc.IncVertexCentricTraversalFunction.TraversalGraphSnapShot; +import org.apache.geaflow.api.graph.function.vc.IncVertexCentricTraversalFunction.TraversalHistoricalGraph; +import org.apache.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.TraversalEdgeQuery; +import org.apache.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.TraversalVertexQuery; +import org.apache.geaflow.api.graph.function.vc.base.IncVertexCentricFunction.TemporaryGraph; +import org.apache.geaflow.api.graph.sampling.SamplingClock; +import org.apache.geaflow.api.graph.sampling.SamplingPhase; +import org.apache.geaflow.api.graph.sampling.SubgraphSamplingSpec; +import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; +import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.RowEdge; +import org.apache.geaflow.dsl.common.data.RowVertex; +import org.apache.geaflow.dsl.common.types.GraphSchema; +import org.apache.geaflow.dsl.common.data.impl.types.ObjectEdge; +import org.apache.geaflow.dsl.common.data.impl.ObjectRow; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.state.sampling.LocalNeighborhood; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class GeaFlowAlgorithmDynamicRuntimeContextTest { + + @Test + public void testSamplingContextMapsRuntimeClock() { + IncVertexCentricTraversalFuncContext traversalContext = mock( + IncVertexCentricTraversalFuncContext.class); + TraversalHistoricalGraph historicalGraph = mock(TraversalHistoricalGraph.class); + TraversalGraphSnapShot snapshot = mock(TraversalGraphSnapShot.class); + RuntimeContext runtimeContext = mock(RuntimeContext.class); + when(traversalContext.getHistoricalGraph()).thenReturn(historicalGraph); + when(historicalGraph.getSnapShot(0L)).thenReturn(snapshot); + when(snapshot.vertex()).thenReturn(mock(TraversalVertexQuery.class)); + when(snapshot.edges()).thenReturn(mock(TraversalEdgeQuery.class)); + when(traversalContext.getRuntimeContext()).thenReturn(runtimeContext); + when(runtimeContext.getWindowId()).thenReturn(7L); + when(traversalContext.getIterationId()).thenReturn(3L); + GeaFlowAlgorithmDynamicRuntimeContext context = new GeaFlowAlgorithmDynamicRuntimeContext( + new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), + mock(AlgorithmUserFunction.class), new Object[0]), traversalContext, + mock(GraphSchema.class)); + SubgraphSamplingSpec spec = new SubgraphSamplingSpec(2, 2, EdgeDirection.OUT, 9L, 5L); + + SamplingClock current = context.getSamplingClock(spec, 11L, 1L); + Assert.assertEquals(current.getHop(), 1); + Assert.assertEquals(current.getPhase(), SamplingPhase.COMMIT_AND_REQUEST); + } + + @Test + public void testSamplingUsesMaterializedSnapshotOnly() { + IncVertexCentricTraversalFuncContext traversalContext = mock( + IncVertexCentricTraversalFuncContext.class); + TraversalHistoricalGraph historicalGraph = mock(TraversalHistoricalGraph.class); + TraversalGraphSnapShot snapshot = mock(TraversalGraphSnapShot.class); + TraversalVertexQuery vertexQuery = mock(TraversalVertexQuery.class); + TraversalEdgeQuery edgeQuery = mock(TraversalEdgeQuery.class); + RuntimeContext runtimeContext = mock(RuntimeContext.class); + TemporaryGraph temporaryGraph = mock(TemporaryGraph.class); + + RowEdge edge = new ObjectEdge(1L, 2L, ObjectRow.create(1.0D)); + edge.setDirect(EdgeDirection.OUT); + when(traversalContext.getHistoricalGraph()).thenReturn(historicalGraph); + when(historicalGraph.getSnapShot(0L)).thenReturn(snapshot); + when(snapshot.vertex()).thenReturn(vertexQuery); + when(snapshot.edges()).thenReturn(edgeQuery); + when(edgeQuery.getOutEdges()).thenReturn(Collections.singletonList(edge)); + when(traversalContext.getRuntimeContext()).thenReturn(runtimeContext); + when(runtimeContext.getWindowId()).thenReturn(7L); + when(traversalContext.getTemporaryGraph()).thenReturn(temporaryGraph); + when(temporaryGraph.getEdges()).thenReturn(Arrays.asList(edge)); + + GeaFlowAlgorithmDynamicRuntimeContext context = new GeaFlowAlgorithmDynamicRuntimeContext( + new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), + mock(AlgorithmUserFunction.class), new Object[0]), traversalContext, mock(GraphSchema.class)); + RowVertex vertex = mock(RowVertex.class); + when(vertex.getId()).thenReturn(1L); + + LocalNeighborhood neighborhood = context.sampleOneHop(vertex, EdgeDirection.OUT, -1); + + Assert.assertEquals(neighborhood.getEdges().size(), 1); + Assert.assertEquals(neighborhood.getSnapshotVersion(), 7L); + verify(temporaryGraph, never()).getEdges(); + } +} From b0b7088952ef76ada471803d14b33430ae73297b Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 12 Aug 2026 13:26:33 +0800 Subject: [PATCH 6/9] refactor(sampling): remove duplicate one-degree sampling entrypoint --- .../geaflow/state/data/OneDegreeGraph.java | 47 -------------- .../DeterministicNeighborSamplerTest.java | 62 ------------------- 2 files changed, 109 deletions(-) diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java index d564aec3b..317c5075d 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java @@ -20,26 +20,15 @@ package org.apache.geaflow.state.data; import java.io.Serializable; -import java.util.Collections; -import java.util.List; -import java.util.Objects; import org.apache.geaflow.common.iterator.CloseableIterator; -import org.apache.geaflow.model.graph.edge.EdgeDirection; import org.apache.geaflow.model.graph.edge.IEdge; import org.apache.geaflow.model.graph.vertex.IVertex; -import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; public class OneDegreeGraph implements Serializable { private IVertex vertex; protected CloseableIterator> edgeIterator; protected K key; - private List> sampledEdges; - private EdgeDirection sampledDirection; - private Integer sampledFanout; - private Long sampledMaxReturnedEdges; - private Long sampledSeed; - private Long sampledVersion; public OneDegreeGraph(K key, IVertex vertex, CloseableIterator> edgeIterator) { this.key = key; @@ -59,40 +48,4 @@ public CloseableIterator> getEdgeIterator() { return edgeIterator; } - /** Samples this vertex's bounded one-hop neighborhood in the state layer. */ - public synchronized List> sampleNeighbors(EdgeDirection direction, int fanout) { - return sampleNeighbors(direction, fanout, - DeterministicNeighborSampler.DEFAULT_MAX_CANDIDATE_EDGES, 0L, 0L); - } - - /** Samples this one-shot edge iterator for one deterministic sampling round. */ - public synchronized List> sampleNeighbors(EdgeDirection direction, int fanout, - long maxReturnedEdges, long seed, - long samplingVersion) { - if (sampledEdges != null) { - if (sampledDirection != direction || !Objects.equals(sampledFanout, fanout) - || !Objects.equals(sampledMaxReturnedEdges, maxReturnedEdges) - || !Objects.equals(sampledSeed, seed) - || !Objects.equals(sampledVersion, samplingVersion)) { - throw new IllegalStateException( - "one-degree edge iterator was already consumed by a different sampling request"); - } - return sampledEdges; - } - try { - Iterable> iterable = () -> edgeIterator; - sampledEdges = Collections.unmodifiableList( - DeterministicNeighborSampler.sample(key, iterable, direction, fanout, - java.util.Comparator.comparing(String::valueOf), maxReturnedEdges, - seed, samplingVersion)); - sampledDirection = direction; - sampledFanout = fanout; - sampledMaxReturnedEdges = maxReturnedEdges; - sampledSeed = seed; - sampledVersion = samplingVersion; - return sampledEdges; - } finally { - edgeIterator.close(); - } - } } diff --git a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java index 6de60ab05..f3504eb05 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java @@ -21,16 +21,13 @@ import java.util.ArrayList; import java.util.Arrays; -import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.stream.Collectors; -import org.apache.geaflow.common.iterator.CloseableIterator; import org.apache.geaflow.model.graph.edge.EdgeDirection; import org.apache.geaflow.model.graph.edge.IEdge; import org.apache.geaflow.model.graph.edge.impl.ValueEdge; import org.apache.geaflow.model.graph.vertex.impl.ValueVertex; -import org.apache.geaflow.state.data.OneDegreeGraph; import org.testng.Assert; import org.testng.annotations.Test; @@ -121,41 +118,6 @@ public void testFanoutCountsNeighborsAndKeepsSelectedParallelEdges() { } } - @Test - public void testOneDegreeStateExposesOneHopSampling() { - TrackingIterator iterator = new TrackingIterator(Arrays.asList( - edge(1L, 2L), edge(1L, 3L), edge(1L, 4L)).iterator()); - OneDegreeGraph oneDegreeGraph = new OneDegreeGraph<>(1L, - new ValueVertex<>(1L, "vertex"), iterator); - - List> sampled = oneDegreeGraph.sampleNeighbors( - EdgeDirection.OUT, 2); - - Assert.assertEquals(sampled.size(), 2); - Assert.assertTrue(iterator.closed); - Assert.assertEquals(oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2), sampled); - } - - @Test(expectedExceptions = IllegalStateException.class) - public void testOneDegreeStateRejectsDifferentRequestAfterIteratorConsumption() { - OneDegreeGraph oneDegreeGraph = new OneDegreeGraph<>(1L, - new ValueVertex<>(1L, "vertex"), new TrackingIterator(Arrays.asList( - edge(1L, 2L), edge(1L, 3L), edge(1L, 4L)).iterator())); - - oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2); - oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 1); - } - - @Test(expectedExceptions = IllegalStateException.class) - public void testOneDegreeStateRejectsDifferentSamplingVersion() { - OneDegreeGraph oneDegreeGraph = new OneDegreeGraph<>(1L, - new ValueVertex<>(1L, "vertex"), new TrackingIterator(Arrays.asList( - edge(1L, 2L), edge(1L, 3L), edge(1L, 4L)).iterator())); - - oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2, 10L, 17L, 1L); - oneDegreeGraph.sampleNeighbors(EdgeDirection.OUT, 2, 10L, 17L, 2L); - } - @Test public void testNeighborhoodMatchesSnapshotAndSamplingVersion() { LocalNeighborhood neighborhood = new LocalNeighborhood<>( @@ -187,28 +149,4 @@ private static List targetIds(List> edges) { return edges.stream().map(IEdge::getTargetId).collect(Collectors.toList()); } - private static class TrackingIterator implements CloseableIterator> { - - private final Iterator> delegate; - private boolean closed; - - private TrackingIterator(Iterator> delegate) { - this.delegate = delegate; - } - - @Override - public void close() { - closed = true; - } - - @Override - public boolean hasNext() { - return delegate.hasNext(); - } - - @Override - public IEdge next() { - return delegate.next(); - } - } } From 3dbb3ebc4b67d531bc81e8de71d4e0112e6944d4 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 19 Aug 2026 13:08:21 +0800 Subject: [PATCH 7/9] refactor(sampling): reduce subgraph sampling implementation --- .../graph/sampling/EmptySamplingRequest.java | 43 --- .../graph/sampling/EmptySamplingResponse.java | 43 --- .../sampling/IterativeSamplingState.java | 103 ------ .../graph/sampling/NeighborStateRequest.java | 51 --- .../graph/sampling/NeighborStateResponse.java | 55 --- .../graph/sampling/PendingSamplingRound.java | 99 ----- .../api/graph/sampling/SampledSubgraph.java | 13 - .../api/graph/sampling/SamplingClock.java | 146 -------- .../api/graph/sampling/SamplingMessage.java | 28 -- .../api/graph/sampling/SamplingPhase.java | 28 -- .../sampling/SamplingResponseCollector.java | 105 ------ .../sampling/IterativeSamplingStateTest.java | 98 ----- .../IterativeSubgraphSamplingE2ETest.java | 342 ------------------ .../LayeredSubgraphAssemblerTest.java | 22 +- .../sampling/PendingSamplingRoundTest.java | 116 ------ .../api/graph/sampling/SamplingClockTest.java | 64 ---- .../algo/AlgorithmSamplingRuntimeContext.java | 53 ++- .../common/algo/AlgorithmUserFunction.java | 4 - .../algo/SubgraphSamplingAlgorithm.java | 24 -- .../GeaFlowAlgorithmAggTraversalFunction.java | 14 +- ...wAlgorithmDynamicAggTraversalFunction.java | 7 - ...GeaFlowAlgorithmDynamicRuntimeContext.java | 46 --- .../GeaFlowAlgorithmRuntimeContext.java | 43 --- ...FlowAlgorithmAggTraversalFunctionTest.java | 43 --- ...orithmDynamicAggTraversalFunctionTest.java | 14 - ...lowAlgorithmDynamicRuntimeContextTest.java | 35 +- .../geaflow/state/data/OneDegreeGraph.java | 2 +- .../DeterministicNeighborSampler.java | 24 +- .../state/sampling/LocalNeighborhood.java | 7 - .../DeterministicNeighborSamplerTest.java | 4 +- 30 files changed, 52 insertions(+), 1624 deletions(-) delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java delete mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java delete mode 100644 geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java delete mode 100644 geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java deleted file mode 100644 index 76eb3f261..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingRequest.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.util.Objects; - -/** First half of the empty-neighborhood self barrier. */ -public final class EmptySamplingRequest implements SamplingMessage { - - private final SamplingClock clock; - private final K vertexId; - - public EmptySamplingRequest(SamplingClock clock, K vertexId) { - this.clock = NeighborStateRequest.requirePhase(clock, SamplingPhase.REQUEST); - this.vertexId = Objects.requireNonNull(vertexId, "vertexId"); - } - - @Override - public SamplingClock getClock() { - return clock; - } - - public K getVertexId() { - return vertexId; - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java deleted file mode 100644 index 63af66ed2..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/EmptySamplingResponse.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.util.Objects; - -/** Second half of the empty-neighborhood self barrier. */ -public final class EmptySamplingResponse implements SamplingMessage { - - private final SamplingClock clock; - private final K vertexId; - - public EmptySamplingResponse(SamplingClock clock, K vertexId) { - this.clock = NeighborStateRequest.requirePhase(clock, SamplingPhase.RESPOND); - this.vertexId = Objects.requireNonNull(vertexId, "vertexId"); - } - - @Override - public SamplingClock getClock() { - return clock; - } - - public K getVertexId() { - return vertexId; - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java deleted file mode 100644 index 33007c602..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingState.java +++ /dev/null @@ -1,103 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.io.Serializable; -import java.util.Objects; - -/** Per-vertex committed payload and bounded in-flight state for one sampling session. */ -public final class IterativeSamplingState implements Serializable { - - private final long snapshotVersion; - private final long sessionId; - private int completedHop; - private P committedPayload; - private PendingSamplingRound pendingRound; - - public IterativeSamplingState(long snapshotVersion, long sessionId, P initialPayload) { - this.snapshotVersion = snapshotVersion; - this.sessionId = sessionId; - this.committedPayload = initialPayload; - } - - public void startRound(PendingSamplingRound pending) { - Objects.requireNonNull(pending, "pending"); - SamplingClock clock = pending.getRequestClock(); - requireSession(clock); - if (pendingRound != null) { - throw new IllegalStateException("a sampling round is already pending"); - } - if (clock.getHop() != completedHop + 1) { - throw new IllegalArgumentException("sampling request hop does not follow committed state"); - } - this.pendingRound = pending; - } - - public NeighborStateResponse respond(K responderId, NeighborStateRequest request) { - Objects.requireNonNull(request, "request"); - requireSession(request.getClock()); - if (request.getClock().getHop() != completedHop + 1) { - throw new IllegalStateException("requested payload is not committed for the preceding hop"); - } - return new NeighborStateResponse<>(request.getClock().responseClock(), - request.getRequesterId(), responderId, committedPayload); - } - - public void commit(SamplingClock commitClock, P payload) { - Objects.requireNonNull(commitClock, "commitClock"); - requireSession(commitClock); - if (commitClock.getPhase() != SamplingPhase.COMMIT_AND_REQUEST - && commitClock.getPhase() != SamplingPhase.COMPLETE) { - throw new IllegalArgumentException("sampling state can only commit in a commit phase"); - } - if (pendingRound == null || !pendingRound.getRequestClock().isSameRound(commitClock)) { - throw new IllegalStateException("sampling commit does not match the pending round"); - } - this.completedHop = commitClock.getHop(); - this.committedPayload = payload; - this.pendingRound = null; - } - - private void requireSession(SamplingClock clock) { - if (clock.getSnapshotVersion() != snapshotVersion || clock.getSessionId() != sessionId) { - throw new IllegalArgumentException("sampling clock does not match vertex session"); - } - } - - public long getSnapshotVersion() { - return snapshotVersion; - } - - public long getSessionId() { - return sessionId; - } - - public int getCompletedHop() { - return completedHop; - } - - public P getCommittedPayload() { - return committedPayload; - } - - public PendingSamplingRound getPendingRound() { - return pendingRound; - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java deleted file mode 100644 index 41529d482..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateRequest.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.util.Objects; - -/** Request for a neighbor's state committed at the preceding hop. */ -public final class NeighborStateRequest implements SamplingMessage { - - private final SamplingClock clock; - private final K requesterId; - - public NeighborStateRequest(SamplingClock clock, K requesterId) { - this.clock = requirePhase(clock, SamplingPhase.REQUEST); - this.requesterId = Objects.requireNonNull(requesterId, "requesterId"); - } - - static SamplingClock requirePhase(SamplingClock clock, SamplingPhase phase) { - Objects.requireNonNull(clock, "clock"); - if (clock.getPhase() != phase) { - throw new IllegalArgumentException("sampling message requires phase " + phase); - } - return clock; - } - - @Override - public SamplingClock getClock() { - return clock; - } - - public K getRequesterId() { - return requesterId; - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java deleted file mode 100644 index 05e41f937..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/NeighborStateResponse.java +++ /dev/null @@ -1,55 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.util.Objects; - -/** State returned by one selected neighbor. */ -public final class NeighborStateResponse implements SamplingMessage { - - private final SamplingClock clock; - private final K requesterId; - private final K responderId; - private final P payload; - - public NeighborStateResponse(SamplingClock clock, K requesterId, K responderId, P payload) { - this.clock = NeighborStateRequest.requirePhase(clock, SamplingPhase.RESPOND); - this.requesterId = Objects.requireNonNull(requesterId, "requesterId"); - this.responderId = Objects.requireNonNull(responderId, "responderId"); - this.payload = payload; - } - - @Override - public SamplingClock getClock() { - return clock; - } - - public K getRequesterId() { - return requesterId; - } - - public K getResponderId() { - return responderId; - } - - public P getPayload() { - return payload; - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java deleted file mode 100644 index 042e8ffb3..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRound.java +++ /dev/null @@ -1,99 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.io.Serializable; -import java.util.ArrayList; -import java.util.Collections; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; -import org.apache.geaflow.model.graph.edge.IEdge; - -/** Bounded requester-side state retained between request and commit supersteps. */ -public final class PendingSamplingRound implements Serializable { - - private final SamplingClock requestClock; - private final K requesterId; - private final Map>> edgesByNeighbor; - - public PendingSamplingRound(SamplingClock requestClock, K requesterId, - Iterable> sampledEdges) { - this.requestClock = NeighborStateRequest.requirePhase(requestClock, SamplingPhase.REQUEST); - this.requesterId = Objects.requireNonNull(requesterId, "requesterId"); - Objects.requireNonNull(sampledEdges, "sampledEdges"); - this.edgesByNeighbor = new LinkedHashMap<>(); - for (IEdge edge : sampledEdges) { - Objects.requireNonNull(edge, "edge"); - K neighborId = neighborId(requesterId, edge); - edgesByNeighbor.computeIfAbsent(neighborId, ignored -> new ArrayList<>()).add(edge); - } - } - - private K neighborId(K vertexId, IEdge edge) { - if (Objects.equals(vertexId, edge.getSrcId())) { - return Objects.requireNonNull(edge.getTargetId(), "neighborId"); - } - if (Objects.equals(vertexId, edge.getTargetId())) { - return Objects.requireNonNull(edge.getSrcId(), "neighborId"); - } - throw new IllegalArgumentException("sampled edge is not incident to requesterId=" + vertexId); - } - - public SamplingClock getRequestClock() { - return requestClock; - } - - public K getRequesterId() { - return requesterId; - } - - public boolean isEmpty() { - return edgesByNeighbor.isEmpty(); - } - - public List getNeighborIds() { - return Collections.unmodifiableList(new ArrayList<>(edgesByNeighbor.keySet())); - } - - public Map>> getEdgesByNeighbor() { - Map>> result = new LinkedHashMap<>(); - for (Map.Entry>> entry : edgesByNeighbor.entrySet()) { - result.put(entry.getKey(), Collections.unmodifiableList(entry.getValue())); - } - return Collections.unmodifiableMap(result); - } - - public Map> createRequests() { - Map> requests = new LinkedHashMap<>(); - for (K neighborId : edgesByNeighbor.keySet()) { - requests.put(neighborId, new NeighborStateRequest<>(requestClock, requesterId)); - } - return Collections.unmodifiableMap(requests); - } - - public EmptySamplingRequest createEmptyRequest() { - if (!isEmpty()) { - throw new IllegalStateException("only an empty sampling round uses an empty request"); - } - return new EmptySamplingRequest<>(requestClock, requesterId); - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java index 9187da687..318282148 100644 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java +++ b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SampledSubgraph.java @@ -81,15 +81,6 @@ public void addVertex(IVertex vertex) { vertices.put(vertexId, vertex); } - public void addLayer(List> edges) { - Objects.requireNonNull(edges, "edges"); - List> layer = new ArrayList<>(); - for (IEdge edge : edges) { - addEdge(layer, edge); - } - edgeLayers.add(layer); - } - public void addNeighborhood(int depth, LocalNeighborhood neighborhood, boolean includeEdges) { if (depth < 0) { @@ -145,10 +136,6 @@ public List>> getEdgeLayers() { return Collections.unmodifiableList(layers); } - public long getEdgeCount() { - return edgeIdentities.size(); - } - public void validateComplete() { for (List> layer : edgeLayers) { for (IEdge edge : layer) { diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java deleted file mode 100644 index c34174d17..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingClock.java +++ /dev/null @@ -1,146 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.io.Serializable; -import java.util.Objects; - -/** Immutable logical clock mapped from the runtime BSP iteration. */ -public final class SamplingClock implements Serializable { - - private final long snapshotVersion; - private final long sessionId; - private final int hop; - private final SamplingPhase phase; - - public SamplingClock(long snapshotVersion, long sessionId, int hop, SamplingPhase phase) { - if (hop < 1) { - throw new IllegalArgumentException("sampling hop must be greater than zero"); - } - this.snapshotVersion = snapshotVersion; - this.sessionId = sessionId; - this.hop = hop; - this.phase = Objects.requireNonNull(phase, "phase"); - } - - public static SamplingClock forIteration(long snapshotVersion, long sessionId, int maxHops, - long startIterationId, long iterationId) { - if (maxHops < 1) { - throw new IllegalArgumentException("maxHops must be greater than zero"); - } - if (iterationId < startIterationId) { - throw new IllegalArgumentException("iteration precedes the sampling session"); - } - long offset = iterationId - startIterationId; - if (offset == 0L) { - return new SamplingClock(snapshotVersion, sessionId, 1, SamplingPhase.REQUEST); - } - if ((offset & 1L) == 1L) { - long responseHop = (offset + 1L) / 2L; - requireHopInRange(responseHop, maxHops, iterationId); - return new SamplingClock(snapshotVersion, sessionId, (int) responseHop, - SamplingPhase.RESPOND); - } - long completedHop = offset / 2L; - requireHopInRange(completedHop, maxHops, iterationId); - SamplingPhase phase = completedHop == maxHops - ? SamplingPhase.COMPLETE : SamplingPhase.COMMIT_AND_REQUEST; - return new SamplingClock(snapshotVersion, sessionId, (int) completedHop, phase); - } - - public static long requiredIterations(int maxHops) { - if (maxHops < 1) { - throw new IllegalArgumentException("maxHops must be greater than zero"); - } - return Math.addExact(Math.multiplyExact((long) maxHops, 2L), 1L); - } - - private static void requireHopInRange(long hop, int maxHops, long iterationId) { - if (hop < 1L || hop > maxHops) { - throw new IllegalArgumentException("iteration is outside the sampling session: " - + iterationId); - } - } - - public SamplingClock responseClock() { - if (phase != SamplingPhase.REQUEST) { - throw new IllegalStateException("only a request clock can create a response clock"); - } - return new SamplingClock(snapshotVersion, sessionId, hop, SamplingPhase.RESPOND); - } - - public SamplingClock nextRequestClock() { - if (phase != SamplingPhase.COMMIT_AND_REQUEST) { - throw new IllegalStateException("current clock does not start another sampling hop"); - } - return new SamplingClock(snapshotVersion, sessionId, Math.addExact(hop, 1), - SamplingPhase.REQUEST); - } - - public boolean isSameRound(SamplingClock other) { - return other != null && snapshotVersion == other.snapshotVersion - && sessionId == other.sessionId && hop == other.hop; - } - - public long getSamplingVersion() { - long value = mix64(snapshotVersion) ^ Long.rotateLeft(mix64(sessionId), 21); - return mix64(value ^ Long.rotateLeft(mix64(hop), 42)); - } - - private static long mix64(long value) { - value = (value ^ (value >>> 30)) * 0xbf58476d1ce4e5b9L; - value = (value ^ (value >>> 27)) * 0x94d049bb133111ebL; - return value ^ (value >>> 31); - } - - public long getSnapshotVersion() { - return snapshotVersion; - } - - public long getSessionId() { - return sessionId; - } - - public int getHop() { - return hop; - } - - public SamplingPhase getPhase() { - return phase; - } - - @Override - public boolean equals(Object other) { - if (this == other) { - return true; - } - if (!(other instanceof SamplingClock)) { - return false; - } - SamplingClock that = (SamplingClock) other; - return snapshotVersion == that.snapshotVersion && sessionId == that.sessionId - && hop == that.hop && phase == that.phase; - } - - @Override - public int hashCode() { - return Objects.hash(snapshotVersion, sessionId, hop, phase); - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java deleted file mode 100644 index 604c58510..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingMessage.java +++ /dev/null @@ -1,28 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.io.Serializable; - -/** Marker for clocked iterative sampling messages. */ -public interface SamplingMessage extends Serializable { - - SamplingClock getClock(); -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java deleted file mode 100644 index 602454fcc..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingPhase.java +++ /dev/null @@ -1,28 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -/** Physical BSP phase for one logical sampling hop. */ -public enum SamplingPhase { - REQUEST, - RESPOND, - COMMIT_AND_REQUEST, - COMPLETE -} diff --git a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java b/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java deleted file mode 100644 index a0d7e9251..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/main/java/org/apache/geaflow/api/graph/sampling/SamplingResponseCollector.java +++ /dev/null @@ -1,105 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.util.ArrayList; -import java.util.Collections; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.Objects; - -/** Validates and orders all responses for one pending sampling round. */ -public final class SamplingResponseCollector { - - private final SamplingClock requestClock; - private final K requesterId; - private final List expectedNeighbors; - private final Map> responses = new LinkedHashMap<>(); - private boolean emptyResponseReceived; - - public SamplingResponseCollector(PendingSamplingRound pending) { - Objects.requireNonNull(pending, "pending"); - this.requestClock = pending.getRequestClock(); - this.requesterId = pending.getRequesterId(); - this.expectedNeighbors = pending.getNeighborIds(); - } - - public void add(NeighborStateResponse response) { - Objects.requireNonNull(response, "response"); - requireResponseRound(response.getClock()); - if (!Objects.equals(requesterId, response.getRequesterId())) { - throw new IllegalArgumentException("sampling response requester does not match pending round"); - } - K responderId = response.getResponderId(); - if (!expectedNeighbors.contains(responderId)) { - throw new IllegalArgumentException("sampling response came from an unrequested neighbor: " - + responderId); - } - if (responses.putIfAbsent(responderId, response) != null) { - throw new IllegalStateException("duplicate sampling response from neighbor: " + responderId); - } - } - - public void addEmpty(EmptySamplingResponse response) { - Objects.requireNonNull(response, "response"); - requireResponseRound(response.getClock()); - if (!expectedNeighbors.isEmpty()) { - throw new IllegalStateException("non-empty sampling round cannot accept an empty response"); - } - if (!Objects.equals(requesterId, response.getVertexId())) { - throw new IllegalArgumentException("empty sampling response vertex does not match requester"); - } - if (emptyResponseReceived) { - throw new IllegalStateException("duplicate empty sampling response"); - } - emptyResponseReceived = true; - } - - private void requireResponseRound(SamplingClock responseClock) { - NeighborStateRequest.requirePhase(responseClock, SamplingPhase.RESPOND); - if (!requestClock.isSameRound(responseClock)) { - throw new IllegalArgumentException("sampling response clock does not match pending round"); - } - } - - public boolean isComplete() { - return expectedNeighbors.isEmpty() ? emptyResponseReceived - : responses.size() == expectedNeighbors.size(); - } - - public void validateComplete() { - if (!isComplete()) { - List missing = new ArrayList<>(expectedNeighbors); - missing.removeAll(responses.keySet()); - throw new IllegalStateException("sampling responses are incomplete, requesterId=" - + requesterId + ", missing=" + missing); - } - } - - public List> getResponses() { - validateComplete(); - List> ordered = new ArrayList<>(expectedNeighbors.size()); - for (K neighborId : expectedNeighbors) { - ordered.add(responses.get(neighborId)); - } - return Collections.unmodifiableList(ordered); - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java deleted file mode 100644 index 462550057..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSamplingStateTest.java +++ /dev/null @@ -1,98 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.util.Arrays; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import org.apache.geaflow.model.graph.edge.EdgeDirection; -import org.apache.geaflow.model.graph.edge.IEdge; -import org.apache.geaflow.model.graph.edge.impl.ValueEdge; -import org.testng.Assert; -import org.testng.annotations.Test; - -public class IterativeSamplingStateTest { - - @Test - public void testAllVerticesCommitTwoHopsWithoutAccumulatingSubgraph() { - Map>> adjacency = new LinkedHashMap<>(); - adjacency.put(1L, Arrays.asList(edge(1L, 2L))); - adjacency.put(2L, Arrays.asList(edge(2L, 1L), edge(2L, 3L))); - adjacency.put(3L, Arrays.asList(edge(3L, 2L))); - Map> states = new LinkedHashMap<>(); - states.put(1L, new IterativeSamplingState<>(7L, 11L, 1)); - states.put(2L, new IterativeSamplingState<>(7L, 11L, 2)); - states.put(3L, new IterativeSamplingState<>(7L, 11L, 3)); - - SamplingClock request = SamplingClock.forIteration(7L, 11L, 2, 1L, 1L); - for (int hop = 1; hop <= 2; hop++) { - Map nextPayloads = new LinkedHashMap<>(); - for (Map.Entry> entry - : states.entrySet()) { - Long requesterId = entry.getKey(); - PendingSamplingRound pending = new PendingSamplingRound<>(request, - requesterId, adjacency.get(requesterId)); - entry.getValue().startRound(pending); - Assert.assertTrue(pending.getNeighborIds().size() <= 2); - - SamplingResponseCollector collector = - new SamplingResponseCollector<>(pending); - for (Map.Entry> outbound - : pending.createRequests().entrySet()) { - collector.add(states.get(outbound.getKey()).respond(outbound.getKey(), - outbound.getValue())); - } - int next = collector.getResponses().stream() - .mapToInt(NeighborStateResponse::getPayload).sum(); - nextPayloads.put(requesterId, next); - } - - SamplingClock commit = SamplingClock.forIteration(7L, 11L, 2, 1L, hop * 2L + 1L); - for (Map.Entry payload : nextPayloads.entrySet()) { - states.get(payload.getKey()).commit(commit, payload.getValue()); - } - if (hop < 2) { - request = commit.nextRequestClock(); - } - } - - Assert.assertEquals(states.get(1L).getCommittedPayload(), Integer.valueOf(4)); - Assert.assertEquals(states.get(2L).getCommittedPayload(), Integer.valueOf(4)); - Assert.assertEquals(states.get(3L).getCommittedPayload(), Integer.valueOf(4)); - for (IterativeSamplingState state : states.values()) { - Assert.assertEquals(state.getCompletedHop(), 2); - Assert.assertNull(state.getPendingRound()); - } - } - - @Test(expectedExceptions = IllegalStateException.class) - public void testCannotServeNextHopBeforePreviousHopCommit() { - IterativeSamplingState state = - new IterativeSamplingState<>(7L, 11L, 1); - NeighborStateRequest request = new NeighborStateRequest<>( - new SamplingClock(7L, 11L, 2, SamplingPhase.REQUEST), 2L); - state.respond(1L, request); - } - - private IEdge edge(long source, long target) { - return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java deleted file mode 100644 index ba5bd83ff..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/IterativeSubgraphSamplingE2ETest.java +++ /dev/null @@ -1,342 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; -import java.util.Comparator; -import java.util.LinkedHashMap; -import java.util.LinkedHashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; -import org.apache.geaflow.model.graph.edge.EdgeDirection; -import org.apache.geaflow.model.graph.edge.IEdge; -import org.apache.geaflow.model.graph.edge.impl.ValueEdge; -import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; -import org.testng.Assert; -import org.testng.annotations.Test; - -public class IterativeSubgraphSamplingE2ETest { - - private static final long SNAPSHOT_VERSION = 7L; - private static final long SESSION_ID = 11L; - private static final long START_ITERATION_ID = 21L; - - @Test - public void testRoutesTwoHopSamplingAcrossAllBspPhases() { - Map>> adjacency = new LinkedHashMap<>(); - adjacency.put(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L))); - adjacency.put(2L, Collections.singletonList(edge(2L, 4L))); - adjacency.put(3L, Collections.singletonList(edge(3L, 5L))); - adjacency.put(4L, Collections.emptyList()); - adjacency.put(5L, Collections.emptyList()); - SubgraphSamplingSpec spec = new SubgraphSamplingSpec( - 2, -1, EdgeDirection.OUT, 100L, 17L); - InMemorySamplingScheduler scheduler = new InMemorySamplingScheduler(adjacency, spec); - - List phases = scheduler.run(); - - Assert.assertEquals(phases, Arrays.asList( - SamplingPhase.REQUEST, - SamplingPhase.RESPOND, - SamplingPhase.COMMIT_AND_REQUEST, - SamplingPhase.RESPOND, - SamplingPhase.COMPLETE)); - Assert.assertEquals(scheduler.getRootPayloads(), Arrays.asList( - vertexIds(1L, 2L, 3L), - vertexIds(1L, 2L, 3L, 4L, 5L))); - Assert.assertEquals(scheduler.getNeighborRequestCount(), 8); - Assert.assertEquals(scheduler.getNeighborResponseCount(), 8); - Assert.assertEquals(scheduler.getEmptyRequestCount(), 4); - Assert.assertEquals(scheduler.getEmptyResponseCount(), 4); - Assert.assertEquals(scheduler.getCommitCount(), 10); - Assert.assertEquals(scheduler.getSamplingCallCount(), 10); - scheduler.assertComplete(); - } - - @Test - public void testDrivesCompleteThreeHopMessageSchedule() { - Map>> adjacency = new LinkedHashMap<>(); - adjacency.put(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L))); - adjacency.put(2L, Collections.singletonList(edge(2L, 4L))); - adjacency.put(3L, Collections.singletonList(edge(3L, 5L))); - adjacency.put(4L, Collections.singletonList(edge(4L, 6L))); - adjacency.put(5L, Collections.singletonList(edge(5L, 7L))); - adjacency.put(6L, Collections.emptyList()); - adjacency.put(7L, Collections.emptyList()); - SubgraphSamplingSpec spec = new SubgraphSamplingSpec( - 3, -1, EdgeDirection.OUT, 100L, 17L); - InMemorySamplingScheduler scheduler = new InMemorySamplingScheduler(adjacency, spec); - - Assert.assertEquals(scheduler.run(), Arrays.asList( - SamplingPhase.REQUEST, - SamplingPhase.RESPOND, - SamplingPhase.COMMIT_AND_REQUEST, - SamplingPhase.RESPOND, - SamplingPhase.COMMIT_AND_REQUEST, - SamplingPhase.RESPOND, - SamplingPhase.COMPLETE)); - Assert.assertEquals(scheduler.getRootPayloads(), Arrays.asList( - vertexIds(1L, 2L, 3L), - vertexIds(1L, 2L, 3L, 4L, 5L), - vertexIds(1L, 2L, 3L, 4L, 5L, 6L, 7L))); - Assert.assertEquals(scheduler.getRootMessageTrace(), Arrays.asList( - "request[1] 1->2", - "request[1] 1->3", - "response[1] 2->1", - "response[1] 3->1", - "request[2] 1->2", - "request[2] 1->3", - "response[2] 2->1", - "response[2] 3->1", - "request[3] 1->2", - "request[3] 1->3", - "response[3] 2->1", - "response[3] 3->1")); - Assert.assertEquals(new LinkedHashSet<>(scheduler.getRootSamplingVersions()).size(), 3, - "each hop must use a distinct sampling version"); - - // Every vertex participates in every hop, including explicit empty rounds for leaves. - Assert.assertEquals(scheduler.getNeighborRequestCount(), 18); - Assert.assertEquals(scheduler.getNeighborResponseCount(), 18); - Assert.assertEquals(scheduler.getEmptyRequestCount(), 6); - Assert.assertEquals(scheduler.getEmptyResponseCount(), 6); - Assert.assertEquals(scheduler.getCommitCount(), 21); - Assert.assertEquals(scheduler.getSamplingCallCount(), 21); - scheduler.assertComplete(); - } - - private static Set vertexIds(Long... ids) { - return new LinkedHashSet<>(Arrays.asList(ids)); - } - - private static IEdge edge(long source, long target) { - return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); - } - - private static final class InMemorySamplingScheduler { - - private final Map>> adjacency; - private final SubgraphSamplingSpec spec; - private final Map>> states = - new LinkedHashMap<>(); - private final List> rootPayloads = new ArrayList<>(); - private final List rootMessageTrace = new ArrayList<>(); - private final List rootSamplingVersions = new ArrayList<>(); - private int neighborRequestCount; - private int neighborResponseCount; - private int emptyRequestCount; - private int emptyResponseCount; - private int commitCount; - private int samplingCallCount; - - private InMemorySamplingScheduler(Map>> adjacency, - SubgraphSamplingSpec spec) { - this.adjacency = adjacency; - this.spec = spec; - for (Long vertexId : adjacency.keySet()) { - states.put(vertexId, new IterativeSamplingState<>( - SNAPSHOT_VERSION, SESSION_ID, vertexIds(vertexId))); - } - } - - private List run() { - List phases = new ArrayList<>(); - Map> inbox = Collections.emptyMap(); - long iterations = SamplingClock.requiredIterations(spec.getHops()); - for (long offset = 0L; offset < iterations; offset++) { - long iterationId = START_ITERATION_ID + offset; - SamplingClock clock = SamplingClock.forIteration(SNAPSHOT_VERSION, SESSION_ID, - spec.getHops(), START_ITERATION_ID, iterationId); - phases.add(clock.getPhase()); - switch (clock.getPhase()) { - case REQUEST: - Assert.assertTrue(inbox.isEmpty()); - inbox = startRounds(clock); - break; - case RESPOND: - inbox = respond(clock, inbox); - break; - case COMMIT_AND_REQUEST: - commit(clock, inbox); - inbox = startRounds(clock.nextRequestClock()); - break; - case COMPLETE: - commit(clock, inbox); - inbox = Collections.emptyMap(); - break; - default: - throw new IllegalStateException("unsupported sampling phase: " - + clock.getPhase()); - } - } - Assert.assertTrue(inbox.isEmpty()); - return phases; - } - - private Map> startRounds(SamplingClock requestClock) { - Map> requests = new LinkedHashMap<>(); - for (Map.Entry>> entry - : states.entrySet()) { - Long requesterId = entry.getKey(); - List> sampled = DeterministicNeighborSampler.sample( - requesterId, adjacency.get(requesterId), spec.getDirection(), spec.getFanout(), - Comparator.naturalOrder(), spec.getMaxReturnedEdges(), spec.getSeed(), - requestClock.getSamplingVersion()); - samplingCallCount++; - PendingSamplingRound pending = new PendingSamplingRound<>( - requestClock, requesterId, sampled); - entry.getValue().startRound(pending); - if (requesterId.equals(1L)) { - rootSamplingVersions.add(requestClock.getSamplingVersion()); - } - if (pending.isEmpty()) { - route(requests, requesterId, pending.createEmptyRequest()); - emptyRequestCount++; - } else { - for (Map.Entry> request - : pending.createRequests().entrySet()) { - route(requests, request.getKey(), request.getValue()); - if (requesterId.equals(1L)) { - rootMessageTrace.add("request[" + requestClock.getHop() + "] " - + requesterId + "->" + request.getKey()); - } - neighborRequestCount++; - } - } - } - return requests; - } - - private Map> respond( - SamplingClock responseClock, Map> requests) { - Map> responses = new LinkedHashMap<>(); - for (Map.Entry> inbox : requests.entrySet()) { - Long responderId = inbox.getKey(); - IterativeSamplingState> responder = states.get(responderId); - Assert.assertNotNull(responder, "message routed to an unknown vertex"); - for (SamplingMessage message : inbox.getValue()) { - Assert.assertTrue(message.getClock().isSameRound(responseClock)); - if (message instanceof NeighborStateRequest) { - NeighborStateRequest request = (NeighborStateRequest) message; - route(responses, request.getRequesterId(), - responder.respond(responderId, request)); - if (request.getRequesterId().equals(1L)) { - rootMessageTrace.add("response[" + responseClock.getHop() + "] " - + responderId + "->" + request.getRequesterId()); - } - neighborResponseCount++; - } else if (message instanceof EmptySamplingRequest) { - EmptySamplingRequest request = (EmptySamplingRequest) message; - Assert.assertEquals(request.getVertexId(), responderId); - route(responses, responderId, new EmptySamplingResponse<>( - request.getClock().responseClock(), responderId)); - emptyResponseCount++; - } else { - throw new IllegalStateException("unexpected sampling request: " + message); - } - } - } - return responses; - } - - private void commit(SamplingClock commitClock, - Map> responses) { - Map> nextPayloads = new LinkedHashMap<>(); - for (Map.Entry>> entry - : states.entrySet()) { - Long requesterId = entry.getKey(); - PendingSamplingRound pending = entry.getValue().getPendingRound(); - SamplingResponseCollector> collector = - new SamplingResponseCollector<>(pending); - for (SamplingMessage message - : responses.getOrDefault(requesterId, Collections.emptyList())) { - if (message instanceof NeighborStateResponse) { - collector.add((NeighborStateResponse>) message); - } else if (message instanceof EmptySamplingResponse) { - collector.addEmpty((EmptySamplingResponse) message); - } else { - throw new IllegalStateException("unexpected sampling response: " + message); - } - } - Set nextPayload = vertexIds(requesterId); - for (NeighborStateResponse> response : collector.getResponses()) { - nextPayload.addAll(response.getPayload()); - } - nextPayloads.put(requesterId, nextPayload); - } - for (Map.Entry> payload : nextPayloads.entrySet()) { - states.get(payload.getKey()).commit(commitClock, payload.getValue()); - commitCount++; - } - rootPayloads.add(new LinkedHashSet<>(states.get(1L).getCommittedPayload())); - } - - private void route(Map> messages, Long destination, - SamplingMessage message) { - messages.computeIfAbsent(destination, ignored -> new ArrayList<>()).add(message); - } - - private void assertComplete() { - for (IterativeSamplingState> state : states.values()) { - Assert.assertEquals(state.getCompletedHop(), spec.getHops()); - Assert.assertNull(state.getPendingRound()); - } - } - - private List> getRootPayloads() { - return rootPayloads; - } - - private List getRootMessageTrace() { - return rootMessageTrace; - } - - private List getRootSamplingVersions() { - return rootSamplingVersions; - } - - private int getNeighborRequestCount() { - return neighborRequestCount; - } - - private int getNeighborResponseCount() { - return neighborResponseCount; - } - - private int getEmptyRequestCount() { - return emptyRequestCount; - } - - private int getEmptyResponseCount() { - return emptyResponseCount; - } - - private int getCommitCount() { - return commitCount; - } - - private int getSamplingCallCount() { - return samplingCallCount; - } - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java index 378f5979e..ac6b95640 100644 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java @@ -48,6 +48,8 @@ public void testAssemblesOneHopPerLayerAndReleasesResult() { SampledSubgraph subgraph = assembler.take(1L); Assert.assertEquals(subgraph.getVertices().size(), 3); + Assert.assertEquals(subgraph.getVertices().get(3L).getValue(), Integer.valueOf(3)); + Assert.assertFalse(subgraph.getVertices().containsKey(4L)); Assert.assertEquals(subgraph.getEdgeLayers().size(), 2); Assert.assertEquals(subgraph.getEdgeLayers().get(0).get(0).getTargetId(), Long.valueOf(2L)); Assert.assertEquals(subgraph.getEdgeLayers().get(1).get(0).getTargetId(), Long.valueOf(3L)); @@ -132,12 +134,14 @@ public void testRejectsMissingResponse() { @Test public void testStructuralEdgeIdentityPreservesCollisionsAndLabels() { SampledSubgraph subgraph = new SampledSubgraph<>("root", 1L); - subgraph.addLayer(java.util.Arrays.asList( + LocalNeighborhood neighborhood = new LocalNeighborhood<>( + new ValueVertex<>("root", 0), java.util.Arrays.asList( new ValueLabelEdge<>("a->b", "c", 1, "first"), new ValueLabelEdge<>("a", "b->c", 1, "first"), - new ValueLabelEdge<>("a", "b->c", 1, "second"))); + new ValueLabelEdge<>("a", "b->c", 1, "second")), 1L); + subgraph.addNeighborhood(0, neighborhood, true); - Assert.assertEquals(subgraph.getEdgeCount(), 3L); + Assert.assertEquals(subgraph.getEdgeLayers().get(0).size(), 3); } @Test @@ -149,19 +153,21 @@ public void testLogicalEdgeIdentityIgnoresReplicaDirectionAndValue() { ValueLabelEdge reciprocal = new ValueLabelEdge<>(2L, 1L, "third", "knows"); SampledSubgraph subgraph = new SampledSubgraph<>(1L, 1L); - subgraph.addLayer(java.util.Arrays.asList(out, inReplica, reciprocal)); + subgraph.addNeighborhood(0, new LocalNeighborhood<>(new ValueVertex<>(1L, 1), + java.util.Arrays.asList(out, inReplica, reciprocal), 1L), true); - Assert.assertEquals(subgraph.getEdgeCount(), 2L); + Assert.assertEquals(subgraph.getEdgeLayers().get(0).size(), 2); } @Test public void testLogicalEdgeIdentityPreservesTemporalParallelEdges() { SampledSubgraph subgraph = new SampledSubgraph<>(1L, 1L); - subgraph.addLayer(java.util.Arrays.asList( + subgraph.addNeighborhood(0, new LocalNeighborhood<>(new ValueVertex<>(1L, 1), + java.util.Arrays.asList( new ValueLabelTimeEdge<>(1L, 2L, "same", "knows", 10L), - new ValueLabelTimeEdge<>(1L, 2L, "same", "knows", 11L))); + new ValueLabelTimeEdge<>(1L, 2L, "same", "knows", 11L)), 1L), true); - Assert.assertEquals(subgraph.getEdgeCount(), 2L); + Assert.assertEquals(subgraph.getEdgeLayers().get(0).size(), 2); } private LocalNeighborhood neighborhood(long source, long target, diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java deleted file mode 100644 index e28be164a..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/PendingSamplingRoundTest.java +++ /dev/null @@ -1,116 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import java.io.ByteArrayOutputStream; -import java.io.ObjectOutputStream; -import java.util.Arrays; -import java.util.Collections; -import java.util.List; -import org.apache.geaflow.model.graph.edge.EdgeDirection; -import org.apache.geaflow.model.graph.edge.IEdge; -import org.apache.geaflow.model.graph.edge.impl.ValueEdge; -import org.testng.Assert; -import org.testng.annotations.Test; - -public class PendingSamplingRoundTest { - - @Test - public void testGroupsParallelEdgesAndOrdersResponsesBySampledNeighbor() { - PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, - Arrays.asList(edge(1L, 2L, "first"), edge(1L, 2L, "parallel"), - edge(1L, 3L, "third"))); - - Assert.assertEquals(pending.getNeighborIds(), Arrays.asList(2L, 3L)); - Assert.assertEquals(pending.getEdgesByNeighbor().get(2L).size(), 2); - Assert.assertEquals(pending.createRequests().keySet(), - new java.util.LinkedHashSet<>(Arrays.asList(2L, 3L))); - - SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); - collector.add(new NeighborStateResponse<>(requestClock().responseClock(), 1L, 3L, "three")); - collector.add(new NeighborStateResponse<>(requestClock().responseClock(), 1L, 2L, "two")); - - List> responses = collector.getResponses(); - Assert.assertEquals(responses.get(0).getResponderId(), Long.valueOf(2L)); - Assert.assertEquals(responses.get(1).getResponderId(), Long.valueOf(3L)); - } - - @Test(expectedExceptions = IllegalStateException.class) - public void testRejectsMissingResponseAtCommit() { - PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, - Arrays.asList(edge(1L, 2L, "first"), edge(1L, 3L, "second"))); - SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); - collector.add(new NeighborStateResponse<>(requestClock().responseClock(), 1L, 2L, "two")); - collector.validateComplete(); - } - - @Test(expectedExceptions = IllegalStateException.class) - public void testRejectsDuplicateResponse() { - PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, - Collections.singletonList(edge(1L, 2L, "first"))); - SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); - NeighborStateResponse response = new NeighborStateResponse<>( - requestClock().responseClock(), 1L, 2L, "two"); - collector.add(response); - collector.add(response); - } - - @Test(expectedExceptions = IllegalArgumentException.class) - public void testRejectsResponseFromAnotherSession() { - PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, - Collections.singletonList(edge(1L, 2L, "first"))); - SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); - SamplingClock stale = new SamplingClock(7L, 12L, 1, SamplingPhase.RESPOND); - collector.add(new NeighborStateResponse<>(stale, 1L, 2L, "two")); - } - - @Test - public void testEmptyRoundUsesTwoPhaseSelfBarrier() { - PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, - Collections.emptyList()); - SamplingResponseCollector collector = new SamplingResponseCollector<>(pending); - - EmptySamplingRequest request = pending.createEmptyRequest(); - Assert.assertEquals(request.getVertexId(), Long.valueOf(1L)); - collector.addEmpty(new EmptySamplingResponse<>(request.getClock().responseClock(), 1L)); - Assert.assertTrue(collector.isComplete()); - Assert.assertTrue(collector.getResponses().isEmpty()); - } - - @Test - public void testProtocolStateIsSerializable() throws Exception { - PendingSamplingRound pending = new PendingSamplingRound<>(requestClock(), 1L, - Collections.singletonList(edge(1L, 2L, "first"))); - ByteArrayOutputStream bytes = new ByteArrayOutputStream(); - try (ObjectOutputStream output = new ObjectOutputStream(bytes)) { - output.writeObject(pending); - output.writeObject(pending.createRequests().get(2L)); - } - Assert.assertTrue(bytes.size() > 0); - } - - private SamplingClock requestClock() { - return new SamplingClock(7L, 11L, 1, SamplingPhase.REQUEST); - } - - private IEdge edge(long source, long target, String value) { - return new ValueEdge<>(source, target, value, EdgeDirection.OUT); - } -} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java deleted file mode 100644 index 24c62bd53..000000000 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SamplingClockTest.java +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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 org.apache.geaflow.api.graph.sampling; - -import org.testng.Assert; -import org.testng.annotations.Test; - -public class SamplingClockTest { - - @Test - public void testMapsTwoHopsToFiveIterations() { - SamplingClock firstRequest = SamplingClock.forIteration(7L, 11L, 2, 1L, 1L); - SamplingClock firstResponse = SamplingClock.forIteration(7L, 11L, 2, 1L, 2L); - SamplingClock firstCommit = SamplingClock.forIteration(7L, 11L, 2, 1L, 3L); - SamplingClock secondResponse = SamplingClock.forIteration(7L, 11L, 2, 1L, 4L); - SamplingClock complete = SamplingClock.forIteration(7L, 11L, 2, 1L, 5L); - - assertClock(firstRequest, 1, SamplingPhase.REQUEST); - assertClock(firstResponse, 1, SamplingPhase.RESPOND); - assertClock(firstCommit, 1, SamplingPhase.COMMIT_AND_REQUEST); - assertClock(firstCommit.nextRequestClock(), 2, SamplingPhase.REQUEST); - assertClock(secondResponse, 2, SamplingPhase.RESPOND); - assertClock(complete, 2, SamplingPhase.COMPLETE); - Assert.assertEquals(SamplingClock.requiredIterations(2), 5L); - } - - @Test - public void testSamplingVersionChangesPerHopButNotPhase() { - SamplingClock request = new SamplingClock(7L, 11L, 1, SamplingPhase.REQUEST); - SamplingClock response = request.responseClock(); - SamplingClock next = new SamplingClock(7L, 11L, 2, SamplingPhase.REQUEST); - - Assert.assertEquals(request.getSamplingVersion(), response.getSamplingVersion()); - Assert.assertNotEquals(request.getSamplingVersion(), next.getSamplingVersion()); - Assert.assertTrue(request.isSameRound(response)); - } - - @Test(expectedExceptions = IllegalArgumentException.class) - public void testRejectsIterationAfterSessionCompletion() { - SamplingClock.forIteration(7L, 11L, 2, 1L, 6L); - } - - private void assertClock(SamplingClock clock, int hop, SamplingPhase phase) { - Assert.assertEquals(clock.getHop(), hop); - Assert.assertEquals(clock.getPhase(), phase); - } -} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java index 2c90440bb..fd0f0e780 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java @@ -19,31 +19,28 @@ package org.apache.geaflow.dsl.common.algo; -import org.apache.geaflow.api.graph.sampling.SamplingClock; -import org.apache.geaflow.api.graph.sampling.SamplingPhase; +import java.util.Comparator; +import java.util.List; import org.apache.geaflow.api.graph.sampling.SubgraphSamplingSpec; +import org.apache.geaflow.common.iterator.CloseableIterator; +import org.apache.geaflow.common.type.IType; import org.apache.geaflow.dsl.common.data.Row; +import org.apache.geaflow.dsl.common.data.RowEdge; import org.apache.geaflow.dsl.common.data.RowVertex; import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; import org.apache.geaflow.state.sampling.LocalNeighborhood; /** Runtime-facing contract for reusable one-hop sampling. */ public interface AlgorithmSamplingRuntimeContext extends AlgorithmRuntimeContext { - LocalNeighborhood sampleOneHop(RowVertex vertex, EdgeDirection direction, - int fanout); - default LocalNeighborhood sampleOneHop(RowVertex vertex, EdgeDirection direction, - int fanout, - long maxCandidateEdges) { - LocalNeighborhood neighborhood = sampleOneHop(vertex, direction, fanout); - if (neighborhood.getEdges().size() > maxCandidateEdges) { - throw new IllegalStateException(String.format( - "one-hop sampling edge limit exceeded, vertexId=%s, actual=%s, limit=%s", - vertex.getId(), neighborhood.getEdges().size(), maxCandidateEdges)); - } - return neighborhood; + int fanout) { + return sampleOneHop(vertex, direction, fanout, + DeterministicNeighborSampler.DEFAULT_MAX_RETURNED_EDGES, 0L, + getSamplingSnapshotVersion()); } default LocalNeighborhood sampleOneHop(RowVertex vertex, @@ -52,26 +49,24 @@ default LocalNeighborhood sampleOneHop(RowVertex vertex, long maxReturnedEdges, long seed, long samplingVersion) { - return sampleOneHop(vertex, direction, fanout, maxReturnedEdges); - } - - default SamplingClock getSamplingClock(SubgraphSamplingSpec spec, long sessionId, - long startIterationId) { - return SamplingClock.forIteration(getSamplingSnapshotVersion(), sessionId, - spec.getHops(), startIterationId, getCurrentIterationId()); + try (CloseableIterator iterator = loadStaticEdgesIterator(direction)) { + Iterable edges = () -> iterator; + Comparator comparator = (left, right) -> + ((IType) getGraphSchema().getIdType()).compare(left, right); + @SuppressWarnings({"unchecked", "rawtypes"}) + List> sampled = (List) DeterministicNeighborSampler.sample( + vertex.getId(), edges, direction, fanout, comparator, maxReturnedEdges, + seed, samplingVersion); + return new LocalNeighborhood<>(vertex, sampled, getSamplingSnapshotVersion(), + samplingVersion); + } } default LocalNeighborhood sampleOneHop(RowVertex vertex, SubgraphSamplingSpec spec, - SamplingClock requestClock) { - if (requestClock.getPhase() != SamplingPhase.REQUEST) { - throw new IllegalArgumentException("one-hop sampling requires a request clock"); - } - if (requestClock.getSnapshotVersion() != getSamplingSnapshotVersion()) { - throw new IllegalArgumentException("sampling clock does not match runtime snapshot"); - } + long samplingVersion) { return sampleOneHop(vertex, spec.getDirection(), spec.getFanout(), - spec.getMaxReturnedEdges(), spec.getSeed(), requestClock.getSamplingVersion()); + spec.getMaxReturnedEdges(), spec.getSeed(), samplingVersion); } long getSamplingSnapshotVersion(); diff --git a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java index e125ad8a8..4058ff6f6 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmUserFunction.java @@ -60,10 +60,6 @@ public interface AlgorithmUserFunction extends Serializable { default void finish() { } - /** Called before an iteration starts processing vertices and messages. */ - default void initIteration(long iterationId) { - } - /** * Finish Iteration method called after each iteration finished. */ diff --git a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java deleted file mode 100644 index 3c6756ff5..000000000 --- a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/SubgraphSamplingAlgorithm.java +++ /dev/null @@ -1,24 +0,0 @@ -/* - * 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 org.apache.geaflow.dsl.common.algo; - -/** Marks an algorithm whose first sampling iteration requires a stable window snapshot. */ -public interface SubgraphSamplingAlgorithm { -} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java index 217409bd2..daca980db 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunction.java @@ -27,7 +27,6 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; -import org.apache.geaflow.api.function.iterator.RichIteratorFunction; import org.apache.geaflow.api.graph.function.vc.VertexCentricAggTraversalFunction; import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; import org.apache.geaflow.dsl.common.data.Row; @@ -44,8 +43,7 @@ import org.apache.geaflow.utils.keygroup.KeyGroupAssignment; public class GeaFlowAlgorithmAggTraversalFunction implements - VertexCentricAggTraversalFunction, - RichIteratorFunction { + VertexCentricAggTraversalFunction { private static final String STATE_SUFFIX = "UpdatedValueState"; @@ -147,16 +145,6 @@ public void close() { algorithmCtx.close(); } - @Override - public void initIteration(long iterationId) { - userFunction.initIteration(iterationId); - } - - @Override - public void finishIteration(long iterationId) { - userFunction.finishIteration(iterationId); - } - @Override public void initContext(VertexCentricAggContext aggContext) { this.algorithmCtx.setAggContext(Objects.requireNonNull(aggContext)); diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java index 72b2d33b5..658793939 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunction.java @@ -33,7 +33,6 @@ import org.apache.geaflow.api.graph.function.vc.IncVertexCentricAggTraversalFunction; import org.apache.geaflow.common.config.keys.FrameworkConfigKeys; import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; -import org.apache.geaflow.dsl.common.algo.SubgraphSamplingAlgorithm; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowVertex; import org.apache.geaflow.dsl.common.types.GraphSchema; @@ -94,11 +93,6 @@ public void open( IncVertexCentricTraversalFuncContext vertexCentricFuncContext) { this.traversalContext = vertexCentricFuncContext; this.materializeInFinish = traversalContext.getRuntimeContext().getConfiguration().getBoolean(FrameworkConfigKeys.UDF_MATERIALIZE_GRAPH_IN_FINISH); - // Sampling must read a stable window snapshot. Apply the window delta before the first - // sampling iteration, then refresh only vertices triggered by that delta. - if (userFunction instanceof SubgraphSamplingAlgorithm) { - this.materializeInFinish = false; - } this.algorithmCtx = new GeaFlowAlgorithmDynamicRuntimeContext(this, traversalContext, graphSchema); this.initVertices = new HashSet<>(); @@ -268,7 +262,6 @@ public void close() { @Override public void initIteration(long iterationId) { - userFunction.initIteration(iterationId); } @Override diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java index 191632cd4..8203a75fd 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContext.java @@ -20,7 +20,6 @@ package org.apache.geaflow.dsl.runtime.engine; import java.util.ArrayList; -import java.util.Comparator; import java.util.List; import java.util.Objects; import org.apache.geaflow.api.graph.function.aggregate.VertexCentricAggContextFunction.VertexCentricAggContext; @@ -33,7 +32,6 @@ import org.apache.geaflow.dsl.common.algo.AlgorithmSamplingRuntimeContext; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowEdge; -import org.apache.geaflow.dsl.common.data.RowVertex; import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; import org.apache.geaflow.dsl.common.types.GraphSchema; import org.apache.geaflow.dsl.runtime.traversal.message.ITraversalAgg; @@ -47,8 +45,6 @@ import org.apache.geaflow.state.pushdown.filter.IFilter; import org.apache.geaflow.state.pushdown.filter.InEdgeFilter; import org.apache.geaflow.state.pushdown.filter.OutEdgeFilter; -import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; -import org.apache.geaflow.state.sampling.LocalNeighborhood; public class GeaFlowAlgorithmDynamicRuntimeContext implements AlgorithmSamplingRuntimeContext { @@ -84,10 +80,6 @@ public void setVertexId(Object vertexId) { this.edgeQuery.withId(vertexId); } - public Object getVertexId() { - return vertexId; - } - public IVertex loadVertex() { return vertexQuery.get(); } @@ -183,44 +175,6 @@ public List loadStaticEdges(EdgeDirection direction) { } } - @Override - public LocalNeighborhood sampleOneHop(RowVertex vertex, - EdgeDirection direction, - int fanout) { - List> sampled = DeterministicNeighborSampler.sample(vertex.getId(), - loadStaticEdges(direction), direction, fanout); - return new LocalNeighborhood<>(vertex, sampled, getSamplingSnapshotVersion()); - } - - @Override - @SuppressWarnings({"unchecked", "rawtypes"}) - public LocalNeighborhood sampleOneHop(RowVertex vertex, - EdgeDirection direction, - int fanout, - long maxCandidateEdges) { - return sampleOneHop(vertex, direction, fanout, maxCandidateEdges, 0L, - getSamplingSnapshotVersion()); - } - - @Override - @SuppressWarnings({"unchecked", "rawtypes"}) - public LocalNeighborhood sampleOneHop(RowVertex vertex, - EdgeDirection direction, - int fanout, - long maxReturnedEdges, - long seed, - long samplingVersion) { - try (CloseableIterator iterator = loadStaticEdgesIterator(direction)) { - Iterable iterable = () -> iterator; - Comparator comparator = (left, right) -> - ((org.apache.geaflow.common.type.IType) graphSchema.getIdType()).compare(left, right); - List> sampled = (List) DeterministicNeighborSampler.sample( - vertex.getId(), iterable, direction, fanout, comparator, maxReturnedEdges, - seed, samplingVersion); - return new LocalNeighborhood<>(vertex, sampled, getSamplingSnapshotVersion(), samplingVersion); - } - } - @Override public long getSamplingSnapshotVersion() { return incVCTraversalCtx.getRuntimeContext().getWindowId(); diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java index 5113b85d4..d3d88b22a 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/main/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmRuntimeContext.java @@ -20,7 +20,6 @@ package org.apache.geaflow.dsl.runtime.engine; import java.util.ArrayList; -import java.util.Comparator; import java.util.List; import java.util.Objects; import org.apache.geaflow.api.graph.function.aggregate.VertexCentricAggContextFunction.VertexCentricAggContext; @@ -32,20 +31,16 @@ import org.apache.geaflow.dsl.common.algo.AlgorithmSamplingRuntimeContext; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowEdge; -import org.apache.geaflow.dsl.common.data.RowVertex; import org.apache.geaflow.dsl.common.exception.GeaFlowDSLException; import org.apache.geaflow.dsl.common.types.GraphSchema; import org.apache.geaflow.dsl.runtime.traversal.message.ITraversalAgg; import org.apache.geaflow.model.graph.edge.EdgeDirection; -import org.apache.geaflow.model.graph.edge.IEdge; import org.apache.geaflow.model.traversal.ITraversalResponse; import org.apache.geaflow.model.traversal.TraversalType.ResponseType; import org.apache.geaflow.state.pushdown.filter.EmptyFilter; import org.apache.geaflow.state.pushdown.filter.IFilter; import org.apache.geaflow.state.pushdown.filter.InEdgeFilter; import org.apache.geaflow.state.pushdown.filter.OutEdgeFilter; -import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; -import org.apache.geaflow.state.sampling.LocalNeighborhood; public class GeaFlowAlgorithmRuntimeContext implements AlgorithmSamplingRuntimeContext { @@ -118,44 +113,6 @@ public List loadStaticEdges(EdgeDirection direction) { return loadEdges(direction); } - @Override - public LocalNeighborhood sampleOneHop(RowVertex vertex, - EdgeDirection direction, - int fanout) { - List> edges = DeterministicNeighborSampler.sample(vertex.getId(), - loadStaticEdges(direction), direction, fanout); - return new LocalNeighborhood<>(vertex, edges, getSamplingSnapshotVersion()); - } - - @Override - @SuppressWarnings({"unchecked", "rawtypes"}) - public LocalNeighborhood sampleOneHop(RowVertex vertex, - EdgeDirection direction, - int fanout, - long maxCandidateEdges) { - return sampleOneHop(vertex, direction, fanout, maxCandidateEdges, 0L, - getSamplingSnapshotVersion()); - } - - @Override - @SuppressWarnings({"unchecked", "rawtypes"}) - public LocalNeighborhood sampleOneHop(RowVertex vertex, - EdgeDirection direction, - int fanout, - long maxReturnedEdges, - long seed, - long samplingVersion) { - try (CloseableIterator iterator = loadStaticEdgesIterator(direction)) { - Iterable iterable = () -> iterator; - Comparator comparator = (left, right) -> - ((org.apache.geaflow.common.type.IType) graphSchema.getIdType()).compare(left, right); - List> edges = (List) DeterministicNeighborSampler.sample( - vertex.getId(), iterable, direction, fanout, comparator, maxReturnedEdges, - seed, samplingVersion); - return new LocalNeighborhood<>(vertex, edges, getSamplingSnapshotVersion(), samplingVersion); - } - } - @Override public long getSamplingSnapshotVersion() { return traversalContext.getRuntimeContext().getWindowId(); diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java deleted file mode 100644 index 7aadb2c90..000000000 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmAggTraversalFunctionTest.java +++ /dev/null @@ -1,43 +0,0 @@ -/* - * 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 org.apache.geaflow.dsl.runtime.engine; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.verify; - -import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; -import org.apache.geaflow.dsl.common.types.GraphSchema; -import org.testng.annotations.Test; - -public class GeaFlowAlgorithmAggTraversalFunctionTest { - - @Test - public void testForwardsIterationLifecycle() { - AlgorithmUserFunction userFunction = mock(AlgorithmUserFunction.class); - GeaFlowAlgorithmAggTraversalFunction function = new GeaFlowAlgorithmAggTraversalFunction( - mock(GraphSchema.class), userFunction, new Object[0]); - - function.initIteration(3L); - function.finishIteration(3L); - - verify(userFunction).initIteration(3L); - verify(userFunction).finishIteration(3L); - } -} diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java index 505201c9c..2c0cd92d8 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicAggTraversalFunctionTest.java @@ -38,20 +38,6 @@ public class GeaFlowAlgorithmDynamicAggTraversalFunctionTest { - @Test - public void testForwardsIterationLifecycle() { - AlgorithmUserFunction userFunction = mock(AlgorithmUserFunction.class); - GeaFlowAlgorithmDynamicAggTraversalFunction function = - new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), - userFunction, new Object[0]); - - function.initIteration(3L); - function.finishIteration(3L); - - verify(userFunction).initIteration(3L); - verify(userFunction).finishIteration(3L); - } - @Test public void testEvolvePersistsNeighborhoodChangeVersion() throws Exception { GeaFlowAlgorithmDynamicAggTraversalFunction function = diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java index e1c382126..4e2ef3de0 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java @@ -33,9 +33,7 @@ import org.apache.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.TraversalEdgeQuery; import org.apache.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.TraversalVertexQuery; import org.apache.geaflow.api.graph.function.vc.base.IncVertexCentricFunction.TemporaryGraph; -import org.apache.geaflow.api.graph.sampling.SamplingClock; -import org.apache.geaflow.api.graph.sampling.SamplingPhase; -import org.apache.geaflow.api.graph.sampling.SubgraphSamplingSpec; +import org.apache.geaflow.common.iterator.CloseableIterator; import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowEdge; @@ -44,37 +42,14 @@ import org.apache.geaflow.dsl.common.data.impl.types.ObjectEdge; import org.apache.geaflow.dsl.common.data.impl.ObjectRow; import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.state.pushdown.filter.OutEdgeFilter; import org.apache.geaflow.state.sampling.LocalNeighborhood; import org.testng.Assert; import org.testng.annotations.Test; public class GeaFlowAlgorithmDynamicRuntimeContextTest { - @Test - public void testSamplingContextMapsRuntimeClock() { - IncVertexCentricTraversalFuncContext traversalContext = mock( - IncVertexCentricTraversalFuncContext.class); - TraversalHistoricalGraph historicalGraph = mock(TraversalHistoricalGraph.class); - TraversalGraphSnapShot snapshot = mock(TraversalGraphSnapShot.class); - RuntimeContext runtimeContext = mock(RuntimeContext.class); - when(traversalContext.getHistoricalGraph()).thenReturn(historicalGraph); - when(historicalGraph.getSnapShot(0L)).thenReturn(snapshot); - when(snapshot.vertex()).thenReturn(mock(TraversalVertexQuery.class)); - when(snapshot.edges()).thenReturn(mock(TraversalEdgeQuery.class)); - when(traversalContext.getRuntimeContext()).thenReturn(runtimeContext); - when(runtimeContext.getWindowId()).thenReturn(7L); - when(traversalContext.getIterationId()).thenReturn(3L); - GeaFlowAlgorithmDynamicRuntimeContext context = new GeaFlowAlgorithmDynamicRuntimeContext( - new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), - mock(AlgorithmUserFunction.class), new Object[0]), traversalContext, - mock(GraphSchema.class)); - SubgraphSamplingSpec spec = new SubgraphSamplingSpec(2, 2, EdgeDirection.OUT, 9L, 5L); - - SamplingClock current = context.getSamplingClock(spec, 11L, 1L); - Assert.assertEquals(current.getHop(), 1); - Assert.assertEquals(current.getPhase(), SamplingPhase.COMMIT_AND_REQUEST); - } - @Test public void testSamplingUsesMaterializedSnapshotOnly() { IncVertexCentricTraversalFuncContext traversalContext = mock( @@ -85,14 +60,18 @@ public void testSamplingUsesMaterializedSnapshotOnly() { TraversalEdgeQuery edgeQuery = mock(TraversalEdgeQuery.class); RuntimeContext runtimeContext = mock(RuntimeContext.class); TemporaryGraph temporaryGraph = mock(TemporaryGraph.class); + CloseableIterator> edgeIterator = mock(CloseableIterator.class); RowEdge edge = new ObjectEdge(1L, 2L, ObjectRow.create(1.0D)); edge.setDirect(EdgeDirection.OUT); + when(edgeIterator.hasNext()).thenReturn(true, false); + when(edgeIterator.next()).thenReturn(edge); when(traversalContext.getHistoricalGraph()).thenReturn(historicalGraph); when(historicalGraph.getSnapShot(0L)).thenReturn(snapshot); when(snapshot.vertex()).thenReturn(vertexQuery); when(snapshot.edges()).thenReturn(edgeQuery); when(edgeQuery.getOutEdges()).thenReturn(Collections.singletonList(edge)); + when(edgeQuery.getEdges(OutEdgeFilter.getInstance())).thenReturn(edgeIterator); when(traversalContext.getRuntimeContext()).thenReturn(runtimeContext); when(runtimeContext.getWindowId()).thenReturn(7L); when(traversalContext.getTemporaryGraph()).thenReturn(temporaryGraph); diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java index 317c5075d..43552c9ca 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/data/OneDegreeGraph.java @@ -47,5 +47,5 @@ public IVertex getVertex() { public CloseableIterator> getEdgeIterator() { return edgeIterator; } - } + diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java index a84632451..8961fc0a2 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java @@ -36,7 +36,7 @@ */ public final class DeterministicNeighborSampler { - public static final long DEFAULT_MAX_CANDIDATE_EDGES = 100000L; + public static final long DEFAULT_MAX_RETURNED_EDGES = 100000L; private DeterministicNeighborSampler() { } @@ -46,16 +46,7 @@ public static List> sample(K vertexId, EdgeDirection direction, int fanout) { return sample(vertexId, edges, direction, fanout, - Comparator.comparing(String::valueOf), DEFAULT_MAX_CANDIDATE_EDGES, 0L, 0L); - } - - public static List> sample(K vertexId, - Iterable> edges, - EdgeDirection direction, - int fanout, - Comparator idComparator, - long maxReturnedEdges) { - return sample(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, 0L, 0L); + Comparator.comparing(String::valueOf), DEFAULT_MAX_RETURNED_EDGES, 0L, 0L); } public static List> sample(K vertexId, @@ -76,16 +67,7 @@ public static List> project(K vertexId, EdgeDirection direction, int fanout) { return project(vertexId, edges, direction, fanout, - Comparator.comparing(String::valueOf), DEFAULT_MAX_CANDIDATE_EDGES, 0L, 0L); - } - - public static List> project(K vertexId, - Iterable> edges, - EdgeDirection direction, - int fanout, - Comparator idComparator, - long maxReturnedEdges) { - return project(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, 0L, 0L); + Comparator.comparing(String::valueOf), DEFAULT_MAX_RETURNED_EDGES, 0L, 0L); } public static List> project(K vertexId, diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java index 5f01c4652..129bed40d 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java @@ -93,13 +93,6 @@ public LocalNeighborhood project(EdgeDirection direction, int fanout) snapshotVersion, samplingVersion); } - public LocalNeighborhood project(EdgeDirection direction, int fanout, - Comparator idComparator, - long maxReturnedEdges) { - return new LocalNeighborhood<>(vertex, DeterministicNeighborSampler.project(vertex.getId(), - edges, direction, fanout, idComparator, maxReturnedEdges), snapshotVersion, samplingVersion); - } - public LocalNeighborhood project(EdgeDirection direction, int fanout, Comparator idComparator, long maxReturnedEdges, diff --git a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java index f3504eb05..3c08ea235 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java @@ -129,10 +129,10 @@ public void testNeighborhoodMatchesSnapshotAndSamplingVersion() { } @Test(expectedExceptions = IllegalStateException.class) - public void testRejectsCandidateEdgeOverflow() { + public void testRejectsReturnedEdgeOverflow() { DeterministicNeighborSampler.sample(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L)), EdgeDirection.OUT, -1, - Long::compare, 1); + Long::compare, 1, 0L, 0L); } private static IEdge edge(long source, long target) { From 23f4089e070390af3caf5bcad06df986d36e90e8 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Wed, 19 Aug 2026 14:45:30 +0800 Subject: [PATCH 8/9] test(sampling): add end-to-end subgraph sampling coverage --- .../LayeredSubgraphAssemblerTest.java | 62 ++++ .../SubgraphSamplingEndToEndTest.java | 303 ++++++++++++++++++ .../sampling/SubgraphSamplingMessageTest.java | 58 ++++ .../sampling/SubgraphSamplingSpecTest.java | 15 + ...lowAlgorithmDynamicRuntimeContextTest.java | 50 +++ .../DeterministicNeighborSamplerTest.java | 31 ++ .../state/sampling/LocalNeighborhoodTest.java | 95 ++++++ 7 files changed, 614 insertions(+) create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java create mode 100644 geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingMessageTest.java create mode 100644 geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java index ac6b95640..de506288e 100644 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/LayeredSubgraphAssemblerTest.java @@ -53,9 +53,71 @@ public void testAssemblesOneHopPerLayerAndReleasesResult() { Assert.assertEquals(subgraph.getEdgeLayers().size(), 2); Assert.assertEquals(subgraph.getEdgeLayers().get(0).get(0).getTargetId(), Long.valueOf(2L)); Assert.assertEquals(subgraph.getEdgeLayers().get(1).get(0).getTargetId(), Long.valueOf(3L)); + Assert.assertTrue(subgraph.getEdgeLayers().stream() + .flatMap(List::stream) + .noneMatch(edge -> Long.valueOf(4L).equals(edge.getTargetId()))); Assert.assertNull(assembler.take(1L)); } + @Test + public void testAcceptsOutOfOrderResponsesAndIgnoresDuplicates() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 1, new LocalNeighborhood<>(new ValueVertex<>(1L, 1), + java.util.Arrays.asList(new ValueEdge<>(1L, 2L, 1, EdgeDirection.OUT), + new ValueEdge<>(1L, 3L, 1, EdgeDirection.OUT)), 7L)); + Assert.assertTrue(assembler.registerRequest(1L, 2L, 1)); + Assert.assertTrue(assembler.registerRequest(1L, 3L, 1)); + + SubgraphSamplingResponse responseForThree = + new SubgraphSamplingResponse<>(1L, 1, neighborhood(3L, 4L, 7L)); + SubgraphSamplingResponse responseForTwo = + new SubgraphSamplingResponse<>(1L, 1, neighborhood(2L, 5L, 7L)); + Assert.assertTrue(assembler.add(responseForThree)); + Assert.assertTrue(assembler.add(responseForTwo)); + Assert.assertFalse(assembler.add(responseForThree)); + + Assert.assertEquals(assembler.take(1L).getVertices().keySet(), + new java.util.LinkedHashSet<>(java.util.Arrays.asList(1L, 2L, 3L))); + } + + @Test + public void testKeepsAssembliesIsolatedAndClearRemovesThem() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 1, neighborhood(1L, 2L, 7L)); + assembler.start(10L, 1, neighborhood(10L, 11L, 7L)); + + Assert.assertTrue(assembler.registerRequest(1L, 2L, 1)); + Assert.assertTrue(assembler.registerRequest(10L, 11L, 1)); + Assert.assertTrue(assembler.add(new SubgraphSamplingResponse<>(1L, 1, + neighborhood(2L, 3L, 7L)))); + Assert.assertTrue(assembler.add(new SubgraphSamplingResponse<>(10L, 1, + neighborhood(11L, 12L, 7L)))); + Assert.assertEquals(assembler.take(1L).getRootId(), Long.valueOf(1L)); + Assert.assertEquals(assembler.take(10L).getRootId(), Long.valueOf(10L)); + + assembler.start(20L, 1, neighborhood(20L, 21L, 7L)); + assembler.clear(); + Assert.assertNull(assembler.take(20L)); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsRequestDepthOutsideAssembly() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 1, neighborhood(1L, 2L, 7L)); + assembler.registerRequest(1L, 2L, 2); + } + + @Test(expectedExceptions = IllegalStateException.class) + public void testRejectsResponseForUnrequestedVertex() { + LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + assembler.start(1L, 1, neighborhood(1L, 2L, 7L)); + assembler.add(new SubgraphSamplingResponse<>(1L, 1, neighborhood(3L, 4L, 7L))); + } + @Test(expectedExceptions = IllegalArgumentException.class) public void testRejectsNeighborhoodFromPreviousWindow() { LayeredSubgraphAssembler assembler = diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java new file mode 100644 index 000000000..f2cea554f --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java @@ -0,0 +1,303 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.apache.geaflow.model.graph.vertex.impl.ValueVertex; +import org.apache.geaflow.state.sampling.DeterministicNeighborSampler; +import org.apache.geaflow.state.sampling.LocalNeighborhood; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class SubgraphSamplingEndToEndTest { + + private static final long SNAPSHOT_VERSION = 7L; + private static final long SAMPLING_VERSION = 42L; + private static final long SEED = 17L; + + @Test + public void testRunsThreeHopSamplingWithOutOfOrderResponses() { + SubgraphSamplingSpec spec = new SubgraphSamplingSpec( + 3, -1, EdgeDirection.OUT, 100L, SEED); + RecordingGraph graph = diamondAndCycleGraph(); + SamplingDriver driver = new SamplingDriver(graph, spec); + + SampledSubgraph result = driver.run(1L); + + Assert.assertEquals(new LinkedHashSet<>(driver.getRequestTrace().subList(0, 2)), + new LinkedHashSet<>(Arrays.asList("1->2@1", "1->3@1"))); + Assert.assertEquals(new LinkedHashSet<>(driver.getRequestTrace().subList(2, 4)), + new LinkedHashSet<>(Arrays.asList("1->4@2", "1->5@2"))); + Assert.assertEquals(driver.getRequestTrace().get(4), "1->6@3"); + Assert.assertNotEquals(driver.getRequestTrace().subList(0, 2), + driver.getEnqueuedRequestTrace().subList(0, 2)); + Assert.assertEquals(new LinkedHashSet<>(driver.getResponseTrace().subList(0, 2)), + new LinkedHashSet<>(Arrays.asList("2->1@1", "3->1@1"))); + Assert.assertEquals(new LinkedHashSet<>(driver.getResponseTrace().subList(2, 4)), + new LinkedHashSet<>(Arrays.asList("4->1@2", "5->1@2"))); + Assert.assertEquals(driver.getResponseTrace().get(4), "6->1@3"); + Assert.assertEquals(driver.getTerminalResponseEdgeCounts(), Collections.singletonList(0)); + Assert.assertEquals(driver.getRequestDepths(), Arrays.asList(1, 1, 2, 2, 3)); + Assert.assertEquals(driver.getResponseCount(), driver.getRequestCount()); + Assert.assertEquals(driver.getRequestCount(), 5); + Assert.assertEquals(result.getVertices().keySet(), new LinkedHashSet<>(Arrays.asList( + 1L, 2L, 3L, 4L, 5L, 6L))); + Assert.assertEquals(result.getEdgeLayers().size(), 3); + Assert.assertEquals(result.getEdgeLayers().get(0).size(), 2); + Assert.assertEquals(result.getEdgeLayers().get(1).size(), 4); + Assert.assertEquals(result.getEdgeLayers().get(2).size(), 3); + Assert.assertEquals(graph.getSampleReads().size(), 5); + Assert.assertEquals(graph.getVertexOnlyReads(), Collections.singletonList("6@3")); + } + + @Test + public void testDoesNotReadEdgesForTerminalDepth() { + SubgraphSamplingSpec spec = new SubgraphSamplingSpec( + 2, -1, EdgeDirection.OUT, 100L, SEED); + RecordingGraph graph = lineGraph(); + SamplingDriver driver = new SamplingDriver(graph, spec); + + SampledSubgraph result = driver.run(1L); + + Assert.assertEquals(graph.getSampleReads(), Arrays.asList("1@0", "2@1")); + Assert.assertEquals(graph.getVertexOnlyReads(), Collections.singletonList("3@2")); + Assert.assertEquals(driver.getTerminalResponseEdgeCounts(), Collections.singletonList(0)); + Assert.assertEquals(result.getVertices().keySet(), new LinkedHashSet<>(Arrays.asList( + 1L, 2L, 3L))); + Assert.assertEquals(result.getEdgeLayers().size(), 2); + Assert.assertEquals(result.getEdgeLayers().get(0).size(), 1); + Assert.assertEquals(result.getEdgeLayers().get(1).size(), 1); + Assert.assertEquals(driver.getRequestDepths(), Arrays.asList(1, 2)); + } + + @Test + public void testEndToEndSamplingHonorsPositiveFanout() { + SubgraphSamplingSpec spec = new SubgraphSamplingSpec( + 1, 1, EdgeDirection.OUT, 100L, SEED); + RecordingGraph graph = fanoutGraph(); + SamplingDriver driver = new SamplingDriver(graph, spec); + + SampledSubgraph result = driver.run(1L); + + Assert.assertEquals(result.getVertices().size(), 2); + Assert.assertEquals(result.getEdgeLayers().get(0).size(), 1); + Assert.assertEquals(driver.getRequestCount(), 1); + Assert.assertEquals(graph.getSampleReads(), Collections.singletonList("1@0")); + Assert.assertEquals(graph.getVertexOnlyReads().size(), 1); + } + + private static RecordingGraph diamondAndCycleGraph() { + Map>> adjacency = new LinkedHashMap<>(); + adjacency.put(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L))); + adjacency.put(2L, Arrays.asList(edge(2L, 4L), edge(2L, 5L))); + adjacency.put(3L, Arrays.asList(edge(3L, 4L), edge(3L, 5L))); + adjacency.put(4L, Arrays.asList(edge(4L, 1L), edge(4L, 6L))); + adjacency.put(5L, Collections.singletonList(edge(5L, 6L))); + adjacency.put(6L, Collections.emptyList()); + return new RecordingGraph(adjacency); + } + + private static RecordingGraph lineGraph() { + Map>> adjacency = new LinkedHashMap<>(); + adjacency.put(1L, Collections.singletonList(edge(1L, 2L))); + adjacency.put(2L, Collections.singletonList(edge(2L, 3L))); + adjacency.put(3L, Collections.emptyList()); + return new RecordingGraph(adjacency); + } + + private static RecordingGraph fanoutGraph() { + Map>> adjacency = new LinkedHashMap<>(); + adjacency.put(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L), edge(1L, 4L))); + adjacency.put(2L, Collections.emptyList()); + adjacency.put(3L, Collections.emptyList()); + adjacency.put(4L, Collections.emptyList()); + return new RecordingGraph(adjacency); + } + + private static IEdge edge(long source, long target) { + return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); + } + + private static final class SamplingDriver { + + private final RecordingGraph graph; + private final SubgraphSamplingSpec spec; + private final LayeredSubgraphAssembler assembler = + new LayeredSubgraphAssembler<>(); + private final List enqueuedRequestTrace = new ArrayList<>(); + private final List requestTrace = new ArrayList<>(); + private final List responseTrace = new ArrayList<>(); + private final List requestDepths = new ArrayList<>(); + private final List terminalResponseEdgeCounts = new ArrayList<>(); + + private SamplingDriver(RecordingGraph graph, SubgraphSamplingSpec spec) { + this.graph = graph; + this.spec = spec; + } + + private SampledSubgraph run(long rootId) { + LocalNeighborhood rootNeighborhood = + graph.sample(rootId, 0, spec); + assembler.start(rootId, spec.getHops(), rootNeighborhood); + Set frontier = neighbors(rootId, rootNeighborhood.getEdges()); + + for (int depth = 1; depth <= spec.getHops(); depth++) { + List requests = new ArrayList<>(); + for (Long vertexId : frontier) { + if (assembler.registerRequest(rootId, vertexId, depth)) { + requests.add(new RoutedRequest(vertexId, + new SubgraphSamplingRequest<>(rootId, depth))); + enqueuedRequestTrace.add(rootId + "->" + vertexId + "@" + depth); + } + } + + // Reverse delivery order to ensure assembly does not depend on transport order. + Collections.reverse(requests); + Set nextFrontier = new LinkedHashSet<>(); + for (RoutedRequest request : requests) { + Long requestRootId = request.request.getRootId(); + int requestDepth = request.request.getDepth(); + Assert.assertEquals(requestRootId, Long.valueOf(rootId)); + Assert.assertEquals(requestDepth, depth); + requestTrace.add(requestRootId + "->" + request.vertexId + "@" + requestDepth); + requestDepths.add(requestDepth); + LocalNeighborhood neighborhood = requestDepth == spec.getHops() + ? graph.vertexOnly(request.vertexId, requestDepth) + : graph.sample(request.vertexId, requestDepth, spec); + responseTrace.add(request.vertexId + "->" + requestRootId + "@" + requestDepth); + if (requestDepth == spec.getHops()) { + terminalResponseEdgeCounts.add(neighborhood.getEdges().size()); + } + SubgraphSamplingResponse response = + new SubgraphSamplingResponse<>(requestRootId, requestDepth, neighborhood); + Assert.assertTrue(assembler.add(response)); + if (requestDepth < spec.getHops()) { + nextFrontier.addAll(neighbors(request.vertexId, neighborhood.getEdges())); + } + } + frontier = nextFrontier; + } + return assembler.take(rootId); + } + + private Set neighbors(long vertexId, List> edges) { + Set neighbors = new LinkedHashSet<>(); + for (IEdge edge : edges) { + if (Long.valueOf(vertexId).equals(edge.getSrcId())) { + neighbors.add(edge.getTargetId()); + } else { + neighbors.add(edge.getSrcId()); + } + } + return neighbors; + } + + private List getRequestTrace() { + return requestTrace; + } + + private List getEnqueuedRequestTrace() { + return enqueuedRequestTrace; + } + + private List getResponseTrace() { + return responseTrace; + } + + private List getRequestDepths() { + return requestDepths; + } + + private int getRequestCount() { + return requestTrace.size(); + } + + private int getResponseCount() { + return responseTrace.size(); + } + + private List getTerminalResponseEdgeCounts() { + return terminalResponseEdgeCounts; + } + } + + private static final class RoutedRequest { + + private final Long vertexId; + private final SubgraphSamplingRequest request; + + private RoutedRequest(Long vertexId, SubgraphSamplingRequest request) { + this.vertexId = vertexId; + this.request = request; + } + } + + private static final class RecordingGraph { + + private final Map>> adjacency; + private final List sampleReads = new ArrayList<>(); + private final List vertexOnlyReads = new ArrayList<>(); + + private RecordingGraph(Map>> adjacency) { + this.adjacency = adjacency; + } + + private LocalNeighborhood sample( + long vertexId, int depth, SubgraphSamplingSpec spec) { + if (depth >= spec.getHops()) { + Assert.fail("terminal depth must not read or sample adjacent edges"); + } + sampleReads.add(vertexId + "@" + depth); + List> sampled = DeterministicNeighborSampler.sample( + vertexId, adjacency.get(vertexId), spec.getDirection(), spec.getFanout(), + Long::compare, spec.getMaxReturnedEdges(), spec.getSeed(), SAMPLING_VERSION); + return neighborhood(vertexId, sampled); + } + + private LocalNeighborhood vertexOnly(long vertexId, int depth) { + vertexOnlyReads.add(vertexId + "@" + depth); + return neighborhood(vertexId, Collections.emptyList()); + } + + private LocalNeighborhood neighborhood( + long vertexId, List> edges) { + return new LocalNeighborhood<>(new ValueVertex<>(vertexId, (int) vertexId), + edges, SNAPSHOT_VERSION, SAMPLING_VERSION); + } + + private List getSampleReads() { + return sampleReads; + } + + private List getVertexOnlyReads() { + return vertexOnlyReads; + } + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingMessageTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingMessageTest.java new file mode 100644 index 000000000..7ae6af8cd --- /dev/null +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingMessageTest.java @@ -0,0 +1,58 @@ +/* + * 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 org.apache.geaflow.api.graph.sampling; + +import java.util.Collections; +import org.apache.geaflow.model.graph.vertex.impl.ValueVertex; +import org.apache.geaflow.state.sampling.LocalNeighborhood; +import org.testng.annotations.Test; + +public class SubgraphSamplingMessageTest { + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsNonPositiveRequestDepth() { + new SubgraphSamplingRequest<>(1L, 0); + } + + @Test(expectedExceptions = NullPointerException.class) + public void testRejectsNullRequestRoot() { + new SubgraphSamplingRequest<>(null, 1); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsNonPositiveResponseDepth() { + new SubgraphSamplingResponse<>(1L, 0, neighborhood(1L)); + } + + @Test(expectedExceptions = NullPointerException.class) + public void testRejectsNullResponseRoot() { + new SubgraphSamplingResponse<>(null, 1, neighborhood(1L)); + } + + @Test(expectedExceptions = NullPointerException.class) + public void testRejectsNullResponseNeighborhood() { + new SubgraphSamplingResponse<>(1L, 1, null); + } + + private LocalNeighborhood neighborhood(long vertexId) { + return new LocalNeighborhood<>(new ValueVertex<>(vertexId, (int) vertexId), + Collections.emptyList(), 7L); + } +} diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java index 02b106c30..603549009 100644 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingSpecTest.java @@ -45,6 +45,16 @@ public void testRejectsZeroFanout() { new SubgraphSamplingSpec(2, 0, EdgeDirection.OUT); } + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsFanoutBelowUnlimitedMarker() { + new SubgraphSamplingSpec(2, -2, EdgeDirection.OUT); + } + + @Test(expectedExceptions = NullPointerException.class) + public void testRejectsNullDirection() { + new SubgraphSamplingSpec(2, 1, null); + } + @Test public void testUnlimitedFanoutKeepsPerVertexEdgeBudget() { SubgraphSamplingSpec spec = new SubgraphSamplingSpec(2, -1, EdgeDirection.BOTH); @@ -57,4 +67,9 @@ public void testUnlimitedFanoutKeepsPerVertexEdgeBudget() { public void testRejectsInvalidReturnedEdgeBudget() { new SubgraphSamplingSpec(2, -1, EdgeDirection.BOTH, 0); } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsNegativeReturnedEdgeBudget() { + new SubgraphSamplingSpec(2, -1, EdgeDirection.BOTH, -1); + } } diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java index 4e2ef3de0..7e37f55bc 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java @@ -19,6 +19,7 @@ package org.apache.geaflow.dsl.runtime.engine; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -27,6 +28,7 @@ import java.util.Arrays; import java.util.Collections; import org.apache.geaflow.api.context.RuntimeContext; +import org.apache.geaflow.api.graph.sampling.SubgraphSamplingSpec; import org.apache.geaflow.api.graph.function.vc.IncVertexCentricTraversalFunction.IncVertexCentricTraversalFuncContext; import org.apache.geaflow.api.graph.function.vc.IncVertexCentricTraversalFunction.TraversalGraphSnapShot; import org.apache.geaflow.api.graph.function.vc.IncVertexCentricTraversalFunction.TraversalHistoricalGraph; @@ -34,6 +36,7 @@ import org.apache.geaflow.api.graph.function.vc.VertexCentricTraversalFunction.TraversalVertexQuery; import org.apache.geaflow.api.graph.function.vc.base.IncVertexCentricFunction.TemporaryGraph; import org.apache.geaflow.common.iterator.CloseableIterator; +import org.apache.geaflow.common.type.primitive.LongType; import org.apache.geaflow.dsl.common.algo.AlgorithmUserFunction; import org.apache.geaflow.dsl.common.data.Row; import org.apache.geaflow.dsl.common.data.RowEdge; @@ -89,4 +92,51 @@ public void testSamplingUsesMaterializedSnapshotOnly() { Assert.assertEquals(neighborhood.getSnapshotVersion(), 7L); verify(temporaryGraph, never()).getEdges(); } + + @Test + public void testSamplingSpecPropagatesVersionAndClosesStaticIterator() { + IncVertexCentricTraversalFuncContext traversalContext = mock( + IncVertexCentricTraversalFuncContext.class); + TraversalHistoricalGraph historicalGraph = mock(TraversalHistoricalGraph.class); + TraversalGraphSnapShot snapshot = mock(TraversalGraphSnapShot.class); + TraversalVertexQuery vertexQuery = mock(TraversalVertexQuery.class); + TraversalEdgeQuery edgeQuery = mock(TraversalEdgeQuery.class); + RuntimeContext runtimeContext = mock(RuntimeContext.class); + CloseableIterator> edgeIterator = mock(CloseableIterator.class); + GraphSchema graphSchema = mock(GraphSchema.class); + + RowEdge first = edge(1L, 2L); + RowEdge second = edge(1L, 3L); + RowEdge third = edge(1L, 4L); + when(edgeIterator.hasNext()).thenReturn(true, true, true, false); + when(edgeIterator.next()).thenReturn(first, second, third); + when(traversalContext.getHistoricalGraph()).thenReturn(historicalGraph); + when(historicalGraph.getSnapShot(0L)).thenReturn(snapshot); + when(snapshot.vertex()).thenReturn(vertexQuery); + when(snapshot.edges()).thenReturn(edgeQuery); + when(edgeQuery.getEdges(OutEdgeFilter.getInstance())).thenReturn(edgeIterator); + when(traversalContext.getRuntimeContext()).thenReturn(runtimeContext); + when(runtimeContext.getWindowId()).thenReturn(7L); + doReturn(LongType.INSTANCE).when(graphSchema).getIdType(); + + GeaFlowAlgorithmDynamicRuntimeContext context = new GeaFlowAlgorithmDynamicRuntimeContext( + new GeaFlowAlgorithmDynamicAggTraversalFunction(graphSchema, + mock(AlgorithmUserFunction.class), new Object[0]), traversalContext, graphSchema); + RowVertex vertex = mock(RowVertex.class); + when(vertex.getId()).thenReturn(1L); + + LocalNeighborhood neighborhood = context.sampleOneHop(vertex, + new SubgraphSamplingSpec(1, 1, EdgeDirection.OUT, 100L, 17L), 9L); + + Assert.assertEquals(neighborhood.getEdges().size(), 1); + Assert.assertEquals(neighborhood.getSnapshotVersion(), 7L); + Assert.assertEquals(neighborhood.getSamplingVersion(), 9L); + verify(edgeIterator).close(); + } + + private RowEdge edge(long source, long target) { + RowEdge edge = new ObjectEdge(source, target, ObjectRow.create(1.0D)); + edge.setDirect(EdgeDirection.OUT); + return edge; + } } diff --git a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java index 3c08ea235..bdbe4e8c3 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java @@ -21,6 +21,7 @@ import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.stream.Collectors; @@ -100,6 +101,36 @@ public void testDirectionAndUnlimitedFanout() { 2); } + @Test + public void testIncomingNormalizationDoesNotMutateInputEdge() { + IEdge incoming = edge(2L, 1L); + incoming.setDirect(EdgeDirection.IN); + + List> sampled = DeterministicNeighborSampler.sample( + 1L, Collections.singletonList(incoming), EdgeDirection.IN, -1); + + Assert.assertEquals(incoming.getSrcId(), Long.valueOf(2L)); + Assert.assertEquals(incoming.getTargetId(), Long.valueOf(1L)); + Assert.assertEquals(incoming.getDirect(), EdgeDirection.IN); + Assert.assertEquals(sampled.get(0).getSrcId(), Long.valueOf(1L)); + Assert.assertEquals(sampled.get(0).getTargetId(), Long.valueOf(2L)); + } + + @Test + public void testComparatorTieUsesStableIdFallback() { + List> edges = Arrays.asList(edge(1L, 3L), edge(1L, 2L)); + + List first = targetIds(DeterministicNeighborSampler.sample( + 1L, edges, EdgeDirection.OUT, 1, (left, right) -> 0, + 100L, 17L, 7L)); + List second = targetIds(DeterministicNeighborSampler.sample( + 1L, Arrays.asList(edges.get(1), edges.get(0)), EdgeDirection.OUT, 1, + (left, right) -> 0, 100L, 17L, 7L)); + + Assert.assertEquals(first, second); + Assert.assertEquals(first.size(), 1); + } + @Test public void testFanoutCountsNeighborsAndKeepsSelectedParallelEdges() { List> edges = Arrays.asList( diff --git a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java new file mode 100644 index 000000000..04da68710 --- /dev/null +++ b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java @@ -0,0 +1,95 @@ +/* + * 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 org.apache.geaflow.state.sampling; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.model.graph.edge.EdgeDirection; +import org.apache.geaflow.model.graph.edge.IEdge; +import org.apache.geaflow.model.graph.edge.impl.ValueEdge; +import org.apache.geaflow.model.graph.vertex.impl.ValueVertex; +import org.testng.Assert; +import org.testng.annotations.Test; + +public class LocalNeighborhoodTest { + + @Test + public void testProjectPreservesVersionsAndBoundsNeighbors() { + LocalNeighborhood neighborhood = neighborhood( + Arrays.asList(edge(1L, 2L), edge(1L, 3L)), 7L, 11L); + + LocalNeighborhood projected = neighborhood.project( + EdgeDirection.OUT, 1, Long::compare, 100L, 17L); + + Assert.assertEquals(projected.getEdges().size(), 1); + Assert.assertEquals(projected.getSnapshotVersion(), 7L); + Assert.assertEquals(projected.getSamplingVersion(), 11L); + Assert.assertTrue(projected.matches(7L, 11L)); + } + + @Test + public void testRevalidatePreservesSamplingVersionOnNewSnapshot() { + LocalNeighborhood neighborhood = neighborhood( + Collections.singletonList(edge(1L, 2L)), 7L, 11L); + + LocalNeighborhood revalidated = neighborhood.revalidate( + new ValueVertex<>(1L, 2), 8L); + + Assert.assertEquals(revalidated.getSnapshotVersion(), 8L); + Assert.assertEquals(revalidated.getSamplingVersion(), 11L); + Assert.assertEquals(revalidated.getEdges().size(), 1); + } + + @Test(expectedExceptions = IllegalArgumentException.class) + public void testRejectsRevalidationToOlderSnapshot() { + LocalNeighborhood neighborhood = neighborhood( + Collections.emptyList(), 7L, 11L); + + neighborhood.revalidate(new ValueVertex<>(1L, 2), 6L); + } + + @Test + public void testCopiesInputEdgesAndProtectsReturnedEdges() { + List> input = new ArrayList<>(); + input.add(edge(1L, 2L)); + LocalNeighborhood neighborhood = neighborhood(input, 7L, 11L); + + input.clear(); + Assert.assertEquals(neighborhood.getEdges().size(), 1); + try { + neighborhood.getEdges().clear(); + Assert.fail("neighborhood edges must be immutable"); + } catch (UnsupportedOperationException expected) { + // Expected defensive view. + } + } + + private LocalNeighborhood neighborhood( + List> edges, long snapshotVersion, long samplingVersion) { + return new LocalNeighborhood<>(new ValueVertex<>(1L, 1), edges, + snapshotVersion, samplingVersion); + } + + private IEdge edge(long source, long target) { + return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); + } +} From 93fde8f22048571e4ecc8d34111a5466ad85024c Mon Sep 17 00:00:00 2001 From: aotenjou Date: Fri, 21 Aug 2026 23:50:28 +0800 Subject: [PATCH 9/9] fix(sampling): stabilize neighbor sampling semantics --- .../SubgraphSamplingEndToEndTest.java | 10 +- .../algo/AlgorithmSamplingRuntimeContext.java | 9 +- ...lowAlgorithmDynamicRuntimeContextTest.java | 6 +- .../DeterministicNeighborSampler.java | 99 ++++++++----- .../state/sampling/LocalNeighborhood.java | 14 +- .../DeterministicNeighborSamplerTest.java | 136 ++++++++++++++++-- .../state/sampling/LocalNeighborhoodTest.java | 9 +- 7 files changed, 226 insertions(+), 57 deletions(-) diff --git a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java index f2cea554f..fde7db764 100644 --- a/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java +++ b/geaflow/geaflow-core/geaflow-api/src/test/java/org/apache/geaflow/api/graph/sampling/SubgraphSamplingEndToEndTest.java @@ -145,6 +145,13 @@ private static IEdge edge(long source, long target) { return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); } + private static byte[] longBytes(long value) { + return new byte[]{ + (byte) (value >>> 56), (byte) (value >>> 48), (byte) (value >>> 40), + (byte) (value >>> 32), (byte) (value >>> 24), (byte) (value >>> 16), + (byte) (value >>> 8), (byte) value}; + } + private static final class SamplingDriver { private final RecordingGraph graph; @@ -277,7 +284,8 @@ private LocalNeighborhood sample( sampleReads.add(vertexId + "@" + depth); List> sampled = DeterministicNeighborSampler.sample( vertexId, adjacency.get(vertexId), spec.getDirection(), spec.getFanout(), - Long::compare, spec.getMaxReturnedEdges(), spec.getSeed(), SAMPLING_VERSION); + Long::compare, spec.getMaxReturnedEdges(), spec.getSeed(), SAMPLING_VERSION, + SubgraphSamplingEndToEndTest::longBytes); return neighborhood(vertexId, sampled); } diff --git a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java index fd0f0e780..1466b52f2 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-common/src/main/java/org/apache/geaflow/dsl/common/algo/AlgorithmSamplingRuntimeContext.java @@ -21,6 +21,7 @@ import java.util.Comparator; import java.util.List; +import java.util.function.Function; import org.apache.geaflow.api.graph.sampling.SubgraphSamplingSpec; import org.apache.geaflow.common.iterator.CloseableIterator; import org.apache.geaflow.common.type.IType; @@ -51,12 +52,14 @@ default LocalNeighborhood sampleOneHop(RowVertex vertex, long samplingVersion) { try (CloseableIterator iterator = loadStaticEdgesIterator(direction)) { Iterable edges = () -> iterator; - Comparator comparator = (left, right) -> - ((IType) getGraphSchema().getIdType()).compare(left, right); + @SuppressWarnings({"unchecked", "rawtypes"}) + IType idType = (IType) getGraphSchema().getIdType(); + Comparator comparator = idType::compare; + Function idEncoder = idType::serialize; @SuppressWarnings({"unchecked", "rawtypes"}) List> sampled = (List) DeterministicNeighborSampler.sample( vertex.getId(), edges, direction, fanout, comparator, maxReturnedEdges, - seed, samplingVersion); + seed, samplingVersion, idEncoder); return new LocalNeighborhood<>(vertex, sampled, getSamplingSnapshotVersion(), samplingVersion); } diff --git a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java index 7e37f55bc..5cb2b699b 100644 --- a/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java +++ b/geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/org/apache/geaflow/dsl/runtime/engine/GeaFlowAlgorithmDynamicRuntimeContextTest.java @@ -64,6 +64,7 @@ public void testSamplingUsesMaterializedSnapshotOnly() { RuntimeContext runtimeContext = mock(RuntimeContext.class); TemporaryGraph temporaryGraph = mock(TemporaryGraph.class); CloseableIterator> edgeIterator = mock(CloseableIterator.class); + GraphSchema graphSchema = mock(GraphSchema.class); RowEdge edge = new ObjectEdge(1L, 2L, ObjectRow.create(1.0D)); edge.setDirect(EdgeDirection.OUT); @@ -79,10 +80,11 @@ public void testSamplingUsesMaterializedSnapshotOnly() { when(runtimeContext.getWindowId()).thenReturn(7L); when(traversalContext.getTemporaryGraph()).thenReturn(temporaryGraph); when(temporaryGraph.getEdges()).thenReturn(Arrays.asList(edge)); + doReturn(LongType.INSTANCE).when(graphSchema).getIdType(); GeaFlowAlgorithmDynamicRuntimeContext context = new GeaFlowAlgorithmDynamicRuntimeContext( - new GeaFlowAlgorithmDynamicAggTraversalFunction(mock(GraphSchema.class), - mock(AlgorithmUserFunction.class), new Object[0]), traversalContext, mock(GraphSchema.class)); + new GeaFlowAlgorithmDynamicAggTraversalFunction(graphSchema, + mock(AlgorithmUserFunction.class), new Object[0]), traversalContext, graphSchema); RowVertex vertex = mock(RowVertex.class); when(vertex.getId()).thenReturn(1L); diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java index 8961fc0a2..58e2ec04a 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java @@ -26,6 +26,7 @@ import java.util.Map; import java.util.Objects; import java.util.PriorityQueue; +import java.util.function.Function; import org.apache.geaflow.model.graph.IGraphElementWithLabelField; import org.apache.geaflow.model.graph.IGraphElementWithTimeField; import org.apache.geaflow.model.graph.edge.EdgeDirection; @@ -33,6 +34,7 @@ /** * Storage-independent seeded one-hop neighbor sampling. + * The supplied ID encoder must return a stable, canonical byte representation across workers. */ public final class DeterministicNeighborSampler { @@ -44,9 +46,11 @@ private DeterministicNeighborSampler() { public static List> sample(K vertexId, Iterable> edges, EdgeDirection direction, - int fanout) { - return sample(vertexId, edges, direction, fanout, - Comparator.comparing(String::valueOf), DEFAULT_MAX_RETURNED_EDGES, 0L, 0L); + int fanout, + Comparator idComparator, + Function idEncoder) { + return sample(vertexId, edges, direction, fanout, idComparator, + DEFAULT_MAX_RETURNED_EDGES, 0L, 0L, idEncoder); } public static List> sample(K vertexId, @@ -56,18 +60,21 @@ public static List> sample(K vertexId, Comparator idComparator, long maxReturnedEdges, long seed, - long samplingVersion) { + long samplingVersion, + Function idEncoder) { return select(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, - seed, samplingVersion, true); + seed, samplingVersion, idEncoder, true); } /** Project an already direction-filtered local neighborhood to a smaller fanout. */ public static List> project(K vertexId, Iterable> edges, EdgeDirection direction, - int fanout) { - return project(vertexId, edges, direction, fanout, - Comparator.comparing(String::valueOf), DEFAULT_MAX_RETURNED_EDGES, 0L, 0L); + int fanout, + Comparator idComparator, + Function idEncoder) { + return project(vertexId, edges, direction, fanout, idComparator, + DEFAULT_MAX_RETURNED_EDGES, 0L, 0L, idEncoder); } public static List> project(K vertexId, @@ -77,9 +84,10 @@ public static List> project(K vertexId, Comparator idComparator, long maxReturnedEdges, long seed, - long samplingVersion) { + long samplingVersion, + Function idEncoder) { return select(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, - seed, samplingVersion, false); + seed, samplingVersion, idEncoder, false); } private static List> select(K vertexId, @@ -90,11 +98,14 @@ private static List> select(K vertexId, long maxReturnedEdges, long seed, long samplingVersion, + Function idEncoder, boolean filterAndNormalize) { - validate(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges); + validate(vertexId, edges, direction, fanout, idComparator, maxReturnedEdges, idEncoder); + long vertexHash = hashId(vertexId, idEncoder); Comparator> rankComparator = (left, right) -> { int result = Long.compareUnsigned(left.score, right.score); - return result != 0 ? result : compareIds(left.neighborId, right.neighborId, idComparator); + return result != 0 ? result + : compareIds(left.neighborId, right.neighborId, idComparator, idEncoder); }; Map> selected = new HashMap<>(); PriorityQueue> worstFirst = @@ -106,7 +117,10 @@ private static List> select(K vertexId, continue; } IEdge edge = filterAndNormalize ? normalize(sourceEdge) : sourceEdge; - K neighborId = Objects.requireNonNull(neighborId(vertexId, edge), "neighborId"); + K neighborId = neighborId(vertexId, edge); + if (neighborId == null) { + continue; + } NeighborGroup group = selected.get(neighborId); if (group != null) { group.edges.add(edge); @@ -114,7 +128,8 @@ private static List> select(K vertexId, } group = new NeighborGroup<>(neighborId, - sampleScore(seed, samplingVersion, vertexId, direction, neighborId), edge); + sampleScore(seed, samplingVersion, vertexHash, direction, + hashId(neighborId, idEncoder)), edge); if (fanout < 0 || selected.size() < fanout) { selected.put(neighborId, group); if (worstFirst != null) { @@ -132,7 +147,7 @@ private static List> select(K vertexId, groups.sort(rankComparator); List> result = new ArrayList<>(); for (NeighborGroup group : groups) { - group.edges.sort((left, right) -> compareEdges(left, right, idComparator)); + group.edges.sort((left, right) -> compareEdges(left, right, idComparator, idEncoder)); result.addAll(group.edges); if (result.size() > maxReturnedEdges) { throw new IllegalStateException(String.format( @@ -144,11 +159,13 @@ private static List> select(K vertexId, } private static void validate(Object vertexId, Iterable edges, EdgeDirection direction, - int fanout, Comparator idComparator, long maxReturnedEdges) { + int fanout, Comparator idComparator, long maxReturnedEdges, + Function idEncoder) { Objects.requireNonNull(vertexId, "vertexId"); Objects.requireNonNull(edges, "edges"); Objects.requireNonNull(direction, "direction"); Objects.requireNonNull(idComparator, "idComparator"); + Objects.requireNonNull(idEncoder, "idEncoder"); if (fanout == 0 || fanout < -1) { throw new IllegalArgumentException("fanout must be -1 or greater than zero"); } @@ -166,8 +183,8 @@ private static IEdge normalize(IEdge edge) { return edge; } IEdge reversed = edge.reverse(); - // Direction remains the sampling-side marker; endpoints are restored to logical order. - reversed.setDirect(edge.getDirect()); + // The reversed endpoints now represent the logical outgoing direction. + reversed.setDirect(EdgeDirection.OUT); return reversed; } @@ -178,23 +195,26 @@ private static K neighborId(K vertexId, IEdge edge) { if (Objects.equals(vertexId, edge.getTargetId())) { return edge.getSrcId(); } - return edge.getTargetId(); + return null; } - private static long sampleScore(long seed, long samplingVersion, Object vertexId, - EdgeDirection direction, Object neighborId) { + private static long sampleScore(long seed, long samplingVersion, long vertexHash, + EdgeDirection direction, long neighborHash) { long value = mix64(seed) ^ Long.rotateLeft(mix64(samplingVersion), 11); - value ^= Long.rotateLeft(stableHash(vertexId), 23); + value ^= Long.rotateLeft(vertexHash, 23); value ^= Long.rotateLeft(mix64(direction.ordinal()), 37); - value ^= Long.rotateLeft(stableHash(neighborId), 47); + value ^= Long.rotateLeft(neighborHash, 47); return mix64(value); } - private static long stableHash(Object value) { - String text = value.getClass().getName() + ':' + value; + private static long hashId(K value, Function idEncoder) { + return stableHash(Objects.requireNonNull(idEncoder.apply(value), "idEncoder result")); + } + + private static long stableHash(byte[] value) { long hash = 0xcbf29ce484222325L; - for (int i = 0; i < text.length(); i++) { - hash ^= text.charAt(i); + for (byte current : value) { + hash ^= current & 0xffL; hash *= 0x100000001b3L; } return mix64(hash); @@ -206,16 +226,20 @@ private static long mix64(long value) { return value ^ (value >>> 31); } - private static int compareIds(K left, K right, Comparator idComparator) { + private static int compareIds(K left, K right, Comparator idComparator, + Function idEncoder) { int result = idComparator.compare(left, right); - return result != 0 ? result : String.valueOf(left).compareTo(String.valueOf(right)); + return result != 0 ? result : compareBytes( + Objects.requireNonNull(idEncoder.apply(left), "idEncoder result"), + Objects.requireNonNull(idEncoder.apply(right), "idEncoder result")); } private static int compareEdges(IEdge left, IEdge right, - Comparator idComparator) { - int result = compareIds(left.getSrcId(), right.getSrcId(), idComparator); + Comparator idComparator, + Function idEncoder) { + int result = compareIds(left.getSrcId(), right.getSrcId(), idComparator, idEncoder); if (result == 0) { - result = compareIds(left.getTargetId(), right.getTargetId(), idComparator); + result = compareIds(left.getTargetId(), right.getTargetId(), idComparator, idEncoder); } if (result == 0) { result = left.getDirect().compareTo(right.getDirect()); @@ -232,6 +256,17 @@ private static int compareEdges(IEdge left, IEdge right, return result; } + private static int compareBytes(byte[] left, byte[] right) { + int length = Math.min(left.length, right.length); + for (int i = 0; i < length; i++) { + int result = Integer.compare(left[i] & 0xff, right[i] & 0xff); + if (result != 0) { + return result; + } + } + return Integer.compare(left.length, right.length); + } + private static String labelOf(IEdge edge) { return edge instanceof IGraphElementWithLabelField ? ((IGraphElementWithLabelField) edge).getLabel() : null; diff --git a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java index 129bed40d..405f6aedf 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/LocalNeighborhood.java @@ -25,6 +25,7 @@ import java.util.Comparator; import java.util.List; import java.util.Objects; +import java.util.function.Function; import org.apache.geaflow.model.graph.edge.EdgeDirection; import org.apache.geaflow.model.graph.edge.IEdge; import org.apache.geaflow.model.graph.vertex.IVertex; @@ -87,18 +88,23 @@ public LocalNeighborhood revalidate(IVertex currentVertex, /** * Create a bounded view of this already direction-filtered neighborhood. */ - public LocalNeighborhood project(EdgeDirection direction, int fanout) { + public LocalNeighborhood project(EdgeDirection direction, int fanout, + Comparator idComparator, + Function idEncoder) { return new LocalNeighborhood<>(vertex, - DeterministicNeighborSampler.project(vertex.getId(), edges, direction, fanout), + DeterministicNeighborSampler.project(vertex.getId(), edges, direction, fanout, + idComparator, idEncoder), snapshotVersion, samplingVersion); } public LocalNeighborhood project(EdgeDirection direction, int fanout, Comparator idComparator, long maxReturnedEdges, - long seed) { + long seed, + Function idEncoder) { return new LocalNeighborhood<>(vertex, DeterministicNeighborSampler.project(vertex.getId(), - edges, direction, fanout, idComparator, maxReturnedEdges, seed, samplingVersion), + edges, direction, fanout, idComparator, maxReturnedEdges, seed, samplingVersion, + idEncoder), snapshotVersion, samplingVersion); } } diff --git a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java index bdbe4e8c3..7e22133db 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/DeterministicNeighborSamplerTest.java @@ -22,8 +22,10 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.Comparator; import java.util.List; import java.util.Map; +import java.util.function.Function; import java.util.stream.Collectors; import org.apache.geaflow.model.graph.edge.EdgeDirection; import org.apache.geaflow.model.graph.edge.IEdge; @@ -42,9 +44,11 @@ public void testSamplingIsBoundedAndIndependentOfInputOrder() { edge(1L, 4L), edge(1L, 2L), edge(1L, 3L)); List> sampledFirst = - DeterministicNeighborSampler.sample(1L, first, EdgeDirection.OUT, 2); + DeterministicNeighborSampler.sample(1L, first, EdgeDirection.OUT, 2, + Long::compare, DeterministicNeighborSamplerTest::longBytes); List> sampledSecond = - DeterministicNeighborSampler.sample(1L, second, EdgeDirection.OUT, 2); + DeterministicNeighborSampler.sample(1L, second, EdgeDirection.OUT, 2, + Long::compare, DeterministicNeighborSamplerTest::longBytes); Assert.assertEquals(sampledFirst.size(), 2); Assert.assertEquals(targetIds(sampledFirst), targetIds(sampledSecond)); @@ -58,7 +62,8 @@ public void testPositiveFanoutDoesNotMaterializeAllCandidates() { } List> sampled = DeterministicNeighborSampler.sample( - 1L, edges, EdgeDirection.OUT, 3, Long::compare, 3L, 17L, 9L); + 1L, edges, EdgeDirection.OUT, 3, Long::compare, 3L, 17L, 9L, + DeterministicNeighborSamplerTest::longBytes); Assert.assertEquals(sampled.stream().map(IEdge::getTargetId).distinct().count(), 3L); Assert.assertEquals(sampled.size(), 3); @@ -74,11 +79,14 @@ public void testSeedAndVersionAreStableAndInputOrderIndependent() { java.util.Collections.reverse(reversed); List firstSample = targetIds(DeterministicNeighborSampler.sample( - 1L, first, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 7L)); + 1L, first, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 7L, + DeterministicNeighborSamplerTest::longBytes)); List reorderedSample = targetIds(DeterministicNeighborSampler.sample( - 1L, reversed, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 7L)); + 1L, reversed, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 7L, + DeterministicNeighborSamplerTest::longBytes)); List nextVersionSample = targetIds(DeterministicNeighborSampler.sample( - 1L, first, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 8L)); + 1L, first, EdgeDirection.OUT, 5, Long::compare, 5L, 123L, 8L, + DeterministicNeighborSamplerTest::longBytes)); Assert.assertEquals(firstSample, reorderedSample); Assert.assertNotEquals(firstSample, nextVersionSample); @@ -91,29 +99,50 @@ public void testDirectionAndUnlimitedFanout() { in.setDirect(EdgeDirection.IN); List> sampledIn = DeterministicNeighborSampler.sample( - 1L, Arrays.asList(out, in), EdgeDirection.IN, -1); + 1L, Arrays.asList(out, in), EdgeDirection.IN, -1, Long::compare, + DeterministicNeighborSamplerTest::longBytes); Assert.assertEquals(sampledIn.size(), 1); Assert.assertEquals(sampledIn.get(0).getSrcId(), Long.valueOf(3L)); Assert.assertEquals(sampledIn.get(0).getTargetId(), Long.valueOf(1L)); - Assert.assertEquals(sampledIn.get(0).getDirect(), EdgeDirection.IN); + Assert.assertEquals(sampledIn.get(0).getDirect(), EdgeDirection.OUT); Assert.assertEquals( - DeterministicNeighborSampler.sample(1L, Arrays.asList(out, in), EdgeDirection.BOTH, -1).size(), + DeterministicNeighborSampler.sample(1L, Arrays.asList(out, in), EdgeDirection.BOTH, + -1, Long::compare, DeterministicNeighborSamplerTest::longBytes).size(), 2); } + @Test + public void testIgnoresUnrelatedEdges() { + IEdge related = edge(1L, 2L); + IEdge unrelated = edge(3L, 4L); + List> edges = Arrays.asList(related, unrelated); + + List> sampled = DeterministicNeighborSampler.sample( + 1L, edges, EdgeDirection.OUT, -1, Long::compare, 1L, 0L, 0L, + DeterministicNeighborSamplerTest::longBytes); + List> projected = DeterministicNeighborSampler.project( + 1L, edges, EdgeDirection.OUT, -1, Long::compare, + DeterministicNeighborSamplerTest::longBytes); + + Assert.assertEquals(sampled, Collections.singletonList(related)); + Assert.assertEquals(projected, Collections.singletonList(related)); + } + @Test public void testIncomingNormalizationDoesNotMutateInputEdge() { IEdge incoming = edge(2L, 1L); incoming.setDirect(EdgeDirection.IN); List> sampled = DeterministicNeighborSampler.sample( - 1L, Collections.singletonList(incoming), EdgeDirection.IN, -1); + 1L, Collections.singletonList(incoming), EdgeDirection.IN, -1, Long::compare, + DeterministicNeighborSamplerTest::longBytes); Assert.assertEquals(incoming.getSrcId(), Long.valueOf(2L)); Assert.assertEquals(incoming.getTargetId(), Long.valueOf(1L)); Assert.assertEquals(incoming.getDirect(), EdgeDirection.IN); Assert.assertEquals(sampled.get(0).getSrcId(), Long.valueOf(1L)); Assert.assertEquals(sampled.get(0).getTargetId(), Long.valueOf(2L)); + Assert.assertEquals(sampled.get(0).getDirect(), EdgeDirection.OUT); } @Test @@ -122,10 +151,11 @@ public void testComparatorTieUsesStableIdFallback() { List first = targetIds(DeterministicNeighborSampler.sample( 1L, edges, EdgeDirection.OUT, 1, (left, right) -> 0, - 100L, 17L, 7L)); + 100L, 17L, 7L, DeterministicNeighborSamplerTest::longBytes)); List second = targetIds(DeterministicNeighborSampler.sample( 1L, Arrays.asList(edges.get(1), edges.get(0)), EdgeDirection.OUT, 1, - (left, right) -> 0, 100L, 17L, 7L)); + (left, right) -> 0, 100L, 17L, 7L, + DeterministicNeighborSamplerTest::longBytes)); Assert.assertEquals(first, second); Assert.assertEquals(first.size(), 1); @@ -137,7 +167,8 @@ public void testFanoutCountsNeighborsAndKeepsSelectedParallelEdges() { edge(1L, 2L), edgeWithValue(1L, 2L, "parallel"), edge(1L, 3L), edge(1L, 4L)); List> sampled = DeterministicNeighborSampler.sample( - 1L, edges, EdgeDirection.OUT, 2); + 1L, edges, EdgeDirection.OUT, 2, Long::compare, + DeterministicNeighborSamplerTest::longBytes); Map allCounts = edges.stream().collect(Collectors.groupingBy( IEdge::getTargetId, Collectors.counting())); Map sampledCounts = sampled.stream().collect(Collectors.groupingBy( @@ -163,7 +194,46 @@ public void testNeighborhoodMatchesSnapshotAndSamplingVersion() { public void testRejectsReturnedEdgeOverflow() { DeterministicNeighborSampler.sample(1L, Arrays.asList(edge(1L, 2L), edge(1L, 3L)), EdgeDirection.OUT, -1, - Long::compare, 1, 0L, 0L); + Long::compare, 1, 0L, 0L, DeterministicNeighborSamplerTest::longBytes); + } + + @Test + public void testSamplesIdsWithoutCallingToString() { + List> edges = Arrays.asList( + stableEdge(1L, 2L, "first"), stableEdge(1L, 2L, "second"), + stableEdge(1L, 3L, "third")); + Comparator comparator = (left, right) -> 0; + + List> first = DeterministicNeighborSampler.sample( + new StableId(1L), edges, EdgeDirection.OUT, -1, comparator, 100L, 17L, 7L, + id -> longBytes(id.value)); + List> second = DeterministicNeighborSampler.sample( + new StableId(1L), Arrays.asList(edges.get(2), edges.get(1), edges.get(0)), + EdgeDirection.OUT, -1, comparator, 100L, 17L, 7L, id -> longBytes(id.value)); + + if (first.size() != second.size()) { + Assert.fail("sampling result sizes differ"); + } + for (int i = 0; i < first.size(); i++) { + if (first.get(i) != second.get(i)) { + Assert.fail("sampling result order differs"); + } + } + } + + @Test(expectedExceptions = NullPointerException.class, + expectedExceptionsMessageRegExp = "idEncoder result") + public void testRejectsNullIdEncoding() { + DeterministicNeighborSampler.sample(1L, Collections.singletonList(edge(1L, 2L)), + EdgeDirection.OUT, -1, Long::compare, 100L, 0L, 0L, id -> null); + } + + @Test(expectedExceptions = NullPointerException.class, + expectedExceptionsMessageRegExp = "idEncoder") + public void testRejectsNullIdEncoder() { + DeterministicNeighborSampler.sample(1L, Collections.singletonList(edge(1L, 2L)), + EdgeDirection.OUT, -1, Long::compare, 100L, 0L, 0L, + (Function) null); } private static IEdge edge(long source, long target) { @@ -180,4 +250,42 @@ private static List targetIds(List> edges) { return edges.stream().map(IEdge::getTargetId).collect(Collectors.toList()); } + private static IEdge stableEdge(long source, long target, String value) { + ValueEdge edge = new ValueEdge<>(new StableId(source), + new StableId(target), value); + edge.setDirect(EdgeDirection.OUT); + return edge; + } + + private static byte[] longBytes(long value) { + return new byte[]{ + (byte) (value >>> 56), (byte) (value >>> 48), (byte) (value >>> 40), + (byte) (value >>> 32), (byte) (value >>> 24), (byte) (value >>> 16), + (byte) (value >>> 8), (byte) value}; + } + + private static final class StableId { + + private final long value; + + private StableId(long value) { + this.value = value; + } + + @Override + public boolean equals(Object other) { + return other instanceof StableId && value == ((StableId) other).value; + } + + @Override + public int hashCode() { + return Long.hashCode(value); + } + + @Override + public String toString() { + throw new AssertionError("sampling must not call StableId.toString()"); + } + } + } diff --git a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java index 04da68710..68f34bff1 100644 --- a/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java +++ b/geaflow/geaflow-state/geaflow-state-common/src/test/java/org/apache/geaflow/state/sampling/LocalNeighborhoodTest.java @@ -38,7 +38,7 @@ public void testProjectPreservesVersionsAndBoundsNeighbors() { Arrays.asList(edge(1L, 2L), edge(1L, 3L)), 7L, 11L); LocalNeighborhood projected = neighborhood.project( - EdgeDirection.OUT, 1, Long::compare, 100L, 17L); + EdgeDirection.OUT, 1, Long::compare, 100L, 17L, LocalNeighborhoodTest::longBytes); Assert.assertEquals(projected.getEdges().size(), 1); Assert.assertEquals(projected.getSnapshotVersion(), 7L); @@ -92,4 +92,11 @@ private LocalNeighborhood neighborhood( private IEdge edge(long source, long target) { return new ValueEdge<>(source, target, 1, EdgeDirection.OUT); } + + private static byte[] longBytes(long value) { + return new byte[]{ + (byte) (value >>> 56), (byte) (value >>> 48), (byte) (value >>> 40), + (byte) (value >>> 32), (byte) (value >>> 24), (byte) (value >>> 16), + (byte) (value >>> 8), (byte) value}; + } }