Add topic, partition, and timestamp column kafka publishing support

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

@@ -504,6 +504,28 @@ public void describeDifferences(@NotNull List<String> differences, @NotNull fina
         }
     }
 
+    /**
+     * Checks if {@link #getDataType() dataType} can be cast to {@code destDataType}. If not, this throws a
+     * {@link ClassCastException}.
+     *
+     * @param destDataType the destination data type
+     */
+    public final void checkCastTo(Class<?> destDataType) {
+        TypeHelper.checkCastTo("[" + name + "]", dataType, destDataType);
+    }
+
+    /**
+     * Checks if {@link #getDataType() dataType} can be cast to {@code destDataType} and checks that
+     * {@link #getComponentType() componentType} can be cast to {@code destComponentType} (both component types must be
+     * present and cast-able, or both must be {@code null}). If not, this throws a {@link ClassCastException}.
+     *
+     * @param destDataType the destination data type
+     * @param destComponentType the destination component type, may be {@code null}
+     */
+    public final void checkCastTo(Class<?> destDataType, @Nullable Class<?> destComponentType) {
+        TypeHelper.checkCastTo("[" + name + "]", dataType, componentType, destDataType, destComponentType);
+    }
+
     public boolean equals(final Object other) {
         if (this == other) {
             return true;

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

@@ -150,7 +150,7 @@ default Object exportElement(T tuple, int elementIndex) {
      * {@code String} data:
      *
      * <pre>
-     *     ColumnSource&lt;String&gt; colSource = table.getColumnSource("MyString").getParameterized(String.class)
+     *     ColumnSource&lt;String&gt; colSource = table.getColumnSource("MyString").cast(String.class)
      * </pre>
      * <p>
      * Due to the nature of type erasure, the JVM will still insert an additional cast to {@code TYPE} when elements are
@@ -163,11 +163,7 @@ default Object exportElement(T tuple, int elementIndex) {
     @FinalDefault
     default <TYPE> ColumnSource<TYPE> cast(Class<? extends TYPE> clazz) {
         Require.neqNull(clazz, "clazz");
-        final Class<?> columnSourceType = getType();
-        if (!clazz.isAssignableFrom(columnSourceType)) {
-            throw new ClassCastException(String.format("Cannot convert column source for type %s to type %s",
-                    columnSourceType.getName(), clazz.getName()));
-        }
+        TypeHelper.checkCastTo("ColumnSource", getType(), clazz);
         // noinspection unchecked
         return (ColumnSource<TYPE>) this;
     }
@@ -184,7 +180,7 @@ default <TYPE> ColumnSource<TYPE> cast(Class<? extends TYPE> clazz) {
      * {@code String} data:
      *
      * <pre>
-     *     ColumnSource&lt;String&gt; colSource = table.getColumnSource("MyString", null).getParameterized(String.class)
+     *     ColumnSource&lt;String&gt; colSource = table.getColumnSource("MyString").cast(String.class, null)
      * </pre>
      * <p>
      * Due to the nature of type erasure, the JVM will still insert an additional cast to {@code TYPE} when elements are
@@ -197,19 +193,10 @@ default <TYPE> ColumnSource<TYPE> cast(Class<? extends TYPE> clazz) {
      */
     @FinalDefault
     default <TYPE> ColumnSource<TYPE> cast(Class<? extends TYPE> clazz, @Nullable Class<?> componentType) {
-        final ColumnSource<TYPE> casted = cast(clazz);
-        final Class<?> columnSourceComponentType = getComponentType();
-        if ((componentType == null && columnSourceComponentType == null) || (componentType != null
-                && columnSourceComponentType != null && componentType.isAssignableFrom(columnSourceComponentType))) {
-            return casted;
-        }
-        final Class<?> columnSourceType = getType();
-        throw new ClassCastException(String.format(
-                "Cannot convert column source componentType for type %s to %s (for %s / %s)",
-                columnSourceComponentType == null ? null : columnSourceComponentType.getName(),
-                componentType == null ? null : componentType.getName(),
-                columnSourceType.getName(),
-                clazz.getName()));
+        Require.neqNull(clazz, "clazz");
+        TypeHelper.checkCastTo("ColumnSource", getType(), getComponentType(), clazz, componentType);
+        // noinspection unchecked
+        return (ColumnSource<TYPE>) this;
     }
 
     /**

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

@@ -228,33 +228,51 @@ public interface Table extends
      * caller expects. This differs from {@link #getColumnSource(String, Class)} which uses the provided {@link Class}
      * object to verify that the data type is a subclass of the expected class.
      *
+     * <p>
+     * The success of this call is equivalent to {@code getDefinition().checkColumn(sourceName)}, which is the preferred
+     * way to check for compatibility in scenarios where the caller does not want the implementation to potentially
+     * invoke {@link #coalesce()}.
+     *
      * @param sourceName The name of the column
      * @param <T> The target type, as a type parameter. Inferred from context.
      * @return The column source for {@code sourceName}, parameterized by {@code T}
+     * @see TableDefinition#checkHasColumn(String)
      */
     <T> ColumnSource<T> getColumnSource(String sourceName);
 
     /**
      * Retrieves a {@code ColumnSource} and {@link ColumnSource#cast(Class) casts} it to the target class {@code clazz}.
      *
+     * <p>
+     * The success of this call is equivalent to {@code getDefinition().checkColumn(sourceName, clazz)}, which is the
+     * preferred way to check for compatibility in scenarios where the caller does not want to the implementation to
+     * potentially invoke {@link #coalesce()}.
+     *
      * @param sourceName The name of the column
      * @param clazz The target type
      * @param <T> The target type, as a type parameter. Intended to be inferred from {@code clazz}.
      * @return The column source for {@code sourceName}, parameterized by {@code T}
      * @see ColumnSource#cast(Class)
+     * @see TableDefinition#checkHasColumn(String, Class)
      */
     <T> ColumnSource<T> getColumnSource(String sourceName, Class<? extends T> clazz);
 
     /**
      * Retrieves a {@code ColumnSource} and {@link ColumnSource#cast(Class, Class)} casts} it to the target class
      * {@code clazz} and {@code componentType}.
      *
+     * <p>
+     * The success of this call is equivalent to {@code getDefinition().checkColumn(sourceName, clazz, componentType)},
+     * which is the preferred way to check for compatibility in scenarios where the caller does not want the
+     * implementation to potentially invoke {@link #coalesce()}.
+     *
      * @param sourceName The name of the column
      * @param clazz The target type
      * @param componentType The target component type, may be null
      * @param <T> The target type, as a type parameter. Intended to be inferred from {@code clazz}.
      * @return The column source for {@code sourceName}, parameterized by {@code T}
      * @see ColumnSource#cast(Class, Class)
+     * @see TableDefinition#checkHasColumn(String, Class, Class)
      */
     <T> ColumnSource<T> getColumnSource(String sourceName, Class<? extends T> clazz, @Nullable Class<?> componentType);
 

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

@@ -313,6 +313,40 @@ public final void checkHasColumn(@NotNull String columnName) {
         NoSuchColumnException.throwIf(getColumnNameSet(), columnName);
     }
 
+    /**
+     * Checks if {@code columnName} exists and if supports {@link ColumnDefinition#checkCastTo(Class)} with
+     * {@code clazz}. Otherwise, throws a {@link NoSuchColumnException} or a {@link ClassCastException}.
+     *
+     * @param columnName the column name
+     * @param clazz the data type
+     * @see ColumnDefinition#checkCastTo(Class)
+     */
+    public final void checkHasColumn(@NotNull String columnName, @NotNull Class<?> clazz) {
+        final ColumnDefinition<?> cd = getColumn(columnName);
+        if (cd == null) {
+            throw new NoSuchColumnException(getColumnNameSet(), columnName);
+        }
+        cd.checkCastTo(clazz);
+    }
+
+    /**
+     * Checks if {@code columnName} exists and if supports {@link ColumnDefinition#checkCastTo(Class, Class)} with
+     * {@code clazz} and {@code componentType}. Otherwise, throws a {@link NoSuchColumnException} or a
+     * {@link ClassCastException}.
+     *
+     * @param columnName the column name
+     * @param clazz the data type
+     * @param componentType the component type
+     * @see ColumnDefinition#checkCastTo(Class, Class)
+     */
+    public final void checkHasColumn(@NotNull String columnName, @NotNull Class<?> clazz, Class<?> componentType) {
+        final ColumnDefinition<?> cd = getColumn(columnName);
+        if (cd == null) {
+            throw new NoSuchColumnException(getColumnNameSet(), columnName);
+        }
+        cd.checkCastTo(clazz, componentType);
+    }
+
     /**
      * Check this definition to ensure that all {@code columns} are present.
      *
@@ -323,6 +357,7 @@ 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/TypeHelper.java

@@ -0,0 +1,36 @@
+/**
+ * Copyright (c) 2016-2023 Deephaven Data Labs and Patent Pending
+ */
+package io.deephaven.engine.table;
+
+class TypeHelper {
+
+    // Could be good to move this to io.deephaven.qst.type.Type layer
+
+    public static void checkCastTo(String context, Class<?> srcType, Class<?> destType) {
+        if (!destType.isAssignableFrom(srcType)) {
+            throw new ClassCastException(String.format("Cannot convert %s of type %s to type %s",
+                    context, srcType.getName(), destType.getName()));
+        }
+    }
+
+    public static void checkCastTo(
+            String prefix,
+            Class<?> srcType,
+            Class<?> srcComponentType,
+            Class<?> destType,
+            Class<?> destComponentType) {
+        checkCastTo(prefix, srcType, destType);
+        if ((srcComponentType == null && destComponentType == null) || (srcComponentType != null
+                && destComponentType != null && destComponentType.isAssignableFrom(srcComponentType))) {
+            return;
+        }
+        throw new ClassCastException(String.format(
+                "Cannot convert %s componentType of type %s to %s (for %s / %s)",
+                prefix,
+                srcComponentType == null ? null : srcComponentType.getName(),
+                destComponentType == null ? null : destComponentType.getName(),
+                srcType.getName(),
+                destType.getName()));
+    }
+}

extensions/kafka/src/main/java/io/deephaven/kafka/KafkaPublishOptions.java

@@ -39,9 +39,10 @@ public static Builder builder() {
     public abstract Table table();
 
     /**
-     * The default kafka topic to publish to.
+     * The default kafka topic to publish to. When {@code null}, {@link #topicColumn()} must be set.
      *
      * @return the default kafka topic
+     * @see #topicColumn()
      */
     @Nullable
     public abstract String topic();
@@ -50,6 +51,7 @@ public static Builder builder() {
      * The default kafka partition to publish to.
      *
      * @return the default kafka partition
+     * @see #partitionColumn()
      */
     public abstract OptionalInt partition();
 
@@ -108,25 +110,33 @@ public boolean publishInitial() {
 
     /**
      * The topic column. When set, uses the the given {@link CharSequence}-compatible column from {@link #table()} as
-     * the first source for setting the kafka record topic.
+     * the first source for setting the kafka record topic. When not present, or the column value is null,
+     * {@link #topic()} will be used.
      *
      * @return the topic column name
      */
     public abstract Optional<ColumnName> topicColumn();
 
     /**
      * The partition column. When set, uses the the given {@code int} column from {@link #table()} as the first source
-     * for setting the kafka record partition.
+     * for setting the kafka record partition. When not present, or the column value is null, {@link #partition()} will
+     * be used if present. If a valid partition number is specified, that partition will be used when sending the
+     * record. If no partition is specified but a key is present, a partition will be chosen using a hash of the key. If
+     * neither key nor partition is present, a partition will be assigned in a round-robin fashion.
      *
      * @return the partition column name
      */
     public abstract Optional<ColumnName> partitionColumn();
 
     /**
      * The timestamp column. When set, uses the the given {@link Instant} column from {@link #table()} as the first
-     * source for setting the kafka record timestamp.
+     * source for setting the kafka record timestamp. When not present, or the column value is null, the producer will
+     * stamp the record with its current time. The timestamp eventually used by Kafka depends on the timestamp type
+     * configured for the topic. If the topic is configured to use CreateTime, the timestamp in the producer record will
+     * be used by the broker. If the topic is configured to use LogAppendTime, the timestamp in the producer record will
+     * be overwritten by the broker with the broker local time when it appends the message to its log.
      *
-     * @return the partition column name
+     * @return the timestamp column name
      */
     public abstract Optional<ColumnName> timestampColumn();
 
@@ -161,21 +171,21 @@ final void checkTopic() {
     @Check
     final void checkTopicColumn() {
         if (topicColumn().isPresent()) {
-            table().getColumnSource(topicColumn().get().name(), CharSequence.class);
+            table().getDefinition().checkHasColumn(topicColumn().get().name(), CharSequence.class);
         }
     }
 
     @Check
     final void checkPartitionColumn() {
         if (partitionColumn().isPresent()) {
-            table().getColumnSource(partitionColumn().get().name(), int.class);
+            table().getDefinition().checkHasColumn(partitionColumn().get().name(), int.class);
         }
     }
 
     @Check
     final void checkTimestampColumn() {
         if (timestampColumn().isPresent()) {
-            table().getColumnSource(timestampColumn().get().name(), Instant.class);
+            table().getDefinition().checkHasColumn(timestampColumn().get().name(), Instant.class);
         }
     }
 

extensions/kafka/src/main/java/io/deephaven/kafka/publish/PublishToKafka.java

@@ -133,7 +133,7 @@ public PublishToKafka(
      */
     public PublishToKafka(
             final Properties props,
-            final Table table,
+            Table table,
             final String defaultTopic,
             final Integer defaultPartition,
             final String[] keyColumns,
@@ -146,7 +146,7 @@ public PublishToKafka(
             final ColumnName partitionColumn,
             final ColumnName timestampColumn,
             final boolean publishInitial) {
-        this.table = table;
+        this.table = (table = table.coalesce());
         this.producer = new KafkaProducer<>(
                 props,
                 Objects.requireNonNull(kafkaKeySerializer),
@@ -201,7 +201,7 @@ private Integer partition(IntChunk<?> partitionChunk, int index) {
             return defaultPartition;
         }
         final int partition = partitionChunk.get(index);
-        return partition == QueryConstants.NULL_INT ? defaultPartition : (Integer) partition;
+        return partition == QueryConstants.NULL_INT ? defaultPartition : Integer.valueOf(partition);
     }
 
     public static Long timestampMillis(LongChunk<?> nanosChunk, int index) {