Drop columns from key table

Author: georgecwan

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

@@ -10,6 +10,7 @@
 import io.deephaven.javascript.proto.dhinternal.io.deephaven.proto.partitionedtable_pb.GetTableRequest;
 import io.deephaven.javascript.proto.dhinternal.io.deephaven.proto.partitionedtable_pb.MergeRequest;
 import io.deephaven.javascript.proto.dhinternal.io.deephaven.proto.partitionedtable_pb.PartitionedTableDescriptor;
+import io.deephaven.javascript.proto.dhinternal.io.deephaven.proto.table_pb.DropColumnsRequest;
 import io.deephaven.javascript.proto.dhinternal.io.deephaven.proto.ticket_pb.TypedTicket;
 import io.deephaven.web.client.api.barrage.WebBarrageUtils;
 import io.deephaven.web.client.api.barrage.def.ColumnDefinition;
@@ -23,6 +24,7 @@
 import io.deephaven.web.shared.data.RangeSet;
 import io.deephaven.web.shared.fu.JsConsumer;
 import jsinterop.annotations.JsIgnore;
+import jsinterop.annotations.JsMethod;
 import jsinterop.annotations.JsProperty;
 import jsinterop.annotations.JsType;
 import jsinterop.base.Js;
@@ -45,10 +47,10 @@ public class JsPartitionedTable extends HasLifecycle implements ServerObject {
     public static final String EVENT_KEYADDED = "keyadded",
             EVENT_DISCONNECT = JsTable.EVENT_DISCONNECT,
             EVENT_RECONNECT = JsTable.EVENT_RECONNECT,
-            /**
-             * Indicates that an error has occurred while communicating with the server.
-             */
-            EVENT_RECONNECTFAILED = JsTable.EVENT_RECONNECTFAILED;
+    /**
+     * Indicates that an error has occurred while communicating with the server.
+     */
+    EVENT_RECONNECTFAILED = JsTable.EVENT_RECONNECTFAILED;
 
     private final WorkerConnection connection;
     private final JsWidget widget;
@@ -167,26 +169,26 @@ private void populateLazyTable(List<Object> key) {
         tables.put(key, JsLazy.of(() -> {
             // If we've entered this lambda, the JsLazy is being used, so we need to go ahead and get the tablehandle
             final ClientTableState entry = connection.newState((c, cts, metadata) -> {
-                // TODO deephaven-core#2529 parallelize this
-                connection.newTable(
-                        descriptor.getKeyColumnNamesList().asArray(new String[0]),
-                        keyColumnTypes.toArray(new String[0]),
-                        key.stream().map(item -> new Object[] {item}).toArray(Object[][]::new),
-                        null,
-                        this)
-                        .then(table -> {
-                            GetTableRequest getTableRequest = new GetTableRequest();
-                            getTableRequest.setPartitionedTable(widget.getTicket());
-                            getTableRequest.setKeyTableTicket(table.getHandle().makeTicket());
-                            getTableRequest.setResultId(cts.getHandle().makeTicket());
-                            connection.partitionedTableServiceClient().getTable(getTableRequest, connection.metadata(),
-                                    (error, success) -> {
-                                        table.close();
-                                        c.apply(error, success);
-                                    });
-                            return null;
-                        });
-            },
+                        // TODO deephaven-core#2529 parallelize this
+                        connection.newTable(
+                                        descriptor.getKeyColumnNamesList().asArray(new String[0]),
+                                        keyColumnTypes.toArray(new String[0]),
+                                        key.stream().map(item -> new Object[]{item}).toArray(Object[][]::new),
+                                        null,
+                                        this)
+                                .then(table -> {
+                                    GetTableRequest getTableRequest = new GetTableRequest();
+                                    getTableRequest.setPartitionedTable(widget.getTicket());
+                                    getTableRequest.setKeyTableTicket(table.getHandle().makeTicket());
+                                    getTableRequest.setResultId(cts.getHandle().makeTicket());
+                                    connection.partitionedTableServiceClient().getTable(getTableRequest, connection.metadata(),
+                                            (error, success) -> {
+                                                table.close();
+                                                c.apply(error, success);
+                                            });
+                                    return null;
+                                });
+                    },
                     "partitioned table key " + key);
 
             // later, when the CTS is released, remove this "table" from the map and replace with an unresolved JsLazy
@@ -202,7 +204,7 @@ private void populateLazyTable(List<Object> key) {
 
     /**
      * Fetch the table with the given key.
-     * 
+     *
      * @param key The key to fetch. An array of values for each key column, in the same order as the key columns are.
      * @return Promise of dh.Table
      */
@@ -225,24 +227,24 @@ public Promise<JsTable> getTable(Object key) {
     /**
      * Open a new table that is the result of merging all constituent tables. See
      * {@link io.deephaven.engine.table.PartitionedTable#merge()} for details.
-     * 
+     *
      * @return A merged representation of the constituent tables.
      */
     public Promise<JsTable> getMergedTable() {
         return connection.newState((c, cts, metadata) -> {
-            MergeRequest requestMessage = new MergeRequest();
-            requestMessage.setPartitionedTable(widget.getTicket());
-            requestMessage.setResultId(cts.getHandle().makeTicket());
-            connection.partitionedTableServiceClient().merge(requestMessage, connection.metadata(), c::apply);
-        }, "partitioned table merged table")
+                    MergeRequest requestMessage = new MergeRequest();
+                    requestMessage.setPartitionedTable(widget.getTicket());
+                    requestMessage.setResultId(cts.getHandle().makeTicket());
+                    connection.partitionedTableServiceClient().merge(requestMessage, connection.metadata(), c::apply);
+                }, "partitioned table merged table")
                 .refetch(this, connection.metadata())
                 .then(cts -> Promise.resolve(new JsTable(cts.getConnection(), cts)));
     }
 
     /**
      * The set of all currently known keys. This is kept up to date, so getting the list after adding an event listener
      * for <b>keyadded</b> will ensure no keys are missed.
-     * 
+     *
      * @return Set of Object
      */
     public JsSet<Object> getKeys() {
@@ -254,7 +256,7 @@ public JsSet<Object> getKeys() {
 
     /**
      * The count of known keys.
-     * 
+     *
      * @return int
      */
     @JsProperty(name = "size")
@@ -284,13 +286,21 @@ public Column[] getColumns() {
     }
 
     /**
-     * A Table object containing all currently known keys used for this partitioned table.
+     * Fetch a table containing all the valid keys of the partitioned table.
      *
-     * @return Table
+     * @return Promise of a Table
      */
-    @JsProperty
-    public JsTable getKeyTable() {
-        return keys;
+    @JsMethod
+    public Promise<JsTable> getKeyTable() {
+        return connection.newState((c, state, metadata) -> {
+                    DropColumnsRequest drop = new DropColumnsRequest();
+                    drop.setColumnNamesList(new String[]{descriptor.getConstituentColumnName()});
+                    drop.setSourceId(keys.state().getHandle().makeTableReference());
+                    drop.setResultId(state.getHandle().makeTicket());
+                    connection.tableServiceClient().dropColumns(drop, metadata, c::apply);
+                }, "drop constituent column")
+                .refetch(this, connection.metadata())
+                .then(state -> Promise.resolve(new JsTable(connection, state)));
     }
 
     /**