diff --git a/geaflow/geaflow-common/src/main/java/org/apache/geaflow/common/config/keys/ExecutionConfigKeys.java b/geaflow/geaflow-common/src/main/java/org/apache/geaflow/common/config/keys/ExecutionConfigKeys.java index d1d3f3035..722e35fbe 100644 --- a/geaflow/geaflow-common/src/main/java/org/apache/geaflow/common/config/keys/ExecutionConfigKeys.java +++ b/geaflow/geaflow-common/src/main/java/org/apache/geaflow/common/config/keys/ExecutionConfigKeys.java @@ -638,4 +638,13 @@ public class ExecutionConfigKeys implements Serializable { .defaultValue(false) .description("if enable detail job metric"); + // ------------------------------------------------------------------------ + // optimizer + // ------------------------------------------------------------------------ + + public static final ConfigKey LOCAL_SHUFFLE_OPTIMIZATION_ENABLE = ConfigKeys + .key("geaflow.local.shuffle.optimization.enable") + .defaultValue(false) + .description("whether to enable local shuffle optimization for graph → sink/map patterns"); + } diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilder.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilder.java index 964b38366..864111d4d 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilder.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilder.java @@ -183,6 +183,10 @@ public ExecutionGraph buildExecutionGraph(Configuration jobConf) { /** * Build execution vertex group. + * + *

This method supports co-location hints from LocalShuffleOptimizer. + * Vertices with the same coLocationGroup will be placed in the same execution group + * to enable automatic local shuffle through LocalInputChannel. */ private Map buildExecutionVertexGroup(Map vertexId2GroupIdMap, Queue pipelineVertexQueue) { @@ -191,6 +195,47 @@ private Map buildExecutionVertexGroup(Map groupedVertices = new HashSet<>(); + // Step 1: Process co-location groups first for local shuffle optimization. + // Collect vertices by coLocationGroup for local shuffle optimization. + Map> coLocationGroupMap = new HashMap<>(); + + for (PipelineVertex vertex : plan.getVertexMap().values()) { + String coLocationGroup = vertex.getCoLocationGroup(); + if (coLocationGroup != null && !coLocationGroup.isEmpty()) { + coLocationGroupMap.computeIfAbsent(coLocationGroup, k -> new ArrayList<>()).add(vertex); + } + } + + // Process co-located vertices first. + for (Map.Entry> entry : coLocationGroupMap.entrySet()) { + List coLocatedVertices = entry.getValue(); + if (coLocatedVertices.isEmpty()) { + continue; + } + + // Create execution group for co-located vertices. + ExecutionVertexGroup vertexGroup = new ExecutionVertexGroup(groupId); + Map currentVertexGroupMap = new HashMap<>(); + + for (PipelineVertex vertex : coLocatedVertices) { + if (!groupedVertices.contains(vertex.getVertexId())) { + ExecutionVertex executionVertex = buildExecutionVertex(vertex); + currentVertexGroupMap.put(vertex.getVertexId(), executionVertex); + groupedVertices.add(vertex.getVertexId()); + vertexId2GroupIdMap.put(vertex.getVertexId(), groupId); + } + } + + if (!currentVertexGroupMap.isEmpty()) { + vertexGroup.getVertexMap().putAll(currentVertexGroupMap); + vertexGroupMap.put(groupId, vertexGroup); + LOGGER.info("Created co-located execution group {} with {} vertices for coLocationGroup '{}'", + groupId, currentVertexGroupMap.size(), entry.getKey()); + groupId++; + } + } + + // Step 2: Process remaining vertices using standard grouping logic. while (!pipelineVertexQueue.isEmpty()) { PipelineVertex pipelineVertex = pipelineVertexQueue.poll(); // Ignore already grouped vertex. diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/PipelinePlanBuilder.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/PipelinePlanBuilder.java index 8ab196b4d..12bf2a9a4 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/PipelinePlanBuilder.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/PipelinePlanBuilder.java @@ -543,7 +543,7 @@ private void optimizePipelinePlan(Configuration pipelineConfig) { LOGGER.info("union optimize: {}", new PlanGraphVisualization(pipelineGraph).getGraphviz()); } - new PipelineGraphOptimizer().optimizePipelineGraph(pipelineGraph); + new PipelineGraphOptimizer().optimizePipelineGraph(pipelineGraph, pipelineConfig); } } diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/graph/PipelineVertex.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/graph/PipelineVertex.java index 8c31ba979..b8f978e2c 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/graph/PipelineVertex.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/graph/PipelineVertex.java @@ -35,6 +35,12 @@ public class PipelineVertex implements Serializable { private AffinityLevel affinity; private boolean duplication; private VertexType chainTailType; + /** + * Co-location group ID for local shuffle optimization. + * Vertices with the same co-location group ID should be deployed on the same node + * to enable automatic local shuffle through LocalInputChannel. + */ + private String coLocationGroup; public PipelineVertex(int vertexId, OP operator, int parallelism) { this.vertexId = vertexId; @@ -139,6 +145,14 @@ public void setChainTailType(VertexType chainTailType) { this.chainTailType = chainTailType; } + public String getCoLocationGroup() { + return coLocationGroup; + } + + public void setCoLocationGroup(String coLocationGroup) { + this.coLocationGroup = coLocationGroup; + } + public String getVertexString() { String operatorStr = operator.toString(); return String.format("%s, p:%d, %s", getVertexName(), parallelism, operatorStr); diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/optimizer/PipelineGraphOptimizer.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/optimizer/PipelineGraphOptimizer.java index ab979332c..47405e897 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/optimizer/PipelineGraphOptimizer.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/optimizer/PipelineGraphOptimizer.java @@ -20,18 +20,30 @@ package org.apache.geaflow.plan.optimizer; import java.io.Serializable; +import org.apache.geaflow.common.config.Configuration; +import org.apache.geaflow.common.config.keys.ExecutionConfigKeys; import org.apache.geaflow.plan.graph.PipelineGraph; import org.apache.geaflow.plan.optimizer.strategy.ChainCombiner; +import org.apache.geaflow.plan.optimizer.strategy.LocalShuffleOptimizer; import org.apache.geaflow.plan.optimizer.strategy.SingleWindowGroupRule; public class PipelineGraphOptimizer implements Serializable { - public void optimizePipelineGraph(PipelineGraph pipelineGraph) { - // Enforce chain combiner opt. + public void optimizePipelineGraph(PipelineGraph pipelineGraph, Configuration config) { + // 1. Enforce chain combiner optimization. + // Merge operators with forward partition into single execution unit. ChainCombiner chainCombiner = new ChainCombiner(); chainCombiner.combineVertex(pipelineGraph); - // Enforce single window rule. + // 2. Local shuffle optimization (disabled by default). + // Mark vertices for co-location to enable automatic local shuffle. + if (config.getBoolean(ExecutionConfigKeys.LOCAL_SHUFFLE_OPTIMIZATION_ENABLE)) { + LocalShuffleOptimizer localShuffleOptimizer = new LocalShuffleOptimizer(); + localShuffleOptimizer.optimize(pipelineGraph); + } + + // 3. Enforce single window rule. + // Disable grouping for single-window batch jobs. SingleWindowGroupRule groupRule = new SingleWindowGroupRule(); groupRule.apply(pipelineGraph); } diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/optimizer/strategy/LocalShuffleOptimizer.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/optimizer/strategy/LocalShuffleOptimizer.java new file mode 100644 index 000000000..0b0e87d03 --- /dev/null +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/main/java/org/apache/geaflow/plan/optimizer/strategy/LocalShuffleOptimizer.java @@ -0,0 +1,300 @@ +/* + * 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.plan.optimizer.strategy; + +import java.io.Serializable; +import java.util.Collection; +import java.util.Map; +import java.util.Set; +import org.apache.geaflow.operator.impl.graph.algo.vc.IGraphVertexCentricAggOp; +import org.apache.geaflow.partitioner.IPartitioner; +import org.apache.geaflow.plan.graph.PipelineEdge; +import org.apache.geaflow.plan.graph.PipelineGraph; +import org.apache.geaflow.plan.graph.PipelineVertex; +import org.apache.geaflow.plan.graph.VertexType; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Local Shuffle Optimizer for graph traversal/computation to sink patterns. + * + *

Optimization scenario: When graph traversal or graph computation operators are followed + * only by Sink or Map operators, mark these vertices for co-location. This enables the runtime + * to automatically use LocalInputChannel for zero-copy memory transfer instead of network shuffle. + * + *

Pattern detected: + * 

+ * GraphTraversal/GraphComputation Operator
+ *     ↓ (Forward Partition)
+ * Sink/Map Operator (single input)
+ *
+ * + *

Optimization conditions: + *

+ * + *

Performance benefits: + *

+ */ +public class LocalShuffleOptimizer implements Serializable { + + private static final Logger LOGGER = LoggerFactory.getLogger(LocalShuffleOptimizer.class); + + /** + * Apply local shuffle optimization to the pipeline graph. + * + * @param pipelineGraph the pipeline graph to optimize + */ + public void optimize(PipelineGraph pipelineGraph) { + Map vertexMap = pipelineGraph.getVertexMap(); + Collection edges = pipelineGraph.getPipelineEdgeList(); + + int optimizedCount = 0; + int skippedCount = 0; + + for (PipelineEdge edge : edges) { + PipelineVertex srcVertex = vertexMap.get(edge.getSrcId()); + PipelineVertex targetVertex = vertexMap.get(edge.getTargetId()); + + if (srcVertex == null || targetVertex == null) { + continue; + } + + // Skip self-loop edges (iteration edges) + if (edge.getSrcId() == edge.getTargetId()) { + continue; + } + + // Check if eligible for local shuffle optimization + if (isEligibleForLocalShuffle(srcVertex, targetVertex, edge, pipelineGraph)) { + // Mark vertices for co-location + markForCoLocation(srcVertex, targetVertex); + optimizedCount++; + + LOGGER.info("LocalShuffleOptimizer: Marked vertices {} -> {} for co-location " + + "(parallelism: {} -> {})", + srcVertex.getVertexId(), targetVertex.getVertexId(), + srcVertex.getParallelism(), targetVertex.getParallelism()); + } else { + skippedCount++; + } + } + + LOGGER.info("LocalShuffleOptimizer: Optimized {} edges, skipped {} edges", + optimizedCount, skippedCount); + } + + /** + * Check if an edge is eligible for local shuffle optimization. + * + * @param srcVertex source vertex + * @param targetVertex target vertex + * @param edge the edge connecting source and target + * @param pipelineGraph the pipeline graph + * @return true if eligible for optimization + */ + private boolean isEligibleForLocalShuffle(PipelineVertex srcVertex, + PipelineVertex targetVertex, + PipelineEdge edge, + PipelineGraph pipelineGraph) { + // Condition 0: Exclude vertices with aggregation requirement + // Vertices implementing IGraphVertexCentricAggOp must stay grouped with their + // aggregation vertex (ID=0) to satisfy SchedulerGraphAggregateProcessor validation + if (hasAggregationRequirement(srcVertex) || hasAggregationRequirement(targetVertex)) { + LOGGER.debug("Skipping co-location for vertex with aggregation requirement: " + + "src={}, target={}", + srcVertex.getVertexId(), targetVertex.getVertexId()); + return false; + } + + // Condition 1: Source vertex must be a graph operator + if (!isGraphOperator(srcVertex)) { + LOGGER.debug("Source vertex {} is not a graph operator, skipping", + srcVertex.getVertexId()); + return false; + } + + // Condition 2: Target vertex must be a sink or map operator + if (!isSinkOrMapOperator(targetVertex)) { + LOGGER.debug("Target vertex {} is not a sink/map operator, skipping", + targetVertex.getVertexId()); + return false; + } + + // Condition 3: Edge partition type must be FORWARD + if (edge.getPartition().getPartitionType() != IPartitioner.PartitionType.forward) { + LOGGER.debug("Edge {}->{} partition type is not FORWARD, skipping", + edge.getSrcId(), edge.getTargetId()); + return false; + } + + // Condition 4: Target vertex must have single input (in-degree = 1) + Set inputEdges = pipelineGraph.getVertexInputEdges(targetVertex.getVertexId()); + if (inputEdges.size() != 1) { + LOGGER.debug("Target vertex {} has {} inputs (expected 1), skipping", + targetVertex.getVertexId(), inputEdges.size()); + return false; + } + + // Condition 5: Parallelism must be compatible + if (!isParallelismCompatible(srcVertex, targetVertex)) { + LOGGER.debug("Parallelism incompatible: src={}, target={}, skipping", + srcVertex.getParallelism(), targetVertex.getParallelism()); + return false; + } + + return true; + } + + /** + * Check if vertex has aggregation requirement. + * + *

Vertices implementing IGraphVertexCentricAggOp must stay grouped with their + * aggregation vertex to ensure proper execution. These vertices should NOT be + * marked for co-location with downstream operators. + * + * @param vertex the vertex to check + * @return true if vertex has aggregation requirement + */ + private boolean hasAggregationRequirement(PipelineVertex vertex) { + if (vertex.getOperator() == null) { + return false; + } + return vertex.getOperator() instanceof IGraphVertexCentricAggOp; + } + + /** + * Check if a vertex is a graph operator (traversal or computation). + * + *

Graph operators include: + *

    + *
  • Graph traversal operators
    • + *
  • Graph algorithm/computation operators
    • + *
+ * + * @param vertex the vertex to check + * @return true if vertex is a graph operator + */ + private boolean isGraphOperator(PipelineVertex vertex) { + if (vertex.getOperator() == null) { + return false; + } + + String className = vertex.getOperator().getClass().getName(); + VertexType type = vertex.getType(); + + // Check for graph-related class names + boolean isGraphClass = className.contains("Graph") + || className.contains("Traversal") + || className.contains("Algorithm"); + + // Some graph traversals may be source vertices + boolean isGraphSource = type == VertexType.source && isGraphClass; + + return isGraphClass || isGraphSource; + } + + /** + * Check if a vertex is a Sink or Map operator. + * + * @param vertex the vertex to check + * @return true if vertex is a sink or map operator + */ + private boolean isSinkOrMapOperator(PipelineVertex vertex) { + if (vertex.getOperator() == null) { + return false; + } + + String className = vertex.getOperator().getClass().getName(); + VertexType type = vertex.getType(); + + // Check vertex type + if (type == VertexType.sink) { + return true; + } + + // Check class name for Map or Sink + return className.contains("Map") || className.contains("Sink"); + } + + /** + * Check if parallelism is compatible between source and target vertices. + * + *

Compatible cases: + *

    + *
  • Exact match: parallelism is equal
    • + *
  • Divisible ratio: source parallelism is multiple of target parallelism (e.g., 8→4)
    • + *
+ * + * @param srcVertex source vertex + * @param targetVertex target vertex + * @return true if parallelism is compatible + */ + private boolean isParallelismCompatible(PipelineVertex srcVertex, PipelineVertex targetVertex) { + int srcParallelism = srcVertex.getParallelism(); + int targetParallelism = targetVertex.getParallelism(); + + // Exact match + if (srcParallelism == targetParallelism) { + return true; + } + + // Allow source parallelism > target parallelism if divisible + // Example: 8 -> 4 (2:1 mapping), 12 -> 4 (3:1 mapping) + if (srcParallelism > targetParallelism && srcParallelism % targetParallelism == 0) { + LOGGER.debug("Parallelism compatible with {}:1 mapping: {} -> {}", + srcParallelism / targetParallelism, srcParallelism, targetParallelism); + return true; + } + + return false; + } + + /** + * Mark vertices for co-location by setting the same co-location group ID. + * + *

The co-location group ID is used by ExecutionGraphBuilder to place + * tasks on the same node, enabling automatic local shuffle through LocalInputChannel. + * + * @param srcVertex source vertex + * @param targetVertex target vertex + */ + private void markForCoLocation(PipelineVertex srcVertex, PipelineVertex targetVertex) { + // Generate co-location group ID + String coLocationGroupId = "local_shuffle_" + srcVertex.getVertexId(); + + // Set co-location markers + srcVertex.setCoLocationGroup(coLocationGroupId); + targetVertex.setCoLocationGroup(coLocationGroupId); + + LOGGER.debug("Marked vertices {} and {} with co-location group '{}'", + srcVertex.getVertexId(), targetVertex.getVertexId(), coLocationGroupId); + } +} diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilderTest.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilderTest.java index f4be83d11..ddf33012e 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilderTest.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/core/graph/builder/ExecutionGraphBuilderTest.java @@ -184,7 +184,7 @@ public void testAllWindowWithReduceTwoAndSinkFourConcurrency() { PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -225,7 +225,7 @@ public void testAllWindowWithSingleConcurrency() { PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -262,7 +262,7 @@ public void testOperatorChain() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -309,7 +309,7 @@ public void flatMap(String value, Collector collector) { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -368,7 +368,7 @@ public void testIncGraphCompute() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -414,7 +414,7 @@ public void testStaticGraphCompute() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -460,7 +460,7 @@ public void testAllWindowStaticGraphCompute() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -507,7 +507,7 @@ public void testWindowGraphTraversal() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -556,7 +556,7 @@ public void testMultiSourceWindowGraphTraversal() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -599,7 +599,7 @@ public void testAllWindowGraphTraversal() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -649,7 +649,7 @@ public void testTwoSourceWithGraphUnion() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -703,7 +703,7 @@ public void testThreeSourceWithGraphUnion() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -789,7 +789,7 @@ public void testTenSourceWithGraphUnion() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -828,7 +828,7 @@ public void testGroupVertexDiamondDependency() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); @@ -878,7 +878,7 @@ public void testMultiGraphTraversal() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); ExecutionGraphBuilder builder = new ExecutionGraphBuilder(pipelineGraph); ExecutionGraph graph = builder.buildExecutionGraph(new Configuration()); diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/PipelinePlanTest.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/PipelinePlanTest.java index 26ff4f355..753e85bdf 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/PipelinePlanTest.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/PipelinePlanTest.java @@ -225,7 +225,7 @@ public int generateId() { PipelinePlanBuilder planBuilder = new PipelinePlanBuilder(); PipelineGraph pipelineGraph = planBuilder.buildPlan(context); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); Map vertexMap = pipelineGraph.getVertexMap(); if (backendType == BackendType.Paimon) { diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/UnionTest.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/UnionTest.java index a3ebd984f..67d219439 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/UnionTest.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/UnionTest.java @@ -70,7 +70,7 @@ public void testUnionPlan() { Assert.assertEquals(vertexMap.size(), 6); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); Assert.assertEquals(vertexMap.size(), 4); } @@ -107,7 +107,7 @@ public void testMultiUnionPlan() { Assert.assertEquals(vertexMap.size(), 7); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); Assert.assertEquals(vertexMap.size(), 5); } @@ -138,7 +138,7 @@ public void testUnionWithKeyByPlan() { Assert.assertEquals(vertexMap.size(), 6); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); Assert.assertEquals(vertexMap.size(), 4); } @@ -170,7 +170,7 @@ public void testWindowUnionWithKeyByPlan() { Assert.assertEquals(vertexMap.size(), 6); PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); Assert.assertEquals(vertexMap.size(), 4); } diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/optimizer/PlanOptimizerTest.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/optimizer/PlanOptimizerTest.java index 17723f7f9..1b5935f13 100644 --- a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/optimizer/PlanOptimizerTest.java +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/optimizer/PlanOptimizerTest.java @@ -150,7 +150,7 @@ public void testOperatorChain() { } PipelineGraphOptimizer optimizer = new PipelineGraphOptimizer(); - optimizer.optimizePipelineGraph(pipelineGraph); + optimizer.optimizePipelineGraph(pipelineGraph, new Configuration()); Assert.assertEquals(pipelineGraph.getVertexMap().size(), 1); PipelineVertex sourceVertex = pipelineGraph.getVertexMap().get(1); Assert.assertEquals(((AbstractOperator) sourceVertex.getOperator()).getNextOperators().size(), 1); diff --git a/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/optimizer/strategy/LocalShuffleOptimizerTest.java b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/optimizer/strategy/LocalShuffleOptimizerTest.java new file mode 100644 index 000000000..dc9759dae --- /dev/null +++ b/geaflow/geaflow-core/geaflow-runtime/geaflow-plan/src/test/java/org/apache/geaflow/plan/optimizer/strategy/LocalShuffleOptimizerTest.java @@ -0,0 +1,458 @@ +/* + * 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.plan.optimizer.strategy; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ImmutableSet; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.geaflow.api.context.RuntimeContext; +import org.apache.geaflow.collector.ICollector; +import org.apache.geaflow.common.encoder.IEncoder; +import org.apache.geaflow.operator.Operator; +import org.apache.geaflow.partitioner.IPartitioner; +import org.apache.geaflow.partitioner.impl.ForwardPartitioner; +import org.apache.geaflow.partitioner.impl.KeyPartitioner; +import org.apache.geaflow.plan.graph.PipelineEdge; +import org.apache.geaflow.plan.graph.PipelineGraph; +import org.apache.geaflow.plan.graph.PipelineVertex; +import org.apache.geaflow.plan.graph.VertexType; +import org.testng.Assert; +import org.testng.annotations.Test; + +/** + * Unit tests for LocalShuffleOptimizer. + * + *

This test class verifies the local shuffle optimization logic which detects + * graph operator → sink/map patterns and marks them for co-location to enable + * automatic local shuffle through LocalInputChannel. + */ +public class LocalShuffleOptimizerTest { + + /** + * Test basic optimization scenario: GraphTraversal → Sink with forward partition. + * + *

Expected behavior: + * - Both vertices should be marked with the same coLocationGroup + * - Optimization should succeed for matching parallelism and forward partition + */ + @Test + public void testGraphToSinkWithForwardPartition() { + // Create mock pipeline graph + PipelineGraph pipelineGraph = createMockPipelineGraph(); + + // Create vertices + PipelineVertex graphVertex = createGraphVertex(1, 4); // Graph operator, parallelism 4 + PipelineVertex sinkVertex = createSinkVertex(2, 4); // Sink operator, parallelism 4 + + // Add vertices to graph + Map vertexMap = new HashMap<>(); + vertexMap.put(1, graphVertex); + vertexMap.put(2, sinkVertex); + when(pipelineGraph.getVertexMap()).thenReturn(vertexMap); + + // Create edge with forward partition + PipelineEdge edge = createEdge(1, 1, 2, new ForwardPartitioner<>()); + List edges = new ArrayList<>(); + edges.add(edge); + when(pipelineGraph.getPipelineEdgeList()).thenReturn(edges); + + // Mock single input for sink vertex + mockSingleInput(pipelineGraph, 2, edge); + + // Apply optimization + LocalShuffleOptimizer optimizer = new LocalShuffleOptimizer(); + optimizer.optimize(pipelineGraph); + + // Verify co-location + Assert.assertNotNull(graphVertex.getCoLocationGroup(), + "Graph vertex should have coLocationGroup"); + Assert.assertNotNull(sinkVertex.getCoLocationGroup(), + "Sink vertex should have coLocationGroup"); + Assert.assertEquals(graphVertex.getCoLocationGroup(), sinkVertex.getCoLocationGroup(), + "Both vertices should have the same coLocationGroup"); + } + + /** + * Test chain scenario: GraphTraversal → Map → Sink. + * + *

Expected behavior: + * - Graph → Map should be optimized (forward partition, single input, graph→sink/map pattern) + * - Map → Sink should NOT be optimized (source is not graph operator) + * - Graph and Map vertices should have co-location group + * - Sink vertex may or may not have co-location (depends on whether map was marked) + */ + @Test + public void testGraphToMapToSinkChain() { + // Create mock pipeline graph + PipelineGraph pipelineGraph = createMockPipelineGraph(); + + // Create vertices + PipelineVertex graphVertex = createGraphVertex(1, 4); + PipelineVertex mapVertex = createMapVertex(2, 4); + PipelineVertex sinkVertex = createSinkVertex(3, 4); + + // Add vertices to graph + Map vertexMap = new HashMap<>(); + vertexMap.put(1, graphVertex); + vertexMap.put(2, mapVertex); + vertexMap.put(3, sinkVertex); + when(pipelineGraph.getVertexMap()).thenReturn(vertexMap); + + // Create edges with forward partition + PipelineEdge edge1 = createEdge(1, 1, 2, new ForwardPartitioner<>()); + PipelineEdge edge2 = createEdge(2, 2, 3, new ForwardPartitioner<>()); + List edges = new ArrayList<>(); + edges.add(edge1); + edges.add(edge2); + when(pipelineGraph.getPipelineEdgeList()).thenReturn(edges); + + // Mock single inputs + mockSingleInput(pipelineGraph, 2, edge1); + mockSingleInput(pipelineGraph, 3, edge2); + + // Apply optimization + LocalShuffleOptimizer optimizer = new LocalShuffleOptimizer(); + optimizer.optimize(pipelineGraph); + + // Verify co-location for graph → map (should be optimized) + Assert.assertNotNull(graphVertex.getCoLocationGroup(), + "Graph vertex should have coLocationGroup"); + Assert.assertNotNull(mapVertex.getCoLocationGroup(), + "Map vertex should have coLocationGroup"); + Assert.assertEquals(graphVertex.getCoLocationGroup(), mapVertex.getCoLocationGroup(), + "Graph and Map vertices should share same coLocationGroup"); + + // Map → Sink should NOT be optimized (map is not a graph operator) + // So Sink may not have a coLocationGroup, or if Map was marked by first edge, + // Sink will not share it because second edge doesn't meet criteria + } + + /** + * Test negative case: Graph → Sink with KEY partition (should not optimize). + * + *

Expected behavior: + * - Optimization should be skipped because partition type is not FORWARD + * - Vertices should NOT have coLocationGroup + */ + @Test + public void testNoOptimizationForKeyPartition() { + // Create mock pipeline graph + PipelineGraph pipelineGraph = createMockPipelineGraph(); + + // Create vertices + PipelineVertex graphVertex = createGraphVertex(1, 4); + PipelineVertex sinkVertex = createSinkVertex(2, 4); + + // Add vertices to graph + Map vertexMap = new HashMap<>(); + vertexMap.put(1, graphVertex); + vertexMap.put(2, sinkVertex); + when(pipelineGraph.getVertexMap()).thenReturn(vertexMap); + + // Create edge with KEY partition (not FORWARD) + PipelineEdge edge = createEdge(1, 1, 2, new KeyPartitioner<>(1)); + List edges = new ArrayList<>(); + edges.add(edge); + when(pipelineGraph.getPipelineEdgeList()).thenReturn(edges); + + // Mock single input for sink vertex + mockSingleInput(pipelineGraph, 2, edge); + + // Apply optimization + LocalShuffleOptimizer optimizer = new LocalShuffleOptimizer(); + optimizer.optimize(pipelineGraph); + + // Verify NO co-location due to key partition + Assert.assertNull(graphVertex.getCoLocationGroup(), + "Graph vertex should NOT have coLocationGroup with key partition"); + Assert.assertNull(sinkVertex.getCoLocationGroup(), + "Sink vertex should NOT have coLocationGroup with key partition"); + } + + /** + * Test negative case: Sink with multiple inputs (should not optimize). + * + *

Expected behavior: + * - Optimization should be skipped because sink has multiple inputs + * - Vertices should NOT have coLocationGroup + */ + @Test + public void testSinkWithMultipleInputs() { + // Create mock pipeline graph + PipelineGraph pipelineGraph = createMockPipelineGraph(); + + // Create vertices + PipelineVertex graphVertex1 = createGraphVertex(1, 4); + PipelineVertex graphVertex2 = createGraphVertex(2, 4); + PipelineVertex sinkVertex = createSinkVertex(3, 4); + + // Add vertices to graph + Map vertexMap = new HashMap<>(); + vertexMap.put(1, graphVertex1); + vertexMap.put(2, graphVertex2); + vertexMap.put(3, sinkVertex); + when(pipelineGraph.getVertexMap()).thenReturn(vertexMap); + + // Create edges - TWO inputs to sink + PipelineEdge edge1 = createEdge(1, 1, 3, new ForwardPartitioner<>()); + PipelineEdge edge2 = createEdge(2, 2, 3, new ForwardPartitioner<>()); + List edges = new ArrayList<>(); + edges.add(edge1); + edges.add(edge2); + when(pipelineGraph.getPipelineEdgeList()).thenReturn(edges); + + // Mock MULTIPLE inputs for sink vertex + mockMultipleInputs(pipelineGraph, 3, edge1, edge2); + + // Apply optimization + LocalShuffleOptimizer optimizer = new LocalShuffleOptimizer(); + optimizer.optimize(pipelineGraph); + + // Verify NO co-location due to multiple inputs + Assert.assertNull(graphVertex1.getCoLocationGroup(), + "Graph vertex 1 should NOT have coLocationGroup with multiple sink inputs"); + Assert.assertNull(graphVertex2.getCoLocationGroup(), + "Graph vertex 2 should NOT have coLocationGroup with multiple sink inputs"); + Assert.assertNull(sinkVertex.getCoLocationGroup(), + "Sink vertex should NOT have coLocationGroup with multiple inputs"); + } + + /** + * Test negative case: Parallelism mismatch (should not optimize). + * + *

Expected behavior: + * - Optimization should be skipped when parallelism doesn't match or divide evenly + * - Example: 8 → 3 is not divisible, should fail + */ + @Test + public void testParallelismMismatch() { + // Create mock pipeline graph + PipelineGraph pipelineGraph = createMockPipelineGraph(); + + // Create vertices with mismatched parallelism (8 → 3, not divisible) + PipelineVertex graphVertex = createGraphVertex(1, 8); + PipelineVertex sinkVertex = createSinkVertex(2, 3); + + // Add vertices to graph + Map vertexMap = new HashMap<>(); + vertexMap.put(1, graphVertex); + vertexMap.put(2, sinkVertex); + when(pipelineGraph.getVertexMap()).thenReturn(vertexMap); + + // Create edge with forward partition + PipelineEdge edge = createEdge(1, 1, 2, new ForwardPartitioner<>()); + List edges = new ArrayList<>(); + edges.add(edge); + when(pipelineGraph.getPipelineEdgeList()).thenReturn(edges); + + // Mock single input for sink vertex + mockSingleInput(pipelineGraph, 2, edge); + + // Apply optimization + LocalShuffleOptimizer optimizer = new LocalShuffleOptimizer(); + optimizer.optimize(pipelineGraph); + + // Verify NO co-location due to parallelism mismatch + Assert.assertNull(graphVertex.getCoLocationGroup(), + "Graph vertex should NOT have coLocationGroup with parallelism mismatch"); + Assert.assertNull(sinkVertex.getCoLocationGroup(), + "Sink vertex should NOT have coLocationGroup with parallelism mismatch"); + } + + /** + * Test positive case: Compatible parallelism ratio (should optimize). + * + *

Expected behavior: + * - Optimization should succeed for divisible parallelism ratios + * - Example: 8 → 4 (2:1 ratio), 12 → 4 (3:1 ratio) should both succeed + */ + @Test + public void testCompatibleParallelismRatio() { + // Test case 1: 8 → 4 (2:1 ratio) + PipelineGraph pipelineGraph1 = createMockPipelineGraph(); + PipelineVertex graphVertex1 = createGraphVertex(1, 8); + PipelineVertex sinkVertex1 = createSinkVertex(2, 4); + + Map vertexMap1 = new HashMap<>(); + vertexMap1.put(1, graphVertex1); + vertexMap1.put(2, sinkVertex1); + when(pipelineGraph1.getVertexMap()).thenReturn(vertexMap1); + + PipelineEdge edge1 = createEdge(1, 1, 2, new ForwardPartitioner<>()); + List edges1 = new ArrayList<>(); + edges1.add(edge1); + when(pipelineGraph1.getPipelineEdgeList()).thenReturn(edges1); + mockSingleInput(pipelineGraph1, 2, edge1); + + LocalShuffleOptimizer optimizer1 = new LocalShuffleOptimizer(); + optimizer1.optimize(pipelineGraph1); + + Assert.assertNotNull(graphVertex1.getCoLocationGroup(), + "Graph vertex should have coLocationGroup with 8→4 parallelism"); + Assert.assertNotNull(sinkVertex1.getCoLocationGroup(), + "Sink vertex should have coLocationGroup with 8→4 parallelism"); + + // Test case 2: 12 → 4 (3:1 ratio) + PipelineGraph pipelineGraph2 = createMockPipelineGraph(); + PipelineVertex graphVertex2 = createGraphVertex(3, 12); + PipelineVertex sinkVertex2 = createSinkVertex(4, 4); + + Map vertexMap2 = new HashMap<>(); + vertexMap2.put(3, graphVertex2); + vertexMap2.put(4, sinkVertex2); + when(pipelineGraph2.getVertexMap()).thenReturn(vertexMap2); + + PipelineEdge edge2 = createEdge(2, 3, 4, new ForwardPartitioner<>()); + List edges2 = new ArrayList<>(); + edges2.add(edge2); + when(pipelineGraph2.getPipelineEdgeList()).thenReturn(edges2); + mockSingleInput(pipelineGraph2, 4, edge2); + + LocalShuffleOptimizer optimizer2 = new LocalShuffleOptimizer(); + optimizer2.optimize(pipelineGraph2); + + Assert.assertNotNull(graphVertex2.getCoLocationGroup(), + "Graph vertex should have coLocationGroup with 12→4 parallelism"); + Assert.assertNotNull(sinkVertex2.getCoLocationGroup(), + "Sink vertex should have coLocationGroup with 12→4 parallelism"); + } + + // ==================== Helper Methods ==================== + + /** + * Create a mock PipelineGraph. + */ + private PipelineGraph createMockPipelineGraph() { + return mock(PipelineGraph.class); + } + + /** + * Create a graph operator vertex. + */ + private PipelineVertex createGraphVertex(int id, int parallelism) { + Operator operator = new MockGraphTraversalOperator(); + PipelineVertex vertex = new PipelineVertex(id, operator, VertexType.source, parallelism); + return vertex; + } + + /** + * Create a sink operator vertex. + */ + private PipelineVertex createSinkVertex(int id, int parallelism) { + Operator operator = new MockSinkOperator(); + PipelineVertex vertex = new PipelineVertex(id, operator, VertexType.sink, parallelism); + return vertex; + } + + /** + * Create a map operator vertex. + */ + private PipelineVertex createMapVertex(int id, int parallelism) { + Operator operator = new MockMapOperator(); + PipelineVertex vertex = new PipelineVertex(id, operator, VertexType.process, parallelism); + return vertex; + } + + /** + * Create a pipeline edge. + */ + private PipelineEdge createEdge(int edgeId, int srcId, int targetId, IPartitioner partitioner) { + IEncoder encoder = mock(IEncoder.class); + return new PipelineEdge(edgeId, srcId, targetId, partitioner, encoder); + } + + /** + * Mock single input for a vertex. + */ + private void mockSingleInput(PipelineGraph graph, int vertexId, PipelineEdge edge) { + Set inputEdges = ImmutableSet.of(edge); + when(graph.getVertexInputEdges(vertexId)).thenReturn(inputEdges); + } + + /** + * Mock multiple inputs for a vertex. + */ + private void mockMultipleInputs(PipelineGraph graph, int vertexId, PipelineEdge... edges) { + Set inputEdges = ImmutableSet.copyOf(edges); + when(graph.getVertexInputEdges(vertexId)).thenReturn(inputEdges); + } + + // ==================== Mock Operator Classes ==================== + + /** + * Mock GraphTraversal operator (name contains "Graph" and "Traversal"). + */ + private static class MockGraphTraversalOperator implements Operator { + + @Override + public void open(OpContext opContext) { + } + + @Override + public void finish() { + } + + @Override + public void close() { + } + } + + /** + * Mock Sink operator (name contains "Sink"). + */ + private static class MockSinkOperator implements Operator { + + @Override + public void open(OpContext opContext) { + } + + @Override + public void finish() { + } + + @Override + public void close() { + } + } + + /** + * Mock Map operator (name contains "Map"). + */ + private static class MockMapOperator implements Operator { + + @Override + public void open(OpContext opContext) { + } + + @Override + public void finish() { + } + + @Override + public void close() { + } + } +}