From db898e10e54ca1910cdfaae798e6232126c60c1d Mon Sep 17 00:00:00 2001
From: Marco polo <salazarm@users.noreply.github.com>
Date: Wed, 8 Jan 2025 15:23:02 -0500
Subject: [PATCH] Fix O(n^2) code (#26951)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary & Motivation
Use a set so that we're not iterating over every partition key


before:

<img width="1435" alt="Screenshot 2025-01-08 at 2 43 54 PM"
src="https://github.com/user-attachments/assets/8e4b1244-79ab-4082-bd71-2971d9e31f7d"
/>


after:

![image](https://github.com/user-attachments/assets/0cd69efa-0d0d-4f12-86ba-092c5e2a304e)
---
 .../ui-core/src/assets/usePartitionHealthData.tsx        | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/usePartitionHealthData.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/usePartitionHealthData.tsx
index 49c8821895cbc..a2aaa5870901c 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/assets/usePartitionHealthData.tsx
+++ b/js_modules/dagster-ui/packages/ui-core/src/assets/usePartitionHealthData.tsx
@@ -416,12 +416,15 @@ function addKeyIndexesToMaterializedRanges(
   }
   if (partitions.__typename === 'DefaultPartitionStatuses') {
     const dim = dimensions[0]!;
+    const materializedPartitionKeys = new Set(partitions.materializedPartitions);
+    const materializingPartitionKeys = new Set(partitions.materializingPartitions);
+    const failedPartitionKeys = new Set(partitions.failedPartitions);
     const spans = assembleIntoSpans(dim.partitionKeys, (key) =>
-      partitions.materializedPartitions.includes(key)
+      materializedPartitionKeys.has(key)
         ? AssetPartitionStatus.MATERIALIZED
-        : partitions.materializingPartitions.includes(key)
+        : materializingPartitionKeys.has(key)
           ? AssetPartitionStatus.MATERIALIZING
-          : partitions.failedPartitions.includes(key)
+          : failedPartitionKeys.has(key)
             ? AssetPartitionStatus.FAILED
             : AssetPartitionStatus.MISSING,
     );