feat: Partitioned Table UI Enhancements (#2110)

Resolves #2079 ( Resolves #2066 , Resolves #2103 , Resolves #2104 ,
Resolves #2105 , Resolves #2106 , Resolves #2107 , Resolves #2108 ,
Resolves #2109 , Resolves #2049 ), Resolves #2120, Resolves #1904

**Changes Implemented:**
- Replaced `keysTable` with `baseTable` 
- Made small UI changes to differentiate between Partition Table and
Partition-aware source table
- Allow double clicking to set partition
- Add context menu action to set partition

**Sample Visuals:**
_Default Partition Table View_
<img width="798" alt="Screenshot 2024-06-26 at 2 57 09 PM"
src="https://github.com/deephaven/web-client-ui/assets/69530774/19059a9f-7e69-470a-818e-c1d6d9ce8118">
_Default Partition-aware source table view:_
<img width="789" alt="Screenshot 2024-06-26 at 2 56 59 PM"
src="https://github.com/deephaven/web-client-ui/assets/69530774/af891a26-3395-4842-963f-c8b2f1c9186d">

packages/iris-grid/src/IrisGrid.tsx

@@ -116,6 +116,7 @@ import {
   IrisGridContextMenuHandler,
   IrisGridCopyCellMouseHandler,
   IrisGridDataSelectMouseHandler,
+  IrisGridPartitionedTableMouseHandler,
   IrisGridFilterMouseHandler,
   IrisGridRowTreeMouseHandler,
   IrisGridSortMouseHandler,
@@ -734,6 +735,7 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
       new IrisGridContextMenuHandler(this, dh),
       new IrisGridDataSelectMouseHandler(this),
       new PendingMouseHandler(this),
+      new IrisGridPartitionedTableMouseHandler(this),
     ];
     if (canCopy) {
       keyHandlers.push(new CopyKeyHandler(this));
@@ -1781,12 +1783,12 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
    * Rebuilds all the current filters. Necessary if something like the time zone has changed.
    */
   rebuildFilters(): void {
-    log.debug('Rebuilding filters');
-
     const { model } = this.props;
     const { advancedFilters, quickFilters } = this.state;
     const { columns, formatter } = model;
 
+    log.debug('Rebuilding filters');
+
     const newAdvancedFilters = new Map();
     const newQuickFilters = new Map();
 
@@ -1812,8 +1814,8 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
         filter: this.makeQuickFilter(column, text, formatter.timeZone),
       });
     });
-
     this.startLoading('Rebuilding filters...', { resetRanges: true });
+
     this.setState({
       quickFilters: newQuickFilters,
       advancedFilters: newAdvancedFilters,
@@ -2024,11 +2026,15 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
 
   async loadPartitionsTable(model: PartitionedGridModel): Promise<void> {
     try {
-      const partitionConfig = await this.getInitialPartitionConfig(model);
-      this.setState(
-        { isSelectingPartition: true, partitionConfig },
-        this.loadTableState
-      );
+      const { partitionConfig } = this.state;
+      if (!partitionConfig) {
+        this.startLoading('Loading partitions...', {
+          loadingCancelShown: false,
+          resetRanges: true,
+        });
+        this.initializePartitionConfig(model);
+      }
+      this.setState({ isSelectingPartition: true }, this.loadTableState);
     } catch (error) {
       if (!PromiseUtils.isCanceled(error)) {
         this.handleTableLoadError(error);
@@ -2037,19 +2043,9 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
   }
 
   /**
-   * Gets the initial partition config for the currently set model.
-   * Sorts the key table and gets the first key.
-   * If the table is ticking, it will wait for the first tick.
+   * Initialize the partition config to the default partition.
    */
-  async getInitialPartitionConfig(
-    model: PartitionedGridModel
-  ): Promise<PartitionConfig> {
-    const { partitionConfig } = this.state;
-    if (partitionConfig !== undefined) {
-      // User already has a partition selected, just use that
-      return partitionConfig;
-    }
-
+  async initializePartitionConfig(model: PartitionedGridModel): Promise<void> {
     const keyTable = await this.pending.add(
       model.partitionKeysTable(),
       resolved => resolved.close()
@@ -2060,37 +2056,63 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
     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];
-            // Core JSAPI returns undefined for null table values, IrisGridPartitionSelector expects null
-            // https://github.com/deephaven/deephaven-core/issues/5400
-            const values = keyTable.columns.map(
-              column => row.get(column) ?? null
-            );
-            const newPartition: PartitionConfig = {
-              partitions: values,
-              mode: 'partition',
-            };
-            keyTable.close();
-            resolve(newPartition);
-          } catch (e) {
-            keyTable.close();
-            reject(e);
+    // 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) {
+            // We received an update and the table is still empty. Stop showing the loading spinner so we know
+            this.stopLoading();
+            return;
           }
+          const row = data.rows[0];
+          const values = keyTable.columns.map(column => row.get(column));
+          const newPartition: PartitionConfig = {
+            partitions: values,
+            mode: model.isPartitionAwareSourceTable ? 'partition' : 'keys',
+          };
+          keyTable.close();
+          this.setState({
+            loadingSpinnerShown: false,
+            partitionConfig: newPartition,
+          });
+        } catch (e) {
+          keyTable.close();
+          log.error('Error getting initial partition config', e);
         }
+      }
+    );
+  }
+
+  /**
+   * Selects a partition key from the table based on the provided row index.
+   *
+   * @param rowIndex The index of the row from which the partition will be selected.
+   * @returns A promise that resolves when the partition key has been successfully selected.
+   */
+  selectPartitionKeyFromTable(rowIndex: GridRangeIndex): void {
+    const { model } = this.props;
+    assertNotNull(rowIndex);
+    if (isPartitionedGridModel(model)) {
+      const data = this.getRowDataMap(rowIndex);
+      const partitionColumnSet = new Set(
+        model.partitionColumns.map(column => column.name)
       );
-    });
+      const values = Object.entries(data)
+        .filter(([key, _]) => partitionColumnSet.has(key))
+        .map(([key, columnData]) => columnData.value);
+      const newPartition: PartitionConfig = {
+        partitions: values,
+        mode: 'partition',
+      };
+      this.setState({
+        isSelectingPartition: true,
+        partitionConfig: newPartition,
+      });
+    }
   }
 
   copyCell(
@@ -2532,7 +2554,10 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
 
   handlePartitionChange(partitionConfig: PartitionConfig): void {
     this.startLoading('Partitioning...');
-    this.setState({ partitionConfig });
+    const { partitionConfig: prevConfig } = this.state;
+    if (prevConfig !== partitionConfig) {
+      this.setState({ partitionConfig });
+    }
   }
 
   handleTableLoadError(error: unknown): void {
@@ -4700,15 +4725,13 @@ class IrisGrid extends Component<IrisGridProps, IrisGridState> {
             unmountOnExit
           >
             <div className="iris-grid-partition-selector-wrapper iris-grid-bar iris-grid-bar-primary">
-              {isPartitionedGridModel(model) &&
-                model.isPartitionRequired &&
-                partitionConfig && (
-                  <IrisGridPartitionSelector
-                    model={model}
-                    partitionConfig={partitionConfig}
-                    onChange={this.handlePartitionChange}
-                  />
-                )}
+              {isPartitionedGridModel(model) && model.isPartitionRequired && (
+                <IrisGridPartitionSelector
+                  model={model}
+                  partitionConfig={partitionConfig}
+                  onChange={this.handlePartitionChange}
+                />
+              )}
             </div>
           </CSSTransition>
           <CSSTransition

packages/iris-grid/src/IrisGridPartitionSelector.scss

@@ -18,8 +18,7 @@
   }
   .partition-button-group {
     display: flex;
-    border: 1px var(--dh-color-hr);
-    border-style: none solid;
+    border-right: 1px solid var(--dh-color-hr);
     padding: 0 $spacer-2;
     gap: $spacer-2;
   }

packages/iris-grid/src/IrisGridPartitionSelector.test.tsx

@@ -17,6 +17,9 @@ function makeModel(
       Promise.resolve(irisGridTestUtils.makeTable())
     ),
     partitionColumns: columns,
+    partitionBaseTable: jest.fn(() =>
+      Promise.resolve(irisGridTestUtils.makeTable())
+    ),
   } as unknown as PartitionedGridModel;
   return model;
 }