Improve unit test coverage for MSQ

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQFaultsTest.java

@@ -0,0 +1,215 @@
+/*
+ * 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.exec;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.msq.indexing.error.TooManyClusteredByColumnsFault;
+import org.apache.druid.msq.indexing.error.TooManyColumnsFault;
+import org.apache.druid.msq.indexing.error.TooManyInputFilesFault;
+import org.apache.druid.msq.indexing.error.TooManyPartitionsFault;
+import org.apache.druid.msq.indexing.error.UnknownFault;
+import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.test.MSQTestFileUtils;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.Collections;
+import java.util.Map;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+public class MSQFaultsTest extends MSQTestBase
+{
+
+  @Test
+  public void testInsertWithTooManySegments() throws IOException
+  {
+    Map<String, Object> context = ImmutableMap.<String, Object>builder()
+                                              .putAll(DEFAULT_MSQ_CONTEXT)
+                                              .put("rowsPerSegment", 1)
+                                              .build();
+
+
+    RowSignature rowSignature = RowSignature.builder()
+                                            .add("__time", ColumnType.LONG)
+                                            .build();
+
+    File file = MSQTestFileUtils.generateTemporaryNdJsonFile(30000, 1);
+    String filePathAsJson = queryFramework().queryJsonMapper().writeValueAsString(file.getAbsolutePath());
+
+    testIngestQuery().setSql(" insert into foo1 SELECT\n"
+                             + "  floor(TIME_PARSE(\"timestamp\") to day) AS __time\n"
+                             + "FROM TABLE(\n"
+                             + "  EXTERN(\n"
+                             + "    '{ \"files\": [" + filePathAsJson + "],\"type\":\"local\"}',\n"
+                             + "    '{\"type\": \"json\"}',\n"
+                             + "    '[{\"name\": \"timestamp\",\"type\":\"string\"}]'\n"
+                             + "  )\n"
+                             + ") PARTITIONED by day")
+                     .setExpectedDataSource("foo1")
+                     .setExpectedRowSignature(rowSignature)
+                     .setQueryContext(context)
+                     .setExpectedMSQFault(new TooManyPartitionsFault(25000))
+                     .verifyResults();
+
+  }
+
+  @Test
+  public void testInsertWithUnsupportedColumnType()
+  {
+    RowSignature dummyRowSignature = RowSignature.builder().add("__time", ColumnType.LONG).build();
+
+    testIngestQuery()
+        .setSql(StringUtils.format(
+            " insert into foo1 SELECT\n"
+            + "  floor(TIME_PARSE(\"timestamp\") to day) AS __time,\n"
+            + " col1\n"
+            + "FROM TABLE(\n"
+            + "  EXTERN(\n"
+            + "    '{ \"files\": [\"ignored\"],\"type\":\"local\"}',\n"
+            + "    '{\"type\": \"json\"}',\n"
+            + "    '[{\"name\": \"timestamp\", \"type\": \"string\"},{\"name\": \"col1\", \"type\": \"long_array\"} ]'\n"
+            + "  )\n"
+            + ") PARTITIONED by day"
+        ))
+        .setExpectedDataSource("foo1")
+        .setExpectedRowSignature(dummyRowSignature)
+        .setExpectedMSQFault(UnknownFault.forMessage(
+            "org.apache.druid.java.util.common.ISE: Cannot create dimension for type [ARRAY<LONG>]"))
+        .verifyResults();
+  }
+
+  @Test
+  public void testInsertWithManyColumns()
+  {
+    RowSignature dummyRowSignature = RowSignature.builder().add("__time", ColumnType.LONG).build();
+
+    final int numColumns = 2000;
+
+    String columnNames = IntStream.range(1, numColumns)
+                                  .mapToObj(i -> "col" + i).collect(Collectors.joining(", "));
+
+    String externSignature = IntStream.range(1, numColumns)
+                                      .mapToObj(i -> StringUtils.format(
+                                          "{\"name\": \"col%d\", \"type\": \"string\"}",
+                                          i
+                                      ))
+                                      .collect(Collectors.joining(", "));
+
+    testIngestQuery()
+        .setSql(StringUtils.format(
+            " insert into foo1 SELECT\n"
+            + "  floor(TIME_PARSE(\"timestamp\") to day) AS __time,\n"
+            + " %s\n"
+            + "FROM TABLE(\n"
+            + "  EXTERN(\n"
+            + "    '{ \"files\": [\"ignored\"],\"type\":\"local\"}',\n"
+            + "    '{\"type\": \"json\"}',\n"
+            + "    '[{\"name\": \"timestamp\", \"type\": \"string\"}, %s]'\n"
+            + "  )\n"
+            + ") PARTITIONED by day",
+            columnNames,
+            externSignature
+        ))
+        .setExpectedDataSource("foo1")
+        .setExpectedRowSignature(dummyRowSignature)
+        .setExpectedMSQFault(new TooManyColumnsFault(numColumns + 2, 2000))
+        .verifyResults();
+  }
+
+  @Test
+  public void testInsertWithHugeClusteringKeys()
+  {
+    RowSignature dummyRowSignature = RowSignature.builder().add("__time", ColumnType.LONG).build();
+
+    final int numColumns = 1700;
+
+    String columnNames = IntStream.range(1, numColumns)
+                                  .mapToObj(i -> "col" + i).collect(Collectors.joining(", "));
+
+    String clusteredByClause = IntStream.range(1, numColumns + 1)
+                                        .mapToObj(String::valueOf)
+                                        .collect(Collectors.joining(", "));
+
+    String externSignature = IntStream.range(1, numColumns)
+                                      .mapToObj(i -> StringUtils.format(
+                                          "{\"name\": \"col%d\", \"type\": \"string\"}",
+                                          i
+                                      ))
+                                      .collect(Collectors.joining(", "));
+
+    testIngestQuery()
+        .setSql(StringUtils.format(
+            " insert into foo1 SELECT\n"
+            + "  floor(TIME_PARSE(\"timestamp\") to day) AS __time,\n"
+            + " %s\n"
+            + "FROM TABLE(\n"
+            + "  EXTERN(\n"
+            + "    '{ \"files\": [\"ignored\"],\"type\":\"local\"}',\n"
+            + "    '{\"type\": \"json\"}',\n"
+            + "    '[{\"name\": \"timestamp\", \"type\": \"string\"}, %s]'\n"
+            + "  )\n"
+            + ") PARTITIONED by day CLUSTERED BY %s",
+            columnNames,
+            externSignature,
+            clusteredByClause
+        ))
+        .setExpectedDataSource("foo1")
+        .setExpectedRowSignature(dummyRowSignature)
+        .setExpectedMSQFault(new TooManyClusteredByColumnsFault(numColumns + 2, 1500, 0))
+        .verifyResults();
+  }
+
+  @Test
+  public void testTooManyInputFiles() throws IOException
+  {
+    RowSignature dummyRowSignature = RowSignature.builder().add("__time", ColumnType.LONG).build();
+
+    final int numFiles = 20000;
+
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
+    final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
+
+    String externalFiles = String.join(", ", Collections.nCopies(numFiles, toReadFileNameAsJson));
+
+    testIngestQuery()
+        .setSql(StringUtils.format(
+            "insert into foo1 SELECT\n"
+            + "  floor(TIME_PARSE(\"timestamp\") to day) AS __time\n"
+            + "FROM TABLE(\n"
+            + "  EXTERN(\n"
+            + "    '{ \"files\": [%s],\"type\":\"local\"}',\n"
+            + "    '{\"type\": \"csv\", \"hasHeaderRow\": true}',\n"
+            + "    '[{\"name\": \"timestamp\", \"type\": \"string\"}]'\n"
+            + "  )\n"
+            + ") PARTITIONED by day",
+            externalFiles
+        ))
+        .setExpectedDataSource("foo1")
+        .setExpectedRowSignature(dummyRowSignature)
+        .setExpectedMSQFault(new TooManyInputFilesFault(numFiles, Limits.MAX_INPUT_FILES_PER_WORKER, 2))
+        .verifyResults();
+  }
+
+}

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQInsertTest.java

@@ -27,13 +27,12 @@
 import org.apache.druid.hll.HyperLogLogCollector;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.Intervals;
-import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.msq.indexing.error.ColumnNameRestrictedFault;
 import org.apache.druid.msq.indexing.error.InsertTimeNullFault;
 import org.apache.druid.msq.indexing.error.RowTooLargeFault;
-import org.apache.druid.msq.indexing.error.TooManyClusteredByColumnsFault;
 import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.test.MSQTestFileUtils;
 import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
 import org.apache.druid.query.aggregation.hyperloglog.HyperUniquesAggregatorFactory;
@@ -55,8 +54,6 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.TreeSet;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 
 public class MSQInsertTest extends MSQTestBase
 {
@@ -83,7 +80,7 @@ public void testInsertOnFoo1()
   @Test
   public void testInsertOnExternalDataSource() throws IOException
   {
-    final File toRead = getResourceAsTemporaryFile("/wikipedia-sampled.json");
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
     final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     RowSignature rowSignature = RowSignature.builder()
@@ -308,7 +305,7 @@ public void testRollUpOnFoo1ComplexCol()
   @Test
   public void testRollUpOnExternalDataSource() throws IOException
   {
-    final File toRead = getResourceAsTemporaryFile("/wikipedia-sampled.json");
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
     final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     RowSignature rowSignature = RowSignature.builder()
@@ -344,7 +341,7 @@ public void testRollUpOnExternalDataSource() throws IOException
   @Test()
   public void testRollUpOnExternalDataSourceWithCompositeKey() throws IOException
   {
-    final File toRead = getResourceAsTemporaryFile("/wikipedia-sampled.json");
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
     final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     RowSignature rowSignature = RowSignature.builder()
@@ -448,48 +445,6 @@ public void testIncorrectInsertQuery()
                      .verifyPlanningErrors();
   }
 
-  @Test
-  public void testInsertWithHugeClusteringKeys()
-  {
-    RowSignature dummyRowSignature = RowSignature.builder().add("__time", ColumnType.LONG).build();
-
-    final int numColumns = 1700;
-
-    String columnNames = IntStream.range(1, numColumns)
-                                  .mapToObj(i -> "col" + i).collect(Collectors.joining(", "));
-
-    String clusteredByClause = IntStream.range(1, numColumns + 1)
-                                        .mapToObj(String::valueOf)
-                                        .collect(Collectors.joining(", "));
-
-    String externSignature = IntStream.range(1, numColumns)
-                                      .mapToObj(i -> StringUtils.format(
-                                          "{\"name\": \"col%d\", \"type\": \"string\"}",
-                                          i
-                                      ))
-                                      .collect(Collectors.joining(", "));
-
-    testIngestQuery()
-        .setSql(StringUtils.format(
-            " insert into foo1 SELECT\n"
-            + "  floor(TIME_PARSE(\"timestamp\") to day) AS __time,\n"
-            + " %s\n"
-            + "FROM TABLE(\n"
-            + "  EXTERN(\n"
-            + "    '{ \"files\": [\"ignored\"],\"type\":\"local\"}',\n"
-            + "    '{\"type\": \"json\"}',\n"
-            + "    '[{\"name\": \"timestamp\", \"type\": \"string\"}, %s]'\n"
-            + "  )\n"
-            + ") PARTITIONED by day CLUSTERED BY %s",
-            columnNames,
-            externSignature,
-            clusteredByClause
-        ))
-        .setExpectedDataSource("foo1")
-        .setExpectedRowSignature(dummyRowSignature)
-        .setExpectedMSQFault(new TooManyClusteredByColumnsFault(numColumns + 2, 1500, 0))
-        .verifyResults();
-  }
 
   @Test
   public void testInsertRestrictedColumns()
@@ -542,7 +497,7 @@ public void testInsertQueryWithInvalidSubtaskCount()
   @Test
   public void testInsertWithTooLargeRowShouldThrowException() throws IOException
   {
-    final File toRead = getResourceAsTemporaryFile("/wikipedia-sampled.json");
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
     final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     Mockito.doReturn(500).when(workerMemoryParameters).getLargeFrameSize();

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQReplaceTest.java

@@ -23,6 +23,7 @@
 import com.google.common.collect.ImmutableSet;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.test.MSQTestFileUtils;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.sql.SqlPlanningException;
@@ -109,7 +110,7 @@ public void testReplaceOnFoo1WithAllExtern() throws IOException
                                             .add("__time", ColumnType.LONG)
                                             .add("cnt", ColumnType.LONG).build();
 
-    final File toRead = getResourceAsTemporaryFile("/wikipedia-sampled.json");
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
     final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     testIngestQuery().setSql(" REPLACE INTO foo1 OVERWRITE ALL SELECT "
@@ -147,7 +148,7 @@ public void testReplaceOnFoo1WithWhereExtern() throws IOException
                                             .add("__time", ColumnType.LONG)
                                             .add("user", ColumnType.STRING).build();
 
-    final File toRead = getResourceAsTemporaryFile("/wikipedia-sampled.json");
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
     final String toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     testIngestQuery().setSql(" REPLACE INTO foo1 OVERWRITE WHERE __time >= TIMESTAMP '2016-06-27 01:00:00.00' AND __time < TIMESTAMP '2016-06-27 02:00:00.00' "

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQSelectTest.java

@@ -33,6 +33,7 @@
 import org.apache.druid.msq.indexing.MSQTuningConfig;
 import org.apache.druid.msq.shuffle.DurableStorageUtils;
 import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.test.MSQTestFileUtils;
 import org.apache.druid.query.InlineDataSource;
 import org.apache.druid.query.QueryDataSource;
 import org.apache.druid.query.TableDataSource;
@@ -732,7 +733,7 @@ public void testGroupByOrderByAggregationWithLimitAndOffset()
   @Test
   public void testExternSelect1() throws IOException
   {
-    final File toRead = getResourceAsTemporaryFile("/wikipedia-sampled.json");
+    final File toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/wikipedia-sampled.json");
     final String toReadAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     RowSignature rowSignature = RowSignature.builder()

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/error/MSQWarningsTest.java

@@ -29,6 +29,7 @@
 import org.apache.druid.msq.indexing.MSQSpec;
 import org.apache.druid.msq.indexing.MSQTuningConfig;
 import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.test.MSQTestFileUtils;
 import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
@@ -65,7 +66,7 @@ public class MSQWarningsTest extends MSQTestBase
   @Before
   public void setUp3() throws IOException
   {
-    toRead = getResourceAsTemporaryFile("/unparseable.gz");
+    toRead = MSQTestFileUtils.getResourceAsTemporaryFile(this, "/unparseable.gz");
     toReadFileNameAsJson = queryFramework().queryJsonMapper().writeValueAsString(toRead.getAbsolutePath());
 
     rowSignature = RowSignature.builder()