fix: Use correct schema when opening transformed partitioned tables

web/client-api/src/main/java/io/deephaven/web/client/api/JsPartitionedTable.java

@@ -83,6 +83,10 @@ public Promise<JsPartitionedTable> refetch() {
         return widget.refetch().then(w -> {
             descriptor = PartitionedTableDescriptor.deserializeBinary(w.getDataAsU8());
 
+            return w.getExportedObjects()[0].fetch();
+        }).then(result -> {
+            keys = (JsTable) result;
+
             keyColumnTypes = new ArrayList<>();
             InitialTableDefinition tableDefinition = WebBarrageUtils.readTableDefinition(
                     WebBarrageUtils.readSchemaMessage(descriptor.getConstituentDefinitionSchema_asU8()));
@@ -97,15 +101,13 @@ public Promise<JsPartitionedTable> refetch() {
             JsArray<String> keyColumnNames = descriptor.getKeyColumnNamesList();
             for (int i = 0; i < keyColumnNames.length; i++) {
                 String name = keyColumnNames.getAt(i);
-                ColumnDefinition columnDefinition = tableDefinition.getColumnsByName().get(false).get(name);
-                keyColumnTypes.add(columnDefinition.getType());
-                keyColumns[keyColumns.length] = columns[columnDefinition.getColumnIndex()];
+                Column keyColumn = keys.findColumn(name);
+                keyColumnTypes.add(keyColumn.getType());
+                keyColumns[keyColumns.length] = keyColumn;
             }
             this.columns = JsObject.freeze(columns);
             this.keyColumns = JsObject.freeze(keyColumns);
-            return w.getExportedObjects()[0].fetch();
-        }).then(result -> {
-            keys = (JsTable) result;
+
             // TODO(deephaven-core#3604) in case of a new session, we should do a full refetch
             keys.addEventListener(JsTable.EVENT_DISCONNECT, event -> fireEvent(EVENT_DISCONNECT));
             keys.addEventListener(JsTable.EVENT_RECONNECT, event -> {

web/client-api/src/main/java/io/deephaven/web/client/ide/IdeSession.java

@@ -138,6 +138,11 @@ public Promise<JsTreeTable> getHierarchicalTable(String name) {
                 .then(connection::getHierarchicalTable);
     }
 
+    public Promise<JsPartitionedTable> getPartitionedTable(String name) {
+        return connection.getVariableDefinition(name, JsVariableType.PARTITIONEDTABLE)
+                .then(connection::getPartitionedTable);
+    }
+
     public Promise<?> getObject(@TsTypeRef(JsVariableDescriptor.class) JsPropertyMap<Object> definitionObject) {
         return connection.getJsObject(definitionObject);
     }

web/client-api/src/test/java/io/deephaven/web/ClientIntegrationTestSuite.java

@@ -6,6 +6,7 @@
 import com.google.gwt.junit.tools.GWTTestSuite;
 import io.deephaven.web.client.api.HierarchicalTableTestGwt;
 import io.deephaven.web.client.api.NullValueTestGwt;
+import io.deephaven.web.client.api.PartitionedTableTestGwt;
 import io.deephaven.web.client.api.subscription.ConcurrentTableTestGwt;
 import io.deephaven.web.client.api.TableManipulationTestGwt;
 import io.deephaven.web.client.api.subscription.ViewportTestGwt;
@@ -26,6 +27,7 @@ public static Test suite() {
         suite.addTestSuite(ConcurrentTableTestGwt.class);
         suite.addTestSuite(NullValueTestGwt.class);
         suite.addTestSuite(HierarchicalTableTestGwt.class);
+        suite.addTestSuite(PartitionedTableTestGwt.class);
 
         // Unfinished:
         // suite.addTestSuite(TotalsTableTestGwt.class);

web/client-api/src/test/java/io/deephaven/web/client/api/AbstractAsyncGwtTestCase.java

@@ -191,6 +191,10 @@ public IThenable.ThenOnFulfilledCallbackFn<IdeSession, JsTreeTable> treeTable(St
         return session -> session.getTreeTable(tableName);
     }
 
+    public IThenable.ThenOnFulfilledCallbackFn<IdeSession, JsPartitionedTable> partitionedTable(String tableName) {
+        return session -> session.getPartitionedTable(tableName);
+    }
+
     /**
      * Utility method to report Promise errors to the unit test framework
      */

web/client-api/src/test/java/io/deephaven/web/client/api/PartitionedTableTestGwt.java

@@ -0,0 +1,173 @@
+//
+// Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
+//
+package io.deephaven.web.client.api;
+
+import elemental2.core.JsArray;
+import elemental2.dom.CustomEvent;
+import io.deephaven.web.client.api.subscription.ViewportData;
+import io.deephaven.web.client.api.tree.JsTreeTable;
+
+public class PartitionedTableTestGwt extends AbstractAsyncGwtTestCase {
+    @Override
+    public String getModuleName() {
+        return "io.deephaven.web.DeephavenIntegrationTest";
+    }
+
+    private final TableSourceBuilder tables = new TableSourceBuilder()
+            .script("from deephaven import empty_table")
+            .script("source = empty_table(100).update(['MyKey=``+i%5', 'x=i'])")
+            .script("partitioned_source = source.partition_by(by=['MyKey'])")
+            .script("partitioned_result = partitioned_source.transform(func=lambda t: t.drop_columns('MyKey'))")
+            .script("constituent_result = partitioned_result.get_constituent(['0'])");
+
+    private final TableSourceBuilder tickingTables = new TableSourceBuilder()
+            .script("from deephaven import time_table")
+            .script("source = time_table('PT0.1s').update(['MyKey=``+i%5', 'x=i'])")
+            .script("partitioned_source = source.partition_by(by=['MyKey'])")
+            .script("partitioned_result = partitioned_source.transform(func=lambda t: t.drop_columns('MyKey'))");
+
+
+    public void testPartitionedTable() {
+        connect(tables)
+                .then(partitionedTable("partitioned_source"))
+                .then(partitionedTable -> {
+                    delayTestFinish(1500);
+                    Column[] keyColumns = partitionedTable.getKeyColumns();
+                    assertEquals(1, keyColumns.length);
+                    assertEquals("MyKey", keyColumns[0].getName());
+
+                    Column[] columns = partitionedTable.getColumns();
+                    assertEquals(2, columns.length);
+                    assertEquals("MyKey", columns[0].getName());
+                    assertEquals("x", columns[1].getName());
+
+                    return partitionedTable.getKeyTable().then(keyTable -> {
+                        System.out.println("KeyTable size: " + keyTable.getSize());
+                        assertEquals(5d, keyTable.getSize());
+
+                        return partitionedTable.getTable("2");
+                    }).then(constituentTable -> {
+                        assertEquals(20d, constituentTable.getSize());
+                        partitionedTable.close();
+
+                        return null;
+                    });
+                })
+                .then(this::finish).catch_(this::report);
+    }
+
+    public void testTransformedPartitionedTable() {
+        connect(tables)
+                .then(partitionedTable("partitioned_result"))
+                .then(partitionedTable -> {
+                    delayTestFinish(1500);
+                    Column[] keyColumns = partitionedTable.getKeyColumns();
+                    assertEquals(1, keyColumns.length);
+                    assertEquals("MyKey", keyColumns[0].getName());
+
+                    Column[] columns = partitionedTable.getColumns();
+                    assertEquals(1, columns.length);
+                    assertEquals("x", columns[0].getName());
+
+                    return partitionedTable.getKeyTable().then(keyTable -> {
+                        assertEquals(5d, keyTable.getSize());
+
+                        return partitionedTable.getTable("2");
+                    }).then(constituentTable -> {
+                        assertEquals(20d, constituentTable.getSize());
+                        constituentTable.close();
+                        partitionedTable.close();
+                        return null;
+                    });
+                })
+                .then(this::finish).catch_(this::report);
+    }
+
+    public void testConstituentResult() {
+        connect(tables)
+                .then(table("constituent_result"))
+                .then(table -> {
+                    delayTestFinish(1500);
+
+                    JsArray<Column> columns = table.getColumns();
+                    assertEquals(1, columns.length);
+                    assertEquals("x", columns.getAt(0).getName());
+
+                    table.close();
+                    return null;
+                })
+                .then(this::finish).catch_(this::report);
+    }
+
+    public void testTickingPartitionedTable() {
+        connect(tickingTables)
+                .then(partitionedTable("partitioned_source"))
+                .then(partitionedTable -> {
+                    delayTestFinish(20_000);
+                    Column[] keyColumns = partitionedTable.getKeyColumns();
+                    assertEquals(1, keyColumns.length);
+                    assertEquals("MyKey", keyColumns[0].getName());
+
+                    Column[] columns = partitionedTable.getColumns();
+                    assertEquals(3, columns.length);
+                    assertEquals("MyKey", columns[1].getName());
+                    assertEquals("x", columns[2].getName());
+
+                    return partitionedTable.getKeyTable().then(keyTable -> {
+                        keyTable.setViewport(0, 99, keyTable.getColumns(), null);
+                        return keyTable.getViewportData().then(data -> {
+                            assertEquals(0d, keyTable.getSize());
+
+                            return waitForEventWhere(keyTable, JsTable.EVENT_UPDATED,
+                                    (CustomEvent<ViewportData> d) -> d.detail.getRows().length == 5, 20004);
+                        });
+                    }).then(event -> partitionedTable.getTable("2")).then(constituentTable -> {
+                        assertEquals(3, constituentTable.getColumns().length);
+                        assertEquals(2d, constituentTable.getSize());
+
+                        constituentTable.close();
+                        partitionedTable.close();
+
+                        return null;
+                    });
+                })
+                .then(this::finish).catch_(this::report);
+    }
+
+    public void testTickingTransformedPartitionedTable() {
+        connect(tickingTables)
+                .then(partitionedTable("partitioned_result"))
+                .then(partitionedTable -> {
+                    delayTestFinish(20_000);
+                    Column[] keyColumns = partitionedTable.getKeyColumns();
+                    assertEquals(1, keyColumns.length);
+                    assertEquals("MyKey", keyColumns[0].getName());
+
+                    Column[] columns = partitionedTable.getColumns();
+                    assertEquals(2, columns.length);
+                    assertEquals("Timestamp", columns[0].getName());
+                    assertEquals("x", columns[1].getName());
+
+                    return partitionedTable.getKeyTable().then(keyTable -> {
+                        keyTable.setViewport(0, 99, keyTable.getColumns(), null);
+                        return keyTable.getViewportData().then(data -> {
+                            assertEquals(0d, keyTable.getSize());
+
+                            return waitForEventWhere(keyTable, JsTable.EVENT_UPDATED,
+                                    (CustomEvent<ViewportData> d) -> d.detail.getRows().length == 5, 20004);
+                        }).then(event -> partitionedTable.getTable("2")).then(constituentTable -> {
+                            assertEquals(2, constituentTable.getColumns().length);
+                            assertEquals(2d, constituentTable.getSize());
+
+                            keyTable.close();
+                            constituentTable.close();
+                            partitionedTable.close();
+
+                            return null;
+                        });
+                    });
+                })
+                .then(this::finish).catch_(this::report);
+    }
+}