Optimize rendering collapsed groups

This change makes sure collapsed groups are not rendered in the DOM. This makes the graph renderer much more performant for large graphs with multiple groups. This optimization only affects nested groups, therefore large graphs without groups do not benefit from this.

PiperOrigin-RevId: 711340032

src/app/demo_page/demo_page.ng.html

@@ -556,6 +556,30 @@
         </select>
       </mat-form-field>
     </div>
+    <mat-divider />
+    <div class="demo-control-group">
+      <p class="demo-control-header">
+        Node selection
+      </p>
+      <mat-form-field
+        appearance="fill"
+        class="demo-select-full-width"
+      >
+        <mat-label>Node/Group</mat-label>
+        <mat-select
+          [(ngModel)]="selectedNode"
+          [compareWith]="selectedNodeComparator"
+          id="node-select"
+        >
+          <mat-option
+            *ngFor="let selectable of allSelectableNodesAndGroups"
+            [value]="selectable"
+          >
+            {{selectable.label}}
+          </mat-option>
+        </mat-select>
+      </mat-form-field>
+    </div>
   </div>
 </div>
 

src/app/demo_page/demo_page.scss

@@ -71,6 +71,9 @@ h2 {
       &.demo-select {
         width: 40%;
       }
+      &.demo-select-full-width {
+        width: 80%;
+      }
     }
 
     .demo-control-group {

src/app/demo_page/demo_page.spec.ts

@@ -113,6 +113,21 @@ describe('Demo Page', () => {
       await screenShot.expectMatch('dark-mode');
     });
   });
+
+  describe('Node selection', () => {
+    it('Selects correct nested iteration loop node (screenshot)', async () => {
+      await harness.getNodeSelectInput().then(
+          select => select.clickOptions(
+              {text: 'Trainer (TensorFlow Training, it-1)'}));
+      await screenShot.expectMatch('select-nested-iteration-loop');
+    });
+
+    it('Selects correct nested group node (screenshot)', async () => {
+      await harness.getNodeSelectInput().then(
+          select => select.clickOptions({text: 'Fake Exec 1 (sub1)'}));
+      await screenShot.expectMatch('select-nested-group');
+    })
+  })
 });
 
 

src/app/demo_page/demo_page.ts

@@ -20,6 +20,7 @@ import {HttpClientModule} from '@angular/common/http';
 import {Component, ElementRef, NgModule, TemplateRef, ViewChild, ViewEncapsulation} from '@angular/core';
 import {FormsModule, ReactiveFormsModule} from '@angular/forms';
 import {MatButtonModule} from '@angular/material/button';
+import {MatOptionModule} from '@angular/material/core';
 import {MatDialog, MatDialogModule} from '@angular/material/dialog';
 import {MatDividerModule} from '@angular/material/divider';
 import {MatFormFieldModule} from '@angular/material/form-field';
@@ -170,6 +171,32 @@ function translateColor(col: string, isBg = false) {
   }
 }
 
+function getAllSelectableNestedNodesAndGroups(
+    params: {nodes: DagNode[], groups: DagGroup[]},
+    path: string[] = []): Array<SelectedNode&{label: string}> {
+  const allNodesAndGroups = [];
+  const pathLabel = path.length ? ` (${path.join(', ')})` : '';
+  for (const subGroup of params.groups) {
+    allNodesAndGroups.push(
+        {node: subGroup, path, label: `${subGroup.id}${pathLabel}`});
+    if (!subGroup.treatAsLoop) {
+      allNodesAndGroups.push(...getAllSelectableNestedNodesAndGroups(
+          subGroup, [...path, subGroup.id]));
+    } else {
+      // We need to skip the iteration node itself, as it is not selectable
+      subGroup.groups.forEach((iterationGroup) => {
+        allNodesAndGroups.push(...getAllSelectableNestedNodesAndGroups(
+            iterationGroup, [...path, subGroup.id, iterationGroup.id]));
+      });
+    }
+  }
+  for (const subNode of params.nodes) {
+    allNodesAndGroups.push(
+        {node: subNode, path, label: `${subNode.id}${pathLabel}`});
+  }
+  return allNodesAndGroups;
+}
+
 /** Demo component for directed acyclic graph view. */
 @Component({
   standalone: false,
@@ -237,6 +264,7 @@ export class DagDemoPage {
   userConfigChange = new Subject<UserConfig>();
   destroy = new Subject<void>();
 
+  allSelectableNodesAndGroups: Array<SelectedNode&{label: string}> = [];
   constructor(public dialog: MatDialog) {
     this.setCurrDataset(DEFAULT_DATASET, true);
     this.resetAll();
@@ -339,6 +367,8 @@ export class DagDemoPage {
     const newGraph = cloneGraph(this.datasets[name]);
     this.calibrateNodes(newGraph);
     this.currDataset = newGraph;
+    this.allSelectableNodesAndGroups =
+        getAllSelectableNestedNodesAndGroups(newGraph);
   }
 
   onDatasetChanged(event: Event) {
@@ -531,6 +561,14 @@ export class DagDemoPage {
     this.destroy.next();
     this.destroy.complete();
   }
+
+  triggerSelection(node: SelectedNode) {
+    this.selectedNode = node;
+  }
+
+  selectedNodeComparator(a: SelectedNode, b: SelectedNode) {
+    return a.node.id === b.node.id && a.path.join('') === b.path.join('');
+  }
 }
 
 @NgModule({
@@ -546,6 +584,7 @@ export class DagDemoPage {
     DagScaffoldModule,
     DagToolbarModule,
     MatSelectModule,
+    MatOptionModule,
     WorkflowGraphIconModule,
     MatDialogModule,
     FormsModule,

src/app/directed_acyclic_graph.ts

@@ -29,7 +29,7 @@ import {ColorThemeLoader} from './color_theme_loader';
 import {DagStateService} from './dag-state.service';
 import {STATE_SERVICE_PROVIDER} from './dag-state.service.provider';
 import {baseColors, BLUE_THEME, clampVal, CLASSIC_THEME, createDAGFeatures, createNewSizeConfig, type DagTheme, DEFAULT_LAYOUT_OPTIONS, DEFAULT_THEME, defaultFeatures, defaultZoomConfig, EdgeStyle, type FeatureToggleOptions, generateTheme, getMargin, isPoint, type LayoutOptions, type Logger, MarkerStyle, type MinimapPosition, nanSafePt, NODE_HEIGHT, NODE_WIDTH, NodeState, OrientationMarginConfig, RankAlignment, RankDirection, RankerAlgorithim, SCROLL_STEP_PER_DELTA, SizeConfig, SVG_ELEMENT_SIZE, type ZoomConfig} from './data_types_internal';
-import {DagRaw, DagRawModule, EnhancedDagGroup, GraphDims} from './directed_acyclic_graph_raw';
+import {DagRaw, DagRawModule, EnhancedDagGroup, GraphDims, setEnhancedGroupSelection} from './directed_acyclic_graph_raw';
 import {DagLogger} from './logger/dag_logger';
 import {Minimap, MinimapModule} from './minimap/minimap';
 import {type DagEdge, DagGroup, DagNode, GraphSpec, GroupIterationRecord, GroupToggleEvent, isDagreInit, NodeMap, type NodeRef, Point, type SelectedNode} from './node_spec';
@@ -577,7 +577,9 @@ export class DirectedAcyclicGraph implements OnInit, OnDestroy {
           const loopGroup = parent?.$groups?.find(g => g.id === segment);
           if (loopGroup) {
             const selectedLoop = path.shift()!;
-            loopGroup.selectedLoopId = selectedLoop;
+            const selectedLoopNode =
+                loopGroup.groups.find(n => n.id === selectedLoop);
+            setEnhancedGroupSelection(loopGroup, selectedLoopNode!);
             // Increase pathDepth as we've consumed 2 path segments
             pathDepth++;
             return true;
@@ -592,6 +594,8 @@ export class DirectedAcyclicGraph implements OnInit, OnDestroy {
         expandedStateChanges++;
         // Detect changes so groups can expand properly
         this.detectChanges();
+        // Detect changes in new dagEl so QueryList is updated properly
+        dagEl.detectChanges();
       }
       pathDepth++;
     }

src/app/directed_acyclic_graph_raw.ng.html

@@ -1,5 +1,5 @@
 <aside
-  *ngIf="!noEmptySpaceAlloc"
+  *ngIf="!noEmptySpaceAlloc && visible"
   aria-hidden="true"
   class="empty-space-alloc"
   [style.width.px]="graphWidth"
@@ -9,6 +9,7 @@
   [ngClass]="['dag-component', toggleClass(collapsed, 'collapse')]"
   [style.width.px]="graphWidth"
   [style.height.px]="graphHeight"
+  *ngIf="visible"
 >
   <svg class="edge canvas" aria-hidden="true">
     <ng-container *ngFor="let e of edges; trackBy: edgeTrack">
@@ -351,6 +352,7 @@
           [groups]="group.groups"
           (graphResize)="storeSubDagDims($event, group)"
           [resolveReference]="resolveReference"
+          [visible]="isGroupExpanded(group)"
         ></ai-dag-raw>
         <ng-container *ngIf="!!(group.treatAsLoop && group._cachedSelection)">
           <ai-dag-raw
@@ -362,6 +364,7 @@
             [nodes]="[group._cachedSelection]"
             (graphResize)="storeSubDagDims($event, group)"
             [resolveReference]="resolveReference"
+            [visible]="isGroupExpanded(group)"
           ></ai-dag-raw>
           <ai-dag-raw
             #subDag
@@ -374,6 +377,7 @@
             [groups]="group._cachedSelection.groups"
             (graphResize)="storeSubDagDims($event, group)"
             [resolveReference]="resolveReference"
+            [visible]="isGroupExpanded(group)"
           ></ai-dag-raw>
         </ng-container>
       </footer>

src/app/directed_acyclic_graph_raw.ts

@@ -131,6 +131,20 @@ function setGroupSizeProps(
   return Object.assign(expandedGroup, {width, height, padY});
 }
 
+/**
+ * Sets the selection state of a group.
+ * _cachedSelection is a property that should be refactored to be clearer. It is
+ * used to determine which node/group is rendered within an iteration group.
+ */
+export function setEnhancedGroupSelection(
+    group: DagGroup,
+    selection: DagGroup|DagNode,
+) {
+  const enhancedGroup = group as EnhancedDagGroup;
+  enhancedGroup._cachedSelection = selection;
+  enhancedGroup.selectedLoopId = selection?.id;
+}
+
 /** Dimension config for a raw DAG */
 export interface GraphDims {
   width: number;
@@ -284,6 +298,12 @@ export class DagRaw implements DoCheck, OnInit, OnDestroy {
 
   @Input() loading = false;
 
+  /**
+   * If set to false, the contents of the DAG will not be rendered to optimize
+   * for performance.
+   */
+  @Input() visible = true;
+
   @Input('nodes')
   set nodes(nodes: DagNode[]) {
     // Avoid pointer/reference stability, so that angular will pick up the
@@ -473,7 +493,7 @@ export class DagRaw implements DoCheck, OnInit, OnDestroy {
       group: DagGroup, iterationNode: GroupIterationRecord['iterationNode']) {
     (group as any)._cachedSelection = iterationNode;
     if (group.selectedLoopId !== iterationNode.id) {
-      group.selectedLoopId = iterationNode.id;
+      setEnhancedGroupSelection(group, iterationNode);
       const iter = {path: this.dagPath, group, iterationNode};
       this.stateService?.setIterationChange(iter);
     }
@@ -940,6 +960,7 @@ export class DagRaw implements DoCheck, OnInit, OnDestroy {
       this.stateService?.setGroupExpandToggled({groupId: id, isExpanded: true});
     }
     this.updateGraphLayout();
+    this.detectChanges();
     return beforeCt !== this.expandedGroups.size;
   }
 

src/app/test_resources/demo_page_harness.ts

@@ -17,6 +17,7 @@
 
 import {ComponentHarness} from '@angular/cdk/testing';
 import {MatNativeSelectHarness} from '@angular/material/input/testing';
+import {MatSelectHarness} from '@angular/material/select/testing';
 
 /** Test harness for the Demo Page. */
 export class DemoPageHarness extends ComponentHarness {
@@ -31,4 +32,8 @@ export class DemoPageHarness extends ComponentHarness {
     return this.locatorFor(
         MatNativeSelectHarness.with({selector: '#color-theme-select'}))();
   }
+
+  getNodeSelectInput(): Promise<MatSelectHarness> {
+    return this.locatorFor(MatSelectHarness.with({selector: '#node-select'}))();
+  }
 }