Add TableDefinition column name helpers

Plot/src/main/java/io/deephaven/plot/util/ArgumentValidations.java

@@ -841,7 +841,7 @@ public static void assertColumnsInTable(final Table t, final PlotInfo plotInfo,
         assertNotNull(t, "t", plotInfo);
         assertNotNull(cols, "cols", plotInfo);
         for (String c : cols) {
-            if (!t.getColumnSourceMap().containsKey(c)) {
+            if (!t.hasColumns(c)) {
                 throw new PlotIllegalArgumentException("Column " + c + " could not be found in table.", plotInfo);
             }
         }

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

@@ -9,13 +9,15 @@
 import io.deephaven.base.log.LogOutput;
 import io.deephaven.base.log.LogOutputAppendable;
 import io.deephaven.base.verify.Assert;
+import io.deephaven.engine.table.impl.NoSuchColumnException;
 import io.deephaven.io.log.impl.LogOutputStringImpl;
 import io.deephaven.qst.column.header.ColumnHeader;
 import org.jetbrains.annotations.NotNull;
 
 import java.util.Map.Entry;
 import java.util.*;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
@@ -204,6 +206,13 @@ public Map<String, ColumnDefinition<?>> getColumnNameMap() {
                 .toMap(ColumnDefinition::getName, Function.identity(), Assert::neverInvoked, LinkedHashMap::new)));
     }
 
+    /**
+     * @return An unmodifiable set of column names
+     */
+    public Set<String> getColumnNameSet() {
+        return getColumnNameMap().keySet();
+    }
+
     /**
      * @return A list of {@link ColumnDefinition column definitions} for all
      *         {@link ColumnDefinition.ColumnType#Partitioning partitioning} columns in the same relative order as the
@@ -295,6 +304,36 @@ public String getColumnNamesAsString() {
         return getColumnStream().map(ColumnDefinition::getName).collect(Collectors.joining(","));
     }
 
+    /**
+     * Check this definition to ensure that {@code columnName} is present.
+     *
+     * @param columnName The column name to check
+     * @throws NoSuchColumnException If {@code columnName} is missing
+     */
+    public final void checkColumn(@NotNull String columnName) {
+        final Set<String> columnNames = getColumnNameSet();
+        if (!columnNames.contains(columnName)) {
+            throw new NoSuchColumnException(columnNames, columnName);
+        }
+    }
+
+    /**
+     * Check this definition to ensure that all {@code columns} are present.
+     *
+     * @param columns The column names to check
+     * @throws NoSuchColumnException If any {@code columns} were missing
+     */
+    public final void checkColumns(@NotNull Collection<String> columns) {
+        final Set<String> columnNames = getColumnNameSet();
+        final List<String> missingColumns = columns
+                .stream()
+                .filter(Predicate.not(columnNames::contains))
+                .collect(Collectors.toList());
+        if (!missingColumns.isEmpty()) {
+            throw new NoSuchColumnException(columnNames, missingColumns);
+        }
+    }
+
     /**
      * Tests mutual-compatibility of {@code this} and {@code other}. To be mutually compatible, they must have the same
      * number of columns, each matched up with {@link ColumnDefinition#isCompatible}. As such, this method has an

engine/table/src/main/java/io/deephaven/engine/table/impl/NoSuchColumnException.java → engine/api/src/main/java/io/deephaven/engine/table/impl/NoSuchColumnException.java

@@ -17,8 +17,9 @@ public class NoSuchColumnException extends IllegalArgumentException {
      * @param requestedColumns the request column names that were not found
      */
     public NoSuchColumnException(Collection<String> presentColumns, Collection<String> requestedColumns) {
-        super("Unknown column names [" + String.join(",", requestedColumns)
-                + "], available column names are [" + String.join(",", presentColumns) + "]");
+        super(String.format("Unknown column names [%s], available column names are [%s]",
+                String.join(",", requestedColumns),
+                String.join(",", presentColumns)));
     }
 
     /**

engine/benchmark/src/benchmark/java/io/deephaven/benchmark/engine/SortBenchmark.java

@@ -138,7 +138,8 @@ public void setupEnv(BenchmarkParams params) {
         mcsWithSortColumn = inputTable.newModifiedColumnSet(sortCol);
         MutableInt ci = new MutableInt();
         final String[] sortColumns = new String[inputTable.numColumns() - 1];
-        inputTable.getColumnSourceMap().keySet().forEach(columnName -> {
+
+        inputTable.getDefinition().getColumnNameSet().forEach(columnName -> {
             if (!columnName.equals(sortCol)) {
                 sortColumns[ci.intValue()] = columnName;
                 ci.increment();

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

@@ -1020,13 +1020,8 @@ private String formatKeyColumns(String... columns) {
     }
 
     @Override
-    public void checkAvailableColumns(@NotNull final Collection<String> columns) {
-        final Map<String, ? extends ColumnSource<?>> sourceMap = getColumnSourceMap();
-        final String[] missingColumns =
-                columns.stream().filter(col -> !sourceMap.containsKey(col)).toArray(String[]::new);
-        if (missingColumns.length > 0) {
-            throw new NoSuchColumnException(sourceMap.keySet(), Arrays.asList(missingColumns));
-        }
+    public final void checkAvailableColumns(@NotNull final Collection<String> columns) {
+        getDefinition().checkColumns(columns);
     }
 
     public void copySortableColumns(
@@ -1109,9 +1104,9 @@ void copySortableColumns(BaseTable<?> destination, SelectColumn[] selectCols) {
         }
 
         // Now go through the other columns in the table and add them if they were unchanged
-        final Map<String, ? extends ColumnSource<?>> sourceMap = destination.getColumnSourceMap();
+        final Set<String> destKeys = destination.getDefinition().getColumnNameSet();
         for (String col : currentSortableSet) {
-            if (sourceMap.containsKey(col)) {
+            if (destKeys.contains(col)) {
                 newSortableSet.add(col);
             }
         }

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

@@ -20,6 +20,7 @@
 import java.time.Instant;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import static io.deephaven.engine.table.impl.MatchPair.matchString;
@@ -41,8 +42,11 @@ class BucketingContext implements SafeCloseable {
 
     BucketingContext(final String listenerPrefix, final QueryTable leftTable, final QueryTable rightTable,
             MatchPair[] columnsToMatch, MatchPair[] columnsToAdd, JoinControl control) {
-        final List<String> conflicts = Arrays.stream(columnsToAdd).map(MatchPair::leftColumn)
-                .filter(cn -> leftTable.getColumnSourceMap().containsKey(cn)).collect(Collectors.toList());
+        final Set<String> leftKeys = leftTable.getDefinition().getColumnNameSet();
+        final List<String> conflicts = Arrays.stream(columnsToAdd)
+                .map(MatchPair::leftColumn)
+                .filter(leftKeys::contains)
+                .collect(Collectors.toList());
         if (!conflicts.isEmpty()) {
             throw new RuntimeException("Conflicting column names " + conflicts);
         }

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

@@ -195,7 +195,7 @@ private static QueryTable internalJoin(
                 jsm.startTrackingPrevValues();
                 final ModifiedColumnSet.Transformer leftTransformer = leftTable.newModifiedColumnSetTransformer(
                         resultTable,
-                        leftTable.getColumnSourceMap().keySet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY));
+                        leftTable.getDefinition().getColumnNameSet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY));
 
                 leftTable.addUpdateListener(new BaseTable.ListenerImpl(bucketingContext.listenerDescription,
                         leftTable, resultTable) {

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

@@ -1772,14 +1772,7 @@ public Table dropColumns(String... columnNames) {
             return memoizeResult(MemoizedOperationKey.dropColumns(columnNames), () -> QueryPerformanceRecorder
                     .withNugget("dropColumns(" + Arrays.toString(columnNames) + ")", sizeForInstrumentation(), () -> {
                         final Mutable<Table> result = new MutableObject<>();
-
-                        final Set<String> existingColumns = new HashSet<>(definition.getColumnNames());
-                        final Set<String> columnNamesToDrop = new HashSet<>(Arrays.asList(columnNames));
-                        if (!existingColumns.containsAll(columnNamesToDrop)) {
-                            columnNamesToDrop.removeAll(existingColumns);
-                            throw new RuntimeException("Unknown columns: " + columnNamesToDrop
-                                    + ", available columns = " + getColumnSourceMap().keySet());
-                        }
+                        definition.checkColumns(Arrays.asList(columnNames));
                         final Map<String, ColumnSource<?>> newColumns = new LinkedHashMap<>(columns);
                         for (String columnName : columnNames) {
                             newColumns.remove(columnName);
@@ -1800,7 +1793,7 @@ public Table dropColumns(String... columnNames) {
                             if (snapshotControl != null) {
                                 final ModifiedColumnSet.Transformer mcsTransformer =
                                         newModifiedColumnSetTransformer(resultTable,
-                                                resultTable.getColumnSourceMap().keySet()
+                                                resultTable.getDefinition().getColumnNameSet()
                                                         .toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY));
                                 final ListenerImpl listener = new ListenerImpl(
                                         "dropColumns(" + Arrays.deepToString(columnNames) + ')', this, resultTable) {
@@ -2400,7 +2393,7 @@ private Table snapshotIncrementalInternal(final Table base, final boolean doInit
 
         // Use the given columns (if specified); otherwise an empty array means all of my columns
         final String[] useStampColumns = stampColumns.length == 0
-                ? getColumnSourceMap().keySet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY)
+                ? definition.getColumnNameSet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY)
                 : stampColumns;
 
         final Map<String, ColumnSource<?>> triggerColumns = new LinkedHashMap<>();

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

@@ -84,15 +84,8 @@ public Table dropColumns(final String... columnNames) {
         if (columnNames == null || columnNames.length == 0) {
             return this;
         }
-
         final Set<String> columnNamesToDrop = new HashSet<>(Arrays.asList(columnNames));
-        final Set<String> existingColumns = new HashSet<>(definition.getColumnNames());
-        if (!existingColumns.containsAll(columnNamesToDrop)) {
-            columnNamesToDrop.removeAll(existingColumns);
-            throw new RuntimeException("Unknown columns: " + columnNamesToDrop.toString() + ", available columns = "
-                    + getColumnSourceMap().keySet());
-        }
-
+        definition.checkColumns(columnNamesToDrop);
         List<ColumnDefinition<?>> resultColumns = new ArrayList<>();
         for (ColumnDefinition<?> cDef : definition.getColumns()) {
             if (!columnNamesToDrop.contains(cDef.getName())) {

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

@@ -98,7 +98,7 @@ default boolean hasColumns(Collection<String> columnNames) {
         if (columnNames == null) {
             throw new IllegalArgumentException("columnNames cannot be null!");
         }
-        return getDefinition().getColumnNameMap().keySet().containsAll(columnNames);
+        return getDefinition().getColumnNameSet().containsAll(columnNames);
     }
 
     @Override

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

@@ -61,11 +61,14 @@ private TableUpdateValidator(final String description, final QueryTable tableToV
         this.description = description == null ? tableToValidate.getDescription() : description;
         this.tableToValidate = tableToValidate;
         this.validationMCS = tableToValidate.newModifiedColumnSet(
-                tableToValidate.getColumnSourceMap().keySet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY));
+                tableToValidate.getDefinition().getColumnStream().map(ColumnDefinition::getName)
+                        .toArray(String[]::new));
         Assert.neq(validationMCS, "validationMCS", ModifiedColumnSet.ALL, "ModifiedColumnSet.ALL");
         Assert.neq(validationMCS, "validationMCS", ModifiedColumnSet.EMPTY, "ModifiedColumnSet.EMPTY");
 
-        columnInfos = tableToValidate.getColumnSourceMap().keySet().stream()
+        columnInfos = tableToValidate.getDefinition()
+                .getColumnStream()
+                .map(ColumnDefinition::getName)
                 .map((name) -> new ColumnInfo(tableToValidate, name))
                 .toArray(ColumnInfo[]::new);
     }

engine/table/src/main/java/io/deephaven/engine/table/impl/by/ChunkedOperatorAggregationHelper.java

@@ -47,9 +47,9 @@
 
 import java.util.*;
 import java.util.function.LongFunction;
-import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.function.UnaryOperator;
+import java.util.stream.Collectors;
 
 import static io.deephaven.engine.table.impl.by.AggregationRowLookup.DEFAULT_UNKNOWN_ROW;
 import static io.deephaven.engine.table.impl.by.AggregationRowLookup.EMPTY_KEY;
@@ -86,34 +86,14 @@ public static QueryTable aggregation(
             final boolean preserveEmpty,
             @Nullable final Table initialKeys,
             @NotNull final Collection<? extends ColumnName> groupByColumns) {
-        final String[] keyNames = groupByColumns.stream().map(ColumnName::name).toArray(String[]::new);
-        if (!input.hasColumns(keyNames)) {
-            final Set<String> colNames = input.getColumnSourceMap().keySet();
-            final String[] missingColumns = Arrays.stream(keyNames)
-                    .filter(Predicate.not(colNames::contains))
-                    .toArray(String[]::new);;
-
-            throw new IllegalArgumentException("aggregation: not all group-by columns " + Arrays.toString(keyNames)
-                    + " are present in input table with columns "
-                    + Arrays.toString(input.getDefinition().getColumnNamesArray()) + ". Missing columns: "
-                    + Arrays.toString(missingColumns));
-        }
+        final List<String> keyNames = groupByColumns.stream().map(ColumnName::name).collect(Collectors.toList());
+        input.checkAvailableColumns(keyNames);
         if (initialKeys != null) {
-            if (keyNames.length == 0) {
+            if (keyNames.isEmpty()) {
                 throw new IllegalArgumentException(
                         "aggregation: initial groups must not be specified if no group-by columns are specified");
             }
-            if (!initialKeys.hasColumns(keyNames)) {
-                final Set<String> colNames = input.getColumnSourceMap().keySet();
-                final String[] missingColumns = Arrays.stream(keyNames)
-                        .filter(Predicate.not(colNames::contains))
-                        .toArray(String[]::new);;
-
-                throw new IllegalArgumentException("aggregation: not all group-by columns " + Arrays.toString(keyNames)
-                        + " are present in initial groups table with columns "
-                        + Arrays.toString(initialKeys.getDefinition().getColumnNamesArray()) + ". Missing columns: "
-                        + Arrays.toString(missingColumns));
-            }
+            initialKeys.getDefinition().checkColumns(keyNames);
             for (final String keyName : keyNames) {
                 final ColumnDefinition<?> inputDef = input.getDefinition().getColumn(keyName);
                 final ColumnDefinition<?> initialKeysDef = initialKeys.getDefinition().getColumn(keyName);
@@ -132,7 +112,7 @@ public static QueryTable aggregation(
                 "by(" + aggregationContextFactory + ", " + groupByColumns + ")", snapshotControl,
                 (usePrev, beforeClockValue) -> {
                     resultHolder.setValue(aggregation(control, snapshotControl, aggregationContextFactory,
-                            input, preserveEmpty, initialKeys, keyNames, usePrev));
+                            input, preserveEmpty, initialKeys, keyNames.toArray(String[]::new), usePrev));
                     return true;
                 });
         return resultHolder.getValue();

engine/table/src/main/java/io/deephaven/engine/table/impl/hierarchical/HierarchicalTableImpl.java

@@ -136,13 +136,8 @@ IFACE_TYPE noopResult() {
     }
 
     @Override
-    protected void checkAvailableColumns(@NotNull final Collection<String> columns) {
-        final Set<String> availableColumns = root.getDefinition().getColumnNameMap().keySet();
-        final List<String> missingColumns =
-                columns.stream().filter(column -> !availableColumns.contains(column)).collect(Collectors.toList());
-        if (!missingColumns.isEmpty()) {
-            throw new NoSuchColumnException(availableColumns, missingColumns);
-        }
+    protected final void checkAvailableColumns(@NotNull final Collection<String> columns) {
+        root.getDefinition().checkColumns(columns);
     }
 
     /**

engine/table/src/main/java/io/deephaven/engine/table/impl/remote/ConstructSnapshot.java

@@ -1403,9 +1403,8 @@ private static boolean serializeAllTable(
         }
 
         LongSizedDataStructure.intSize("construct snapshot", snapshot.rowsIncluded.size());
-
-        final Map<String, ? extends ColumnSource<?>> sourceMap = table.getColumnSourceMap();
-        final String[] columnSources = sourceMap.keySet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY);
+        final String[] columnSources =
+                table.getDefinition().getColumnNameSet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY);
 
         snapshot.dataColumns = new Object[columnSources.length];
         try (final SharedContext sharedContext =
@@ -1480,8 +1479,8 @@ private static boolean serializeAllTable(
             snapshot.rowsIncluded = snapshot.rowsAdded.copy();
         }
 
-        final Map<String, ? extends ColumnSource<?>> sourceMap = table.getColumnSourceMap();
-        final String[] columnSources = sourceMap.keySet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY);
+        final String[] columnSources =
+                table.getDefinition().getColumnNameSet().toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY);
 
         try (final SharedContext sharedContext =
                 (columnSources.length > 1) ? SharedContext.makeSharedContext() : null) {

engine/table/src/main/java/io/deephaven/engine/table/impl/snapshot/SnapshotInternalListener.java

@@ -33,7 +33,7 @@ public SnapshotInternalListener(QueryTable triggerTable,
             Map<String, SingleValueColumnSource<?>> resultTriggerColumns,
             Map<String, WritableColumnSource<?>> resultBaseColumns,
             TrackingWritableRowSet resultRowSet) {
-        super("snapshot " + result.getColumnSourceMap().keySet(), triggerTable, result);
+        super("snapshot " + result.getDefinition().getColumnNameSet(), triggerTable, result);
         this.triggerTable = triggerTable;
         this.result = result;
         this.lazySnapshot = lazySnapshot;

engine/table/src/main/java/io/deephaven/engine/table/impl/updateby/UpdateBy.java

@@ -47,6 +47,7 @@
 import java.util.concurrent.atomic.AtomicIntegerArray;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
 /**
@@ -1198,7 +1199,10 @@ public static Table updateBy(@NotNull final QueryTable source,
 
         final MutableObject<String> timestampColumnName = new MutableObject<>(null);
         // create an initial set of all source columns
-        final Set<String> preservedColumnSet = new LinkedHashSet<>(source.getColumnSourceMap().keySet());
+        final LinkedHashSet<String> preservedColumnSet = source.getDefinition()
+                .getColumnStream()
+                .map(ColumnDefinition::getName)
+                .collect(Collectors.toCollection(LinkedHashSet::new));
 
         final Set<String> problems = new LinkedHashSet<>();
         final Map<String, ColumnSource<?>> opResultSources = new LinkedHashMap<>();

engine/table/src/main/java/io/deephaven/engine/table/impl/util/DynamicTableWriter.java

@@ -778,7 +778,7 @@ private DynamicTableRow() {
         public PermissiveRowSetter<?> getSetter(final String name) {
             final PermissiveRowSetter<?> rowSetter = columnToSetter.get(name);
             if (rowSetter == null) {
-                if (table.getColumnSourceMap().containsKey(name)) {
+                if (table.hasColumns(name)) {
                     throw new RuntimeException("Column has a constant value, can not get setter " + name);
                 } else {
                     throw new RuntimeException("Unknown column name " + name);