From b8f31c940babf437adaf532fcae3335891a325b3 Mon Sep 17 00:00:00 2001
From: Vlad Babich <vladbabich@deephaven.io>
Date: Mon, 22 Apr 2024 03:08:18 -0600
Subject: [PATCH] Fix null partition filter

---
 packages/iris-grid/src/IrisGrid.tsx | 40 +++++++++--------------------
 1 file changed, 12 insertions(+), 28 deletions(-)

diff --git a/packages/iris-grid/src/IrisGrid.tsx b/packages/iris-grid/src/IrisGrid.tsx
index 3bdf4edce..6b8478bed 100644
--- a/packages/iris-grid/src/IrisGrid.tsx
+++ b/packages/iris-grid/src/IrisGrid.tsx
@@ -2043,39 +2043,23 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
       model.partitionKeysTable(),
       resolved => resolved.close()
     );
-    const { dh } = model;
 
     const sorts = keyTable.columns.map(column => column.sort().desc());
     keyTable.applySort(sorts);
     keyTable.setViewport(0, 0);
 
-    return new Promise((resolve, reject) => {
-      // We want to wait for the first UPDATED event instead of just getting viewport data here
-      // It's possible that the key table does not have any rows of data yet, so just wait until it does have one
-      keyTable.addEventListener(
-        dh.Table.EVENT_UPDATED,
-        (event: CustomEvent<DhType.ViewportData>) => {
-          try {
-            const { detail: data } = event;
-            if (data.rows.length === 0) {
-              // Table is empty, wait for the next updated event
-              return;
-            }
-            const row = data.rows[0];
-            const values = keyTable.columns.map(column => row.get(column));
-            const newPartition: PartitionConfig = {
-              partitions: values,
-              mode: 'partition',
-            };
-            keyTable.close();
-            resolve(newPartition);
-          } catch (e) {
-            keyTable.close();
-            reject(e);
-          }
-        }
-      );
-    });
+    const { rows } = await keyTable.getViewportData();
+    let values = [];
+    if (rows.length > 0) {
+      const row = rows[0];
+      values = keyTable.columns.map(column => row.get(column));
+    }
+    const newPartition: PartitionConfig = {
+      partitions: values,
+      mode: 'partition',
+    };
+    keyTable.close();
+    return newPartition;
   }
 
   copyCell(