From ecf312a3d0a4136c7d6af13fddfb2dd1de9e2535 Mon Sep 17 00:00:00 2001
From: Clint Wylie <cwylie@apache.org>
Date: Mon, 10 Apr 2023 20:38:58 -0700
Subject: [PATCH] fix bug in nested v4 format merger from refactoring (#14053)

---
 .../QueryableIndexIndexableAdapter.java       |  5 ++-
 .../druid/query/NestedDataTestUtils.java      | 12 +++----
 .../apache/druid/segment/IndexBuilder.java    | 34 +++++++++++++++----
 3 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/segment/QueryableIndexIndexableAdapter.java b/processing/src/main/java/org/apache/druid/segment/QueryableIndexIndexableAdapter.java
index 28e26c02ded7..cf7b352dbff0 100644
--- a/processing/src/main/java/org/apache/druid/segment/QueryableIndexIndexableAdapter.java
+++ b/processing/src/main/java/org/apache/druid/segment/QueryableIndexIndexableAdapter.java
@@ -35,6 +35,7 @@
 import org.apache.druid.segment.data.ImmutableBitmapValues;
 import org.apache.druid.segment.data.IndexedIterable;
 import org.apache.druid.segment.nested.NestedCommonFormatColumn;
+import org.apache.druid.segment.nested.NestedDataComplexTypeSerde;
 import org.apache.druid.segment.nested.SortedValueDictionary;
 import org.apache.druid.segment.selector.settable.SettableColumnValueSelector;
 import org.apache.druid.segment.selector.settable.SettableLongColumnValueSelector;
@@ -174,7 +175,9 @@ public NestedColumnMergable getNestedColumnMergeables(String columnName)
     if (columnHolder == null) {
       return null;
     }
-    if (!(columnHolder.getColumnFormat() instanceof NestedCommonFormatColumn.Format)) {
+    final ColumnFormat format = columnHolder.getColumnFormat();
+    if (!(format instanceof NestedCommonFormatColumn.Format
+          || format instanceof NestedDataComplexTypeSerde.NestedColumnFormatV4)) {
       return null;
     }
 
diff --git a/processing/src/test/java/org/apache/druid/query/NestedDataTestUtils.java b/processing/src/test/java/org/apache/druid/query/NestedDataTestUtils.java
index c6b77c211cca..9a43c77b9c21 100644
--- a/processing/src/test/java/org/apache/druid/query/NestedDataTestUtils.java
+++ b/processing/src/test/java/org/apache/druid/query/NestedDataTestUtils.java
@@ -40,10 +40,10 @@
 import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
 import org.apache.druid.query.expression.TestExprMacroTable;
-import org.apache.druid.segment.AutoTypeColumnSchema;
 import org.apache.druid.segment.IncrementalIndexSegment;
 import org.apache.druid.segment.IndexBuilder;
 import org.apache.druid.segment.IndexSpec;
+import org.apache.druid.segment.NestedDataDimensionSchema;
 import org.apache.druid.segment.QueryableIndexSegment;
 import org.apache.druid.segment.Segment;
 import org.apache.druid.segment.TestHelper;
@@ -96,11 +96,11 @@ public class NestedDataTestUtils
       DimensionsSpec.builder()
                     .setDimensions(
                         Arrays.asList(
-                            new AutoTypeColumnSchema("dim"),
-                            new AutoTypeColumnSchema("nest_json"),
-                            new AutoTypeColumnSchema("nester_json"),
-                            new AutoTypeColumnSchema("variant_json"),
-                            new AutoTypeColumnSchema("list_json")
+                            new NestedDataDimensionSchema("dim"),
+                            new NestedDataDimensionSchema("nest_json"),
+                            new NestedDataDimensionSchema("nester_json"),
+                            new NestedDataDimensionSchema("variant_json"),
+                            new NestedDataDimensionSchema("list_json")
                         )
                     )
                     .build();
diff --git a/processing/src/test/java/org/apache/druid/segment/IndexBuilder.java b/processing/src/test/java/org/apache/druid/segment/IndexBuilder.java
index 90a4411c87b1..95a863dbb017 100644
--- a/processing/src/test/java/org/apache/druid/segment/IndexBuilder.java
+++ b/processing/src/test/java/org/apache/druid/segment/IndexBuilder.java
@@ -53,6 +53,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.ThreadLocalRandom;
@@ -222,15 +223,34 @@ public QueryableIndex buildMMappedIndex()
     Preconditions.checkNotNull(indexMerger, "indexMerger");
     Preconditions.checkNotNull(tmpDir, "tmpDir");
     try (final IncrementalIndex incrementalIndex = buildIncrementalIndex()) {
+      List<IndexableAdapter> adapters = Collections.singletonList(
+          new QueryableIndexIndexableAdapter(
+              indexIO.loadIndex(
+                  indexMerger.persist(
+                      incrementalIndex,
+                      new File(
+                          tmpDir,
+                          StringUtils.format("testIndex-%s", ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE))
+                      ),
+                      indexSpec,
+                      null
+                  )
+              )
+          )
+      );
+      // Do a 'merge' of the persisted segment even though there is only one; this time it will be reading from the
+      // queryable index instead of the incremental index, which also mimics the behavior of real ingestion tasks
+      // which persist incremental indexes as intermediate segments and then merges all the intermediate segments to
+      // publish
       return indexIO.loadIndex(
-          indexMerger.persist(
-              incrementalIndex,
-              new File(
-                  tmpDir,
-                  StringUtils.format("testIndex-%s", ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE))
-              ),
+          indexMerger.merge(
+              adapters,
+              schema.isRollup(),
+              schema.getMetrics(),
+              tmpDir,
+              schema.getDimensionsSpec(),
               indexSpec,
-              null
+              Integer.MAX_VALUE
           )
       );
     }