From e1dc3f8258fae03ef632a9d34f2eb00103651c48 Mon Sep 17 00:00:00 2001
From: David Liu <d.liu@dagsterlabs.com>
Date: Wed, 18 Dec 2024 17:09:22 -0500
Subject: [PATCH] fix searching for multi partitions

---
 .../ui-core/src/assets/AssetPartitions.tsx    | 24 +++++++++++++------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/AssetPartitions.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/AssetPartitions.tsx
index eaa05536efbaa..933284f41ceae 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/assets/AssetPartitions.tsx
+++ b/js_modules/dagster-ui/packages/ui-core/src/assets/AssetPartitions.tsx
@@ -94,10 +94,20 @@ export const AssetPartitions = ({
         .filter((s: AssetPartitionStatus) => DISPLAYED_STATUSES.includes(s)),
   });
 
-  const [searchValue, setSearchValue] = useQueryPersistedState<string>({
-    queryKey: 'search',
-    defaults: {search: ''},
-  });
+  const [searchValues, setSearchValues] = useState<string[]>([]);
+  const updateSearchValue = (idx: number, value: string) => {
+    setSearchValues((prev) => {
+      const next = [...prev];
+
+      // add empty strings for missing indices
+      while (next.length <= idx) {
+        next.push('');
+      }
+
+      next[idx] = value;
+      return next;
+    });
+  };
 
   // Determine which axis we will show at the top of the page, if any.
   const timeDimensionIdx = selections.findIndex((s) => isTimeseriesDimension(s.dimension));
@@ -148,7 +158,7 @@ export const AssetPartitions = ({
     const sortType = getSort(sortTypes, idx, selections[idx]!.dimension.type);
 
     // Apply the search filter
-    const searchLower = searchValue.toLocaleLowerCase().trim();
+    const searchLower = searchValues?.[idx]?.toLocaleLowerCase().trim() || '';
     const filteredKeys = allKeys.filter((key) => key.toLowerCase().includes(searchLower));
 
     const getSelectionKeys = () =>
@@ -269,8 +279,8 @@ export const AssetPartitions = ({
                   <TextInput
                     fill
                     icon="search"
-                    value={searchValue}
-                    onChange={(e) => setSearchValue(e.target.value)}
+                    value={searchValues[idx] || ''}
+                    onChange={(e) => updateSearchValue(idx, e.target.value)}
                     placeholder="Filter by name…"
                   />
                 </Box>