Fix time-ordered scan queries on realtime segments

processing/src/main/java/org/apache/druid/query/SinkQueryRunners.java

@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query;
+
+import org.apache.druid.java.util.common.Pair;
+import org.joda.time.Interval;
+
+import java.util.Iterator;
+
+public class SinkQueryRunners<T> implements Iterable<QueryRunner<T>>
+{
+  Iterable<Pair<Interval, QueryRunner<T>>> runners;
+
+  public SinkQueryRunners(Iterable<Pair<Interval, QueryRunner<T>>> runners)
+  {
+    this.runners = runners;
+  }
+
+  public Iterator<Pair<Interval, QueryRunner<T>>> runnerIntervalMappingIterator()
+  {
+    return runners.iterator();
+  }
+
+  @Override
+  public Iterator<QueryRunner<T>> iterator()
+  {
+    Iterator<Pair<Interval, QueryRunner<T>>> runnerIntervalIterator = runners.iterator();
+    return new Iterator<QueryRunner<T>>()
+    {
+      @Override
+      public boolean hasNext()
+      {
+        return runnerIntervalIterator.hasNext();
+      }
+
+      @Override
+      public QueryRunner<T> next()
+      {
+        return runnerIntervalIterator.next().rhs;
+      }
+    };
+  }
+}

processing/src/main/java/org/apache/druid/query/scan/ScanQueryRunnerFactory.java

@@ -20,7 +20,6 @@
 package org.apache.druid.query.scan;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Ordering;
 import com.google.inject.Inject;
@@ -39,6 +38,7 @@
 import org.apache.druid.query.QueryRunnerFactory;
 import org.apache.druid.query.QueryToolChest;
 import org.apache.druid.query.SegmentDescriptor;
+import org.apache.druid.query.SinkQueryRunners;
 import org.apache.druid.query.spec.MultipleSpecificSegmentSpec;
 import org.apache.druid.query.spec.QuerySegmentSpec;
 import org.apache.druid.query.spec.SpecificSegmentSpec;
@@ -114,11 +114,11 @@ public QueryRunner<ScanResultValue> mergeRunners(
           return returnedRows;
         }
       } else {
-        List<SegmentDescriptor> descriptorsOrdered = getSegmentDescriptorsFromSpecificQuerySpec(query.getQuerySegmentSpec());
+        List<Interval> intervalsOrdered = getIntervalsFromSpecificQuerySpec(query.getQuerySegmentSpec());
         List<QueryRunner<ScanResultValue>> queryRunnersOrdered = Lists.newArrayList(queryRunners);
 
         if (query.getOrder().equals(ScanQuery.Order.DESCENDING)) {
-          descriptorsOrdered = Lists.reverse(descriptorsOrdered);
+          intervalsOrdered = Lists.reverse(intervalsOrdered);
           queryRunnersOrdered = Lists.reverse(queryRunnersOrdered);
         }
         int maxRowsQueuedForOrdering = (query.getMaxRowsQueuedForOrdering() == null
@@ -132,28 +132,29 @@ public QueryRunner<ScanResultValue> mergeRunners(
                   input -> input.run(queryPlus, responseContext)
               )),
               query,
-              descriptorsOrdered
+              intervalsOrdered
           );
         } else {
-          Preconditions.checkState(
-              descriptorsOrdered.size() == queryRunnersOrdered.size(),
-              "Number of segment descriptors does not equal number of "
-              + "query runners...something went wrong!"
-          );
-
-          // Combine the two lists of segment descriptors and query runners into a single list of
-          // segment descriptors - query runner pairs.  This makes it easier to use stream operators.
-          List<Pair<SegmentDescriptor, QueryRunner<ScanResultValue>>> descriptorsAndRunnersOrdered = new ArrayList<>();
-          for (int i = 0; i < queryRunnersOrdered.size(); i++) {
-            descriptorsAndRunnersOrdered.add(new Pair<>(descriptorsOrdered.get(i), queryRunnersOrdered.get(i)));
+          // Use n-way merge strategy
+          List<Pair<Interval, QueryRunner<ScanResultValue>>> intervalsAndRunnersOrdered = new ArrayList<>();
+          if (intervalsOrdered.size() == queryRunnersOrdered.size()) {
+            for (int i = 0; i < queryRunnersOrdered.size(); i++) {
+              intervalsAndRunnersOrdered.add(new Pair<>(intervalsOrdered.get(i), queryRunnersOrdered.get(i)));
+            }
+          } else if (queryRunners instanceof SinkQueryRunners) {
+            ((SinkQueryRunners<ScanResultValue>) queryRunners).runnerIntervalMappingIterator()
+                                                              .forEachRemaining(intervalsAndRunnersOrdered::add);
+          } else {
+            throw new ISE("Number of segment descriptors does not equal number of "
+                          + "query runners...something went wrong!");
           }
 
           // Group the list of pairs by interval.  The LinkedHashMap will have an interval paired with a list of all the
           // query runners for that segment
-          LinkedHashMap<Interval, List<Pair<SegmentDescriptor, QueryRunner<ScanResultValue>>>> partitionsGroupedByInterval =
-              descriptorsAndRunnersOrdered.stream()
+          LinkedHashMap<Interval, List<Pair<Interval, QueryRunner<ScanResultValue>>>> partitionsGroupedByInterval =
+              intervalsAndRunnersOrdered.stream()
                                           .collect(Collectors.groupingBy(
-                                              x -> x.lhs.getInterval(),
+                                              x -> x.lhs,
                                               LinkedHashMap::new,
                                               Collectors.toList()
                                           ));
@@ -167,9 +168,9 @@ public QueryRunner<ScanResultValue> mergeRunners(
                                          .max(Comparator.comparing(Integer::valueOf))
                                          .get();
 
-          int segmentPartitionLimit = (query.getMaxSegmentPartitionsOrderedInMemory() == null
-                                       ? scanQueryConfig.getMaxSegmentPartitionsOrderedInMemory()
-                                       : query.getMaxSegmentPartitionsOrderedInMemory());
+          int segmentPartitionLimit = query.getMaxSegmentPartitionsOrderedInMemory() == null
+                                      ? scanQueryConfig.getMaxSegmentPartitionsOrderedInMemory()
+                                      : query.getMaxSegmentPartitionsOrderedInMemory();
           if (maxNumPartitionsInSegment <= segmentPartitionLimit) {
             // Use n-way merge strategy
 
@@ -205,7 +206,7 @@ public QueryRunner<ScanResultValue> mergeRunners(
   Sequence<ScanResultValue> priorityQueueSortAndLimit(
       Sequence<ScanResultValue> inputSequence,
       ScanQuery scanQuery,
-      List<SegmentDescriptor> descriptorsOrdered
+      List<Interval> intervalsOrdered
   )
   {
     Comparator<ScanResultValue> priorityQComparator = new ScanResultValueTimestampComparator(scanQuery);
@@ -254,9 +255,9 @@ public ScanResultValue accumulate(ScanResultValue accumulated, ScanResultValue i
         // Finish scanning the interval containing the limit row
         if (numRowsScanned > limit && finalInterval == null) {
           long timestampOfLimitRow = srv.getFirstEventTimestamp(scanQuery.getResultFormat());
-          for (SegmentDescriptor descriptor : descriptorsOrdered) {
-            if (descriptor.getInterval().contains(timestampOfLimitRow)) {
-              finalInterval = descriptor.getInterval();
+          for (Interval interval : intervalsOrdered) {
+            if (interval.contains(timestampOfLimitRow)) {
+              finalInterval = interval;
             }
           }
           if (finalInterval == null) {
@@ -280,23 +281,28 @@ public ScanResultValue accumulate(ScanResultValue accumulated, ScanResultValue i
   }
 
   @VisibleForTesting
-  List<SegmentDescriptor> getSegmentDescriptorsFromSpecificQuerySpec(QuerySegmentSpec spec)
+  List<Interval> getIntervalsFromSpecificQuerySpec(QuerySegmentSpec spec)
   {
     // Query segment spec must be an instance of MultipleSpecificSegmentSpec or SpecificSegmentSpec because
     // segment descriptors need to be present for a 1:1 matching of intervals with query runners.
     // The other types of segment spec condense the intervals (i.e. merge neighbouring intervals), eliminating
     // the 1:1 relationship between intervals and query runners.
-    List<SegmentDescriptor> descriptorsOrdered;
+    List<Interval> descriptorsOrdered;
 
     if (spec instanceof MultipleSpecificSegmentSpec) {
       // Ascending time order for both descriptors and query runners by default
-      descriptorsOrdered = ((MultipleSpecificSegmentSpec) spec).getDescriptors();
+      descriptorsOrdered = ((MultipleSpecificSegmentSpec) spec).getDescriptors()
+                                                               .stream()
+                                                               .map(SegmentDescriptor::getInterval)
+                                                               .collect(Collectors.toList());
     } else if (spec instanceof SpecificSegmentSpec) {
-      descriptorsOrdered = Collections.singletonList(((SpecificSegmentSpec) spec).getDescriptor());
+      descriptorsOrdered = Collections.singletonList(((SpecificSegmentSpec) spec).getDescriptor().getInterval());
     } else {
-      throw new UOE("Time-ordering on scan queries is only supported for queries with segment specs"
-                    + "of type MultipleSpecificSegmentSpec or SpecificSegmentSpec...a [%s] was received instead.",
-                    spec.getClass().getSimpleName());
+      throw new UOE(
+          "Time-ordering on scan queries is only supported for queries with segment specs"
+          + "of type MultipleSpecificSegmentSpec or SpecificSegmentSpec...a [%s] was received instead.",
+          spec.getClass().getSimpleName()
+      );
     }
     return descriptorsOrdered;
   }

processing/src/main/java/org/apache/druid/query/scan/ScanResultValue.java

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.UOE;
+import org.apache.druid.segment.DimensionHandlerUtils;
 import org.apache.druid.segment.column.ColumnHolder;
 
 import javax.annotation.Nullable;
@@ -79,18 +80,18 @@ public Object getEvents()
   public long getFirstEventTimestamp(ScanQuery.ResultFormat resultFormat)
   {
     if (resultFormat.equals(ScanQuery.ResultFormat.RESULT_FORMAT_LIST)) {
-      Long timestamp = (Long) ((Map<String, Object>) ((List<Object>) this.getEvents()).get(0)).get(ColumnHolder.TIME_COLUMN_NAME);
-      if (timestamp == null) {
+      Object timestampObj = ((Map<String, Object>) ((List<Object>) this.getEvents()).get(0)).get(ColumnHolder.TIME_COLUMN_NAME);
+      if (timestampObj == null) {
         throw new ISE("Unable to compare timestamp for rows without a time column");
       }
-      return timestamp;
+      return DimensionHandlerUtils.convertObjectToLong(timestampObj);
     } else if (resultFormat.equals(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)) {
       int timeColumnIndex = this.getColumns().indexOf(ColumnHolder.TIME_COLUMN_NAME);
       if (timeColumnIndex == -1) {
         throw new ISE("Unable to compare timestamp for rows without a time column");
       }
       List<Object> firstEvent = (List<Object>) ((List<Object>) this.getEvents()).get(0);
-      return (Long) firstEvent.get(timeColumnIndex);
+      return DimensionHandlerUtils.convertObjectToLong(firstEvent.get(timeColumnIndex));
     }
     throw new UOE("Unable to get first event timestamp using result format of [%s]", resultFormat.toString());
   }
@@ -105,7 +106,6 @@ public List<ScanResultValue> toSingleEventScanResultValues()
     return singleEventScanResultValues;
   }
 
-
   @Override
   public boolean equals(Object o)
   {

processing/src/test/java/org/apache/druid/query/scan/ScanQueryRunnerFactoryTest.java

@@ -140,10 +140,10 @@ public void testSortAndLimitScanResultValues()
         List<ScanResultValue> output = factory.priorityQueueSortAndLimit(
             inputSequence,
             query,
-            ImmutableList.of(new SegmentDescriptor(new Interval(
+            ImmutableList.of(new Interval(
                 DateTimes.of("2010-01-01"),
                 DateTimes.of("2019-01-01").plusHours(1)
-            ), "1", 0))
+            ))
         ).toList();
         if (query.getLimit() > Integer.MAX_VALUE) {
           Assert.fail("Unsupported exception should have been thrown due to high limit");
@@ -275,7 +275,7 @@ public static class ScanQueryRunnerFactoryNonParameterizedTest
     ), "1", 0);
 
     @Test
-    public void testGetValidSegmentDescriptorsFromSpec()
+    public void testGetValidIntervalsFromSpec()
     {
       QuerySegmentSpec multiSpecificSpec = new MultipleSpecificSegmentSpec(
           Collections.singletonList(
@@ -284,13 +284,13 @@ public void testGetValidSegmentDescriptorsFromSpec()
       );
       QuerySegmentSpec singleSpecificSpec = new SpecificSegmentSpec(descriptor);
 
-      List<SegmentDescriptor> descriptors = factory.getSegmentDescriptorsFromSpecificQuerySpec(multiSpecificSpec);
-      Assert.assertEquals(1, descriptors.size());
-      Assert.assertEquals(descriptor, descriptors.get(0));
+      List<Interval> intervals = factory.getIntervalsFromSpecificQuerySpec(multiSpecificSpec);
+      Assert.assertEquals(1, intervals.size());
+      Assert.assertEquals(descriptor.getInterval(), intervals.get(0));
 
-      descriptors = factory.getSegmentDescriptorsFromSpecificQuerySpec(singleSpecificSpec);
-      Assert.assertEquals(1, descriptors.size());
-      Assert.assertEquals(descriptor, descriptors.get(0));
+      intervals = factory.getIntervalsFromSpecificQuerySpec(singleSpecificSpec);
+      Assert.assertEquals(1, intervals.size());
+      Assert.assertEquals(descriptor.getInterval(), intervals.get(0));
     }
 
     @Test(expected = UOE.class)
@@ -304,7 +304,7 @@ public void testGetSegmentDescriptorsFromInvalidIntervalSpec()
               )
           )
       );
-      factory.getSegmentDescriptorsFromSpecificQuerySpec(multiIntervalSpec);
+      factory.getIntervalsFromSpecificQuerySpec(multiIntervalSpec);
     }
 
     @Test(expected = UOE.class)
@@ -316,7 +316,7 @@ public void testGetSegmentDescriptorsFromInvalidLegacySpec()
               DateTimes.of("2019-01-01").plusHours(1)
           )
       );
-      factory.getSegmentDescriptorsFromSpecificQuerySpec(legacySpec);
+      factory.getIntervalsFromSpecificQuerySpec(legacySpec);
     }
   }
 }