Preserve empty and key initializer table support for aggregations

engine/api/src/main/java/io/deephaven/engine/table/ColumnSource.java

@@ -110,11 +110,6 @@ default void releaseCachedResources() {
     <ALTERNATE_DATA_TYPE> ColumnSource<ALTERNATE_DATA_TYPE> reinterpret(
             @NotNull final Class<ALTERNATE_DATA_TYPE> alternateDataType) throws IllegalArgumentException;
 
-    @Override
-    default List<ColumnSource> getColumnSources() {
-        return Collections.singletonList(this);
-    }
-
     @Override
     default T createTuple(final long rowKey) {
         return get(rowKey);

engine/api/src/main/java/io/deephaven/engine/table/PartitionedTable.java

@@ -9,6 +9,7 @@
 import io.deephaven.base.log.LogOutputAppendable;
 import io.deephaven.engine.liveness.LivenessNode;
 import io.deephaven.engine.liveness.LivenessReferent;
+import io.deephaven.engine.updategraph.ConcurrentMethod;
 import io.deephaven.util.annotations.FinalDefault;
 import org.jetbrains.annotations.NotNull;
 
@@ -62,6 +63,7 @@ interface Proxy extends TableOperations<Proxy, TableOperations<?, ?>> {
      *
      * @return The underlying {@link Table partitioned table}
      */
+    @ConcurrentMethod
     Table table();
 
     /**
@@ -70,6 +72,7 @@ interface Proxy extends TableOperations<Proxy, TableOperations<?, ?>> {
      *
      * @return The key column names
      */
+    @ConcurrentMethod
     Set<String> keyColumnNames();
 
     /**
@@ -82,13 +85,15 @@ interface Proxy extends TableOperations<Proxy, TableOperations<?, ?>> {
      *
      * @return Whether the keys in the underlying partitioned table are unique
      */
+    @ConcurrentMethod
     boolean uniqueKeys();
 
     /**
      * Get the name of the "constituent" column of {@link Table tables}.
      *
      * @return The constituent column name
      */
+    @ConcurrentMethod
     String constituentColumnName();
 
     /**
@@ -98,6 +103,7 @@ interface Proxy extends TableOperations<Proxy, TableOperations<?, ?>> {
      *
      * @return The constituent definition
      */
+    @ConcurrentMethod
     TableDefinition constituentDefinition();
 
     /**
@@ -115,6 +121,7 @@ interface Proxy extends TableOperations<Proxy, TableOperations<?, ?>> {
      *
      * @return Whether the constituents of the underlying partitioned table can change
      */
+    @ConcurrentMethod
     boolean constituentChangesPermitted();
 
     /**
@@ -125,6 +132,7 @@ interface Proxy extends TableOperations<Proxy, TableOperations<?, ?>> {
      * @see #proxy(boolean, boolean)
      */
     @FinalDefault
+    @ConcurrentMethod
     default Proxy proxy() {
         return proxy(true, true);
     }
@@ -147,6 +155,7 @@ default Proxy proxy() {
      * @return A proxy that allows {@link TableOperations table operations} to be applied to the constituent tables of
      *         this PartitionedTable
      */
+    @ConcurrentMethod
     Proxy proxy(boolean requireMatchingKeys, boolean sanityCheckJoinOperations);
 
     /**
@@ -169,6 +178,7 @@ default Proxy proxy() {
      * @param filters The filters to apply. Must not reference the constituent column.
      * @return The filtered PartitionedTable
      */
+    @ConcurrentMethod
     PartitionedTable filter(Collection<? extends Filter> filters);
 
     /**
@@ -180,6 +190,7 @@ default Proxy proxy() {
      * @param sortColumns The columns to sort by. Must not reference the constituent column.
      * @return The sorted PartitionedTable
      */
+    @ConcurrentMethod
     PartitionedTable sort(Collection<SortColumn> sortColumns);
 
     /**
@@ -242,6 +253,7 @@ PartitionedTable partitionedTransform(
      * @return The {@link Table constituent} at the single row in {@link #table()} matching the {@code keyColumnValues},
      *         or {@code null} if no matches were found
      */
+    @ConcurrentMethod
     Table constituentFor(@NotNull Object... keyColumnValues);
 
     /**
@@ -255,5 +267,6 @@ PartitionedTable partitionedTransform(
      *
      * @return An array of all current {@link Table constituents}
      */
+    @ConcurrentMethod
     Table[] constituents();
 }

engine/api/src/main/java/io/deephaven/engine/table/Table.java

@@ -142,8 +142,10 @@ public interface Table extends
      * all rows:
      * <ol>
      * <li>{@link #groupBy} is unsupported
-     * <li>{@link #aggBy} is unsupported if {@link Aggregation#AggGroup(String...)} is used
      * <li>{@link #partitionBy} is unsupported</li>
+     * <li>{@link #partitionedAggBy(Collection, boolean, Table, String...) partitionedAggBy} is unsupported</li>
+     * <li>{@link #aggBy} is unsupported if either of {@link io.deephaven.api.agg.spec.AggSpecGroup group} or
+     * {@link io.deephaven.api.agg.Partition partition} are used</li>
      * <li>{@link #rollup(Collection, boolean, ColumnName...) rollup()} is unsupported if
      * {@code includeConstituents == true}</li>
      * <li>{@link #treeTable(String, String) treeTable()} is unsupported</li>
@@ -995,6 +997,14 @@ Table join(Table rightTable, Collection<? extends JoinMatch> columnsToMatch,
     @ConcurrentMethod
     Table aggBy(Aggregation aggregation);
 
+    @Override
+    @ConcurrentMethod
+    Table aggBy(Collection<? extends Aggregation> aggregations);
+
+    @Override
+    @ConcurrentMethod
+    Table aggBy(Collection<? extends Aggregation> aggregations, boolean preserveEmpty);
+
     @Override
     @ConcurrentMethod
     Table aggBy(Aggregation aggregation, String... groupByColumns);
@@ -1005,15 +1015,16 @@ Table join(Table rightTable, Collection<? extends JoinMatch> columnsToMatch,
 
     @Override
     @ConcurrentMethod
-    Table aggBy(Collection<? extends Aggregation> aggregations, Collection<? extends ColumnName> groupByColumns);
+    Table aggBy(Collection<? extends Aggregation> aggregations, String... groupByColumns);
 
     @Override
     @ConcurrentMethod
-    Table aggBy(Collection<? extends Aggregation> aggregations, String... groupByColumns);
+    Table aggBy(Collection<? extends Aggregation> aggregations, Collection<? extends ColumnName> groupByColumns);
 
     @Override
     @ConcurrentMethod
-    Table aggBy(Collection<? extends Aggregation> aggregations);
+    Table aggBy(Collection<? extends Aggregation> aggregations, boolean preserveEmpty, Table initialGroups,
+            Collection<? extends ColumnName> groupByColumns);
 
     Table headBy(long nRows, Collection<String> groupByColumnNames);
 
@@ -1602,6 +1613,29 @@ Table join(Table rightTable, Collection<? extends JoinMatch> columnsToMatch,
     @ConcurrentMethod
     PartitionedTable partitionBy(String... keyColumnNames);
 
+    /**
+     * Convenience method that performs an {@link #aggBy(Collection, boolean, Table, Collection)} and wraps the result
+     * in a {@link PartitionedTable}. If {@code aggregations} does not include a {@link io.deephaven.api.agg.Partition
+     * partition}, one will be added automatically with the default constituent column name and behavior used in
+     * {@link #partitionBy(String...)}.
+     *
+     * @param aggregations The {@link Aggregation aggregations} to apply
+     * @param preserveEmpty Whether to keep result rows for groups that are initially empty or become empty as a result
+     *        of updates. Each aggregation operator defines its own value for empty groups.
+     * @param initialGroups A table whose distinct combinations of values for the {@code groupByColumns} should be used
+     *        to create an initial set of aggregation groups. All other columns are ignored. This is useful in
+     *        combination with {@code preserveEmpty == true} to ensure that particular groups appear in the result
+     *        table, or with {@code preserveEmpty == false} to control the encounter order for a collection of groups
+     *        and thus their relative order in the result. Changes to {@code initialGroups} are not expected or handled;
+     *        if {@code initialGroups} is a refreshing table, only its contents at instantiation time will be used. If
+     *        {@code initialGroups == null}, the result will be the same as if a table with no rows was supplied.
+     * @param keyColumnNames The names of the key columns to aggregate by
+     * @return A {@link PartitionedTable} keyed by {@code keyColumnNames}
+     */
+    @ConcurrentMethod
+    PartitionedTable partitionedAggBy(Collection<? extends Aggregation> aggregations, boolean preserveEmpty,
+            Table initialGroups, String... keyColumnNames);
+
     // -----------------------------------------------------------------------------------------------------------------
     // Hierarchical table operations (rollup and treeTable).
     // -----------------------------------------------------------------------------------------------------------------

engine/table/src/main/java/io/deephaven/engine/table/impl/AbstractColumnSource.java

@@ -26,6 +26,7 @@
 
 import java.util.Collections;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 
 public abstract class AbstractColumnSource<T> implements
@@ -42,6 +43,7 @@ public abstract class AbstractColumnSource<T> implements
     protected final Class<?> componentType;
 
     protected volatile Map<T, RowSet> groupToRange;
+    protected volatile List<ColumnSource> rowSetIndexerKey;
 
     protected AbstractColumnSource(@NotNull final Class<T> type) {
         this(type, Object.class);
@@ -102,6 +104,19 @@ public ColumnSource<T> getPrevSource() {
         return new PrevColumnSource<>(this);
     }
 
+    @Override
+    public List<ColumnSource> getColumnSources() {
+        List<ColumnSource> localRowSetIndexerKey;
+        if ((localRowSetIndexerKey = rowSetIndexerKey) == null) {
+            synchronized (this) {
+                if ((localRowSetIndexerKey = rowSetIndexerKey) == null) {
+                    rowSetIndexerKey = localRowSetIndexerKey = Collections.singletonList(this);
+                }
+            }
+        }
+        return localRowSetIndexerKey;
+    }
+
     @Override
     public Map<T, RowSet> getGroupToRange() {
         return groupToRange;

engine/table/src/main/java/io/deephaven/engine/table/impl/GroupingUtils.java

@@ -10,8 +10,10 @@
 import io.deephaven.engine.rowset.TrackingWritableRowSet;
 import io.deephaven.engine.rowset.WritableRowSet;
 import io.deephaven.engine.table.ColumnSource;
+import io.deephaven.engine.table.WritableColumnSource;
 import io.deephaven.engine.table.impl.indexer.RowSetIndexer;
 import io.deephaven.engine.table.impl.sources.ArrayBackedColumnSource;
+import io.deephaven.engine.table.impl.sources.InMemoryColumnSource;
 import io.deephaven.engine.table.impl.sources.ObjectArraySource;
 import org.apache.commons.lang3.mutable.MutableInt;
 import org.jetbrains.annotations.NotNull;
@@ -158,4 +160,25 @@ public static <TYPE> Pair<ArrayBackedColumnSource<TYPE>, ObjectArraySource<Track
 
         return new Pair<>(resultKeyColumnSource, resultIndexColumnSource);
     }
+
+    /**
+     * Convert a group-to-RowSet map to a flat, immutable, in-memory column of keys.
+     *
+     * @param originalKeyColumnSource The key column source whose contents are reflected by the group-to-RowSet map
+     *        (used for typing, only)
+     * @param groupToRowSet The group-to-RowSet map to convert
+     * @return A flat, immutable, in-memory column of keys
+     */
+    public static <TYPE> WritableColumnSource<TYPE> groupingKeysToImmutableFlatSource(
+            @NotNull final ColumnSource<TYPE> originalKeyColumnSource,
+            @NotNull final Map<TYPE, RowSet> groupToRowSet) {
+        final WritableColumnSource<TYPE> destination = InMemoryColumnSource.makeImmutableSource(
+                originalKeyColumnSource.getType(), originalKeyColumnSource.getComponentType());
+        destination.ensureCapacity(groupToRowSet.size());
+        int ri = 0;
+        for (final TYPE key : groupToRowSet.keySet()) {
+            destination.set(ri++, key);
+        }
+        return destination;
+    }
 }

engine/table/src/main/java/io/deephaven/engine/table/impl/HierarchicalTable.java

@@ -160,11 +160,6 @@ public Table join(Table rightTable, MatchPair[] columnsToMatch, MatchPair[] colu
         return throwUnsupported("join()");
     }
 
-    @Override
-    public Table countBy(String countColumnName, ColumnName... groupByColumns) {
-        return throwUnsupported("countBy()");
-    }
-
     @Override
     public Table ungroup(boolean nullFill, String... columnsToUngroup) {
         return throwUnsupported("ungroup()");
@@ -186,7 +181,7 @@ public Table aggAllBy(AggSpec spec, ColumnName... groupByColumns) {
     }
 
     @Override
-    public Table aggBy(Collection<? extends Aggregation> aggregations,
+    public Table aggBy(Collection<? extends Aggregation> aggregations, boolean preserveEmpty, Table initialGroups,
             Collection<? extends ColumnName> groupByColumns) {
         return throwUnsupported("aggBy()");
     }
@@ -251,6 +246,12 @@ public PartitionedTable partitionBy(boolean dropKeys, String... keyColumnNames)
         return throwUnsupported("partitionBy()");
     }
 
+    @Override
+    public PartitionedTable partitionedAggBy(Collection<? extends Aggregation> aggregations, boolean preserveEmpty,
+            Table initialGroups, String... keyColumnNames) {
+        return throwUnsupported("partitionedAggBy()");
+    }
+
     @Override
     public Table rollup(Collection<? extends Aggregation> aggregations, boolean includeConstituents,
             ColumnName... groupByColumns) {