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

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,15 +231,24 @@ 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`
  // In the other modes, we disable the dropdowns anyway
  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}

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)
  ),
  };
  }

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);
  }
  );