Re: [I] Support subgraph sampling for GNN [geaflow]

["kaori-seasons (via GitHub)" <gi...@apache.org>]

kaori-seasons commented on issue #502:
URL: https://github.com/apache/geaflow/issues/502#issuecomment-3509623766

   @Loognqiang Hello, I have designed a GNN sampling scheme. You can take a look.
   
   ### 1.1 Core Capabilities Already Available
   
   GeaFlow already provides solid infrastructure:
   
   **Edge Iterator Mechanism**: The system has implemented `EdgeScanIterator`, supporting iterative edge access and filtering [1](#1-0) .
   
   **Edge Limit Support**: The `IEdgeLimit` interface supports limiting the number of incoming and outgoing edges per vertex<cite/>.
   
   **Filter Pushdown**: `StatePushDown` supports pushing filter conditions and limits down to the storage layer<cite/>.
   
   **Vertex-Grouped Edge Iteration**: `EdgeListScanIterator` has implemented edge iteration grouped by source vertex ID [2](#1-1) .
   
   **Algorithm Runtime Context**: `AlgorithmRuntimeContext` provides iterator interfaces for loading edges<cite/>.
   
   **K-hop Algorithm Example**: The existing K-hop implementation demonstrates multi-layer neighbor traversal patterns<cite/>.
   
   ### 1.2 Capabilities Requiring Extension
   
   Current architecture **limitations**:
   1. `LimitFilter` only supports sequential limiting, not random sampling<cite/>
   2. No dedicated layered sampling API for GNN scenarios<cite/>
   3. Lack of flexible sampling strategy configuration (random sampling, importance sampling, etc.)<cite/>
   
   ## II. Core Design Plan
   
   ### 2.1 New GNN Sampling API Design
   
   #### 2.1.1 Sampling Configuration Interface
   
   ```java
   /**
    * GNN subgraph sampling configuration
    */
   public interface IGNNSamplingConfig extends Serializable {
       
       // Sample size configuration per layer
       List<Integer> getLayerSampleSizes();
       
       // Sampling depth (K)
       int getSamplingDepth();
       
       // Sampling strategy
       SamplingStrategy getStrategy();
       
       // Edge direction
       EdgeDirection getDirection();
       
       // Include vertex properties
       boolean includeVertexProperties();
       
       // Include edge properties
       boolean includeEdgeProperties();
   }
   ```
   
   #### 2.1.2 Sampling Strategy Enum
   
   ```java
   public enum SamplingStrategy {
       RANDOM,           // Random sampling
       TOP_K,            // Top K
       WEIGHTED_RANDOM,  // Weighted random sampling
       RESERVOIR         // Reservoir sampling
   }
   ```
   
   #### 2.1.3 Extend AlgorithmRuntimeContext
   
   Add new methods to the `AlgorithmRuntimeContext` interface:
   
   ```java
   /**
    * Load edge iterator with sampling configuration
    */
   CloseableIterator<RowEdge> loadEdgesWithSampling(
       EdgeDirection direction, 
       IGNNSamplingConfig samplingConfig
   );
   
   /**
    * Load K-layer neighbor subgraph
    */
   GNNSubgraph loadKLayerSubgraph(
       IGNNSamplingConfig samplingConfig
   );
   ```
   
   ### 2.2 Random Sampling Filter Implementation
   
   #### 2.2.1 RandomSamplingEdgeLimit
   
   Extend the existing `IEdgeLimit` interface:
   
   ```java
   public class RandomSamplingEdgeLimit implements IEdgeLimit {
       private final long sampleSize;
       private final long totalEdgeCount;
       private final Random random;
       private final EdgeDirection direction;
       
       @Override
       public long inEdgeLimit() {
           return direction == EdgeDirection.IN ? sampleSize : Long.MAX_VALUE;
       }
       
       @Override
       public long outEdgeLimit() {
           return direction == EdgeDirection.OUT ? sampleSize : Long.MAX_VALUE;
       }
       
       @Override
       public LimitType limitType() {
           return LimitType.COMPOSED;
       }
   }
   ```
   
   #### 2.2.2 RandomSamplingFilter
   
   Implement random sampling filter based on reservoir algorithm:
   
   ```java
   public class RandomSamplingFilter extends BaseGraphFilter {
       private final IGraphFilter baseFilter;
       private final int sampleSize;
       private final Random random;
       private List<IEdge> reservoir;
       private int edgeCount;
       
       @Override
       public boolean filterEdge(IEdge edge) {
           if (!baseFilter.filterEdge(edge)) {
               return false;
           }
           
           // Reservoir sampling algorithm
           if (edgeCount < sampleSize) {
               reservoir.add(edge);
               edgeCount++;
               return true;
           } else {
               int replaceIndex = random.nextInt(edgeCount + 1);
               if (replaceIndex < sampleSize) {
                   reservoir.set(replaceIndex, edge);
               }
               edgeCount++;
               return false; // Return via reservoir later
           }
       }
   }
   ```
   
   ### 2.3 Layered Sampling Iterator
   
   #### 2.3.1 LayeredSamplingIterator
   
   ```java
   public class LayeredSamplingIterator<K, EV> 
       implements CloseableIterator<LayeredEdges<K, EV>> {
       
       private final GraphState<K, ?, EV> graphState;
       private final IGNNSamplingConfig config;
       private final Queue<LayerTask<K>> taskQueue;
       private final Map<K, Integer> visitedVertices; // Track visit layer
       
       private int currentLayer = 0;
       
       @Override
       public boolean hasNext() {
           return !taskQueue.isEmpty() && currentLayer < config.getSamplingDepth();
       }
       
       @Override
       public LayeredEdges<K, EV> next() {
           List<IEdge<K, EV>> edges = new ArrayList<>();
           int layerSampleSize = config.getLayerSampleSizes().get(currentLayer);
           
           while (!taskQueue.isEmpty()) {
               LayerTask<K> task = taskQueue.poll();
               
               // Configure sampling using StatePushDown
               StatePushDown pushDown = StatePushDown.of()
                   .withEdgeLimit(new RandomSamplingEdgeLimit(
                       layerSampleSize, 
                       config.getDirection()
                   ));
                   
               CloseableIterator<IEdge<K, EV>> edgeIter = 
                   graphState.staticGraph().getEdgeIterator(
                       Collections.singletonList(task.vertexId), 
                       pushDown
                   );
                   
               // Sample edges and add target vertices to next layer queue
               while (edgeIter.hasNext()) {
                   IEdge<K, EV> edge = edgeIter.next();
                   edges.add(edge);
                   
                   K targetId = edge.getTargetId();
                   if (!visitedVertices.containsKey(targetId)) {
                       taskQueue.offer(new LayerTask<>(targetId, currentLayer + 1));
                       visitedVertices.put(targetId, currentLayer + 1);
                   }
               }
           }
           
           currentLayer++;
           return new LayeredEdges<>(currentLayer - 1, edges);
       }
   }
   ```
   
   ### 2.4 GNN Subgraph Data Structure
   
   #### 2.4.1 GNNSubgraph
   
   ```java
   public class GNNSubgraph<K, VV, EV> implements Serializable {
       
       private final K rootVertexId;
       private final Map<Integer, Set<K>> layeredVertices; // Layer -> vertex set
       private final Map<Integer, List<IEdge<K, EV>>> layeredEdges; // Layer -> edge list
       private final Map<K, IVertex<K, VV>> vertexMap; // Vertex properties
       
       public int getDepth() {
           return layeredVertices.size();
       }
       
       public Set<K> getVerticesAtLayer(int layer) {
           return layeredVertices.getOrDefault(layer, Collections.emptySet());
       }
       
       public List<IEdge<K, EV>> getEdgesAtLayer(int layer) {
           return layeredEdges.getOrDefault(layer, Collections.emptyList());
       }
       
       public IVertex<K, VV> getVertex(K vertexId) {
           return vertexMap.get(vertexId);
       }
   }
   ```
   
   
   ## III. Usage Examples
   
   ### 3.1 Basic GNN Sampling UDF
   
   ```java
   @Description(name = "gnn_sampling", description = "GNN subgraph sampling")
   public class GNNSamplingAlgorithm implements AlgorithmUserFunction<Object, GNNMessage> {
       
       private AlgorithmRuntimeContext<Object, GNNMessage> context;
       private IGNNSamplingConfig samplingConfig;
       
       @Override
       public void init(AlgorithmRuntimeContext<Object, GNNMessage> context, Object[] params) {
           this.context = context;
           
           // Configuration: sample 10 neighbors per layer, depth of 2
           this.samplingConfig = GNNSamplingConfig.builder()
               .layerSampleSizes(Arrays.asList(10, 10))
               .samplingDepth(2)
               .strategy(SamplingStrategy.RANDOM)
               .direction(EdgeDirection.OUT)
               .includeVertexProperties(true)
               .includeEdgeProperties(true)
               .build();
       }
       
       @Override
       public void process(RowVertex vertex, Optional<Row> updatedValues, 
                          Iterator<GNNMessage> messages) {
           if (context.getCurrentIterationId() == 1L) {
               // First iteration: load K-layer subgraph
               GNNSubgraph subgraph = context.loadKLayerSubgraph(samplingConfig);
               
               // Process subgraph, send to GNN training component
               context.take(ObjectRow.create(
                   vertex.getId(), 
                   subgraph.serialize()
               ));
           }
       }
       
       @Override
       public void finish(RowVertex vertex, Optional<Row> newValue) {
           // Cleanup resources
       }
       
       @Override
       public StructType getOutputType(GraphSchema graphSchema) {
           return new StructType(
               new TableField("vertex_id", graphSchema.getIdType(), false),
               new TableField("subgraph", StringType.INSTANCE, false)
           );
       }
   }
   ```
   
   ### 3.2 Progressive Sampling Using Iterator
   
   ```java
   @Override
   public void process(RowVertex vertex, Optional<Row> updatedValues, 
                      Iterator<GNNMessage> messages) {
       
       // Load layer by layer using iterator to control memory usage
       for (int layer = 0; layer < samplingConfig.getSamplingDepth(); layer++) {
           
           // Current layer sample size
           int sampleSize = samplingConfig.getLayerSampleSizes().get(layer);
           
           // Load edges using iterator
           CloseableIterator<RowEdge> edgeIter = context.loadEdgesWithSampling(
               EdgeDirection.OUT,
               samplingConfig.withLayerSampleSize(sampleSize)
           );
           
           List<RowEdge> sampledEdges = new ArrayList<>();
           int count = 0;
           
           // Control load count in UDF
           while (edgeIter.hasNext() && count < sampleSize) {
               sampledEdges.add(edgeIter.next());
               count++;
           }
           
           // Process current layer sampling results
           processLayerEdges(layer, sampledEdges);
           
           edgeIter.close();
       }
   }
   ```


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: issues-unsubscribe@geaflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@geaflow.apache.org
For additional commands, e-mail: issues-help@geaflow.apache.org