From bd9c04244ffbf3e9ca07c7c8885e619f1c072c8f Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Thu, 4 Feb 2016 17:35:47 -0800
Subject: [PATCH] SegmentMetadataQuery: Retain segment id when merging, if
 possible.

This is helpful on realtime nodes, where two analyses from two different hydrants
are merged together but they are actually from the same segment.
---
 .../SegmentMetadataQueryQueryToolChest.java   | 10 ++++-
 .../metadata/SegmentMetadataQueryTest.java    | 44 ++++++++++++-------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/processing/src/main/java/io/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java b/processing/src/main/java/io/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
index 1e640dacb30..d6981b3b025 100644
--- a/processing/src/main/java/io/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
+++ b/processing/src/main/java/io/druid/query/metadata/SegmentMetadataQueryQueryToolChest.java
@@ -329,8 +329,16 @@ public class SegmentMetadataQueryQueryToolChest extends QueryToolChest<SegmentAn
       }
     }
 
+    final String mergedId;
+
+    if (arg1.getId() != null && arg2.getId() != null && arg1.getId().equals(arg2.getId())) {
+      mergedId = arg1.getId();
+    } else {
+      mergedId = "merged";
+    }
+
     return new SegmentAnalysis(
-        "merged",
+        mergedId,
         newIntervals,
         columns,
         arg1.getSize() + arg2.getSize(),
diff --git a/processing/src/test/java/io/druid/query/metadata/SegmentMetadataQueryTest.java b/processing/src/test/java/io/druid/query/metadata/SegmentMetadataQueryTest.java
index 4aec1c0eca2..3b199f15edb 100644
--- a/processing/src/test/java/io/druid/query/metadata/SegmentMetadataQueryTest.java
+++ b/processing/src/test/java/io/druid/query/metadata/SegmentMetadataQueryTest.java
@@ -75,23 +75,27 @@ public class SegmentMetadataQueryTest
 
   @SuppressWarnings("unchecked")
   public static QueryRunner makeMMappedQueryRunner(
+      String segmentId,
       QueryRunnerFactory factory
   )
   {
     return QueryRunnerTestHelper.makeQueryRunner(
         factory,
-        new QueryableIndexSegment(QueryRunnerTestHelper.segmentId, TestIndex.getMMappedTestIndex())
+        segmentId,
+        new QueryableIndexSegment(segmentId, TestIndex.getMMappedTestIndex())
     );
   }
 
   @SuppressWarnings("unchecked")
   public static QueryRunner makeIncrementalIndexQueryRunner(
+      String segmentId,
       QueryRunnerFactory factory
   )
   {
     return QueryRunnerTestHelper.makeQueryRunner(
         factory,
-        new IncrementalIndexSegment(TestIndex.getIncrementalTestIndex(), QueryRunnerTestHelper.segmentId)
+        segmentId,
+        new IncrementalIndexSegment(TestIndex.getIncrementalTestIndex(), segmentId)
     );
   }
 
@@ -99,30 +103,36 @@ public class SegmentMetadataQueryTest
   private final QueryRunner runner2;
   private final boolean mmap1;
   private final boolean mmap2;
+  private final boolean differentIds;
   private final SegmentMetadataQuery testQuery;
   private final SegmentAnalysis expectedSegmentAnalysis1;
   private final SegmentAnalysis expectedSegmentAnalysis2;
 
-  @Parameterized.Parameters(name = "mmap1 = {0}, mmap2 = {1}")
+  @Parameterized.Parameters(name = "mmap1 = {0}, mmap2 = {1}, differentIds = {2}")
   public static Collection<Object[]> constructorFeeder()
   {
     return ImmutableList.of(
-        new Object[]{true, true},
-        new Object[]{true, false},
-        new Object[]{false, true},
-        new Object[]{false, false}
+        new Object[]{true, true, false},
+        new Object[]{true, false, false},
+        new Object[]{false, true, false},
+        new Object[]{false, false, false},
+        new Object[]{false, false, true}
     );
   }
 
   public SegmentMetadataQueryTest(
       boolean mmap1,
-      boolean mmap2
+      boolean mmap2,
+      boolean differentIds
   )
   {
-    this.runner1 = mmap1 ? makeMMappedQueryRunner(FACTORY) : makeIncrementalIndexQueryRunner(FACTORY);
-    this.runner2 = mmap2 ? makeMMappedQueryRunner(FACTORY) : makeIncrementalIndexQueryRunner(FACTORY);
+    final String id1 = differentIds ? "testSegment1" : "testSegment";
+    final String id2 = differentIds ? "testSegment2" : "testSegment";
+    this.runner1 = mmap1 ? makeMMappedQueryRunner(id1, FACTORY) : makeIncrementalIndexQueryRunner(id1, FACTORY);
+    this.runner2 = mmap2 ? makeMMappedQueryRunner(id2, FACTORY) : makeIncrementalIndexQueryRunner(id2, FACTORY);
     this.mmap1 = mmap1;
     this.mmap2 = mmap2;
+    this.differentIds = differentIds;
     testQuery = Druids.newSegmentMetadataQueryBuilder()
                       .dataSource("testing")
                       .intervals("2013/2014")
@@ -132,7 +142,7 @@ public class SegmentMetadataQueryTest
                       .build();
 
     expectedSegmentAnalysis1 = new SegmentAnalysis(
-        "testSegment",
+        id1,
         ImmutableList.of(
             new Interval("2011-01-12T00:00:00.000Z/2011-04-15T00:00:00.001Z")
         ),
@@ -166,7 +176,7 @@ public class SegmentMetadataQueryTest
         null
     );
     expectedSegmentAnalysis2 = new SegmentAnalysis(
-        "testSegment",
+        id2,
         ImmutableList.of(
             new Interval("2011-01-12T00:00:00.000Z/2011-04-15T00:00:00.001Z")
         ),
@@ -217,7 +227,7 @@ public class SegmentMetadataQueryTest
   public void testSegmentMetadataQueryWithHasMultipleValuesMerge()
   {
     SegmentAnalysis mergedSegmentAnalysis = new SegmentAnalysis(
-        "merged",
+        differentIds ? "merged" : "testSegment",
         null,
         ImmutableMap.of(
             "placement",
@@ -279,7 +289,7 @@ public class SegmentMetadataQueryTest
   public void testSegmentMetadataQueryWithComplexColumnMerge()
   {
     SegmentAnalysis mergedSegmentAnalysis = new SegmentAnalysis(
-        "merged",
+        differentIds ? "merged" : "testSegment",
         null,
         ImmutableMap.of(
             "placement",
@@ -341,7 +351,7 @@ public class SegmentMetadataQueryTest
   public void testSegmentMetadataQueryWithDefaultAnalysisMerge()
   {
     SegmentAnalysis mergedSegmentAnalysis = new SegmentAnalysis(
-        "merged",
+        differentIds ? "merged" : "testSegment",
         ImmutableList.of(expectedSegmentAnalysis1.getIntervals().get(0)),
         ImmutableMap.of(
             "__time",
@@ -405,7 +415,7 @@ public class SegmentMetadataQueryTest
   public void testSegmentMetadataQueryWithNoAnalysisTypesMerge()
   {
     SegmentAnalysis mergedSegmentAnalysis = new SegmentAnalysis(
-        "merged",
+        differentIds ? "merged" : "testSegment",
         null,
         ImmutableMap.of(
             "placement",
@@ -463,7 +473,7 @@ public class SegmentMetadataQueryTest
       expectedAggregators.put(agg.getName(), agg.getCombiningFactory());
     }
     SegmentAnalysis mergedSegmentAnalysis = new SegmentAnalysis(
-        "merged",
+        differentIds ? "merged" : "testSegment",
         null,
         ImmutableMap.of(
             "placement",