From a9021e4cd7d65f524ec21fbce615f55dc16de3a9 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Wed, 4 Oct 2023 02:41:41 -0700
Subject: [PATCH] Fix NPE with lenient aggregators merging in segmentMetadata.
 (#15078)

When merging analyses, lenient merging sets unmergeable aggregators
to null. Merging such a null aggregator record into a nonnull record
would potentially lead to NPE in getMergingFactory.

The new code only calls getMergingFactory if both the old and new
aggregators are nonnull; else, if either is null, then the merged
aggregator is also set to null.
---
 .../SegmentMetadataQueryQueryToolChest.java   | 25 ++++++++----
 ...egmentMetadataQueryQueryToolChestTest.java | 38 +++++++++++++++++++
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java b/processing/src/main/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
index 4bb73d845ec..912ecb1ac32 100644
--- a/processing/src/main/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
+++ b/processing/src/main/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
@@ -333,17 +333,26 @@ public class SegmentMetadataQueryQueryToolChest extends QueryToolChest<SegmentAn
           for (Map.Entry<String, AggregatorFactory> entry : analysis.getAggregators().entrySet()) {
             final String aggregatorName = entry.getKey();
             final AggregatorFactory aggregator = entry.getValue();
-            AggregatorFactory merged = aggregators.get(aggregatorName);
-            if (merged != null) {
-              try {
-                merged = merged.getMergingFactory(aggregator);
-              }
-              catch (AggregatorFactoryNotMergeableException e) {
+            final boolean isMergedYet = aggregators.containsKey(aggregatorName);
+            AggregatorFactory merged;
+
+            if (!isMergedYet) {
+              merged = aggregator;
+            } else {
+              merged = aggregators.get(aggregatorName);
+
+              if (merged != null && aggregator != null) {
+                try {
+                  merged = merged.getMergingFactory(aggregator);
+                }
+                catch (AggregatorFactoryNotMergeableException e) {
+                  merged = null;
+                }
+              } else {
                 merged = null;
               }
-            } else {
-              merged = aggregator;
             }
+
             aggregators.put(aggregatorName, merged);
           }
         }
diff --git a/processing/src/test/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChestTest.java b/processing/src/test/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChestTest.java
index 1926e14eae5..4afb202b433 100644
--- a/processing/src/test/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChestTest.java
+++ b/processing/src/test/java/org/apache/druid/query/metadata/SegmentMetadataQueryQueryToolChestTest.java
@@ -634,6 +634,44 @@ public class SegmentMetadataQueryQueryToolChestTest
         )
     );
 
+    // Simulate multi-level lenient merge (unmerged first)
+    Assert.assertEquals(
+        new SegmentAnalysis(
+            "dummy_2021-01-01T00:00:00.000Z_2021-01-02T00:00:00.000Z_merged",
+            null,
+            new LinkedHashMap<>(),
+            0,
+            0,
+            expectedLenient,
+            null,
+            null,
+            null
+        ),
+        mergeLenient(
+            analysis1,
+            mergeLenient(analysis1, analysis2)
+        )
+    );
+
+    // Simulate multi-level lenient merge (unmerged second)
+    Assert.assertEquals(
+        new SegmentAnalysis(
+            "dummy_2021-01-01T00:00:00.000Z_2021-01-02T00:00:00.000Z_merged",
+            null,
+            new LinkedHashMap<>(),
+            0,
+            0,
+            expectedLenient,
+            null,
+            null,
+            null
+        ),
+        mergeLenient(
+            mergeLenient(analysis1, analysis2),
+            analysis1
+        )
+    );
+
     Assert.assertEquals(
         new SegmentAnalysis(
             "dummy_2021-01-01T00:00:00.000Z_2021-01-02T00:00:00.000Z_merged",