Re: [PR] [BEAM-5090] Use topological sort during ProcessBundle in Java SDKHarness [beam]

[via GitHub]

gemini-code-assist[bot] commented on code in PR #36169:
URL: https://github.com/apache/beam/pull/36169#discussion_r2351505872


##########
sdks/java/harness/src/main/java/org/apache/beam/fn/harness/control/ProcessBundleHandler.java:
##########
@@ -843,32 +807,31 @@ public void afterBundleCommit(Instant callbackExpiry, 
Callback callback) {
             bundleFinalizationCallbackRegistrations,
             runnerCapabilities);
 
-    // Create a BeamFnStateClient
-    for (Map.Entry<String, PTransform> entry : 
bundleDescriptor.getTransformsMap().entrySet()) {
-
-      // Skip anything which isn't a root.
-      // Also force data output transforms to be unconditionally instantiated 
(see BEAM-10450).
-      // TODO: Remove source as a root and have it be triggered by the Runner.
-      if (!DATA_INPUT_URN.equals(entry.getValue().getSpec().getUrn())
-          && !DATA_OUTPUT_URN.equals(entry.getValue().getSpec().getUrn())
-          && !JAVA_SOURCE_URN.equals(entry.getValue().getSpec().getUrn())
-          && !PTransformTranslation.READ_TRANSFORM_URN.equals(
-              entry.getValue().getSpec().getUrn())) {
-        continue;
-      }
-
-      RunnerApi.Components components =
-          RunnerApi.Components.newBuilder()
-              .putAllCoders(bundleDescriptor.getCodersMap())
-              .putAllPcollections(bundleDescriptor.getPcollectionsMap())
-              
.putAllWindowingStrategies(bundleDescriptor.getWindowingStrategiesMap())
-              .build();
-
-      addRunnerAndConsumersForPTransformRecursively(
+    RunnerApi.Components components =
+        RunnerApi.Components.newBuilder()
+            .putAllCoders(bundleDescriptor.getCodersMap())
+            .putAllPcollections(bundleDescriptor.getPcollectionsMap())
+            .putAllTransforms(bundleDescriptor.getTransformsMap())
+            
.putAllWindowingStrategies(bundleDescriptor.getWindowingStrategiesMap())
+            .build();
+
+    QueryablePipeline queryablePipeline =
+        
QueryablePipeline.forTransforms(bundleDescriptor.getTransformsMap().keySet(), 
components);
+
+    // Get the transforms in topological order, reverse and build the DAG.
+    List<PipelineNode.PTransformNode> reverseTopologicallyOrderedTransforms =
+        StreamSupport.stream(
+                
queryablePipeline.getTopologicallyOrderedTransforms().spliterator(), false)
+            .collect(
+                Collectors.collectingAndThen(
+                    Collectors.toCollection(ArrayList::new), Lists::reverse));

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   This is a great simplification. To improve readability, you could use 
Guava's `Lists.newArrayList` and `Lists.reverse` directly, which makes the 
intent clearer than using a stream with `collectingAndThen`.
   
   ```java
       List<PipelineNode.PTransformNode> reverseTopologicallyOrderedTransforms =
           
Lists.reverse(Lists.newArrayList(queryablePipeline.getTopologicallyOrderedTransforms()));
   ```



-- 
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: github-unsubscr...@beam.apache.org

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