tests

codestyle/druid-forbidden-apis.txt

@@ -54,7 +54,7 @@ org.apache.calcite.sql.type.OperandTypes#NULLABLE_LITERAL @ Create an instance o
 org.apache.commons.io.FileUtils#getTempDirectory() @ Use org.junit.rules.TemporaryFolder for tests instead
 org.apache.commons.io.FileUtils#deleteDirectory(java.io.File) @ Use org.apache.druid.java.util.common.FileUtils#deleteDirectory()
 org.apache.commons.io.FileUtils#forceMkdir(java.io.File) @ Use org.apache.druid.java.util.common.FileUtils.mkdirp instead
-org.apache.datasketches.memory.Memory#wrap(byte[], int, int, java.nio.ByteOrder) @ The implementation isn't correct in datasketches-memory-2.2.0. Wrap the array in a ByteBuffer before wrapping it up in Memory
+org.apache.datasketches.memory.Memory#wrap(byte[], int, int, java.nio.ByteOrder) @ The implementation isn't correct in datasketches-memory-2.2.0. Please refer to https://github.com/apache/datasketches-memory/issues/178. Use wrap(byte[]) and modify the offset by the callers instead
 java.lang.Class#getCanonicalName() @ Class.getCanonicalName can return null for anonymous types, use Class.getName instead.
 java.util.concurrent.Executors#newFixedThreadPool(int) @ Executor is non-daemon and can prevent JVM shutdown, use org.apache.druid.java.util.common.concurrent.Execs#multiThreaded(int, java.lang.String) instead.
 

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

@@ -31,20 +31,24 @@
 import org.apache.druid.msq.indexing.MSQTuningConfig;
 import org.apache.druid.msq.indexing.destination.TaskReportMSQDestination;
 import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.query.DataSource;
 import org.apache.druid.query.NestedDataTestUtils;
+import org.apache.druid.query.aggregation.LongSumAggregatorFactory;
+import org.apache.druid.query.groupby.GroupByQueryConfig;
+import org.apache.druid.query.scan.ScanQuery;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.segment.nested.StructuredData;
 import org.apache.druid.sql.calcite.external.ExternalDataSource;
 import org.apache.druid.sql.calcite.filtration.Filtration;
 import org.apache.druid.sql.calcite.planner.ColumnMapping;
 import org.apache.druid.sql.calcite.planner.ColumnMappings;
-import org.apache.druid.sql.calcite.util.CalciteTests;
 import org.apache.druid.timeline.SegmentId;
 import org.apache.druid.utils.CompressionUtils;
 import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
 
 import java.io.File;
 import java.io.IOException;
@@ -54,6 +58,8 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 
 public class MSQComplexGroupByTest extends MSQTestBase
 {
@@ -65,7 +71,6 @@
   private String dataFileSignatureJsonString;
   private DataSource dataFileExternalDataSource;
 
-
   public static Collection<Object[]> data()
   {
     Object[][] data = new Object[][]{
@@ -95,8 +100,6 @@
     RowSignature dataFileSignature = RowSignature.builder()
                                                  .add("timestamp", ColumnType.STRING)
                                                  .add("obj", ColumnType.NESTED_DATA)
-                                                 .add("arrayObject", ColumnType.STRING_ARRAY)
-                                                 .add("arrayNestedLong", ColumnType.LONG_ARRAY)
                                                  .build();
     dataFileSignatureJsonString = queryFramework().queryJsonMapper().writeValueAsString(dataFileSignature);
 
@@ -109,8 +112,9 @@
     objectMapper.registerModules(NestedDataModule.getJacksonModulesList());
   }
 
-  @Test
-  public void testInsertWithRollupOnNestedData()
+  @MethodSource("data")
+  @ParameterizedTest(name = "{index}:with context {0}")
+  public void testInsertWithoutRollupOnNestedData(String contextName, Map<String, Object> context)
   {
     testIngestQuery().setSql("INSERT INTO foo1 SELECT\n"
                              + " obj,\n"
@@ -124,7 +128,7 @@
                              + " )\n"
                              + " GROUP BY 1\n"
                              + " PARTITIONED BY ALL")
-                     .setQueryContext(ImmutableMap.of())
+                     .setQueryContext(context)
                      .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
                      .setExpectedDataSource("foo1")
                      .setExpectedRowSignature(RowSignature.builder()
@@ -228,13 +232,148 @@
                              1L
                          }
                      ))
+                     .setQueryContext(context)
                      .verifyResults();
 
   }
 
-  @Test
-  public void testSortingOnNestedData()
+  @MethodSource("data")
+  @ParameterizedTest(name = "{index}:with context {0}")
+  public void testInsertWithRollupOnNestedData(String contextName, Map<String, Object> context)
   {
+    final Map<String, Object> adjustedContext = new HashMap<>(context);
+    adjustedContext.put(GroupByQueryConfig.CTX_KEY_ENABLE_MULTI_VALUE_UNNESTING, false);
+    adjustedContext.put(MultiStageQueryContext.CTX_FINALIZE_AGGREGATIONS, false);
+    testIngestQuery().setSql("INSERT INTO foo1 SELECT\n"
+                             + " obj,\n"
+                             + " COUNT(*) as cnt\n"
+                             + "FROM TABLE(\n"
+                             + "  EXTERN(\n"
+                             + "    '{ \"files\": [" + dataFileNameJsonString + "],\"type\":\"local\"}',\n"
+                             + "    '{\"type\": \"json\"}',\n"
+                             + "    '[{\"name\": \"timestamp\", \"type\": \"STRING\"}, {\"name\": \"obj\", \"type\": \"COMPLEX<json>\"}]'\n"
+                             + "   )\n"
+                             + " )\n"
+                             + " GROUP BY 1\n"
+                             + " PARTITIONED BY ALL")
+                     .setQueryContext(adjustedContext)
+                     .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
+                     .setExpectedDataSource("foo1")
+                     .setExpectedRowSignature(RowSignature.builder()
+                                                          .add("__time", ColumnType.LONG)
+                                                          .add("obj", ColumnType.NESTED_DATA)
+                                                          .add("cnt", ColumnType.LONG)
+                                                          .build())
+                     .addExpectedAggregatorFactory(new LongSumAggregatorFactory("cnt", "cnt"))
+                     .setExpectedResultRows(ImmutableList.of(
+                         new Object[]{
+                             0L,
+                             StructuredData.wrap(
+                                 ImmutableMap.of(
+                                     "a", 500,
+                                     "b", ImmutableMap.of(
+                                         "x", "e",
+                                         "z", ImmutableList.of(1, 2, 3, 4)
+                                     ),
+                                     "v", "a"
+                                 )
+                             ),
+                             1L
+                         },
+                         new Object[]{
+                             0L,
+                             StructuredData.wrap(
+                                 ImmutableMap.of(
+                                     "a", 100,
+                                     "b", ImmutableMap.of(
+                                         "x", "a",
+                                         "y", 1.1,
+                                         "z", ImmutableList.of(1, 2, 3, 4)
+                                     ),
+                                     "v", Collections.emptyList()
+                                 )
+                             ),
+                             2L
+                         },
+                         new Object[]{
+                             0L,
+                             StructuredData.wrap(
+                                 ImmutableMap.of(
+                                     "a", 700,
+                                     "b", ImmutableMap.of(
+                                         "x", "g",
+                                         "y", 1.1,
+                                         "z", Arrays.asList(9, null, 9, 9)
+                                     ),
+                                     "v", Collections.emptyList()
+                                 )
+                             ),
+                             1L
+                         },
+                         new Object[]{
+                             0L,
+                             StructuredData.wrap(
+                                 ImmutableMap.of(
+                                     "a", 200,
+                                     "b", ImmutableMap.of(
+                                         "x", "b",
+                                         "y", 1.1,
+                                         "z", ImmutableList.of(2, 4, 6)
+                                     ),
+                                     "v", Collections.emptyList()
+                                 )
+                             ),
+                             1L
+                         },
+                         new Object[]{
+                             0L,
+                             StructuredData.wrap(
+                                 ImmutableMap.of(
+                                     "a", 600,
+                                     "b", ImmutableMap.of(
+                                         "x", "f",
+                                         "y", 1.1,
+                                         "z", ImmutableList.of(6, 7, 8, 9)
+                                     ),
+                                     "v", "b"
+                                 )
+                             ),
+                             1L
+                         },
+                         new Object[]{
+                             0L,
+                             StructuredData.wrap(
+                                 ImmutableMap.of(
+                                     "a", 400,
+                                     "b", ImmutableMap.of(
+                                         "x", "d",
+                                         "y", 1.1,
+                                         "z", ImmutableList.of(3, 4)
+                                     ),
+                                     "v", Collections.emptyList()
+                                 )
+                             ),
+                             1L
+                         },
+                         new Object[]{
+                             0L,
+                             StructuredData.wrap(ImmutableMap.of("a", 300)),
+                             1L
+                         }
+                     ))
+                     .setExpectedRollUp(true)
+                     .setQueryContext(adjustedContext)
+                     .verifyResults();
+
+  }
+
+  @MethodSource("data")
+  @ParameterizedTest(name = "{index}:with context {0}")
+  public void testSortingOnNestedData(String contextName, Map<String, Object> context)
+  {
+    RowSignature rowSignature = RowSignature.builder()
+                                            .add("obj", ColumnType.NESTED_DATA)
+                                            .build();
     testSelectQuery().setSql("SELECT\n"
                              + " obj\n"
                              + "FROM TABLE(\n"
@@ -250,9 +389,11 @@
                      .setExpectedMSQSpec(MSQSpec
                                              .builder()
                                              .query(newScanQueryBuilder()
-                                                        .dataSource(CalciteTests.DATASOURCE1)
+                                                        .dataSource(dataFileExternalDataSource)
                                                         .intervals(querySegmentSpec(Filtration.eternity()))
-                                                        .columns("cnt", "dim1")
+                                                        .columns("obj")
+                                                        .context(defaultScanQueryContext(context, rowSignature))
+                                                        .orderBy(Collections.singletonList(new ScanQuery.OrderBy("obj", ScanQuery.Order.ASCENDING)))
                                                         .build()
                                              )
                                              .columnMappings(new ColumnMappings(ImmutableList.of(
@@ -262,9 +403,8 @@
                                              .destination(TaskReportMSQDestination.INSTANCE)
                                              .build()
                      )
-                     .setExpectedRowSignature(RowSignature.builder()
-                                                          .add("obj", ColumnType.NESTED_DATA)
-                                                          .build())
+                     .setExpectedRowSignature(rowSignature)
+                     .setQueryContext(context)
                      .setExpectedResultRows(ImmutableList.of(
                          new Object[]{"{\"a\":500,\"b\":{\"x\":\"e\",\"z\":[1,2,3,4]},\"v\":\"a\"}"},
                          new Object[]{"{\"a\":100,\"b\":{\"x\":\"a\",\"y\":1.1,\"z\":[1,2,3,4]},\"v\":[]}"},