Modified sequence limit to accept longs and added test for long limits

core/src/main/java/org/apache/druid/java/util/common/guava/LimitedSequence.java

@@ -30,11 +30,11 @@ import java.io.IOException;
 final class LimitedSequence<T> extends YieldingSequenceBase<T>
 {
   private final Sequence<T> baseSequence;
-  private final int limit;
+  private final long limit;
 
   LimitedSequence(
       Sequence<T> baseSequence,
-      int limit
+      long limit
   )
   {
     Preconditions.checkNotNull(baseSequence);
@@ -106,7 +106,7 @@ final class LimitedSequence<T> extends YieldingSequenceBase<T>
 
   private class LimitedYieldingAccumulator<OutType, T> extends DelegatingYieldingAccumulator<OutType, T>
   {
-    int count;
+    long count;
     boolean interruptYield = false;
 
     LimitedYieldingAccumulator(YieldingAccumulator<OutType, T> accumulator)

core/src/main/java/org/apache/druid/java/util/common/guava/Sequence.java

@@ -82,7 +82,7 @@ public interface Sequence<T>
     return accumulate(new ArrayList<>(), Accumulators.list());
   }
 
-  default Sequence<T> limit(int limit)
+  default Sequence<T> limit(long limit)
   {
     return new LimitedSequence<>(this, limit);
   }

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

@@ -207,6 +207,12 @@ public class ScanQueryRunnerFactory implements QueryRunnerFactory<ScanResultValu
   {
     Comparator<ScanResultValue> priorityQComparator = new ScanResultValueTimestampComparator(scanQuery);
 
+    if (scanQuery.getLimit() > Integer.MAX_VALUE) {
+      throw new UOE("Limit of %,d rows not supported for priority queue strategy of time-ordering scan results",
+                    scanQuery.getLimit()
+      );
+    }
+
     // Converting the limit from long to int could theoretically throw an ArithmeticException but this branch
     // only runs if limit < MAX_LIMIT_FOR_IN_MEMORY_TIME_ORDERING (which should be < Integer.MAX_VALUE)
     int limit = Math.toIntExact(scanQuery.getLimit());
@@ -282,7 +288,6 @@ public class ScanQueryRunnerFactory implements QueryRunnerFactory<ScanResultValu
     // (3) Create a sequence of results from each runner in the group and flatmerge based on timestamp
     // (4) Create a sequence of results from each runner group
     // (5) Join all the results into a single sequence
-
     return Sequences.concat(
         Sequences.map(
             Sequences.simple(groupedRunners),
@@ -303,7 +308,7 @@ public class ScanQueryRunnerFactory implements QueryRunnerFactory<ScanResultValu
                 )
         )
     ).limit(
-        Math.toIntExact(((ScanQuery) (queryPlus.getQuery())).getLimit())
+        ((ScanQuery) (queryPlus.getQuery())).getLimit()
     );
   }
 

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

@@ -22,6 +22,7 @@ package org.apache.druid.query.scan;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.java.util.common.UOE;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Sequences;
 import org.apache.druid.query.DefaultGenericQueryMetricsFactory;
@@ -83,7 +84,7 @@ public class ScanQueryRunnerFactoryTest
   {
     List<Integer> numsElements = ImmutableList.of(0, 10, 100);
     List<Integer> batchSizes = ImmutableList.of(1, 100);
-    List<Long> limits = ImmutableList.of(3L, 1000L);
+    List<Long> limits = ImmutableList.of(3L, 1000L, Long.MAX_VALUE);
     List<ScanQuery.ResultFormat> resultFormats = ImmutableList.of(
         ScanQuery.ResultFormat.RESULT_FORMAT_LIST,
         ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST
@@ -125,17 +126,25 @@ public class ScanQueryRunnerFactoryTest
       return retVal;
     });
     Sequence<ScanResultValue> inputSequence = Sequences.simple(srvs);
-    List<ScanResultValue> output =
-        factory.priorityQueueSortAndLimit(
-            inputSequence,
-            query,
-            ImmutableList.of(new SegmentDescriptor(new Interval(
-                DateTimes.of("2010-01-01"),
-                DateTimes.of("2019-01-01").plusHours(1)
-            ), "1", 0))
-        ).toList();
-
-    validateSortedOutput(output, expectedEventTimestamps);
+    try {
+      List<ScanResultValue> output = factory.priorityQueueSortAndLimit(
+          inputSequence,
+          query,
+          ImmutableList.of(new SegmentDescriptor(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");
+      }
+      validateSortedOutput(output, expectedEventTimestamps);
+    }
+    catch (UOE e) {
+      if (query.getLimit() <= Integer.MAX_VALUE) {
+        Assert.fail("Unsupported operation exception should not have been thrown here");
+      }
+    }
   }
 
   @Test