fix: DH-18803: Add qst Type find support for componentType (#6667)

This adds the needed, centralized logic for adapting to qst Type when componentType is not null. This allows downstream callers, such as the Sql adapter layer, to correctly construct the appropriate Types.
deephaven · Mar 4, 2025 · 68d7bdc · 68d7bdc
1 parent fb6c6f4
commit 68d7bdc
Show file tree

Hide file tree

Showing 11 changed files with 294 additions and 22 deletions.
diff --git a/engine/sql/src/main/java/io/deephaven/engine/sql/Sql.java b/engine/sql/src/main/java/io/deephaven/engine/sql/Sql.java
@@ -109,7 +109,7 @@ private static TableHeader adapt(TableDefinition tableDef) {
     }
 
     private static ColumnHeader<?> adapt(ColumnDefinition<?> columnDef) {
-        return ColumnHeader.of(columnDef.getName(), Type.find(columnDef.getDataType()));
+        return ColumnHeader.of(columnDef.getName(), Type.find(columnDef.getDataType(), columnDef.getComponentType()));
     }
 
     private enum ToGraphvizDot implements ObjFormatter<TableSpec> {

diff --git a/...ne/table/src/main/java/io/deephaven/engine/table/impl/locations/util/PartitionParser.java b/...ne/table/src/main/java/io/deephaven/engine/table/impl/locations/util/PartitionParser.java
@@ -195,9 +195,11 @@ public Comparable<?> parse(@NotNull final String stringValue) {
     @Nullable
     public static PartitionParser lookup(@NotNull final Class<?> dataType, @Nullable final Class<?> componentType) {
         if (componentType != null) {
+            // This is a short-circuit since we know that Resolver does not support ArrayType.
             return null;
         }
-        return lookup(Type.find(dataType));
+        // noinspection ConstantValue
+        return lookup(Type.find(dataType, componentType));
     }
 
     /**
@@ -289,6 +291,8 @@ public PartitionParser visit(@NotNull final InstantType instantType) {
 
         @Override
         public PartitionParser visit(@NotNull final ArrayType<?, ?> arrayType) {
+            // If the partition parser ever supports ArrayTypes, make sure the short-circuit in
+            // PartitionParser.lookup(java.lang.Class<?>, java.lang.Class<?>) is removed.
             return null;
         }
 

diff --git a/engine/table/src/test/java/io/deephaven/qst/type/GenericVectorTest.java b/engine/table/src/test/java/io/deephaven/qst/type/GenericVectorTest.java
@@ -0,0 +1,28 @@
+//
+// Copyright (c) 2016-2025 Deephaven Data Labs and Patent Pending
+//
+package io.deephaven.qst.type;
+
+import io.deephaven.vector.ObjectVector;
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class GenericVectorTest {
+
+    @Test
+    public void stringType() throws ClassNotFoundException {
+        testConstruction(GenericVectorType.of(ObjectVector.class, Type.stringType()));
+    }
+
+    @Test
+    public void instantType() throws ClassNotFoundException {
+        testConstruction(GenericVectorType.of(ObjectVector.class, Type.instantType()));
+    }
+
+    private static <ComponentType> void testConstruction(GenericVectorType<?, ComponentType> vectorType)
+            throws ClassNotFoundException {
+        assertThat(Type.find(vectorType.clazz(), vectorType.componentType().clazz())).isEqualTo(vectorType);
+        assertThat(GenericVectorType.of(vectorType.clazz(), vectorType.componentType())).isEqualTo(vectorType);
+    }
+}
diff --git a/engine/table/src/test/java/io/deephaven/qst/type/PrimitiveVectorTest.java b/engine/table/src/test/java/io/deephaven/qst/type/PrimitiveVectorTest.java
@@ -3,13 +3,20 @@
 //
 package io.deephaven.qst.type;
 
-import static org.assertj.core.api.Assertions.assertThat;
-
-import io.deephaven.vector.*;
+import io.deephaven.vector.ByteVector;
+import io.deephaven.vector.CharVector;
+import io.deephaven.vector.DoubleVector;
+import io.deephaven.vector.FloatVector;
+import io.deephaven.vector.IntVector;
+import io.deephaven.vector.LongVector;
+import io.deephaven.vector.ShortVector;
+import org.junit.Test;
 
 import java.lang.reflect.InvocationTargetException;
+import java.util.stream.Collectors;
 
-import org.junit.Test;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.failBecauseExceptionWasNotThrown;
 
 public class PrimitiveVectorTest {
 
@@ -25,4 +32,77 @@ public void types()
                 FloatVector.type(),
                 DoubleVector.type());
     }
+
+    @Test
+    public void byteVector() throws ClassNotFoundException {
+        testConstruction(ByteVector.type());
+    }
+
+    @Test
+    public void charVector() throws ClassNotFoundException {
+        testConstruction(CharVector.type());
+    }
+
+    @Test
+    public void shortVector() throws ClassNotFoundException {
+        testConstruction(ShortVector.type());
+    }
+
+    @Test
+    public void intVector() throws ClassNotFoundException {
+        testConstruction(IntVector.type());
+    }
+
+    @Test
+    public void longVector() throws ClassNotFoundException {
+        testConstruction(LongVector.type());
+    }
+
+    @Test
+    public void floatVector() throws ClassNotFoundException {
+        testConstruction(FloatVector.type());
+    }
+
+    @Test
+    public void doubleVector() throws ClassNotFoundException {
+        testConstruction(DoubleVector.type());
+    }
+
+    private static <T, ComponentType> void testConstruction(PrimitiveVectorType<T, ComponentType> vectorType)
+            throws ClassNotFoundException {
+        assertThat(Type.find(vectorType.clazz())).isEqualTo(vectorType);
+        assertThat(Type.find(vectorType.clazz(), vectorType.componentType().clazz())).isEqualTo(vectorType);
+        assertThat(PrimitiveVectorType.of(vectorType.clazz(), vectorType.componentType())).isEqualTo(vectorType);
+        // fail if component type is bad
+        for (PrimitiveType<?> badComponent : PrimitiveType.instances().collect(Collectors.toList())) {
+            if (badComponent.equals(vectorType.componentType())) {
+                continue;
+            }
+            fail(vectorType.clazz(), badComponent);
+        }
+        // fail if data type is bad
+        fail(Object.class, vectorType.componentType());
+        for (PrimitiveVectorType<?, ?> primitiveVectorType : TypeHelper.primitiveVectorTypes()
+                .collect(Collectors.toList())) {
+            if (primitiveVectorType.equals(vectorType)) {
+                continue;
+            }
+            fail(primitiveVectorType.clazz(), vectorType.componentType());
+        }
+    }
+
+    public static void fail(Class<?> clazz, PrimitiveType<?> ct) {
+        try {
+            Type.find(clazz, ct.clazz());
+            failBecauseExceptionWasNotThrown(IllegalArgumentException.class);
+        } catch (IllegalArgumentException e) {
+            assertThat(e).hasMessageContaining("Invalid PrimitiveVectorType");
+        }
+        try {
+            PrimitiveVectorType.of(clazz, ct);
+            failBecauseExceptionWasNotThrown(IllegalArgumentException.class);
+        } catch (IllegalArgumentException e) {
+            assertThat(e).hasMessageContaining("Invalid PrimitiveVectorType");
+        }
+    }
 }
diff --git a/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTest.java b/extensions/flight-sql/src/test/java/io/deephaven/server/flightsql/FlightSqlTest.java
@@ -17,10 +17,20 @@
 import io.deephaven.engine.table.TableDefinition;
 import io.deephaven.engine.util.TableTools;
 import io.deephaven.proto.backplane.grpc.WrappedAuthenticationRequest;
+import io.deephaven.qst.type.GenericVectorType;
+import io.deephaven.qst.type.Type;
 import io.deephaven.server.auth.AuthorizationProvider;
 import io.deephaven.server.config.ServerConfig;
 import io.deephaven.server.runner.DeephavenApiServerTestBase;
 import io.deephaven.server.runner.DeephavenApiServerTestBase.TestComponent.Builder;
+import io.deephaven.vector.ByteVector;
+import io.deephaven.vector.CharVector;
+import io.deephaven.vector.DoubleVector;
+import io.deephaven.vector.FloatVector;
+import io.deephaven.vector.IntVector;
+import io.deephaven.vector.LongVector;
+import io.deephaven.vector.ObjectVector;
+import io.deephaven.vector.ShortVector;
 import io.grpc.ManagedChannel;
 import org.apache.arrow.flight.Action;
 import org.apache.arrow.flight.ActionType;
@@ -588,6 +598,31 @@ public void badSqlQuery() {
         queryError("this is not SQL", FlightStatusCode.INVALID_ARGUMENT, "Flight SQL: query can't be parsed");
     }
 
+    @Test
+    public void testDh18803() throws Exception {
+        // https://deephaven.atlassian.net/browse/DH-18803: Sql fails to adapt Vector types
+        final TableDefinition td = TableDefinition.of(
+                ColumnDefinition.ofVector("ByteVector", ByteVector.class),
+                ColumnDefinition.ofVector("CharVector", CharVector.class),
+                ColumnDefinition.ofVector("ShortVector", ShortVector.class),
+                ColumnDefinition.ofVector("IntVector", IntVector.class),
+                ColumnDefinition.ofVector("LongVector", LongVector.class),
+                ColumnDefinition.ofVector("FloatVector", FloatVector.class),
+                ColumnDefinition.ofVector("DoubleVector", DoubleVector.class),
+                ColumnDefinition.of("StringVector", GenericVectorType.of(ObjectVector.class, Type.stringType())));
+        final Table emptyTable = TableTools.newTable(td);
+        ExecutionContext.getContext().getQueryScope().putParam("MyTable", emptyTable);
+        {
+            final SchemaResult schema = flightSqlClient.getExecuteSchema("SELECT * FROM MyTable");
+            assertThat(schema.getSchema().getFields()).hasSize(8);
+        }
+        {
+            final FlightInfo info = flightSqlClient.execute("SELECT * FROM MyTable");
+            assertThat(info.getSchema().getFields()).hasSize(8);
+            consume(info, 0, 0, false);
+        }
+    }
+
     @Test
     public void executeSubstrait() {
         getSchemaUnimplemented(() -> flightSqlClient.getExecuteSubstraitSchema(fakePlan()),

diff --git a/extensions/kafka/src/main/java/io/deephaven/kafka/SimpleImpl.java b/extensions/kafka/src/main/java/io/deephaven/kafka/SimpleImpl.java
@@ -233,13 +233,14 @@ public Optional<SchemaProvider> getSchemaProvider() {
 
         @Override
         Serializer<?> getSerializer(SchemaRegistryClient schemaRegistryClient, TableDefinition definition) {
-            final Class<?> dataType = definition.getColumn(columnName).getDataType();
-            final Serializer<?> serializer = serializer(Type.find(dataType)).orElse(null);
+            final ColumnDefinition<?> cd = definition.getColumn(columnName);
+            final Type<?> type = Type.find(cd.getDataType(), cd.getComponentType());
+            final Serializer<?> serializer = serializer(type).orElse(null);
             if (serializer != null) {
                 return serializer;
             }
             throw new UncheckedDeephavenException(
-                    String.format("Serializer not found for column %s, type %s", columnName, dataType.getName()));
+                    String.format("Serializer not found for column %s, type %s", columnName, type));
         }
 
         @Override

diff --git a/qst/type/build.gradle b/qst/type/build.gradle
@@ -9,6 +9,8 @@ dependencies {
     compileOnly project(':util-immutables')
     annotationProcessor libs.immutables.value
 
+    compileOnly libs.jetbrains.annotations
+
     testImplementation libs.assertj
     testImplementation platform(libs.junit.bom)
     testImplementation libs.junit.jupiter

diff --git a/qst/type/src/main/java/io/deephaven/qst/type/GenericVectorType.java b/qst/type/src/main/java/io/deephaven/qst/type/GenericVectorType.java
@@ -4,13 +4,16 @@
 package io.deephaven.qst.type;
 
 import io.deephaven.annotations.SimpleStyle;
+import org.immutables.value.Value.Check;
 import org.immutables.value.Value.Immutable;
 import org.immutables.value.Value.Parameter;
 
 @Immutable
 @SimpleStyle
 public abstract class GenericVectorType<T, ComponentType> extends ArrayTypeBase<T, ComponentType> {
 
+    private static final String OBJECT_VECTOR = "io.deephaven.vector.ObjectVector";
+
     public static <T, ComponentType> GenericVectorType<T, ComponentType> of(
             Class<T> clazz,
             GenericType<ComponentType> componentType) {
@@ -27,4 +30,12 @@ public static <T, ComponentType> GenericVectorType<T, ComponentType> of(
     public final <R> R walk(ArrayType.Visitor<R> visitor) {
         return visitor.visit(this);
     }
+
+    @Check
+    final void checkClazz() {
+        if (!OBJECT_VECTOR.equals(clazz().getName())) {
+            throw new IllegalArgumentException(String.format("Invalid GenericVectorType. clazz=%s, componentType=%s",
+                    clazz().getName(), componentType()));
+        }
+    }
 }
diff --git a/qst/type/src/main/java/io/deephaven/qst/type/PrimitiveVectorType.java b/qst/type/src/main/java/io/deephaven/qst/type/PrimitiveVectorType.java
@@ -73,10 +73,55 @@ public final <R> R walk(ArrayType.Visitor<R> visitor) {
     }
 
     @Check
-    final void checkClazz() {
-        if (!VALID_CLASSES.contains(clazz().getName())) {
-            throw new IllegalArgumentException(String.format("Class '%s' is not a valid '%s'",
-                    clazz(), PrimitiveVectorType.class));
+    final void checkPairing() {
+        final String vectorClassNameFromComponent = componentType().walk(VectorClassName.INSTANCE);
+        if (!clazz().getName().equals(vectorClassNameFromComponent)) {
+            throw new IllegalArgumentException(String.format("Invalid PrimitiveVectorType. clazz=%s, componentType=%s",
+                    clazz().getName(), componentType()));
+        }
+    }
+
+    private enum VectorClassName implements PrimitiveType.Visitor<String> {
+        INSTANCE;
+
+        @Override
+        public String visit(BooleanType booleanType) {
+            return null;
+        }
+
+        @Override
+        public String visit(ByteType byteType) {
+            return BYTE_VECTOR;
+        }
+
+        @Override
+        public String visit(CharType charType) {
+            return CHAR_VECTOR;
+        }
+
+        @Override
+        public String visit(ShortType shortType) {
+            return SHORT_VECTOR;
+        }
+
+        @Override
+        public String visit(IntType intType) {
+            return INT_VECTOR;
+        }
+
+        @Override
+        public String visit(LongType longType) {
+            return LONG_VECTOR;
+        }
+
+        @Override
+        public String visit(FloatType floatType) {
+            return FLOAT_VECTOR;
+        }
+
+        @Override
+        public String visit(DoubleType doubleType) {
+            return DOUBLE_VECTOR;
         }
     }
 }
diff --git a/qst/type/src/main/java/io/deephaven/qst/type/Type.java b/qst/type/src/main/java/io/deephaven/qst/type/Type.java
@@ -3,6 +3,8 @@
 //
 package io.deephaven.qst.type;
 
+import org.jetbrains.annotations.Nullable;
+
 import java.util.List;
 import java.util.Optional;
 
@@ -19,19 +21,42 @@ public interface Type<T> {
      * Finds the {@link #knownTypes() known type}, or else creates the relevant {@link NativeArrayType native array
      * type} or {@link CustomType custom type}.
      *
-     * @param clazz the class
-     * @param <T> the generic type of {@code clazz}
+     * @param dataType the data type
+     * @param <T> the generic type of {@code dataType}
      * @return the type
      */
-    static <T> Type<T> find(Class<T> clazz) {
-        Optional<Type<T>> found = TypeHelper.findStatic(clazz);
+    static <T> Type<T> find(Class<T> dataType) {
+        Optional<Type<T>> found = TypeHelper.findStatic(dataType);
         if (found.isPresent()) {
             return found.get();
         }
-        if (clazz.isArray()) {
-            return NativeArrayType.of(clazz, find(clazz.getComponentType()));
+        if (dataType.isArray()) {
+            return NativeArrayType.of(dataType, find(dataType.getComponentType()));
         }
-        return CustomType.of(clazz);
+        return CustomType.of(dataType);
+    }
+
+    /**
+     * If {@code componentType} is not {@code null}, this will find the appropriate {@link ArrayType}. Otherwise, this
+     * is equivalent to {@link #find(Class)}.
+     *
+     * @param dataType the data type
+     * @param componentType the component type
+     * @return the type
+     * @param <T> the generic type of {@code dataType}
+     */
+    static <T> Type<T> find(final Class<T> dataType, @Nullable final Class<?> componentType) {
+        if (componentType == null) {
+            return find(dataType);
+        }
+        final Type<?> ct = find(componentType);
+        if (dataType.isArray()) {
+            return NativeArrayType.of(dataType, ct);
+        }
+        if (componentType.isPrimitive()) {
+            return PrimitiveVectorType.of(dataType, (PrimitiveType<?>) ct);
+        }
+        return GenericVectorType.of(dataType, (GenericType<?>) ct);
     }
 
     /**