Add selectDistinct / countBy to table API

Partial deephaven#1617
devinrsmith · Dec 3, 2021 · d9d8fb1 · d9d8fb1
1 parent 9ef33fe
commit d9d8fb1
Show file tree

Hide file tree

Showing 12 changed files with 270 additions and 5 deletions.
diff --git a/...lient/flight-examples/src/main/java/io/deephaven/client/examples/AggregateAllExample.java b/...lient/flight-examples/src/main/java/io/deephaven/client/examples/AggregateAllExample.java
@@ -88,6 +88,9 @@ protected void execute(FlightSession flight) throws Exception {
         final TableSpec spec = InMemoryKeyBackedInputTable.of(TableHeader.of(inputKey, groupKey, u, v, w, x, y, z),
                 Collections.singletonList(inputKey.name()));
         builder.putMap("base", spec);
+        builder.putMap("countBy", spec.countBy("Count", groupKey.name()).sort(groupKey.name()));
+        builder.putMap("selectDistinct", spec.selectDistinct(groupKey.name()).sort(groupKey.name()));
+
         // todo: countBy
         for (AggSpec aggSpec : aggSpecs()) {
             final String name = aggSpec.toString()
@@ -96,7 +99,7 @@ protected void execute(FlightSession flight) throws Exception {
                     .replace('}', '_')
                     .replace(' ', '_')
                     .replace('.', '_');
-            final TableSpec tableSpec = spec.aggAllBy(aggSpec, groupKey.name());
+            final TableSpec tableSpec = spec.aggAllBy(aggSpec, groupKey.name()).sort(groupKey.name());
             builder.putMap(name, tableSpec);
         }
         final LabeledTables tables = builder.build();

diff --git a/java-client/session/src/main/java/io/deephaven/client/impl/BatchTableRequestBuilder.java b/java-client/session/src/main/java/io/deephaven/client/impl/BatchTableRequestBuilder.java
@@ -74,6 +74,7 @@
 import io.deephaven.proto.backplane.grpc.NotCondition;
 import io.deephaven.proto.backplane.grpc.OrCondition;
 import io.deephaven.proto.backplane.grpc.Reference;
+import io.deephaven.proto.backplane.grpc.SelectDistinctRequest;
 import io.deephaven.proto.backplane.grpc.SelectOrUpdateRequest;
 import io.deephaven.proto.backplane.grpc.SnapshotTableRequest;
 import io.deephaven.proto.backplane.grpc.SortDescriptor;
@@ -87,6 +88,7 @@
 import io.deephaven.qst.table.AggregationTable;
 import io.deephaven.qst.table.AsOfJoinTable;
 import io.deephaven.qst.table.ByTableBase;
+import io.deephaven.qst.table.CountByTable;
 import io.deephaven.qst.table.EmptyTable;
 import io.deephaven.qst.table.ExactJoinTable;
 import io.deephaven.qst.table.HeadTable;
@@ -99,6 +101,7 @@
 import io.deephaven.qst.table.NewTable;
 import io.deephaven.qst.table.ReverseAsOfJoinTable;
 import io.deephaven.qst.table.ReverseTable;
+import io.deephaven.qst.table.SelectDistinctTable;
 import io.deephaven.qst.table.SelectTable;
 import io.deephaven.qst.table.SingleParentTable;
 import io.deephaven.qst.table.SnapshotTable;
@@ -451,6 +454,16 @@ public void visit(InMemoryKeyBackedInputTable inMemoryKeyBacked) {
             out = op(Builder::setCreateInputTable, builder);
         }
 
+        @Override
+        public void visit(SelectDistinctTable selectDistinctTable) {
+            out = op(Builder::setSelectDistinct, selectDistinct(selectDistinctTable));
+        }
+
+        @Override
+        public void visit(CountByTable countByTable) {
+            out = op(Builder::setComboAggregate, countBy(countByTable));
+        }
+
         private SelectOrUpdateRequest selectOrUpdate(SingleParentTable x,
                 Collection<Selectable> columns) {
             SelectOrUpdateRequest.Builder builder =
@@ -485,6 +498,24 @@ private ComboAggregateRequest.Builder groupByColumns(ByTableBase base) {
             }
             return builder;
         }
+
+        private SelectDistinctRequest selectDistinct(SelectDistinctTable selectDistinctTable) {
+            SelectDistinctRequest.Builder builder = SelectDistinctRequest.newBuilder()
+                    .setResultId(ticket)
+                    .setSourceId(ref(selectDistinctTable.parent()));
+            for (Selectable column : selectDistinctTable.groupByColumns()) {
+                builder.addColumnNames(Strings.of(column));
+            }
+            return builder.build();
+        }
+
+        private ComboAggregateRequest countBy(CountByTable countByTable) {
+            final Aggregate aggregate = Aggregate.newBuilder()
+                    .setType(AggType.COUNT)
+                    .setColumnName(countByTable.countName().name())
+                    .build();
+            return groupByColumns(countByTable).addAggregates(aggregate).build();
+        }
     }
 
     private static class AggregationAdapter implements Aggregation.Visitor {

diff --git a/qst/graphviz/src/main/java/io/deephaven/graphviz/LabelBuilder.java b/qst/graphviz/src/main/java/io/deephaven/graphviz/LabelBuilder.java
@@ -4,6 +4,7 @@
 import io.deephaven.qst.table.AggregateAllByTable;
 import io.deephaven.qst.table.AggregationTable;
 import io.deephaven.qst.table.AsOfJoinTable;
+import io.deephaven.qst.table.CountByTable;
 import io.deephaven.qst.table.EmptyTable;
 import io.deephaven.qst.table.ExactJoinTable;
 import io.deephaven.qst.table.HeadTable;
@@ -14,6 +15,7 @@
 import io.deephaven.qst.table.JoinTable;
 import io.deephaven.qst.table.NaturalJoinTable;
 import io.deephaven.qst.table.ReverseAsOfJoinTable;
+import io.deephaven.qst.table.SelectDistinctTable;
 import io.deephaven.qst.table.SelectTable;
 import io.deephaven.qst.table.SelectableTable;
 import io.deephaven.qst.table.TableSpec;
@@ -164,6 +166,20 @@ public void visit(InMemoryKeyBackedInputTable inMemoryKeyBacked) {
         });
     }
 
+    @Override
+    public void visit(SelectDistinctTable selectDistinctTable) {
+        sb.append("selectDistinct(");
+        append(Strings::of, selectDistinctTable.groupByColumns(), sb);
+        sb.append(')');
+    }
+
+    @Override
+    public void visit(CountByTable countByTable) {
+        sb.append("countBy(").append(countByTable.countName()).append(',');
+        append(Strings::of, countByTable.groupByColumns(), sb);
+        sb.append(')');
+    }
+
     private void join(String name, Join j) {
         sb.append(name).append("([");
         append(Strings::of, j.matches(), sb);

diff --git a/qst/src/main/java/io/deephaven/qst/TableAdapterImpl.java b/qst/src/main/java/io/deephaven/qst/TableAdapterImpl.java
@@ -7,6 +7,7 @@
 import io.deephaven.qst.table.AggregateAllByTable;
 import io.deephaven.qst.table.AggregationTable;
 import io.deephaven.qst.table.AsOfJoinTable;
+import io.deephaven.qst.table.CountByTable;
 import io.deephaven.qst.table.EmptyTable;
 import io.deephaven.qst.table.ExactJoinTable;
 import io.deephaven.qst.table.HeadTable;
@@ -18,6 +19,7 @@
 import io.deephaven.qst.table.ParentsVisitor;
 import io.deephaven.qst.table.ReverseAsOfJoinTable;
 import io.deephaven.qst.table.ReverseTable;
+import io.deephaven.qst.table.SelectDistinctTable;
 import io.deephaven.qst.table.SelectTable;
 import io.deephaven.qst.table.SingleParentTable;
 import io.deephaven.qst.table.SnapshotTable;
@@ -230,14 +232,22 @@ public void visit(ReverseAsOfJoinTable raj) {
     @Override
     public void visit(AggregateAllByTable aggAllByTable) {
         final AggSpec spec = aggAllByTable.spec();
-        final Selectable[] groupByColumns = aggAllByTable.groupByColumns().toArray(new Selectable[0]);
-        addOp(aggAllByTable, parentOps(aggAllByTable).aggAllBy(spec, groupByColumns));
+        if (aggAllByTable.groupByColumns().isEmpty()) {
+            addOp(aggAllByTable, parentOps(aggAllByTable).aggAllBy(spec));
+        } else {
+            final Selectable[] groupByColumns = aggAllByTable.groupByColumns().toArray(new Selectable[0]);
+            addOp(aggAllByTable, parentOps(aggAllByTable).aggAllBy(spec, groupByColumns));
+        }
     }
 
     @Override
     public void visit(AggregationTable aggregationTable) {
-        addOp(aggregationTable,
-                parentOps(aggregationTable).aggBy(aggregationTable.aggregations(), aggregationTable.groupByColumns()));
+        if (aggregationTable.groupByColumns().isEmpty()) {
+            addOp(aggregationTable, parentOps(aggregationTable).aggBy(aggregationTable.aggregations()));
+        } else {
+            addOp(aggregationTable, parentOps(aggregationTable).aggBy(aggregationTable.aggregations(),
+                    aggregationTable.groupByColumns()));
+        }
     }
 
     @Override
@@ -250,6 +260,26 @@ public void visit(InputTable inputTable) {
         addTable(inputTable, tableCreation.of(inputTable));
     }
 
+    @Override
+    public void visit(SelectDistinctTable selectDistinctTable) {
+        if (selectDistinctTable.groupByColumns().isEmpty()) {
+            addOp(selectDistinctTable, parentOps(selectDistinctTable).selectDistinct());
+        } else {
+            addOp(selectDistinctTable,
+                    parentOps(selectDistinctTable).selectDistinct(selectDistinctTable.groupByColumns()));
+        }
+    }
+
+    @Override
+    public void visit(CountByTable countByTable) {
+        if (countByTable.groupByColumns().isEmpty()) {
+            addOp(countByTable, parentOps(countByTable).countBy(countByTable.countName().name()));
+        } else {
+            addOp(countByTable, parentOps(countByTable).countBy(countByTable.countName().name(),
+                    countByTable.groupByColumns().toArray(new Selectable[0])));
+        }
+    }
+
     private final class OutputTable implements Output<TOPS, TABLE> {
         private final TABLE table;
 

diff --git a/qst/src/main/java/io/deephaven/qst/table/CountByTable.java b/qst/src/main/java/io/deephaven/qst/table/CountByTable.java
@@ -0,0 +1,27 @@
+package io.deephaven.qst.table;
+
+import io.deephaven.annotations.NodeStyle;
+import io.deephaven.api.ColumnName;
+import org.immutables.value.Value.Immutable;
+
+@Immutable
+@NodeStyle
+public abstract class CountByTable extends ByTableBase {
+
+    public static Builder builder() {
+        return ImmutableCountByTable.builder();
+    }
+
+    public abstract ColumnName countName();
+
+    @Override
+    public final <V extends Visitor> V walk(V visitor) {
+        visitor.visit(this);
+        return visitor;
+    }
+
+    public interface Builder extends ByTableBase.Builder<CountByTable, CountByTable.Builder> {
+
+        Builder countName(ColumnName countName);
+    }
+}
diff --git a/qst/src/main/java/io/deephaven/qst/table/ParentsVisitor.java b/qst/src/main/java/io/deephaven/qst/table/ParentsVisitor.java
@@ -269,6 +269,16 @@ public void visit(TableHeader header) {
         });
     }
 
+    @Override
+    public void visit(SelectDistinctTable selectDistinctTable) {
+        out = single(selectDistinctTable);
+    }
+
+    @Override
+    public void visit(CountByTable countByTable) {
+        out = single(countByTable);
+    }
+
     private static class Search {
 
         private final Predicate<TableSpec> excludePaths;

diff --git a/qst/src/main/java/io/deephaven/qst/table/SelectDistinctTable.java b/qst/src/main/java/io/deephaven/qst/table/SelectDistinctTable.java
@@ -0,0 +1,23 @@
+package io.deephaven.qst.table;
+
+import io.deephaven.annotations.NodeStyle;
+import org.immutables.value.Value.Immutable;
+
+@Immutable
+@NodeStyle
+public abstract class SelectDistinctTable extends ByTableBase {
+
+    public static Builder builder() {
+        return ImmutableSelectDistinctTable.builder();
+    }
+
+    @Override
+    public final <V extends Visitor> V walk(V visitor) {
+        visitor.visit(this);
+        return visitor;
+    }
+
+    public interface Builder extends ByTableBase.Builder<SelectDistinctTable, SelectDistinctTable.Builder> {
+
+    }
+}
diff --git a/qst/src/main/java/io/deephaven/qst/table/TableBase.java b/qst/src/main/java/io/deephaven/qst/table/TableBase.java
@@ -451,6 +451,55 @@ public final AggregationTable aggBy(Collection<? extends Aggregation> aggregatio
                 .addAllAggregations(aggregations).build();
     }
 
+    @Override
+    public final SelectDistinctTable selectDistinct() {
+        return SelectDistinctTable.builder().parent(this).build();
+    }
+
+    @Override
+    public final SelectDistinctTable selectDistinct(String... groupByColumns) {
+        final SelectDistinctTable.Builder builder = SelectDistinctTable.builder().parent(this);
+        for (String groupByColumn : groupByColumns) {
+            builder.addGroupByColumns(Selectable.parse(groupByColumn));
+        }
+        return builder.build();
+    }
+
+    @Override
+    public final SelectDistinctTable selectDistinct(Selectable... groupByColumns) {
+        return selectDistinct(Arrays.asList(groupByColumns));
+    }
+
+    @Override
+    public final SelectDistinctTable selectDistinct(Collection<? extends Selectable> groupByColumns) {
+        return SelectDistinctTable.builder().parent(this).addAllGroupByColumns(groupByColumns).build();
+    }
+
+    @Override
+    public final CountByTable countBy(String countColumnName) {
+        return CountByTable.builder().parent(this).countName(ColumnName.of(countColumnName)).build();
+    }
+
+    @Override
+    public final CountByTable countBy(String countColumnName, String... groupByColumns) {
+        return countBy(countColumnName, Arrays.asList(groupByColumns));
+    }
+
+    @Override
+    public final CountByTable countBy(String countColumnName, Selectable... groupByColumns) {
+        return CountByTable.builder().parent(this).countName(ColumnName.of(countColumnName))
+                .addGroupByColumns(groupByColumns).build();
+    }
+
+    @Override
+    public final CountByTable countBy(String countColumnName, Collection<String> groupByColumns) {
+        CountByTable.Builder builder = CountByTable.builder().parent(this).countName(ColumnName.of(countColumnName));
+        for (String groupByColumn : groupByColumns) {
+            builder.addGroupByColumns(Selectable.parse(groupByColumn));
+        }
+        return builder.build();
+    }
+
     @Override
     public final <V extends TableSchema.Visitor> V walk(V visitor) {
         visitor.visit(this);

diff --git a/qst/src/main/java/io/deephaven/qst/table/TableSpec.java b/qst/src/main/java/io/deephaven/qst/table/TableSpec.java
@@ -146,5 +146,9 @@ interface Visitor {
         void visit(TicketTable ticketTable);
 
         void visit(InputTable inputTable);
+
+        void visit(SelectDistinctTable selectDistinctTable);
+
+        void visit(CountByTable countByTable);
     }
 }
diff --git a/qst/src/main/java/io/deephaven/qst/table/TableVisitorGeneric.java b/qst/src/main/java/io/deephaven/qst/table/TableVisitorGeneric.java
@@ -128,4 +128,14 @@ public void visit(TicketTable ticketTable) {
     public void visit(InputTable inputTable) {
         accept(inputTable);
     }
+
+    @Override
+    public void visit(SelectDistinctTable selectDistinctTable) {
+        accept(selectDistinctTable);
+    }
+
+    @Override
+    public void visit(CountByTable countByTable) {
+        accept(countByTable);
+    }
 }
diff --git a/table-api/src/main/java/io/deephaven/api/TableOperations.java b/table-api/src/main/java/io/deephaven/api/TableOperations.java
@@ -511,4 +511,26 @@ TOPS raj(TABLE rightTable, Collection<? extends JoinMatch> columnsToMatch,
     TOPS aggBy(Collection<? extends Aggregation> aggregations, String... groupByColumns);
 
     TOPS aggBy(Collection<? extends Aggregation> aggregations, Collection<? extends Selectable> groupByColumns);
+
+    // -------------------------------------------------------------------------------------------
+
+    TOPS selectDistinct();
+
+    TOPS selectDistinct(String... groupByColumns);
+
+    TOPS selectDistinct(Selectable... groupByColumns);
+
+    TOPS selectDistinct(Collection<? extends Selectable> groupByColumns);
+
+    // -------------------------------------------------------------------------------------------
+
+    TOPS countBy(String countColumnName);
+
+    TOPS countBy(String countColumnName, String... groupByColumns);
+
+    TOPS countBy(String countColumnName, Selectable... groupByColumns);
+
+    TOPS countBy(String countColumnName, Collection<String> groupByColumns);
+
+    // -------------------------------------------------------------------------------------------
 }