Fix MSQ compaction state and native interval locking, add test coverage #18950

kfaraz · 2026-02-05T07:51:02Z

Did centralized schema also run into an issue similar to batch allocation?

it's not supported by msq engine, schema is not saved

kfaraz · 2026-02-05T05:58:48Z

Just curious, have we added new tests which need this?

ah i think i had some issues with batch segment allocation with segment lock, and some pending segments cache was not removed properly or something (maybe i should have written it down), and have decided to disable the batch allocation to avoid this issue

yea so it's set to false in CompactionTaskRunBase

Oh, if that is the case, we would need to dig deeper into it since it sounds like a bug and batch segment allocation is the default. It should work for all cases.

it's only for segment lock, and i wasn't sure it's due to test setup or anything.

Okay, in that case, let's skip the test case for segment lock and add a comment or add a test case and mark it disabled. That way, we would be aware that there is a bug and can address it later.

i could not reproduce the test cases any more, maybe it disappeared after we switched to use the widen interval. anyways, so i added the segmentQueue test param.

Oh, okay. So do the tests now use batch segment allocation or not?

kfaraz · 2026-02-05T07:50:44Z

Why do we want to impose this limit? It should be okay to be able to change the config between tests.
The before method would be invoked before every test, right?

the before method was called before @test, wont that mean we can't change the config between tests, but in setup only?

well this is just to prevent ppl from accidently calling it inside test method

Why do we need to protect devs against that? This is test-only code.

wont that mean we can't change the config between tests, but in setup only?

If the setup method (i.e. before) is called before every @Test, then each test can potentially have its own setup with a different set of configs if needed. Changing the config while a @Test method is in progress will have no effect anyway.

-Original file line number
+Diff line change
@@ Expand Up / @@ -102,7 +102,6 @@ @@
     import java.util.concurrent.CountDownLatch;
     import java.util.concurrent.ExecutorService;
     import java.util.concurrent.TimeUnit;
-    import java.util.function.BiConsumer;
     import java.util.function.Consumer;
     import java.util.function.Function;
     import java.util.function.Predicate;
@@ Expand Down Expand Up @@
         }
       }
-      /**
-       * Attempts to acquire a lock that covers certain segments.
-       * <p>
-       * Will look at {@link Tasks#FORCE_TIME_CHUNK_LOCK_KEY} to decide whether to acquire a time chunk or segment lock.
-       * <p>
-       * This method will initialize {@link #taskLockHelper} as a side effect.
-       *
-       * @return whether the lock was acquired
-       */
-      boolean determineLockGranularityAndTryLockWithSegments(
-          TaskActionClient client,
-          List<DataSegment> segments,
-          BiConsumer<LockGranularity, List<DataSegment>> segmentCheckFunction
-      ) throws IOException
-      {
-        final boolean forceTimeChunkLock = getContextValue(
-            Tasks.FORCE_TIME_CHUNK_LOCK_KEY,
-            Tasks.DEFAULT_FORCE_TIME_CHUNK_LOCK
-        );
-        if (forceTimeChunkLock) {
-          log.info("[%s] is set to true in task context. Use timeChunk lock", Tasks.FORCE_TIME_CHUNK_LOCK_KEY);
-          taskLockHelper = createLockHelper(LockGranularity.TIME_CHUNK);
-          segmentCheckFunction.accept(LockGranularity.TIME_CHUNK, segments);
-          return tryTimeChunkLock(
-              client,
-              new ArrayList<>(segments.stream().map(DataSegment::getInterval).collect(Collectors.toSet()))
-          );
-        } else {
-          final LockGranularityDetermineResult result = determineSegmentGranularity(segments);
-          taskLockHelper = createLockHelper(result.lockGranularity);
-          segmentCheckFunction.accept(result.lockGranularity, segments);
-          return tryLockWithDetermineResult(client, result);
-        }
-      }
       private LockGranularityDetermineResult determineSegmentGranularity(TaskActionClient client, List<Interval> intervals)
           throws IOException
       {
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -45,7 +45,7 @@ public class CompactionIOConfig implements IOConfig @@
       @JsonCreator
       public CompactionIOConfig(
           @JsonProperty("inputSpec") CompactionInputSpec inputSpec,
-          @JsonProperty("allowNonAlignedInterval") boolean allowNonAlignedInterval,
+          @Deprecated @JsonProperty("allowNonAlignedInterval") boolean allowNonAlignedInterval,
           @JsonProperty("dropExisting") @Nullable Boolean dropExisting
       )
       {
@@ Expand All / @@ -60,6 +60,10 @@ public CompactionInputSpec getInputSpec() @@
         return inputSpec;
       }
+      /**
+       * @deprecated Provided for backwards compatibility. Can lead to data being accidentally overshadowed.
+       */
+      @Deprecated
       @JsonProperty("allowNonAlignedInterval")
       @JsonInclude(JsonInclude.Include.NON_DEFAULT)
       public boolean isAllowNonAlignedInterval()
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -441,7 +441,7 @@ public int getPriority() @@
       public boolean isReady(TaskActionClient taskActionClient) throws Exception
       {
         final List<DataSegment> segments = segmentProvider.findSegments(taskActionClient);
-        return determineLockGranularityAndTryLockWithSegments(taskActionClient, segments, segmentProvider::checkSegments);
+        return determineLockGranularityAndTryLock(taskActionClient, List.of(umbrellaInterval(segments, segmentProvider)));
       }
       @Override
@@ Expand Down Expand Up / @@ -474,6 +474,7 @@ public boolean isPerfectRollup() @@
        * <li> Rollup is done on a multi-valued string dimension or an unknown dimension
        * (since MSQ requires multi-valued string dimensions to be converted to arrays for rollup) </li>
        * </ul>
+       *
        * @return false for native engine, true for MSQ engine only when partitioning or rollup is done on a multi-valued
        * string or unknown dimension.
        */
@@ Expand Down Expand Up @@
               toolbox.getEmitter(),
               metricBuilder,
               segmentProvider.dataSource,
-              JodaUtils.umbrellaInterval(
-                  Iterables.transform(
-                      timelineSegments,
-                      DataSegment::getInterval
-                  )
-              ),
+              umbrellaInterval(timelineSegments, segmentProvider),
               lazyFetchSegments(
                   timelineSegments,
                   toolbox.getSegmentCacheManager()
@@ Expand Down Expand Up @@
         return VersionedIntervalTimeline.getAllObjects(timelineSegments);
       }
+      private static Interval umbrellaInterval(Iterable<DataSegment> segments, SegmentProvider segmentProvider)
+      {
+        return JodaUtils.umbrellaInterval(
+            List.of(
+                JodaUtils.umbrellaInterval(Iterables.transform(segments, DataSegment::getInterval)),
+                segmentProvider.interval
+            ));
+      }
       private static DataSchema createDataSchema(
           ServiceEmitter emitter,
           ServiceMetricEvent.Builder metricBuilder,
@@ Expand Down Expand Up / @@ -964,9 +969,9 @@ public DimensionsSpec getDimensionsSpec() @@
           // Store forceSegmentSortByTime only if false, for compatibility with legacy compaction states.
           final Boolean forceSegmentSortByTime = includeTimeAsDimension ? false : null;
           return DimensionsSpec.builder()
-              .setDimensions(dimensionSchemas)
-              .setForceSegmentSortByTime(forceSegmentSortByTime)
-              .build();
+                               .setDimensions(dimensionSchemas)
+                               .setForceSegmentSortByTime(forceSegmentSortByTime)
+                               .build();
         }
         public AggregatorFactory[] getMetricsSpec()
@@ Expand Down Expand Up / @@ -1016,7 +1021,7 @@ public Set<String> getMultiValuedDimensions() @@
          * Sort {@link #segmentsIterable} in order, such that we look at later segments prior to earlier ones. Useful when
          * analyzing dimensions, as it allows us to take the latest value we see, and therefore prefer types from more
          * recent segments, if there was a change.
-         *
+         * <p>
          * Returns a List copy of the original Iterable.
          */
         private List<Pair<DataSegment, Supplier<ResourceHolder<QueryableIndex>>>> sortSegmentsListNewestFirst()
@@ Expand Down Expand Up @@
               final QueryableIndex projectionIndex = Preconditions.checkNotNull(
                   index.getProjectionQueryableIndex(schema.getName())
               );
-              final List<DimensionSchema> columnSchemas = Lists.newArrayListWithExpectedSize(schema.getGroupingColumns().size());
+              final List<DimensionSchema> columnSchemas = Lists.newArrayListWithExpectedSize(
+                  schema.getGroupingColumns().size()
+              );
               for (String groupingColumn : schema.getGroupingColumns()) {
                 if (groupingColumn.equals(schema.getTimeColumnName())) {
                   columnSchemas.add(
@@ Expand Down Expand Up / @@ -1291,6 +1298,11 @@ public Builder interval(Interval interval) @@
           return inputSpec(new CompactionIntervalSpec(interval, null));
         }
+        public Builder interval(Interval interval, boolean dropExisting)
+        {
+          return inputSpec(new CompactionIntervalSpec(interval, null), dropExisting);
+        }
         public Builder segments(List<DataSegment> segments)
         {
           return inputSpec(SpecificSegmentsSpec.fromSegments(segments));
@@ Expand Down Expand Up / @@ -1394,7 +1406,7 @@ public CompactionTask build() @@
       /**
        * Compcation Task Tuning Config.
-       *
+       * <p>
        * An extension of ParallelIndexTuningConfig. As of now, all this TuningConfig
        * does is fail if the TuningConfig contains
        * `awaitSegmentAvailabilityTimeoutMillis` that is != 0 since it is not
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -181,6 +181,7 @@ public Set<DataSegment> computeTombstoneSegmentsForReplace( @@
       /**
        * See the method body for an example and an indepth explanation as to how the replace interval is created
+       *
        * @param intervalsToDrop    Empty intervals in the query that need to be dropped. They should be aligned with the
        *                           replaceGranularity
        * @param intervalsToReplace Intervals in the query which are eligible for replacement with new data.
@@ Expand Down Expand Up / @@ -278,7 +279,8 @@ public DataSegment createTombstoneForTimeChunkInterval( @@
                        .version(version)
                        .shardSpec(shardSpec)
                        .loadSpec(tombstoneLoadSpec) // load spec is special for tombstone
-                       .size(1); // in case coordinator segment balancing chokes with zero size
+                       .size(1) // in case coordinator segment balancing chokes with zero size
+                       .totalRows(0);
         return dataSegmentBuilder.build();
@@ Expand All / @@ -291,8 +293,8 @@ public DataSegment createTombstoneForTimeChunkInterval( @@
        * For a datasource having segments for 2020-01-01/2020-12-31 and 2022-01-01/2022-12-31, this method would return
        * the segment 2020-01-01/2020-12-31 if the input intervals asked for the segment between 2019 and 2021.
        *
-       * @param inputIntervals   Intervals corresponding to the task
-       * @param dataSource       Datasource corresponding to the task
+       * @param inputIntervals Intervals corresponding to the task
+       * @param dataSource     Datasource corresponding to the task
        * @return Intervals corresponding to used segments that overlap with any of the spec's input intervals
        * @throws IOException If used segments cannot be retrieved
        */
@@ Expand Down @@

-Original file line number
+Diff line change
@@ Expand Up / @@ -20,6 +20,7 @@ @@
     package org.apache.druid.indexing.common.actions;
     import com.fasterxml.jackson.databind.ObjectMapper;
+    import com.google.common.base.Preconditions;
     import com.google.common.base.Suppliers;
     import org.apache.druid.indexing.common.TestUtils;
     import org.apache.druid.indexing.common.config.TaskStorageConfig;
@@ Expand Down Expand Up / @@ -48,10 +49,11 @@ @@
     import org.apache.druid.server.coordinator.simulate.BlockingExecutorService;
     import org.apache.druid.server.coordinator.simulate.TestDruidLeaderSelector;
     import org.apache.druid.server.coordinator.simulate.WrappingScheduledExecutorService;
-    import org.easymock.EasyMock;
     import org.joda.time.Period;
     import org.junit.rules.ExternalResource;
+    import java.util.concurrent.atomic.AtomicBoolean;
     public class TaskActionTestKit extends ExternalResource
     {
       private final MetadataStorageTablesConfig metadataStorageTablesConfig = MetadataStorageTablesConfig.fromBase("druid");
@@ Expand All / @@ -66,7 +68,46 @@ public class TaskActionTestKit extends ExternalResource @@
       private BlockingExecutorService metadataCachePollExec;
       private boolean useSegmentMetadataCache = false;
+      private boolean useCentralizedDatasourceSchema = false;
+      private boolean batchSegmentAllocation = true;
       private boolean skipSegmentPayloadFetchForAllocation = new TaskLockConfig().isBatchAllocationReduceMetadataIO();
+      private AtomicBoolean configFinalized = new AtomicBoolean();
+      public TaskActionTestKit setUseSegmentMetadataCache(boolean useSegmentMetadataCache)
+      {
+        if (configFinalized.get()) {
+          throw new IllegalStateException("Test config already finalized");
+        }
+        this.useSegmentMetadataCache = useSegmentMetadataCache;
+        return this;
+      }
+      public TaskActionTestKit setUseCentralizedDatasourceSchema(boolean useCentralizedDatasourceSchema)
+      {
+        if (configFinalized.get()) {
+          throw new IllegalStateException("Test config already finalized");
+        }
+        this.useCentralizedDatasourceSchema = useCentralizedDatasourceSchema;
+        return this;
+      }
+      public TaskActionTestKit setBatchSegmentAllocation(boolean batchSegmentAllocation)
+      {
+        if (configFinalized.get()) {
+          throw new IllegalStateException("Test config already finalized");
+        }
+        this.batchSegmentAllocation = batchSegmentAllocation;
+        return this;
+      }
+      public TaskActionTestKit setSkipSegmentPayloadFetchForAllocation(boolean skipSegmentPayloadFetchForAllocation)
+      {
+        if (configFinalized.get()) {
+          throw new IllegalStateException("Test config already finalized");
+        }
+        this.skipSegmentPayloadFetchForAllocation = skipSegmentPayloadFetchForAllocation;
+        return this;
+      }
       public StubServiceEmitter getServiceEmitter()
       {
@@ Expand All / @@ -88,6 +129,11 @@ public GlobalTaskLockbox getTaskLockbox() @@
         return taskLockbox;
       }
+      public TaskStorage getTaskStorage()
+      {
+        return taskStorage;
+      }
       public IndexerMetadataStorageCoordinator getMetadataStorageCoordinator()
       {
         return metadataStorageCoordinator;
@@ Expand All / @@ -98,16 +144,6 @@ public TaskActionToolbox getTaskActionToolbox() @@
         return taskActionToolbox;
       }
-      public void setSkipSegmentPayloadFetchForAllocation(boolean skipSegmentPayloadFetchForAllocation)
-      {
-        this.skipSegmentPayloadFetchForAllocation = skipSegmentPayloadFetchForAllocation;
-      }
-      public void setUseSegmentMetadataCache(boolean useSegmentMetadataCache)
-      {
-        this.useSegmentMetadataCache = useSegmentMetadataCache;
-      }
       public void syncSegmentMetadataCache()
       {
         metadataCachePollExec.finishNextPendingTasks(4);
@@ Expand All / @@ -116,11 +152,13 @@ public void syncSegmentMetadataCache() @@
       @Override
       public void before()
       {
+        Preconditions.checkState(configFinalized.compareAndSet(false, true));
         emitter = new StubServiceEmitter();
         taskStorage = new HeapMemoryTaskStorage(new TaskStorageConfig(new Period("PT24H")));
         testDerbyConnector = new TestDerbyConnector(
             new MetadataStorageConnectorConfig(),
-            metadataStorageTablesConfig
+            metadataStorageTablesConfig,
+            CentralizedDatasourceSchemaConfig.enabled(useCentralizedDatasourceSchema)
         );
         final ObjectMapper objectMapper = new TestUtils().getTestObjectMapper();
         final SegmentSchemaManager segmentSchemaManager = new SegmentSchemaManager(
@@ Expand All / @@ -136,13 +174,19 @@ public void before() @@
             metadataStorageTablesConfig,
             testDerbyConnector,
             segmentSchemaManager,
-            CentralizedDatasourceSchemaConfig.create(),
+            CentralizedDatasourceSchemaConfig.enabled(useCentralizedDatasourceSchema),
             new HeapMemoryIndexingStateStorage()
         );
         taskLockbox = new GlobalTaskLockbox(taskStorage, metadataStorageCoordinator);
         taskLockbox.syncFromStorage();
         final TaskLockConfig taskLockConfig = new TaskLockConfig()
         {
+          @Override
+          public boolean isBatchSegmentAllocation()
+          {
+            return batchSegmentAllocation;
+          }
           @Override
           public long getBatchAllocationWaitTime()
           {
@@ Expand All / @@ -156,22 +200,28 @@ public boolean isBatchAllocationReduceMetadataIO() @@
           }
         };
+        SupervisorManager supervisorManager = new SupervisorManager(objectMapper, null);
+        SegmentAllocationQueue segmentAllocationQueue = new SegmentAllocationQueue(
+            taskLockbox,
+            taskLockConfig,
+            metadataStorageCoordinator,
+            emitter,
+            ScheduledExecutors::fixed
+        );
+        if (segmentAllocationQueue.isEnabled()) {
+          segmentAllocationQueue.becomeLeader();
+        }
         taskActionToolbox = new TaskActionToolbox(
             taskLockbox,
             taskStorage,
             metadataStorageCoordinator,
-            new SegmentAllocationQueue(
-                taskLockbox,
-                taskLockConfig,
-                metadataStorageCoordinator,
-                emitter,
-                ScheduledExecutors::fixed
-            ),
+            segmentAllocationQueue,
             emitter,
-            EasyMock.createMock(SupervisorManager.class),
+            supervisorManager,
             objectMapper
         );
         testDerbyConnector.createDataSourceTable();
+        testDerbyConnector.createUpgradeSegmentsTable();
         testDerbyConnector.createPendingSegmentsTable();
         testDerbyConnector.createSegmentSchemasTable();
         testDerbyConnector.createSegmentTable();
@@ Expand Down @@

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix MSQ compaction state and native interval locking, add test coverage #18950

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

kfaraz Feb 5, 2026

Uh oh!

cecemei Feb 5, 2026

Uh oh!

kfaraz Feb 5, 2026

Uh oh!

cecemei Feb 5, 2026

Uh oh!

cecemei Feb 5, 2026

Uh oh!

kfaraz Feb 5, 2026

Uh oh!

cecemei Feb 5, 2026

Uh oh!

kfaraz Feb 5, 2026

Uh oh!

cecemei Feb 5, 2026

Uh oh!

kfaraz Feb 6, 2026

Uh oh!

kfaraz Feb 5, 2026

Uh oh!

cecemei Feb 5, 2026

Uh oh!

cecemei Feb 5, 2026

Uh oh!

kfaraz Feb 6, 2026

Uh oh!

Uh oh!

Uh oh!

-Original file line number
+Diff line change
@@ Expand Up @@
                   Granularities.HOUR,
                   Granularities.MINUTE,
                   true,
-                  // Umbrella interval for all segments, since CompactionTasks generated a single granularitySpec.
-                  ImmutableList.of(Intervals.of("2014-01-01/2014-01-01T03:00:00"))
+                  ImmutableList.of(INTERVAL_TO_INDEX)
               ),
               null
           );
@@ Expand Down @@

Fix MSQ compaction state and native interval locking, add test coverage #18950

Uh oh!

Fix MSQ compaction state and native interval locking, add test coverage #18950

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!