test comments

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

@@ -27,12 +27,20 @@
 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.indexing.MSQSpec;
+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.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.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;
@@ -45,6 +53,7 @@
 import java.nio.file.StandardCopyOption;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 
 public class MSQComplexGroupByTest extends MSQTestBase
 {
@@ -119,10 +128,106 @@ public void testInsertWithRollupOnNestedData()
                      .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())
-                     .setExpectedResultRows(ImmutableList.of())
+                     .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
+                         }
+                     ))
                      .verifyResults();
 
   }
@@ -142,11 +247,34 @@ public void testSortingOnNestedData()
                              + " )\n"
                              + " ORDER BY 1")
                      .setQueryContext(ImmutableMap.of())
+                     .setExpectedMSQSpec(MSQSpec
+                                             .builder()
+                                             .query(newScanQueryBuilder()
+                                                        .dataSource(CalciteTests.DATASOURCE1)
+                                                        .intervals(querySegmentSpec(Filtration.eternity()))
+                                                        .columns("cnt", "dim1")
+                                                        .build()
+                                             )
+                                             .columnMappings(new ColumnMappings(ImmutableList.of(
+                                                 new ColumnMapping("obj", "obj")
+                                             )))
+                                             .tuningConfig(MSQTuningConfig.defaultConfig())
+                                             .destination(TaskReportMSQDestination.INSTANCE)
+                                             .build()
+                     )
                      .setExpectedRowSignature(RowSignature.builder()
                                                           .add("obj", ColumnType.NESTED_DATA)
                                                           .build())
-                     .setExpectedResultRows(ImmutableList.of())
+                     .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\":[]}"},
+                         new Object[]{"{\"a\":100,\"b\":{\"x\":\"a\",\"y\":1.1,\"z\":[1,2,3,4]},\"v\":[]}"},
+                         new Object[]{"{\"a\":700,\"b\":{\"x\":\"g\",\"y\":1.1,\"z\":[9,null,9,9]},\"v\":[]}"},
+                         new Object[]{"{\"a\":200,\"b\":{\"x\":\"b\",\"y\":1.1,\"z\":[2,4,6]},\"v\":[]}"},
+                         new Object[]{"{\"a\":600,\"b\":{\"x\":\"f\",\"y\":1.1,\"z\":[6,7,8,9]},\"v\":\"b\"}"},
+                         new Object[]{"{\"a\":400,\"b\":{\"x\":\"d\",\"y\":1.1,\"z\":[3,4]},\"v\":[]}"},
+                         new Object[]{"{\"a\":300}"}
+                     ))
                      .verifyResults();
   }
-
 }

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

@@ -1234,7 +1234,7 @@ public void verifyResults()
                                                                        .collect(Collectors.toList());
 
         log.info(
-            "Found rows which are sorted forcefully %s",
+            "Found rows which are sorted forcefully\n%s",
             transformedOutputRows.stream().map(Arrays::deepToString).collect(Collectors.joining("\n"))
         );
 

processing/src/main/java/org/apache/druid/frame/key/ByteRowKeyComparator.java

@@ -110,6 +110,12 @@ public int compare(final byte[] keyArray1, final byte[] keyArray2)
         continue;
       }
 
+      // Index of the next field that will get considered. Excludes the last field of the current run length that is being
+      // compared in this iteration
+      final int nextField = fieldsComparedTillNow + runLengthEntry.getRunLength();
+      final int currentRunEndPosition1 = RowKeyReader.fieldEndPosition(keyArray1, nextField - 1);
+      final int currentRunEndPosition2 = RowKeyReader.fieldEndPosition(keyArray2, nextField - 1);
+
       if (!runLengthEntry.isByteComparable()) {
         // Only complex types are not byte comparable. Nested arrays aren't supported in MSQ
         assert runLengthEntry.getRunLength() == 1;
@@ -128,10 +134,6 @@ public int compare(final byte[] keyArray1, final byte[] keyArray2)
             complexTypeName
         );
 
-        // Index of the next field that will get considered. Excludes the current field that we are comparing right now
-        final int nextField = fieldsComparedTillNow + 1;
-        final int currentRunEndPosition1 = RowKeyReader.fieldEndPosition(keyArray1, nextField - 1);
-        final int currentRunEndPosition2 = RowKeyReader.fieldEndPosition(keyArray2, nextField - 1);
 
         int cmp = FrameReaderUtils.compareComplexTypes(
             keyArray1,

processing/src/main/java/org/apache/druid/frame/key/FrameComparisonWidgetImpl.java

@@ -37,6 +37,7 @@
 import org.apache.druid.segment.serde.ComplexMetricSerde;
 import org.apache.druid.segment.serde.ComplexMetrics;
 
+import javax.annotation.concurrent.NotThreadSafe;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -47,6 +48,7 @@
  * Comparison logic in this class is very similar to {@link RowKeyComparator}, but is different because it works
  * on Frames instead of byte[].
  */
+@NotThreadSafe
 public class FrameComparisonWidgetImpl implements FrameComparisonWidget
 {
   private final Frame frame;
@@ -58,6 +60,9 @@ public class FrameComparisonWidgetImpl implements FrameComparisonWidget
   private final List<FieldReader> keyFieldReaders;
   private final int firstFieldPosition;
   private final RowKeyComparisonRunLengths rowKeyComparisonRunLengths;
+  // We memoize the serde instead of fetching it everytime from the global map, since that is thread-safe and is guarded by a
+  // ConcurrentHashMap, while we only access FrameComparisonWidget from a single thread.
+  private final Map<String, ComplexMetricSerde> serdeMap = new HashMap<>();
 
   private FrameComparisonWidgetImpl(
       final Frame frame,
@@ -288,9 +293,6 @@ public int compare(final int row, final FrameComparisonWidget otherWidget, final
     // Number of fields compared till now, which is equivalent to the index of the field to compare next
     int fieldsComparedTillNow = 0;
 
-    // Memoize the serde instead of calling ComplexMetrics.getSerdeForType everytime
-    final Map<String, ComplexMetricSerde> serdeMap = new HashMap<>();
-
     for (RowKeyComparisonRunLengths.RunLengthEntry runLengthEntry : rowKeyComparisonRunLengths.getRunLengthEntries()) {
 
       if (runLengthEntry.getRunLength() <= 0) {
@@ -322,10 +324,11 @@ public int compare(final int row, final FrameComparisonWidget otherWidget, final
         Preconditions.checkArgument(columnType1.equals(columnType2), "Different complex types cannot be compared");
 
         // Use serde for the current implementation.
-        ComplexMetricSerde serde = Preconditions.checkNotNull(
-            ComplexMetrics.getSerdeForType(complexTypeName),
-            "serde for type [%s] not present",
-            complexTypeName
+        ComplexMetricSerde serde = serdeMap.computeIfAbsent(
+            complexTypeName, name ->
+                Preconditions.checkNotNull(
+                    ComplexMetrics.getSerdeForType(name), "serde for type [%s] not present", complexTypeName
+                )
         );
 
         int nextField = fieldsComparedTillNow + 1;