Support new selection syntax in Run GanttChart. (#26632)

## Summary & Motivation

Adds support for the new selection syntax in the Gantt chart.

## How I Tested These Changes

Tested:

1. Retrying runs using the selection syntax to select the steps.
2. Selecting / Unselecting multiple steps by shift clicking on steps.
3. Using status and name filters.
4. Making sure the logs are filtered to the correct steps matching the
selection.

<img width="677" alt="Screenshot 2024-12-20 at 10 30 52 AM"
src="https://github.com/user-attachments/assets/9f94cc19-bc7a-4a5f-900a-5c351a5a2677"
/>
<img width="625" alt="Screenshot 2024-12-20 at 10 30 46 AM"
src="https://github.com/user-attachments/assets/2bb45583-52c5-4695-ac54-0b5f931fb2ae"
/>
<img width="1724" alt="Screenshot 2024-12-20 at 10 30 40 AM"
src="https://github.com/user-attachments/assets/bd7507a7-77d4-402d-9b5b-4d3f0386445b"
/>
<img width="1726" alt="Screenshot 2024-12-20 at 10 30 18 AM"
src="https://github.com/user-attachments/assets/bacb4060-b879-4251-bff6-f7a3e085be6a"
/>
<img width="1463" alt="Screenshot 2024-12-20 at 10 30 00 AM"
src="https://github.com/user-attachments/assets/6a87c64f-0c0a-4543-aaac-d03218cab19c"
/>

js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts

@@ -9,6 +9,7 @@ import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 
 import {AntlrAssetSelectionVisitor} from './AntlrAssetSelectionVisitor';
 import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData';
+import {weakMapMemoize} from '../util/weakMapMemoize';
 import {AssetSelectionLexer} from './generated/AssetSelectionLexer';
 import {AssetSelectionParser} from './generated/AssetSelectionParser';
 import {featureEnabled} from '../app/Flags';
@@ -65,17 +66,17 @@ export const parseAssetSelectionQuery = (
   }
 };
 
-export const filterAssetSelectionByQuery = (
-  all_assets: AssetGraphQueryItem[],
-  query: string,
-): AssetSelectionQueryResult => {
-  if (featureEnabled(FeatureFlag.flagAssetSelectionSyntax)) {
-    const result = parseAssetSelectionQuery(all_assets, query);
-    if (result instanceof Error) {
-      // fall back to old behavior
-      return filterByQuery(all_assets, query);
+export const filterAssetSelectionByQuery = weakMapMemoize(
+  (all_assets: AssetGraphQueryItem[], query: string): AssetSelectionQueryResult => {
+    if (featureEnabled(FeatureFlag.flagAssetSelectionSyntax)) {
+      const result = parseAssetSelectionQuery(all_assets, query);
+      if (result instanceof Error) {
+        // fall back to old behavior
+        return filterByQuery(all_assets, query);
+      }
+      return result;
     }
-    return result;
-  }
-  return filterByQuery(all_assets, query);
-};
+    return filterByQuery(all_assets, query);
+  },
+  {maxEntries: 20},
+);

js_modules/dagster-ui/packages/ui-core/src/assets/VirtualizedSimpleAssetKeyList.tsx

@@ -13,7 +13,7 @@ export const VirtualizedSimpleAssetKeyList = ({
   style,
 }: {
   assetKeys: AssetKeyInput[];
-  style: CSSProperties;
+  style?: CSSProperties;
 }) => {
   const parentRef = useRef<HTMLDivElement>(null);
   const rowVirtualizer = useVirtualizer({

js_modules/dagster-ui/packages/ui-core/src/gantt/GanttChart.tsx

@@ -15,6 +15,7 @@ import isEqual from 'lodash/isEqual';
 import * as React from 'react';
 import {useMemo} from 'react';
 import {Link} from 'react-router-dom';
+import {FeatureFlag} from 'shared/app/FeatureFlags.oss';
 import styled from 'styled-components';
 
 import {
@@ -48,13 +49,15 @@ import {
   interestingQueriesFor,
 } from './GanttChartLayout';
 import {GanttChartModeControl} from './GanttChartModeControl';
+import {GanttChartSelectionInput} from './GanttChartSelectionInput';
 import {GanttChartTimescale} from './GanttChartTimescale';
 import {GanttStatusPanel} from './GanttStatusPanel';
 import {OptionsContainer, OptionsSpacer} from './VizComponents';
 import {ZoomSlider} from './ZoomSlider';
 import {RunGraphQueryItem} from './toGraphQueryItems';
 import {useGanttChartMode} from './useGanttChartMode';
 import {AppContext} from '../app/AppContext';
+import {featureEnabled} from '../app/Flags';
 import {GraphQueryItem} from '../app/GraphQueryImpl';
 import {withMiddleTruncation} from '../app/Util';
 import {WebSocketContext} from '../app/WebSocketProvider';
@@ -411,14 +414,22 @@ const GanttChartInner = React.memo((props: GanttChartInnerProps) => {
           </WebsocketWarning>
         ) : null}
         <FilterInputsBackgroundBox flex={{direction: 'row', alignItems: 'center', gap: 12}}>
-          <GraphQueryInput
-            items={props.graph}
-            value={props.selection.query}
-            placeholder="Type a step subset"
-            onChange={props.onUpdateQuery}
-            presets={presets}
-            className={selection.keys.length > 0 ? 'has-step' : ''}
-          />
+          {featureEnabled(FeatureFlag.flagRunSelectionSyntax) ? (
+            <GanttChartSelectionInput
+              items={props.graph}
+              value={props.selection.query}
+              onChange={props.onUpdateQuery}
+            />
+          ) : (
+            <GraphQueryInput
+              items={props.graph}
+              value={props.selection.query}
+              placeholder="Type a step subset"
+              onChange={props.onUpdateQuery}
+              presets={presets}
+              className={selection.keys.length > 0 ? 'has-step' : ''}
+            />
+          )}
           <Checkbox
             checked={options.hideUnselectedSteps}
             label="Hide unselected steps"

js_modules/dagster-ui/packages/ui-core/src/gantt/GanttChartSelectionInput.tsx

@@ -0,0 +1,61 @@
+import {useMemo} from 'react';
+import styled from 'styled-components';
+
+import {RunGraphQueryItem} from './toGraphQueryItems';
+import {NO_STATE} from '../run-selection/AntlrRunSelectionVisitor';
+import {RunSelectionLexer} from '../run-selection/generated/RunSelectionLexer';
+import {RunSelectionParser} from '../run-selection/generated/RunSelectionParser';
+import {InputDiv, SelectionAutoCompleteInput} from '../selection/SelectionAutoCompleteInput';
+import {createSelectionLinter} from '../selection/createSelectionLinter';
+import {weakMapMemoize} from '../util/weakMapMemoize';
+
+export const GanttChartSelectionInput = ({
+  items,
+  value,
+  onChange,
+}: {
+  items: RunGraphQueryItem[];
+  value: string;
+  onChange: (value: string) => void;
+}) => {
+  const attributesMap = useMemo(() => {
+    const statuses = new Set<string>();
+    const names = new Set<string>();
+
+    items.forEach((item) => {
+      if (item.metadata?.state) {
+        statuses.add(item.metadata.state);
+      } else {
+        statuses.add(NO_STATE);
+      }
+      names.add(item.name);
+    });
+    return {name: Array.from(names), status: Array.from(statuses)};
+  }, [items]);
+
+  return (
+    <Wrapper>
+      <SelectionAutoCompleteInput
+        nameBase="name"
+        attributesMap={attributesMap}
+        placeholder="Type a step subset"
+        functions={FUNCTIONS}
+        linter={getLinter()}
+        value={value}
+        onChange={onChange}
+      />
+    </Wrapper>
+  );
+};
+
+const getLinter = weakMapMemoize(() =>
+  createSelectionLinter({Lexer: RunSelectionLexer, Parser: RunSelectionParser}),
+);
+
+const FUNCTIONS = ['sinks', 'roots'];
+
+const Wrapper = styled.div`
+  ${InputDiv} {
+    width: 24vw;
+  }
+`;

js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelection.ts

@@ -8,8 +8,9 @@ import {RunSelectionLexer} from './generated/RunSelectionLexer';
 import {RunSelectionParser} from './generated/RunSelectionParser';
 import {featureEnabled} from '../app/Flags';
 import {filterByQuery} from '../app/GraphQueryImpl';
+import {weakMapMemoize} from '../util/weakMapMemoize';
 
-type RunSelectionQueryResult = {
+export type RunSelectionQueryResult = {
   all: RunGraphQueryItem[];
   focus: RunGraphQueryItem[];
 };
@@ -44,17 +45,17 @@ export const parseRunSelectionQuery = (
   }
 };
 
-export const filterRunSelectionByQuery = (
-  all_runs: RunGraphQueryItem[],
-  query: string,
-): RunSelectionQueryResult => {
-  if (featureEnabled(FeatureFlag.flagRunSelectionSyntax)) {
-    const result = parseRunSelectionQuery(all_runs, query);
-    if (result instanceof Error) {
-      // fall back to old behavior
-      return filterByQuery(all_runs, query);
+export const filterRunSelectionByQuery = weakMapMemoize(
+  (all_runs: RunGraphQueryItem[], query: string): RunSelectionQueryResult => {
+    if (featureEnabled(FeatureFlag.flagRunSelectionSyntax)) {
+      const result = parseRunSelectionQuery(all_runs, query);
+      if (result instanceof Error) {
+        // fall back to old behavior
+        return filterByQuery(all_runs, query);
+      }
+      return result;
     }
-    return result;
-  }
-  return filterByQuery(all_runs, query);
-};
+    return filterByQuery(all_runs, query);
+  },
+  {maxEntries: 20},
+);

js_modules/dagster-ui/packages/ui-core/src/run-selection/AntlrRunSelectionVisitor.ts

@@ -160,6 +160,12 @@ export class AntlrRunSelectionVisitor
 
   visitStatusAttributeExpr(ctx: StatusAttributeExprContext) {
     const state: string = getValue(ctx.value()).toLowerCase();
-    return new Set([...this.all_runs].filter((i) => i.metadata?.state === state));
+    return new Set(
+      [...this.all_runs].filter(
+        (i) => i.metadata?.state === state || (state === NO_STATE && !i.metadata?.state),
+      ),
+    );
   }
 }
+
+export const NO_STATE = 'none';