lower max bucket for tests and fixup count

apache · Nov 3, 2023 · 8bb2c38 · 8bb2c38
1 parent bdba1bf
commit 8bb2c38
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 15 deletions.
diff --git a/...e/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/TombstoneHelper.java b/...e/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/TombstoneHelper.java
@@ -52,12 +52,10 @@
 
 public class TombstoneHelper
 {
-
   private final TaskActionClient taskActionClient;
 
   public TombstoneHelper(TaskActionClient taskActionClient)
   {
-
     this.taskActionClient = Preconditions.checkNotNull(taskActionClient, "taskActionClient");
   }
 
@@ -216,18 +214,15 @@ public Set<Interval> computeTombstoneIntervalsForReplace(
         if (overlap == null) {
           continue;
         }
-        if (buckets > maxBuckets) {
-          throw new TooManyBucketsException(maxBuckets);
-        }
 
         if (Intervals.ETERNITY.getStart().equals(overlap.getStart())) {
           // Generate a tombstone interval covering the negative eternity interval.
+          buckets = validateAndGetBuckets(buckets + 1, maxBuckets);
           retVal.add(new Interval(overlap.getStart(), replaceGranularity.bucketStart(overlap.getEnd())));
-          buckets += 1;
         } else if ((Intervals.ETERNITY.getEnd()).equals(overlap.getEnd())) {
           // Generate a tombstone interval covering the positive eternity interval.
+          buckets = validateAndGetBuckets(buckets + 1, maxBuckets);
           retVal.add(new Interval(replaceGranularity.bucketStart(overlap.getStart()), overlap.getEnd()));
-          buckets += 1;
         } else {
           // Overlap might not be aligned with the granularity if the used interval is not aligned with the granularity
           // However when fetching from the iterator, the first interval is found using the bucketStart, which
@@ -250,8 +245,8 @@ public Set<Interval> computeTombstoneIntervalsForReplace(
           // Helps in deduplication if required. Since all the intervals are uniformly granular, there should be no
           // no overlap post deduplication
           HashSet<Interval> intervals = Sets.newHashSet(intervalsToDropByGranularity.granularityIntervalsIterator());
+          buckets = validateAndGetBuckets(buckets + intervals.size(), maxBuckets);
           retVal.addAll(intervals);
-          buckets += intervals.size();
         }
       }
     }
@@ -327,4 +322,12 @@ private List<Interval> getExistingNonEmptyIntervalsOfDatasource(
     return JodaUtils.condenseIntervals(retVal);
   }
 
+  private int validateAndGetBuckets(final int buckets, final int maxBuckets)
+  {
+    if (buckets >= maxBuckets) {
+      throw new TooManyBucketsException(maxBuckets);
+    }
+    return buckets;
+  }
+
 }
diff --git a/...c/test/java/org/apache/druid/indexing/common/task/batch/parallel/TombstoneHelperTest.java b/...c/test/java/org/apache/druid/indexing/common/task/batch/parallel/TombstoneHelperTest.java
@@ -57,7 +57,7 @@ public class TombstoneHelperTest
 
   private final TaskActionClient taskActionClient = Mockito.mock(TaskActionClient.class);
 
-  private final int MAX_BUCKETS = 5_000;
+  private final int MAX_BUCKETS = 100;
 
   @Test
   public void noTombstonesWhenNoDataInInputIntervalAndNoExistingSegments() throws Exception
@@ -552,8 +552,8 @@ public void testTombstoneIntervalsForReplaceOverLargeFiniteInterval() throws IOE
     Interval usedInterval = Intervals.of("1000-01-01/9000-12-31");
     Interval replaceInterval = Intervals.of("1000-01-01/9000-12-31");
     List<Interval> dropIntervals = ImmutableList.of(
-        Intervals.of("1000-01-01/5000-12-31"),
-        Intervals.of("6000-01-01/9000-12-31")
+        Intervals.of("1000-01-01/1001-01-01"),
+        Intervals.of("6000-01-01/6001-01-01")
     );
     Granularity replaceGranularity = Granularities.DAY;
 
@@ -587,8 +587,8 @@ public void testTombstoneIntervalsForReplaceOverLargeFiniteIntervalAndMaxBucket(
     Interval usedInterval = Intervals.of("1000-01-01/9000-12-31");
     Interval replaceInterval = Intervals.of("1000-01-01/9000-12-31");
     List<Interval> dropIntervals = ImmutableList.of(
-        Intervals.of("1000-01-01/5000-12-31"),
-        Intervals.of("6000-01-01/9000-12-31")
+        Intervals.of("1000-01-01/1001-01-01"),
+        Intervals.of("6000-01-01/6001-01-01")
     );
     Granularity replaceGranularity = Granularities.DAY;
 
@@ -609,10 +609,10 @@ public void testTombstoneIntervalsForReplaceOverLargeFiniteIntervalAndMaxBucket(
         ImmutableList.of(replaceInterval),
         "test",
         replaceGranularity,
-        3_000_000
+        800
     );
 
-    // ((5000 - 1000) * 365) + ((9000 - 6000) * 365 * 24) ~= 2_557_426 day intervals
+    // (365 * 2) ~= 730 day intervals
     Assert.assertEquals(
         dropIntervals.stream()
                      .mapToLong(interval -> interval.toDuration().getStandardDays())