insert test case

LakshSingla · Oct 4, 2023 · c84c6db · c84c6db
1 parent d41d1b0
commit c84c6db
Show file tree

Hide file tree

Showing 3 changed files with 192 additions and 1 deletion.
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
@@ -75,6 +75,10 @@ public static DimensionSchema createDimensionSchema(
           switch (type.getElementType().getType()) {
             case STRING:
               return new StringDimensionSchema(column, DimensionSchema.MultiValueHandling.ARRAY, null);
+            case LONG:
+            case FLOAT:
+            case DOUBLE:
+              return new AutoTypeColumnSchema(column);
             default:
               throw new ISE("Cannot create dimension for type [%s]", type.toString());
           }

diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java
@@ -63,6 +63,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -1463,6 +1464,193 @@ public void testInsertArraysAutoType() throws IOException
                      .verifyResults();
   }
 
+  @Test
+  public void testInsertArrays() throws IOException
+  {
+    List<Object[]> expectedRows = Arrays.asList(
+        new Object[]{
+            1672531200000L,
+            null,
+            null,
+            new Object[]{1L, 2L, 3L},
+            new Object[]{},
+            new Object[]{1.1d, 2.2d, 3.3d},
+            null
+        },
+        new Object[]{
+            1672531200000L,
+            null,
+            Arrays.asList("a", "b"),
+            null,
+            new Object[]{2L, 3L},
+            null,
+            new Object[]{null}
+        },
+        new Object[]{
+            1672531200000L,
+            Arrays.asList("a", "b"),
+            null,
+            null,
+            new Object[]{null, 2L, 9L},
+            null,
+            new Object[]{999.0d, 5.5d, null}
+        },
+        new Object[]{
+            1672531200000L,
+            Arrays.asList("a", "b"),
+            Arrays.asList("a", "b"),
+            new Object[]{1L, 2L, 3L},
+            new Object[]{1L, null, 3L},
+            new Object[]{1.1d, 2.2d, 3.3d},
+            new Object[]{1.1d, 2.2d, null}
+        },
+        new Object[]{
+            1672531200000L,
+            Arrays.asList("a", "b", "c"),
+            Arrays.asList(null, "b"),
+            new Object[]{2L, 3L},
+            null,
+            new Object[]{3.3d, 4.4d, 5.5d},
+            new Object[]{999.0d, null, 5.5d}
+        },
+        new Object[]{
+            1672531200000L,
+            Arrays.asList("b", "c"),
+            Arrays.asList("d", null, "b"),
+            new Object[]{1L, 2L, 3L, 4L},
+            new Object[]{1L, 2L, 3L},
+            new Object[]{1.1d, 3.3d},
+            new Object[]{null, 2.2d, null}
+        },
+        new Object[]{
+            1672531200000L,
+            Arrays.asList("d", "e"),
+            Arrays.asList("b", "b"),
+            new Object[]{1L, 4L},
+            new Object[]{1L},
+            new Object[]{2.2d, 3.3d, 4.0d},
+            null
+        },
+        new Object[]{
+            1672617600000L,
+            null,
+            null,
+            new Object[]{1L, 2L, 3L},
+            null,
+            new Object[]{1.1d, 2.2d, 3.3d},
+            new Object[]{}
+        },
+        new Object[]{
+            1672617600000L,
+            null,
+            Arrays.asList("a", "b"),
+            null,
+            new Object[]{2L, 3L},
+            null,
+            new Object[]{null, 1.1d}
+        },
+        new Object[]{
+            1672617600000L,
+            Arrays.asList("a", "b"),
+            null,
+            null,
+            new Object[]{null, 2L, 9L},
+            null,
+            new Object[]{999.0d, 5.5d, null}
+        },
+        new Object[]{
+            1672617600000L,
+            Arrays.asList("a", "b"),
+            Collections.emptyList(),
+            new Object[]{1L, 2L, 3L},
+            new Object[]{1L, null, 3L},
+            new Object[]{1.1d, 2.2d, 3.3d},
+            new Object[]{1.1d, 2.2d, null}
+        },
+        new Object[]{
+            1672617600000L,
+            Arrays.asList("a", "b", "c"),
+            Arrays.asList(null, "b"),
+            new Object[]{2L, 3L},
+            null,
+            new Object[]{3.3d, 4.4d, 5.5d},
+            new Object[]{999.0d, null, 5.5d}
+        },
+        new Object[]{
+            1672617600000L,
+            Arrays.asList("b", "c"),
+            Arrays.asList("d", null, "b"),
+            new Object[]{1L, 2L, 3L, 4L},
+            new Object[]{1L, 2L, 3L},
+            new Object[]{1.1d, 3.3d},
+            new Object[]{null, 2.2d, null}
+        },
+        new Object[]{
+            1672617600000L,
+            Arrays.asList("d", "e"),
+            Arrays.asList("b", "b"),
+            new Object[]{1L, 4L},
+            new Object[]{null},
+            new Object[]{2.2d, 3.3d, 4.0},
+            null
+        }
+    );
+
+    RowSignature rowSignatureWithoutTimeAndStringColumns =
+        RowSignature.builder()
+                    .add("arrayLong", ColumnType.LONG_ARRAY)
+                    .add("arrayLongNulls", ColumnType.LONG_ARRAY)
+                    .add("arrayDouble", ColumnType.DOUBLE_ARRAY)
+                    .add("arrayDoubleNulls", ColumnType.DOUBLE_ARRAY)
+                    .build();
+
+
+    RowSignature fileSignature = RowSignature.builder()
+                                             .add("timestamp", ColumnType.STRING)
+                                             .add("arrayString", ColumnType.STRING_ARRAY)
+                                             .add("arrayStringNulls", ColumnType.STRING_ARRAY)
+                                             .addAll(rowSignatureWithoutTimeAndStringColumns)
+                                             .build();
+
+    // MSQ writes strings instead of string arrays
+    RowSignature rowSignature = RowSignature.builder()
+                                            .add("__time", ColumnType.LONG)
+                                            .add("arrayString", ColumnType.STRING)
+                                            .add("arrayStringNulls", ColumnType.STRING)
+                                            .addAll(rowSignatureWithoutTimeAndStringColumns)
+                                            .build();
+
+    final Map<String, Object> adjustedContext = new HashMap<>(context);
+    final File tmpFile = temporaryFolder.newFile();
+    final InputStream resourceStream = NestedDataTestUtils.class.getClassLoader().getResourceAsStream(NestedDataTestUtils.ARRAY_TYPES_DATA_FILE);
+    final InputStream decompressing = CompressionUtils.decompress(resourceStream, NestedDataTestUtils.ARRAY_TYPES_DATA_FILE);
+    Files.copy(decompressing, tmpFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
+    decompressing.close();
+
+    final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(tmpFile);
+
+    testIngestQuery().setSql(" INSERT INTO foo1 SELECT\n"
+                             + "  TIME_PARSE(\"timestamp\") as __time,\n"
+                             + "  arrayString,\n"
+                             + "  arrayStringNulls,\n"
+                             + "  arrayLong,\n"
+                             + "  arrayLongNulls,\n"
+                             + "  arrayDouble,\n"
+                             + "  arrayDoubleNulls\n"
+                             + "FROM TABLE(\n"
+                             + "  EXTERN(\n"
+                             + "    '{ \"files\": [" + toReadFileNameAsJson + "],\"type\":\"local\"}',\n"
+                             + "    '{\"type\": \"json\"}',\n"
+                             + "    '" + queryFramework().queryJsonMapper().writeValueAsString(fileSignature) + "'\n"
+                             + "  )\n"
+                             + ") PARTITIONED BY day")
+                     .setQueryContext(adjustedContext)
+                     .setExpectedResultRows(expectedRows)
+                     .setExpectedDataSource("foo1")
+                     .setExpectedRowSignature(rowSignature)
+                     .verifyResults();
+  }
+
   @Nonnull
   private List<Object[]> expectedFooRows()
   {

diff --git a/processing/src/main/java/org/apache/druid/frame/field/NumericArrayFieldReader.java b/processing/src/main/java/org/apache/druid/frame/field/NumericArrayFieldReader.java
@@ -56,5 +56,4 @@ public boolean isComparable()
   {
     return true;
   }
-
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -56,5 +56,4 @@ public boolean isComparable() @@
       {
         return true;
       }
     }