tests, checkstyle

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/common/SortMergeJoinFrameProcessorFactory.java

@@ -214,7 +214,7 @@ public static List<List<KeyColumn>> toKeyColumns(final JoinConditionAnalysis con
           "leftExpr#getBindingIfIdentifier"
       );
 
-      retVal.get(0).add(new KeyColumn(leftColumn,  KeyOrder.ASCENDING));
+      retVal.get(0).add(new KeyColumn(leftColumn, KeyOrder.ASCENDING));
       retVal.get(1).add(new KeyColumn(equiCondition.getRightColumn(), KeyOrder.ASCENDING));
     }
 

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/statistics/KeyCollectors.java

@@ -36,7 +36,7 @@ private KeyCollectors()
       final ClusterBy clusterBy,
       final boolean aggregate,
       final RowSignature rowSignature
-      )
+  )
   {
     final KeyCollectorFactory<?, ?> baseFactory;
 

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

@@ -0,0 +1,152 @@
+/*
+ * 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.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import org.apache.druid.data.input.impl.JsonInputFormat;
+import org.apache.druid.data.input.impl.LocalInputSource;
+import org.apache.druid.data.input.impl.systemfield.SystemFields;
+import org.apache.druid.guice.NestedDataModule;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.msq.test.MSQTestBase;
+import org.apache.druid.query.DataSource;
+import org.apache.druid.query.NestedDataTestUtils;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.sql.calcite.external.ExternalDataSource;
+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 java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.StandardCopyOption;
+import java.util.Arrays;
+import java.util.Collection;
+
+public class MSQComplexGroupByTest extends MSQTestBase
+{
+  static {
+    NestedDataModule.registerHandlersAndSerde();
+  }
+
+  private String dataFileNameJsonString;
+  private String dataFileSignatureJsonString;
+  private DataSource dataFileExternalDataSource;
+
+
+  public static Collection<Object[]> data()
+  {
+    Object[][] data = new Object[][]{
+        {DEFAULT, DEFAULT_MSQ_CONTEXT},
+        {DURABLE_STORAGE, DURABLE_STORAGE_MSQ_CONTEXT},
+        {FAULT_TOLERANCE, FAULT_TOLERANCE_MSQ_CONTEXT},
+        {PARALLEL_MERGE, PARALLEL_MERGE_MSQ_CONTEXT}
+    };
+    return Arrays.asList(data);
+  }
+
+  @BeforeEach
+  public void setup() throws IOException
+  {
+    File dataFile = newTempFile("dataFile");
+    final InputStream resourceStream = this.getClass().getClassLoader()
+                                                                .getResourceAsStream(NestedDataTestUtils.ALL_TYPES_TEST_DATA_FILE);
+    final InputStream decompressing = CompressionUtils.decompress(
+        resourceStream,
+        "nested-all-types-test-data.json"
+    );
+    Files.copy(decompressing, dataFile.toPath(), StandardCopyOption.REPLACE_EXISTING);
+    decompressing.close();
+
+    dataFileNameJsonString = queryFramework().queryJsonMapper().writeValueAsString(dataFile);
+
+    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);
+
+    dataFileExternalDataSource = new ExternalDataSource(
+        new LocalInputSource(null, null, ImmutableList.of(dataFile), SystemFields.none()),
+        new JsonInputFormat(null, null, null, null, null),
+        dataFileSignature
+    );
+
+    objectMapper.registerModules(NestedDataModule.getJacksonModulesList());
+  }
+
+  @Test
+  public void testInsertWithRollupOnNestedData()
+  {
+    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(ImmutableMap.of())
+                     .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
+                     .setExpectedDataSource("foo1")
+                     .setExpectedRowSignature(RowSignature.builder()
+                                                          .add("obj", ColumnType.NESTED_DATA)
+                                                          .add("cnt", ColumnType.LONG)
+                                                          .build())
+                     .setExpectedResultRows(ImmutableList.of())
+                     .verifyResults();
+
+  }
+
+  @Test
+  public void testSortingOnNestedData()
+  {
+    testSelectQuery().setSql("SELECT\n"
+                             + " obj\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"
+                             + " ORDER BY 1")
+                     .setQueryContext(ImmutableMap.of())
+                     .setExpectedRowSignature(RowSignature.builder()
+                                                          .add("obj", ColumnType.NESTED_DATA)
+                                                          .build())
+                     .setExpectedResultRows(ImmutableList.of())
+                     .verifyResults();
+  }
+
+}

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

@@ -77,6 +77,7 @@
 import org.apache.druid.server.SegmentManager;
 import org.apache.druid.server.coordination.DataSegmentAnnouncer;
 import org.apache.druid.server.coordination.NoopDataSegmentAnnouncer;
+import org.apache.druid.sql.calcite.CalciteNestedDataQueryTest;
 import org.apache.druid.sql.calcite.util.CalciteTests;
 import org.apache.druid.sql.calcite.util.TestDataBuilder;
 import org.apache.druid.timeline.DataSegment;
@@ -292,6 +293,101 @@ private static Supplier<ResourceHolder<Segment>> getSupplierForSegment(Function<
                             .inputTmpDir(tempFolderProducer.apply("tmpDir"))
                             .buildMMappedIndex();
         break;
+      case CalciteNestedDataQueryTest.DATA_SOURCE:
+      case CalciteNestedDataQueryTest.DATA_SOURCE_MIXED:
+        if (segmentId.getPartitionNum() == 0) {
+          index = IndexBuilder.create()
+                              .tmpDir(tempFolderProducer.apply("tmpDir"))
+                              .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+                              .schema(
+                                  new IncrementalIndexSchema.Builder()
+                                      .withMetrics(
+                                          new CountAggregatorFactory("cnt")
+                                      )
+                                      .withDimensionsSpec(CalciteNestedDataQueryTest.ALL_JSON_COLUMNS.getDimensionsSpec())
+                                      .withRollup(false)
+                                      .build()
+                              )
+                              .rows(CalciteNestedDataQueryTest.ROWS)
+                              .buildMMappedIndex();
+        } else if (segmentId.getPartitionNum() == 1) {
+          index = IndexBuilder.create()
+                              .tmpDir(tempFolderProducer.apply("tmpDir"))
+                              .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+                              .schema(
+                                  new IncrementalIndexSchema.Builder()
+                                      .withMetrics(
+                                          new CountAggregatorFactory("cnt")
+                                      )
+                                      .withDimensionsSpec(CalciteNestedDataQueryTest.JSON_AND_SCALAR_MIX.getDimensionsSpec())
+                                      .withRollup(false)
+                                      .build()
+                              )
+                              .rows(CalciteNestedDataQueryTest.ROWS_MIX)
+                              .buildMMappedIndex();
+        } else {
+          throw new ISE("Cannot query segment %s in test runner", segmentId);
+        }
+        break;
+      case CalciteNestedDataQueryTest.DATA_SOURCE_MIXED_2:
+        if (segmentId.getPartitionNum() == 0) {
+          index = IndexBuilder.create()
+                              .tmpDir(tempFolderProducer.apply("tmpDir"))
+                              .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+                              .schema(
+                                  new IncrementalIndexSchema.Builder()
+                                      .withMetrics(
+                                          new CountAggregatorFactory("cnt")
+                                      )
+                                      .withDimensionsSpec(CalciteNestedDataQueryTest.JSON_AND_SCALAR_MIX.getDimensionsSpec())
+                                      .withRollup(false)
+                                      .build()
+                              )
+                              .rows(CalciteNestedDataQueryTest.ROWS_MIX)
+                              .buildMMappedIndex();
+        } else if (segmentId.getPartitionNum() == 1) {
+          index = IndexBuilder.create()
+                      .tmpDir(tempFolderProducer.apply("tmpDir"))
+                      .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+                      .schema(
+                          new IncrementalIndexSchema.Builder()
+                              .withMetrics(
+                                  new CountAggregatorFactory("cnt")
+                              )
+                              .withDimensionsSpec(CalciteNestedDataQueryTest.ALL_JSON_COLUMNS.getDimensionsSpec())
+                              .withRollup(false)
+                              .build()
+                      )
+                      .rows(CalciteNestedDataQueryTest.ROWS)
+                      .buildMMappedIndex();
+        } else {
+          throw new ISE("Cannot query segment %s in test runner", segmentId);
+        }
+        break;
+      case CalciteNestedDataQueryTest.DATA_SOURCE_ALL:
+        index = IndexBuilder.create()
+                            .tmpDir(tempFolderProducer.apply("tmpDir"))
+                            .segmentWriteOutMediumFactory(OffHeapMemorySegmentWriteOutMediumFactory.instance())
+                            .schema(
+                                new IncrementalIndexSchema.Builder()
+                                    .withTimestampSpec(NestedDataTestUtils.AUTO_SCHEMA.getTimestampSpec())
+                                    .withDimensionsSpec(NestedDataTestUtils.AUTO_SCHEMA.getDimensionsSpec())
+                                    .withMetrics(
+                                        new CountAggregatorFactory("cnt")
+                                    )
+                                    .withRollup(false)
+                                    .build()
+                            )
+                            .inputSource(
+                                ResourceInputSource.of(
+                                    NestedDataTestUtils.class.getClassLoader(),
+                                    NestedDataTestUtils.ALL_TYPES_TEST_DATA_FILE
+                                )
+                            )
+                            .inputFormat(TestDataBuilder.DEFAULT_JSON_INPUT_FORMAT)
+                            .inputTmpDir(tempFolderProducer.apply("tmpDir"))
+                            .buildMMappedIndex();
+        break;
       case CalciteTests.WIKIPEDIA_FIRST_LAST:
         index = TestDataBuilder.makeWikipediaIndexWithAggregation(tempFolderProducer.apply("tmpDir"));
         break;

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

@@ -0,0 +1,95 @@
+/*
+ * 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.test;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
+import com.google.inject.Injector;
+import com.google.inject.Module;
+import org.apache.druid.guice.DruidInjectorBuilder;
+import org.apache.druid.msq.exec.WorkerMemoryParameters;
+import org.apache.druid.msq.sql.MSQTaskSqlEngine;
+import org.apache.druid.query.groupby.TestGroupByBuffers;
+import org.apache.druid.server.QueryLifecycleFactory;
+import org.apache.druid.sql.calcite.CalciteNestedDataQueryTest;
+import org.apache.druid.sql.calcite.QueryTestBuilder;
+import org.apache.druid.sql.calcite.TempDirProducer;
+import org.apache.druid.sql.calcite.run.SqlEngine;
+import org.apache.druid.sql.calcite.util.SqlTestFramework;
+
+/**
+ * Runs {@link CalciteNestedDataQueryTest} but with MSQ engine
+ */
+@SqlTestFramework.SqlTestFrameWorkModule(CalciteNestedDataQueryMSQTest.NestedDataQueryMSQComponentSupplier.class)
+public class CalciteNestedDataQueryMSQTest extends CalciteNestedDataQueryTest
+{
+
+  public static class NestedDataQueryMSQComponentSupplier extends NestedComponentSupplier
+  {
+    public NestedDataQueryMSQComponentSupplier(TempDirProducer tempFolderProducer)
+    {
+      super(tempFolderProducer);
+    }
+
+    @Override
+    public void configureGuice(DruidInjectorBuilder builder)
+    {
+      super.configureGuice(builder);
+      builder.addModules(
+          CalciteMSQTestsHelper.fetchModules(tempDirProducer::newTempFolder, TestGroupByBuffers.createDefault()).toArray(new Module[0])
+      );
+    }
+
+    @Override
+    public SqlEngine createEngine(
+        QueryLifecycleFactory qlf,
+        ObjectMapper queryJsonMapper,
+        Injector injector
+    )
+    {
+      final WorkerMemoryParameters workerMemoryParameters =
+          WorkerMemoryParameters.createInstance(
+              WorkerMemoryParameters.PROCESSING_MINIMUM_BYTES * 50,
+              2,
+              10,
+              2,
+              0,
+              0
+          );
+      final MSQTestOverlordServiceClient indexingServiceClient = new MSQTestOverlordServiceClient(
+          queryJsonMapper,
+          injector,
+          new MSQTestTaskActionClient(queryJsonMapper, injector),
+          workerMemoryParameters,
+          ImmutableList.of()
+      );
+      return new MSQTaskSqlEngine(indexingServiceClient, queryJsonMapper);
+    }
+  }
+
+  @Override
+  protected QueryTestBuilder testBuilder()
+  {
+    return new QueryTestBuilder(new CalciteTestConfig(true))
+        .addCustomRunner(new ExtractResultsFactory(() -> (MSQTestOverlordServiceClient) ((MSQTaskSqlEngine) queryFramework().engine()).overlordClient()))
+        .skipVectorize(true)
+        .verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate());
+  }
+}

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

@@ -1202,7 +1202,7 @@ public void verifyResults()
           // assert rollup
           Assert.assertEquals(expectedRollUp, queryableIndex.getMetadata().isRollup());
 
-          // asset query granularity
+          // assert query granularity
           Assert.assertEquals(expectedQueryGranularity, queryableIndex.getMetadata().getQueryGranularity());
 
           // assert aggregator factories
@@ -1467,9 +1467,9 @@ public Pair<MSQSpec, Pair<List<MSQResultsReport.ColumnAndType>, List<Object[]>>>
     public void verifyResults()
     {
       if (expectedMSQFault == null) {
-//        Preconditions.checkArgument(expectedResultRows != null, "Result rows cannot be null");
-//        Preconditions.checkArgument(expectedRowSignature != null, "Row signature cannot be null");
-//        Preconditions.checkArgument(expectedMSQSpec != null, "MultiStageQuery Query spec cannot be null ");
+        Preconditions.checkArgument(expectedResultRows != null, "Result rows cannot be null");
+        Preconditions.checkArgument(expectedRowSignature != null, "Row signature cannot be null");
+        Preconditions.checkArgument(expectedMSQSpec != null, "MultiStageQuery Query spec cannot be null ");
       }
       Pair<MSQSpec, Pair<List<MSQResultsReport.ColumnAndType>, List<Object[]>>> specAndResults = runQueryWithResult();