Skip to content

feat: Add subgraph sampling as reusable multi-hop subgraph sampling foundation - #831

Open
aotenjou wants to merge 9 commits into
apache:masterfrom
aotenjou:subgraph-sampling
Open

feat: Add subgraph sampling as reusable multi-hop subgraph sampling foundation#831
aotenjou wants to merge 9 commits into
apache:masterfrom
aotenjou:subgraph-sampling

Conversation

@aotenjou

@aotenjou aotenjou commented Aug 12, 2026

Copy link
Copy Markdown

What changes were proposed in this pull request?

This PR introduces a model-agnostic subgraph sampling foundation for graph learning algorithms such as GCN and GNN and an inner implementation of #502.

Changes

  • Add fanout-based one-hop neighbor sampling

    • fanout > 0 selects at most the requested number of neighbors
    • fanout = -1 returns all neighbors matching the requested direction
    • Parallel edges belonging to the same neighbor are retained as a group
    • Sampling supports seeds, sampling versions, and returned-edge limits
  • Add layered sampled-subgraph assembly

    • Assemble vertices and edges by sampling depth
    • At the terminal depth, return only vertices and their features without reading or returning adjacent edges
    • Enforce sampled-node and sampled-edge limits
    • Reject neighborhoods from inconsistent snapshot versions
  • Add stable logical edge identities

    • Identify logical edges using sourceId + targetId + label + time
    • Deduplicate incoming and outgoing storage replicas of the same logical edge
    • Preserve parallel edges with different labels or timestamps
  • Add an iterative multi-hop BSP protocol

    • Define REQUEST, RESPOND, COMMIT_AND_REQUEST, and COMPLETE phases
    • Add sampling clocks, request/response messages, response collectors, and per-round state
    • Support empty request/response barriers
    • Cover complete two-hop and three-hop message schedules
  • Integrate sampling with the DSL algorithm runtime

    • Provide one-hop sampling contexts for static and dynamic graphs
    • Forward initIteration and finishIteration lifecycle callbacks
    • Read from a stable window snapshot during sampling
    • Track neighborhood change versions to determine when cached state must be refreshed
    • Register both source and target vertices when processing edge changes

Testing

The following local validation has passed:

  • 453 tests across the four affected modules under JDK 8
  • 47 focused sampling tests after splitting the changes into atomic commits

Coverage includes:

  • Fanout, direction, seed, and sampling-version behavior
  • Unlimited sampling with fanout = -1
  • Parallel-edge grouping and logical-edge deduplication
  • Empty edge lists at the terminal depth
  • Snapshot consistency and capacity limits
  • Complete two-hop and three-hop BSP message schedules
  • Dynamic neighborhood change-version persistence
  • Static and dynamic runtime lifecycle forwarding

How was this PR tested?

  • Tests have Added for the changes
  • Production environment verified

@aotenjou aotenjou changed the title feat: Subgraph reusable multi-hop subgraph sampling foundation feat: Add subgraph sampling as reusable multi-hop subgraph sampling foundation Aug 12, 2026
@aotenjou
aotenjou marked this pull request as draft August 12, 2026 05:37
@aotenjou
aotenjou marked this pull request as ready for review August 19, 2026 07:03
return reversed;
}

private static <K, EV> K neighborId(K vertexId, IEdge<K, EV> edge) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have identified an interesting issue: in DeterministicNeighborSampler, the neighborId method returns targetId even when the edge is not connected to the specified vertex. This causes edges unrelated to the vertex to be included in the sampling as neighbors, leading to incorrect sampling results or "neighbor contamination."

geaflow/geaflow-state/geaflow-state-common/src/main/java/org/apache/geaflow/state/sampling/DeterministicNeighborSampler.java
Method: private static <K,EV> K neighborId(K vertexId, IEdge<K,EV> edge)

Current code snippet (problematic code)

Original implementation (last three lines):

private static <K, EV> K neighborId(K vertexId, IEdge<K, EV> edge) {
if (Objects.equals(vertexId, edge.getSrcId())) {
return edge.getTargetId();
}
if (Objects.equals(vertexId, edge.getTargetId())) {
return edge.getSrcId();
}
return edge.getTargetId();
}

Issue: When neither the edge's source nor target matches vertexId, the method still returns edge.getTargetId(), thereby treating an "unrelated edge" as a neighbor of the vertex.
Risks and consequences

Unrelated edges are included in neighbor groups, leading to incorrect neighbor grouping or fanout selection. This violates sampling semantics, potentially breaches budget constraints, and causes downstream assembly processes to produce "dangling edges" or incorrect subgraphs.
Suggested fix (implementation example)

Change the logic to return null (or throw an exception) when the edge is not incident to the vertex, and skip the edge at the call site. Here is an example:

private static <K, EV> K neighborId(K vertexId, IEdge<K, EV> edge) {
if (Objects.equals(vertexId, edge.getSrcId())) {
return edge.getTargetId();
}
if (Objects.equals(vertexId, edge.getTargetId())) {
return edge.getSrcId();
}
// edge not incident on vertex: signal caller to skip
return null;
}

Change the call site within select(...) (previously Objects.requireNonNull(neighborId(...), "neighborId")) to perform a defensive skip:

K neighborId = neighborId(vertexId, edge);
if (neighborId == null) {
// skip edges that are not incident on vertexId
continue;
}

In my opinion, this approach is more robust and prevents unrelated edges from being included in the sampling. Although the select loop already performs null checks and directional filtering on the edges, this defensive check is still warranted to handle cases where the underlying storage or iterator returns anomalous or non-normalized data.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching this. I agree that the current fallback in neighborId(...) is unsafe.

Although select(...) filters by edge direction, it does not guarantee that every returned edge is incident to the requested vertexId. With the current fallback, an unrelated edge can be assigned to edge.getTargetId() and incorrectly enter the neighbor groups, which may consume fanout capacity and contaminate the sampled subgraph.

I will update the implementation to return null when neither endpoint matches vertexId, and defensively skip such edges in select(...) rather than throwing an exception. This keeps the sampler resilient to anomalous or non-normalized edges returned by the underlying storage.

I will also add regression coverage for both sample(...) and project(...) to ensure unrelated edges are ignored, including cases where they pass the direction check.

return direction == EdgeDirection.BOTH || edge.getDirect() == direction;
}

private static <K, EV> IEdge<K, EV> normalize(IEdge<K, EV> edge) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have identified a potential semantic inconsistency regarding normalize and direction: retaining the original direct value after reversing the endpoints could confuse upstream logic.

private static <K, EV> IEdge<K, EV> normalize(IEdge<K, EV> edge) {
if (edge.getDirect() != EdgeDirection.IN) {
return edge;
}
IEdge<K, EV> reversed = edge.reverse();
// Direction remains the sampling-side marker; endpoints are restored to logical order.
reversed.setDirect(edge.getDirect());
return reversed;
}

Potential Issues

The comment states that "Direction remains the sampling-side marker," but setting the reversed edge's direct to the original edge.getDirect() (e.g., IN) could confuse subsequent logic—such as matching, sorting, or serialization—that relies on edge.direct. Since the endpoints have been reversed but the direct flag remains IN, this increases cognitive load and the risk of errors (e.g., if comparison logic in compareEdges or LogicalEdgeId depends on direct or endpoint order).
Even if the code passes current tests, this semantic ambiguity creates a hidden pitfall for future extensions or other modules that might use normalize and make decisions based on edge.direct.

Recommendations

Clarify the design: The purpose of normalize is to return an edge in "logical order" (src -> target), and edge.getDirect() should reflect this logical direction. (If you need to distinguish between the original storage direction and the logical direction from a sampling perspective, use an additional field or wrapper rather than overloading the semantics of IEdge.direct.)
Specific fix (simple and safe): Set the direction to OUT after reversing (indicating the current state is logical src -> target):

IEdge<K, EV> reversed = edge.reverse();
reversed.setDirect(EdgeDirection.OUT);
return reversed;

Alternatively, be more explicit:
Add a method or field to IEdge to distinguish between "storageDirection" and "logicalDirection," or
Clearly state in the normalize comments—and assert in documentation or code—that the normalized edge will not have an IN direction.
Rationale: Reduce ambiguity and ensure consistency in subsequent processing based on endpoints and direction (e.g., neighborId, compareEdges, etc.).

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for pointing this out. I agree that preserving IN after reversing the edge endpoints creates an inconsistent representation: the edge is normalized to logical src -> target order, but its direction still describes the original storage-side view.

I will update normalize(...) so that an incoming edge is reversed and returned with direct = OUT. Direction filtering will still occur before normalization, using the original edge direction, so sample(..., EdgeDirection.IN, ...) will continue to select incoming storage edges correctly. The returned sampled edge will then consistently represent its normalized logical orientation.

The original input edge will remain unchanged. I will also update the regression tests to verify that:

  • an input IN edge is still selected for IN sampling;
  • the returned normalized edge has reversed endpoints and direct = OUT;
  • the original input edge retains its original endpoints and IN direction.

return mix64(value);
}

private static long stableHash(Object value) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sampleScore / stableHash generates a stable hash by combining value.toString() with the class name—an approach that may entail performance and stability issues.

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);
}

Potential Issues and Risks

  • It relies on value.toString() semantics (where the string representation of an object may vary across implementations, JVMs, or versions), which can incur performance overhead for complex or large objects.
  • It may conflict with serialization consistency requirements (e.g., if toString() behavior differs across nodes in a distributed environment, it could compromise the cross-process reproducibility of deterministic sampling).
  • Stability risks are a concern in distributed deployments (a more robust ID hashing strategy is recommended).

Recommendations

  • Switch to a more stable and efficient hashing strategy, such as:
  • If K is serializable or has a stable ID string (e.g., Long or String), prioritize using the bytes of these native IDs.
  • Alternatively, combine K.hashCode() with the hash of the class name (while hashCode() is not guaranteed to be stable across implementations, it is usually acceptable for common ID types).
  • Ideally, use a known 64-bit non-cryptographic hash function (like MurmurHash3 or CityHash) based on the object's bytes.
  • Simple alternative example (if K is a common primitive or string):
private static long stableHash(Object value) {
String s = String.valueOf(value);
return mix64(s.hashCode() * 31L ^ value.getClass().getName().hashCode());
}

Rationale: Improves cross-node reproducibility and performance (by avoiding the frequent construction of long strings).

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for raising this. I agree that deriving the sampling score from value.toString() and the runtime class name is not a reliable distributed- systems contract. It can add avoidable allocation and make deterministic sampling depend on presentation-level behavior rather than the canonical
vertex ID representation.

I will replace this with an explicit stable ID byte encoder in the sampling and projection APIs. The sampler will hash those canonical bytes directly using the existing 64-bit mixing flow, and use the same byte representation as the deterministic tie-breaker when the supplied ID comparator returns equality.

For the DSL runtime, the encoder will come from the graph schema’s IType.serialize(...), so it matches the representation already used for graph IDs. There will be no implicit toString() or hashCode() fallback: callers with custom ID types must provide a stable cross-worker encoding explicitly.

I will also add regression coverage for IDs whose toString() must not be called, plus validation for missing or null ID encodings.

@aotenjou

Copy link
Copy Markdown
Author

@kitalkuyo-gita I have resolved the above-mentioned issue and ran ci locally.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants