union query fix and a broken test

common/src/test/java/io/druid/collections/OrderedMergeSequenceTest.java

@@ -23,6 +23,7 @@ import com.google.common.base.Function;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Ordering;
 import com.metamx.common.guava.BaseSequence;
+import com.metamx.common.guava.MergeSequence;
 import com.metamx.common.guava.Sequence;
 import com.metamx.common.guava.SequenceTestHelper;
 import com.metamx.common.guava.Sequences;
@@ -230,17 +231,65 @@ public class OrderedMergeSequenceTest
         ordering,
         Sequences.simple(
             Lists.transform( // OMG WTF, the java type system really annoys me at times...
-                seqs,
+                             seqs,
-                new Function<TestSequence<T>, Sequence<T>>()
+                             new Function<TestSequence<T>, Sequence<T>>()
-                {
+                             {
-                  @Override
+                               @Override
-                  public Sequence<T> apply(@Nullable TestSequence<T> input)
+                               public Sequence<T> apply(@Nullable TestSequence<T> input)
-                  {
+                               {
-                    return input;
+                                 return input;
-                  }
+                               }
-                }
+                             }
             )
         )
     );
   }
+
+  private <T> MergeSequence<T> makeUnorderedMergedSequence(
+      Ordering<T> ordering,
+      List<TestSequence<T>> seqs
+  )
+  {
+    return new MergeSequence<T>(
+        ordering,
+        Sequences.simple(
+            Lists.transform( // OMG WTF, the java type system really annoys me at times...
+                             seqs,
+                             new Function<TestSequence<T>, Sequence<T>>()
+                             {
+                               @Override
+                               public Sequence<T> apply(@Nullable TestSequence<T> input)
+                               {
+                                 return input;
+                               }
+                             }
+            )
+        )
+    );
+  }
+
+  @Test
+  public void testHierarchicalMerge() throws Exception
+  {
+    final Sequence<Integer> seq1 = makeUnorderedMergedSequence(
+        Ordering.<Integer>natural(), Lists.newArrayList(
+        TestSequence.create(1)
+    )
+    );
+
+
+    final Sequence<Integer> seq2 = makeUnorderedMergedSequence(
+        Ordering.<Integer>natural(), Lists.newArrayList(
+        TestSequence.create(1)
+    )
+    );
+    final OrderedMergeSequence<Integer> finalMerged = new OrderedMergeSequence<Integer>(
+        Ordering.<Integer>natural(),
+        Sequences.simple(
+            Lists.<Sequence<Integer>>newArrayList(seq1, seq2)
+        )
+    );
+
+    SequenceTestHelper.testAll(finalMerged, Arrays.asList(1, 1));
+  }
 }

processing/src/main/java/io/druid/query/UnionQueryRunner.java

@@ -22,17 +22,24 @@
 package io.druid.query;
 
 import com.google.common.base.Function;
-import com.google.common.collect.Iterables;
+import com.google.common.collect.Lists;
 import com.metamx.common.guava.Sequence;
 import com.metamx.common.guava.Sequences;
 
+import java.util.ArrayList;
+
 public class UnionQueryRunner<T> implements QueryRunner<T>
 {
   private final QueryRunner<T> baseRunner;
+  private final QueryToolChest<T, Query<T>> toolChest;
 
-  public UnionQueryRunner(QueryRunner<T> baseRunner)
+  public UnionQueryRunner(
+      QueryRunner<T> baseRunner,
+      QueryToolChest<T, Query<T>> toolChest
+  )
   {
     this.baseRunner = baseRunner;
+    this.toolChest = toolChest;
   }
 
   @Override
@@ -40,19 +47,21 @@ public class UnionQueryRunner<T> implements QueryRunner<T>
   {
     DataSource dataSource = query.getDataSource();
     if (dataSource instanceof UnionDataSource) {
-      return Sequences.concat(
+      return toolChest.mergeSequences(
-          Iterables.transform(
+          Sequences.simple(
-              ((UnionDataSource) dataSource).getDataSources(),
+              Lists.transform(
-              new Function<DataSource, Sequence<T>>()
+                  ((UnionDataSource) dataSource).getDataSources(),
-              {
+                  new Function<DataSource, Sequence<T>>()
-                @Override
+                  {
-                public Sequence<T> apply(DataSource singleSource)
+                    @Override
-                {
+                    public Sequence<T> apply(DataSource singleSource)
-                  return baseRunner.run(
+                    {
-                      query.withDataSource(singleSource)
+                      return baseRunner.run(
-                  );
+                          query.withDataSource(singleSource)
-                }
+                      );
-              }
+                    }
+                  }
+              )
           )
       );
     } else {

server/src/main/java/io/druid/server/ClientQuerySegmentWalker.java

@@ -87,7 +87,8 @@ public class ClientQuerySegmentWalker implements QuerySegmentWalker
                           }
                         },
                         toolChest.preMergeQueryDecoration(baseClient)
-                    ).withWaitMeasuredFromNow()
+                    ).withWaitMeasuredFromNow(),
+                    toolChest
                 )
             )
         ),