add optional 'castToType' parameter to 'auto' column schema (apache#1…

…5417) * auto but.. with an expected type
yashdeep97 · Dec 1, 2023 · 083d7dc · 083d7dc
1 parent 1f274be
commit 083d7dc
Show file tree

Hide file tree

Showing 37 changed files with 973 additions and 194 deletions.
diff --git a/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlNestedDataBenchmark.java b/benchmarks/src/test/java/org/apache/druid/benchmark/query/SqlNestedDataBenchmark.java
@@ -298,7 +298,7 @@ public void setup()
     );
     List<DimensionSchema> dims = ImmutableList.<DimensionSchema>builder()
                                               .addAll(schemaInfo.getDimensionsSpec().getDimensions())
-                                              .add(new AutoTypeColumnSchema("nested"))
+                                              .add(new AutoTypeColumnSchema("nested", null))
                                               .build();
     DimensionsSpec dimsSpec = new DimensionsSpec(dims);
 

diff --git a/.../avro-extensions/src/test/java/org/apache/druid/data/input/AvroStreamInputFormatTest.java b/.../avro-extensions/src/test/java/org/apache/druid/data/input/AvroStreamInputFormatTest.java
@@ -303,15 +303,15 @@ public void testParseTransformNested() throws SchemaValidationException, IOExcep
 
     DimensionsSpec dimensionsSpec = new DimensionsSpec(
         ImmutableList.of(
-            new AutoTypeColumnSchema("someIntValueMap"),
-            new AutoTypeColumnSchema("someStringValueMap"),
-            new AutoTypeColumnSchema("someRecord"),
-            new AutoTypeColumnSchema("someRecordArray"),
+            new AutoTypeColumnSchema("someIntValueMap", null),
+            new AutoTypeColumnSchema("someStringValueMap", null),
+            new AutoTypeColumnSchema("someRecord", null),
+            new AutoTypeColumnSchema("someRecordArray", null),
             new LongDimensionSchema("tSomeIntValueMap8"),
             new LongDimensionSchema("tSomeIntValueMap8_2"),
             new StringDimensionSchema("tSomeStringValueMap8"),
             new LongDimensionSchema("tSomeRecordSubLong"),
-            new AutoTypeColumnSchema("tSomeRecordArray0"),
+            new AutoTypeColumnSchema("tSomeRecordArray0", null),
             new StringDimensionSchema("tSomeRecordArray0nestedString")
         )
     );

diff --git a/...query/src/main/java/org/apache/druid/msq/querykit/scan/ExternalColumnSelectorFactory.java b/...query/src/main/java/org/apache/druid/msq/querykit/scan/ExternalColumnSelectorFactory.java
@@ -107,7 +107,7 @@ public Object getObject()
           if (expressionType == null) {
             return delegateDimensionSelector.getObject();
           }
-          return ExprEval.ofType(expressionType, delegateDimensionSelector.getObject()).value();
+          return ExprEval.bestEffortOf(delegateDimensionSelector.getObject()).castTo(expressionType).value();
         }
         catch (Exception e) {
           throw createException(e, dimensionSpec.getDimension(), inputSource, offset);
@@ -211,7 +211,7 @@ public Object getObject()
           if (expressionType == null) {
             return delegateColumnValueSelector.getObject();
           }
-          return ExprEval.ofType(expressionType, delegateColumnValueSelector.getObject()).value();
+          return ExprEval.bestEffortOf(delegateColumnValueSelector.getObject()).castTo(expressionType).value();
         }
         catch (Exception e) {
           throw createException(e, columnName, inputSource, offset);

diff --git a/...-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java b/...-core/multi-stage-query/src/main/java/org/apache/druid/msq/util/DimensionSchemaUtils.java
@@ -73,7 +73,10 @@ public static DimensionSchema createDimensionSchema(
                                     .getDimensionSchema(capabilities);
       }
 
-      return new AutoTypeColumnSchema(column);
+      if (type != null && (type.isPrimitive() || type.isPrimitiveArray())) {
+        return new AutoTypeColumnSchema(column, type);
+      }
+      return new AutoTypeColumnSchema(column, null);
     } else {
       // if schema information is not available, create a string dimension
       if (type == null) {
@@ -102,12 +105,12 @@ public static DimensionSchema createDimensionSchema(
             return new StringDimensionSchema(column, DimensionSchema.MultiValueHandling.ARRAY, null);
           } else {
             // arrayIngestMode == ArrayIngestMode.ARRAY would be true
-            return new AutoTypeColumnSchema(column);
+            return new AutoTypeColumnSchema(column, type);
           }
         } else if (elementType.isNumeric()) {
           // ValueType == LONG || ValueType == FLOAT || ValueType == DOUBLE
           if (arrayIngestMode == ArrayIngestMode.ARRAY) {
-            return new AutoTypeColumnSchema(column);
+            return new AutoTypeColumnSchema(column, type);
           } else {
             throw InvalidInput.exception(
                 "Numeric arrays can only be ingested when '%s' is set to 'array' in the MSQ query's context. "

diff --git a/...e/multi-stage-query/src/test/java/org/apache/druid/msq/util/DimensionSchemaUtilsTest.java b/...e/multi-stage-query/src/test/java/org/apache/druid/msq/util/DimensionSchemaUtilsTest.java
@@ -0,0 +1,236 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.util;
+
+import org.apache.druid.data.input.impl.DimensionSchema;
+import org.apache.druid.data.input.impl.DoubleDimensionSchema;
+import org.apache.druid.data.input.impl.FloatDimensionSchema;
+import org.apache.druid.data.input.impl.LongDimensionSchema;
+import org.apache.druid.data.input.impl.StringDimensionSchema;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.segment.AutoTypeColumnSchema;
+import org.apache.druid.segment.column.ColumnType;
+import org.junit.Assert;
+import org.junit.Test;
+
+public class DimensionSchemaUtilsTest
+{
+
+  @Test
+  public void testSchemaScalars()
+  {
+    for (ArrayIngestMode mode : ArrayIngestMode.values()) {
+      DimensionSchema dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.LONG,
+          false,
+          mode
+      );
+      DimensionSchema expected = new LongDimensionSchema("x");
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.DOUBLE,
+          false,
+          mode
+      );
+      expected = new DoubleDimensionSchema("x");
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.FLOAT,
+          false,
+          mode
+      );
+      expected = new FloatDimensionSchema("x");
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.STRING,
+          false,
+          mode
+      );
+      expected = new StringDimensionSchema("x");
+      Assert.assertEquals(expected, dimensionSchema);
+    }
+  }
+
+  @Test
+  public void testSchemaForceAuto()
+  {
+    for (ArrayIngestMode mode : ArrayIngestMode.values()) {
+      DimensionSchema dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.LONG,
+          true,
+          mode
+      );
+      DimensionSchema expected = new AutoTypeColumnSchema("x", ColumnType.LONG);
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.DOUBLE,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", ColumnType.DOUBLE);
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.FLOAT,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", ColumnType.FLOAT);
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.STRING,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", ColumnType.STRING);
+      Assert.assertEquals(expected, dimensionSchema);
+
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.LONG_ARRAY,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", ColumnType.LONG_ARRAY);
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.DOUBLE_ARRAY,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", ColumnType.DOUBLE_ARRAY);
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.FLOAT_ARRAY,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", ColumnType.FLOAT_ARRAY);
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.STRING_ARRAY,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", ColumnType.STRING_ARRAY);
+      Assert.assertEquals(expected, dimensionSchema);
+
+      dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+          "x",
+          ColumnType.NESTED_DATA,
+          true,
+          mode
+      );
+      expected = new AutoTypeColumnSchema("x", null);
+      Assert.assertEquals(expected, dimensionSchema);
+    }
+  }
+
+  @Test
+  public void testSchemaMvdMode()
+  {
+    DimensionSchema dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+        "x",
+        ColumnType.STRING_ARRAY,
+        false,
+        ArrayIngestMode.MVD
+    );
+    DimensionSchema expected = new StringDimensionSchema("x", DimensionSchema.MultiValueHandling.ARRAY, null);
+    Assert.assertEquals(expected, dimensionSchema);
+
+    Throwable t = Assert.assertThrows(
+        DruidException.class,
+        () -> DimensionSchemaUtils.createDimensionSchema("x", ColumnType.LONG_ARRAY, false, ArrayIngestMode.MVD)
+    );
+    Assert.assertEquals("Numeric arrays can only be ingested when 'arrayIngestMode' is set to 'array' in the MSQ query's context. Current value of the parameter [mvd]", t.getMessage());
+
+    t = Assert.assertThrows(
+        DruidException.class,
+        () -> DimensionSchemaUtils.createDimensionSchema("x", ColumnType.DOUBLE_ARRAY, false, ArrayIngestMode.MVD)
+    );
+    Assert.assertEquals("Numeric arrays can only be ingested when 'arrayIngestMode' is set to 'array' in the MSQ query's context. Current value of the parameter [mvd]", t.getMessage());
+
+    t = Assert.assertThrows(
+        DruidException.class,
+        () -> DimensionSchemaUtils.createDimensionSchema("x", ColumnType.FLOAT_ARRAY, false, ArrayIngestMode.MVD)
+    );
+    Assert.assertEquals("Numeric arrays can only be ingested when 'arrayIngestMode' is set to 'array' in the MSQ query's context. Current value of the parameter [mvd]", t.getMessage());
+  }
+
+  @Test
+  public void testSchemaArrayMode()
+  {
+    DimensionSchema dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+        "x",
+        ColumnType.STRING_ARRAY,
+        false,
+        ArrayIngestMode.ARRAY
+    );
+    DimensionSchema expected = new AutoTypeColumnSchema("x", ColumnType.STRING_ARRAY);
+    Assert.assertEquals(expected, dimensionSchema);
+
+    dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+        "x",
+        ColumnType.LONG_ARRAY,
+        false,
+        ArrayIngestMode.ARRAY
+    );
+    expected = new AutoTypeColumnSchema("x", ColumnType.LONG_ARRAY);
+    Assert.assertEquals(expected, dimensionSchema);
+
+    dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+        "x",
+        ColumnType.DOUBLE_ARRAY,
+        false,
+        ArrayIngestMode.ARRAY
+    );
+    expected = new AutoTypeColumnSchema("x", ColumnType.DOUBLE_ARRAY);
+    Assert.assertEquals(expected, dimensionSchema);
+
+    dimensionSchema = DimensionSchemaUtils.createDimensionSchema(
+        "x",
+        ColumnType.FLOAT_ARRAY,
+        false,
+        ArrayIngestMode.ARRAY
+    );
+    expected = new AutoTypeColumnSchema("x", ColumnType.FLOAT_ARRAY);
+    Assert.assertEquals(expected, dimensionSchema);
+  }
+}
diff --git a/...ions-core/orc-extensions/src/test/java/org/apache/druid/data/input/orc/OrcReaderTest.java b/...ions-core/orc-extensions/src/test/java/org/apache/druid/data/input/orc/OrcReaderTest.java
@@ -333,9 +333,9 @@ public void testNestedColumn() throws IOException
         new TimestampSpec("ts", "millis", null),
         new DimensionsSpec(
             ImmutableList.of(
-                new AutoTypeColumnSchema("middle"),
-                new AutoTypeColumnSchema("list"),
-                new AutoTypeColumnSchema("map")
+                new AutoTypeColumnSchema("middle", null),
+                new AutoTypeColumnSchema("list", null),
+                new AutoTypeColumnSchema("map", null)
             )
         ),
         inputFormat,
@@ -542,8 +542,8 @@ public void testListMap() throws IOException
         new TimestampSpec("timestamp", "auto", null),
         new DimensionsSpec(
             ImmutableList.of(
-                new AutoTypeColumnSchema("a"),
-                new AutoTypeColumnSchema("b")
+                new AutoTypeColumnSchema("a", null),
+                new AutoTypeColumnSchema("b", null)
             )
         ),
         inputFormat,
@@ -608,11 +608,11 @@ public void testNestedArray() throws IOException
         new TimestampSpec("timestamp", "auto", null),
         new DimensionsSpec(
             ImmutableList.of(
-                new AutoTypeColumnSchema("a"),
-                new AutoTypeColumnSchema("b"),
-                new AutoTypeColumnSchema("c"),
-                new AutoTypeColumnSchema("d"),
-                new AutoTypeColumnSchema("t_d_0")
+                new AutoTypeColumnSchema("a", null),
+                new AutoTypeColumnSchema("b", null),
+                new AutoTypeColumnSchema("c", null),
+                new AutoTypeColumnSchema("d", null),
+                new AutoTypeColumnSchema("t_d_0", null)
             )
         ),
         inputFormat,

diff --git a/...ions/src/test/java/org/apache/druid/data/input/parquet/NestedColumnParquetReaderTest.java b/...ions/src/test/java/org/apache/druid/data/input/parquet/NestedColumnParquetReaderTest.java
@@ -51,8 +51,8 @@ public void testNestedColumnTransformsNestedTestFile() throws IOException
         new TimestampSpec("timestamp", "auto", null),
         new DimensionsSpec(
             ImmutableList.of(
-                new AutoTypeColumnSchema("nestedData"),
-                new AutoTypeColumnSchema("t_nestedData_listDim"),
+                new AutoTypeColumnSchema("nestedData", null),
+                new AutoTypeColumnSchema("t_nestedData_listDim", null),
                 new StringDimensionSchema("t_nestedData_listDim_string"),
                 new StringDimensionSchema("t_nestedData_dim2"),
                 new LongDimensionSchema("t_nestedData_dim3"),
@@ -105,10 +105,10 @@ public void testNestedColumnTransformsNestedNullableListFile() throws IOExceptio
         new TimestampSpec("timestamp", "auto", null),
         new DimensionsSpec(
             ImmutableList.of(
-                new AutoTypeColumnSchema("a1"),
-                new AutoTypeColumnSchema("a2"),
-                new AutoTypeColumnSchema("t_a2"),
-                new AutoTypeColumnSchema("t_a1_b1"),
+                new AutoTypeColumnSchema("a1", null),
+                new AutoTypeColumnSchema("a2", null),
+                new AutoTypeColumnSchema("t_a2", null),
+                new AutoTypeColumnSchema("t_a1_b1", null),
                 new LongDimensionSchema("t_a1_b1_c1"),
                 new LongDimensionSchema("t_e2_0_b1"),
                 new LongDimensionSchema("tt_a2_0_b1")