From 02542dd8e7a94172a7240f93615f8e816b3fd638 Mon Sep 17 00:00:00 2001
From: mikebender <mikebender@deephaven.io>
Date: Tue, 16 Jan 2024 15:24:11 -0500
Subject: [PATCH] Fix the columns used in the partition selector

- Was referencing the columns from the original table instead of the actual table
- Was causing a mismatch when fetching the viewport, where it would just stall instead of throwing an error
---
 .../src/IrisGridPartitionSelector.tsx         | 34 ++++++++++++-------
 packages/iris-grid/src/IrisGridUtils.ts       | 31 ++++++++++++-----
 .../jsapi-components/src/TableDropdown.tsx    |  9 ++---
 3 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/packages/iris-grid/src/IrisGridPartitionSelector.tsx b/packages/iris-grid/src/IrisGridPartitionSelector.tsx
index f1b4b60d3f..4b1dc380e8 100644
--- a/packages/iris-grid/src/IrisGridPartitionSelector.tsx
+++ b/packages/iris-grid/src/IrisGridPartitionSelector.tsx
@@ -77,8 +77,12 @@ class IrisGridPartitionSelector extends Component<
         )
       );
 
-      this.setState({ isLoading: false, keysTable, partitionTables }, () => {
-        this.updatePartitionFilters();
+      const partitionFilters = this.getPartitionFilters(partitionTables);
+      this.setState({
+        isLoading: false,
+        keysTable,
+        partitionFilters,
+        partitionTables,
       });
     } catch (e) {
       if (!PromiseUtils.isCanceled(e)) {
@@ -167,7 +171,7 @@ class IrisGridPartitionSelector extends Component<
       const partitionFilters = newPartitions
         .slice(0, index + 1)
         .map((partition, i) => {
-          const partitionColumn = model.partitionColumns[i];
+          const partitionColumn = t.columns[i];
           return partitionColumn
             .filter()
             .eq(
@@ -180,14 +184,11 @@ class IrisGridPartitionSelector extends Component<
       t.applyFilter(partitionFilters);
       t.setViewport(0, 0, t.columns);
       const data = await this.pending.add(t.getViewportData());
-      t.close();
-
       const newConfig: PartitionConfig = {
-        partitions: model.partitionColumns.map(column =>
-          data.rows[0].get(column)
-        ),
+        partitions: t.columns.map(column => data.rows[0].get(column)),
         mode: 'partition',
       };
+      t.close();
       this.sendUpdate(newConfig);
     } catch (e) {
       if (!PromiseUtils.isCanceled(e)) {
@@ -230,8 +231,8 @@ class IrisGridPartitionSelector extends Component<
     const { partitionTables } = this.state;
     assertNotNull(partitionTables);
 
-    const { model, partitionConfig } = this.props;
-    const { mode, partitions } = partitionConfig;
+    const { partitionConfig } = this.props;
+    const { mode } = partitionConfig;
     log.debug('updatePartitionFilters', partitionConfig);
     if (mode !== 'partition') {
       // We only need to update the filters if the mode is `partitions`
@@ -239,6 +240,15 @@ class IrisGridPartitionSelector extends Component<
       return;
     }
 
+    const partitionFilters = this.getPartitionFilters(partitionTables);
+    this.setState({ partitionFilters });
+  }
+
+  getPartitionFilters(partitionTables: Table[]): FilterCondition[][] {
+    const { model, partitionConfig } = this.props;
+    const { partitions } = partitionConfig;
+    log.debug('getPartitionFilters', partitionConfig);
+
     if (partitions.length !== partitionTables.length) {
       throw new Error(
         `Invalid partition config set. Expected ${partitionTables.length} partitions, but got ${partitions.length}`
@@ -269,7 +279,7 @@ class IrisGridPartitionSelector extends Component<
         partitionFilters.push(partitionFilter);
       }
     }
-    this.setState({ partitionFilters });
+    return partitionFilters;
   }
 
   getCachedChangeCallback = memoizee(
@@ -293,7 +303,7 @@ class IrisGridPartitionSelector extends Component<
         <TableDropdown
           className="custom-select-sm"
           table={partitionTables?.[index]}
-          column={column}
+          column={partitionTables?.[index]?.columns[index]}
           filter={partitionFilters?.[index]}
           onChange={this.getCachedChangeCallback(index)}
           selectedValue={mode === 'partition' ? partitions[index] : undefined}
diff --git a/packages/iris-grid/src/IrisGridUtils.ts b/packages/iris-grid/src/IrisGridUtils.ts
index c601703a7a..868f4f5add 100644
--- a/packages/iris-grid/src/IrisGridUtils.ts
+++ b/packages/iris-grid/src/IrisGridUtils.ts
@@ -49,7 +49,10 @@ import IrisGridModel from './IrisGridModel';
 import type AdvancedSettingsType from './sidebar/AdvancedSettingsType';
 import AdvancedSettings from './sidebar/AdvancedSettings';
 import ColumnHeaderGroup from './ColumnHeaderGroup';
-import { PartitionConfig } from './PartitionedGridModel';
+import {
+  isPartitionedGridModelProvider,
+  PartitionConfig,
+} from './PartitionedGridModel';
 
 const log = Log.module('IrisGridUtils');
 
@@ -1177,6 +1180,10 @@ class IrisGridUtils {
     assertNotNull(metrics);
     const { userColumnWidths, userRowHeights } = metrics;
     const { columns } = model;
+    const partitionColumns = isPartitionedGridModelProvider(model)
+      ? model.partitionColumns
+      : [];
+
     // Return value will be serialized, should not contain undefined
     return {
       advancedFilters: this.dehydrateAdvancedFilters(columns, advancedFilters),
@@ -1210,7 +1217,10 @@ class IrisGridUtils {
         children: item.children,
         color: item.color,
       })),
-      partitionConfig: this.dehydratePartitionConfig(columns, partitionConfig),
+      partitionConfig: this.dehydratePartitionConfig(
+        partitionColumns,
+        partitionConfig
+      ),
     };
   }
 
@@ -1250,7 +1260,9 @@ class IrisGridUtils {
       partitionConfig,
     } = irisGridState;
     const { columns, formatter } = model;
-
+    const partitionColumns = isPartitionedGridModelProvider(model)
+      ? model.partitionColumns
+      : [];
     return {
       advancedFilters: this.hydrateAdvancedFilters(
         columns,
@@ -1306,7 +1318,10 @@ class IrisGridUtils {
         model,
         columnHeaderGroups ?? model.layoutHints?.columnGroups ?? []
       ).groups,
-      partitionConfig: this.hydratePartitionConfig(columns, partitionConfig),
+      partitionConfig: this.hydratePartitionConfig(
+        partitionColumns,
+        partitionConfig
+      ),
     };
   }
 
@@ -1501,7 +1516,7 @@ class IrisGridUtils {
   }
 
   dehydratePartitionConfig(
-    columns: readonly Column[],
+    partitionColumns: readonly Column[],
     partitionConfig: PartitionConfig | undefined
   ): PartitionConfig | undefined {
     if (partitionConfig == null) {
@@ -1511,13 +1526,13 @@ class IrisGridUtils {
     return {
       ...partitionConfig,
       partitions: partitionConfig.partitions.map((partition, index) =>
-        this.dehydrateValue(partition, columns[index].type)
+        this.dehydrateValue(partition, partitionColumns[index].type)
       ),
     };
   }
 
   hydratePartitionConfig(
-    columns: readonly Column[],
+    partitionColumns: readonly Column[],
     partitionConfig: PartitionConfig | undefined
   ): PartitionConfig | undefined {
     if (partitionConfig == null) {
@@ -1527,7 +1542,7 @@ class IrisGridUtils {
     return {
       ...partitionConfig,
       partitions: partitionConfig.partitions.map((partition, index) =>
-        this.hydrateValue(partition, columns[index].type)
+        this.hydrateValue(partition, partitionColumns[index].type)
       ),
     };
   }
diff --git a/packages/jsapi-components/src/TableDropdown.tsx b/packages/jsapi-components/src/TableDropdown.tsx
index a2f53e8854..67cbe0a390 100644
--- a/packages/jsapi-components/src/TableDropdown.tsx
+++ b/packages/jsapi-components/src/TableDropdown.tsx
@@ -30,8 +30,8 @@ export type TableDropdownProps = {
   /** Table to use as the source of data. Does not own the table, does not close it on unmount. */
   table?: Table;
 
-  /** Column to read data from the table */
-  column: Column;
+  /** Column to read data from the table. Defaults to the first column in the table if it's not provided. */
+  column?: Column;
 
   /** Triggered when the dropdown selection has changed */
   onChange: (value: unknown) => void;
@@ -78,15 +78,16 @@ export function TableDropdown({
       return undefined;
     }
 
+    const tableColumn = column ?? table.columns[0];
     // Need to set a viewport on the table and start listening to get the values to populate the dropdown
     table.applyFilter(filter as FilterCondition[]);
-    const subscription = table.setViewport(0, maxSize, [column]);
+    const subscription = table.setViewport(0, maxSize, [tableColumn]);
 
     subscription.addEventListener(
       dh.Table.EVENT_UPDATED,
       (event: CustomEvent<ViewportData>) => {
         const { detail } = event;
-        const newValues = detail.rows.map(row => row.get(column));
+        const newValues = detail.rows.map(row => row.get(tableColumn));
         setValues(newValues);
       }
     );