feat: Add subgraph sampling as reusable multi-hop subgraph sampling foundation - #831
feat: Add subgraph sampling as reusable multi-hop subgraph sampling foundation#831aotenjou wants to merge 9 commits into
Conversation
50f3e6b to
052cb3c
Compare
| return reversed; | ||
| } | ||
|
|
||
| private static <K, EV> K neighborId(K vertexId, IEdge<K, EV> edge) { |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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) { |
There was a problem hiding this comment.
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.).
There was a problem hiding this comment.
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) { |
There was a problem hiding this comment.
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
Kis serializable or has a stable ID string (e.g.,LongorString), prioritize using the bytes of these native IDs. - Alternatively, combine
K.hashCode()with the hash of the class name (whilehashCode()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
Kis 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).
There was a problem hiding this comment.
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.
|
@kitalkuyo-gita I have resolved the above-mentioned issue and ran ci locally. |
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 > 0selects at most the requested number of neighborsfanout = -1returns all neighbors matching the requested directionAdd layered sampled-subgraph assembly
Add stable logical edge identities
sourceId + targetId + label + timeAdd an iterative multi-hop BSP protocol
REQUEST,RESPOND,COMMIT_AND_REQUEST, andCOMPLETEphasesIntegrate sampling with the DSL algorithm runtime
initIterationandfinishIterationlifecycle callbacksTesting
The following local validation has passed:
Coverage includes:
fanout = -1How was this PR tested?