Add TableDefinition column name helpers (#4813)

Adds some helper methods for on `TableDefinition#checkHasColumn`, `TableDefinition#checkHasColumns`, and `TableDefinition#getColumnNameSet`.

Additionally, fixes up call sites that were (ab)using `Table#getColumnSourceMap` to simply get the keySet. This invokes a potentially extraneous Table#coalesce which can be avoided in these cases.

In support of common scaffolding so #4771 won't need to call `Table#getColumnSource` for validation purposes.

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,6 +9,7 @@
 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;
@@ -204,6 +205,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 +303,26 @@ 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 checkHasColumn(@NotNull String columnName) {
+        NoSuchColumnException.throwIf(getColumnNameSet(), 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 checkHasColumns(@NotNull Collection<String> columns) {
+        NoSuchColumnException.throwIf(getColumnNameSet(), columns);
+    }
+
     /**
      * 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/api/src/main/java/io/deephaven/engine/table/impl/NoSuchColumnException.java

@@ -0,0 +1,121 @@
+/**
+ * Copyright (c) 2016-2022 Deephaven Data Labs and Patent Pending
+ */
+package io.deephaven.engine.table.impl;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+/**
+ * Exception thrown when a column is not found.
+ */
+public class NoSuchColumnException extends IllegalArgumentException {
+
+    public static final String DELIMITER = ", ";
+
+    public static final String DEFAULT_FORMAT_STR = "Unknown column names [%s], available column names are [%s]";
+
+    public enum Type {
+        MISSING, AVAILABLE, REQUESTED
+    }
+
+    /**
+     * Equivalent to {@code throwIf(available, Collections.singleton(requested))}.
+     *
+     * @param available the available columns
+     * @param requested the requested columns
+     * @see #throwIf(Set, Collection)
+     */
+    public static void throwIf(Set<String> available, String requested) {
+        throwIf(available, Collections.singleton(requested));
+    }
+
+    /**
+     * Equivalent to {@code throwIf(available, requested, DEFAULT_FORMAT_STR, Type.MISSING, Type.AVAILABLE)} where
+     * {@code DEFAULT_FORMAT_STR} is {@value DEFAULT_FORMAT_STR}.
+     *
+     * @param available the available columns
+     * @param requested the requested columns
+     * @see #throwIf(Set, Collection, String, Type...)
+     */
+    public static void throwIf(Set<String> available, Collection<String> requested) {
+        throwIf(available, requested, DEFAULT_FORMAT_STR, Type.MISSING, Type.AVAILABLE);
+    }
+
+    /**
+     * Throws a {@link NoSuchColumnException} if any name from {@code requested} is not in {@code available}. The
+     * message will be constructed by {@link String#join(CharSequence, Iterable) joining} the respective collection with
+     * {@value DELIMITER} and presenting them to {@link String#format(String, Object...) format} in {@code types} order.
+     *
+     * @param available the available columns
+     * @param requested the requested columns
+     * @param formatStr the format string
+     * @param types the collection types order for formatting
+     */
+    public static void throwIf(Set<String> available, Collection<String> requested, String formatStr, Type... types) {
+        final List<String> missing = requested
+                .stream()
+                .filter(Predicate.not(available::contains))
+                .collect(Collectors.toList());
+        if (!missing.isEmpty()) {
+            final Object[] formatArgs = new Object[types.length];
+            for (int i = 0; i < types.length; ++i) {
+                switch (types[i]) {
+                    case MISSING:
+                        formatArgs[i] = String.join(DELIMITER, missing);
+                        break;
+                    case AVAILABLE:
+                        formatArgs[i] = String.join(DELIMITER, available);
+                        break;
+                    case REQUESTED:
+                        formatArgs[i] = String.join(DELIMITER, requested);
+                        break;
+                    default:
+                        throw new IllegalStateException("Unexpected case " + types[i]);
+                }
+            }
+            throw new NoSuchColumnException(String.format(formatStr, formatArgs));
+        }
+    }
+
+    /**
+     * Thrown when an operation can not find a required column(s).
+     *
+     * <p>
+     * Callers may prefer to use {@link #throwIf(Set, Collection, String, Type...)} when applicable.
+     *
+     * @param message the message
+     */
+    public NoSuchColumnException(String message) {
+        super(message);
+    }
+
+    /**
+     * Thrown when an operation can not find a required column(s).
+     *
+     * <p>
+     * Callers may prefer to use {@link #throwIf(Set, Collection)} when applicable.
+     *
+     * @param presentColumns the column names present in the table
+     * @param missingColumns the request column names that were not found
+     */
+    public NoSuchColumnException(Collection<String> presentColumns, Collection<String> missingColumns) {
+        this(String.format(DEFAULT_FORMAT_STR,
+                String.join(DELIMITER, missingColumns),
+                String.join(DELIMITER, presentColumns)));
+    }
+
+    /**
+     * Thrown when an operation can not find a required column.
+     *
+     * @param presentColumns the column names present in the table
+     * @param missingColumn the request column name that was not found
+     */
+    public NoSuchColumnException(Collection<String> presentColumns, String missingColumn) {
+        this(presentColumns, Collections.singleton(missingColumn));
+    }
+}

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().checkHasColumns(columns);
     }
 
     public void copySortableColumns(
@@ -1063,7 +1058,7 @@ void copySortableColumns(BaseTable<?> destination, MatchPair[] renamedColumns) {
         // Process the original set of sortable columns, adding them to the new set if one of the below
         // 1) The column exists in the new table and was not renamed in any way but the Identity (C1 = C1)
         // 2) The column does not exist in the new table, but was renamed to another (C2 = C1)
-        final Set<String> resultColumnNames = destination.getDefinition().getColumnNameMap().keySet();
+        final Set<String> resultColumnNames = destination.getDefinition().getColumnNameSet();
         for (final String columnName : currentSortableColumns) {
             // Only add it to the set of sortable columns if it hasn't changed in an unknown way
             final String maybeRenamedColumn = columnMapping.get(columnName);
@@ -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().getColumnNamesArray());
 
                 leftTable.addUpdateListener(new BaseTable.ListenerImpl(bucketingContext.listenerDescription,
                         leftTable, resultTable) {

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

@@ -355,7 +355,7 @@ public long size() {
     public <T> ColumnSource<T> getColumnSource(String sourceName) {
         final ColumnSource<?> columnSource = columns.get(sourceName);
         if (columnSource == null) {
-            throw new NoSuchColumnException(columns.keySet(), Collections.singletonList(sourceName));
+            throw new NoSuchColumnException(columns.keySet(), sourceName);
         }
         // noinspection unchecked
         return (ColumnSource<T>) columnSource;
@@ -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.checkHasColumns(Arrays.asList(columnNames));
                         final Map<String, ColumnSource<?>> newColumns = new LinkedHashMap<>(columns);
                         for (String columnName : columnNames) {
                             newColumns.remove(columnName);
@@ -1794,14 +1787,13 @@ public Table dropColumns(String... columnNames) {
 
                             copyAttributes(resultTable, CopyAttributeOperation.DropColumns);
                             copySortableColumns(resultTable,
-                                    resultTable.getDefinition().getColumnNameMap()::containsKey);
+                                    resultTable.getDefinition().getColumnNameSet()::contains);
                             maybeCopyColumnDescriptions(resultTable);
 
                             if (snapshotControl != null) {
                                 final ModifiedColumnSet.Transformer mcsTransformer =
                                         newModifiedColumnSetTransformer(resultTable,
-                                                resultTable.getColumnSourceMap().keySet()
-                                                        .toArray(CollectionUtil.ZERO_LENGTH_STRING_ARRAY));
+                                                resultTable.getDefinition().getColumnNamesArray());
                                 final ListenerImpl listener = new ListenerImpl(
                                         "dropColumns(" + Arrays.deepToString(columnNames) + ')', this, resultTable) {
                                     @Override
@@ -2400,7 +2392,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.getColumnNamesArray()
                 : 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.checkHasColumns(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,13 @@ 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().getColumnNamesArray());
         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);
     }