MSQ: Use VSF, give stages control over inputs. #18873

gianm · 2025-12-29T18:23:44Z

This patch integrates MSQ with virtual storage. It also refactors how MSQ reads inputs to give stages more control over how inputs are read and merged. In particular, stages are now able to fully control merging logic.

The main changes:

Integrate with virtual storage. Removed DataSegmentProvider, replaced it with direct usage of SegmentManager by SegmentsInputSliceReader. The SegmentManager reference ends up wrapped into RegularLoadableSegment, which provides methods acquire() and acquireIfCached() to the query logic.
Give stages control over input merging: rework InputSliceReader to return ReadablePartitions directly, without embedding any merging logic. Break out StandardPartitionReader as a separate class.

Other changes:

Move ReadableInput to the querykit package. It is no longer specific to the MSQ framework.
Remove StandardStageProcessor, refactoring dependent code to not require it.
Remove ExternalColumnSelectorFactory wrapper. Type casting is now handled directly by RowBasedColumnSelectorFactory.
Include full query context in worker context, rather than just a subset.

The purpose of this config is to enable using SegmentLocalCacheManager for loading segments on MSQ worker tasks, where segments are not assigned by load/drop rules, and where there is not generally a specific maxSize configured for the local cache. We need to evict segments immediately so local disks don't fill up. The main changes: 1) In StorageLocation, update the releaseHold runnable to check for evictImmediately. If it is set, unmount the cache entry if all holds have been released. 2) In SegmentLocalCacheManager, when evictImmediately is set, "mount" sets an onUnmount handler to delete the info file.

This patch integrates MSQ with virtual storage. It also refactors how MSQ reads inputs to give stages more control over how inputs are read and merged. In particular, stages are now able to fully control merging logic. The main changes: 1) Integrate with virtual storage: merge the two DataSegmentProvider impls (Dart and Task) into DataSegmentProviderImpl that relies on SegmentManager. 2) Give stages control over input merging: rework InputSliceReader to return ReadablePartitions directly, without embedding any merging logic. Break out StandardPartitionReader as a separate class. Other changes: 1) Move ReadableInput to the querykit package. It is no longer specific to the MSQ framework. 2) Remove StandardStageProcessor, refactoring dependent code to not require it. 3) Remove ExternalColumnSelectorFactory wrapper. Type casting is now handled directly by RowBasedColumnSelectorFactory. 4) Include full query context in worker context, rather than just a subset. Includes apache#18871.

multi-stage-query/src/main/java/org/apache/druid/msq/shuffle/output/StageOutputHolder.java

clintropolis · 2026-01-09T19:57:13Z

multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java

+      for (AcquireSegmentAction acquireSegmentAction : loadingSegments) {
+        CloseableUtils.closeAndSuppressExceptions(
+            acquireSegmentAction,
+            e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction)


this looks legit, maybe we should decorate AcquireSegmentAction to have a segmentId or something and implement toString

I thought about that but didn't want to add a field to AcquireSegmentAction just for this. I was going to leave it the way it is.

I'll add a comment about this though.

multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java

multi-stage-query/src/test/java/org/apache/druid/msq/input/RegularLoadableSegmentTest.java

+      final int partitionNum = i % 2;
+
+      segments.add(
+          DataSegment.builder()


server/src/test/java/org/apache/druid/test/utils/TestSegmentManager.java

+   */
+  public static DataSegment createDataSegmentForTest(final SegmentId segmentId)
+  {
+    return DataSegment.builder()


clintropolis

overall lgtm 🤘

clintropolis · 2026-01-07T09:33:26Z

indexing-service/src/main/java/org/apache/druid/indexing/common/SegmentCacheManagerFactory.java

-    );
+    final SegmentLoaderConfig loaderConfig =
+        new SegmentLoaderConfig()
+            .setLocations(Collections.singletonList(new StorageLocationConfig(storageDir, null, null)))


i know this isn't new, but should this have a size so like... there is a guard rail to not just wreck things if something tries to load too much? There is tmpStorageBytesPerTask already on the task config, but that is used for tmp sorter output it looks like so.. maybe those files should be managed with StorageLocation too. This doesn't need to be done in this PR, just thinking

I was hoping to manage more stuff through the storage locations as time goes on and that would be a good time to unify things. The other thing that isn't currently tracked is temporary storage of external data that has been downloaded.

clintropolis · 2026-01-09T19:23:24Z

...-stage-query/src/main/java/org/apache/druid/msq/input/external/ExternalInputSliceReader.java

+      // The LoadableSegment generated here does not acquire a real hold, and ends up loading the external data in a
+      // processing thread (when the cursor is created). Ideally, this would be better integrated with the virtual
+      // storage system, giving us storage holds and the ability to load data outside of a processing thread.


👍 it feels like this should be possible, StorageLocation is pretty generic, but agree we can worry about this later

clintropolis · 2026-01-09T19:53:10Z

processing/src/main/java/org/apache/druid/segment/SegmentReference.java

+   * throws an exception. Regardless of success or failure of this method, the old reference should be discarded.
+   * Do not call {@link #close()} on the old reference, only call it on the new one.
+   */
+  public SegmentReference map(final SegmentMapFunction segmentMapFn)


could you update the javadoc at the top of this class to indicate that callers can call this if they don't have a mapFn to apply to the segment at the time of creation of the SegmentReference?

Updated to:

/** * Wrapper for a {@link SegmentDescriptor} and {@link Optional<Segment>}. For regular segments, the latter is typically * created by a {@link ReferenceCountedSegmentProvider} that may have a {@link SegmentMapFunction} applied to it. * * If the {@link SegmentMapFunction} you want to apply is not available at the time the {@link SegmentReference} * is created, use {@link #map(SegmentMapFunction)} to apply it. * * Closing this object closes both the {@link #getSegmentReference()} and any closeables attached from the process of * creating this object, such as from {@link AcquireSegmentAction}. The object from {@link #getSegmentReference()} * should not be closed directly by callers. */

clintropolis · 2026-01-09T19:57:13Z

multi-stage-query/src/main/java/org/apache/druid/msq/querykit/ReadableInputQueue.java

+      for (AcquireSegmentAction acquireSegmentAction : loadingSegments) {
+        CloseableUtils.closeAndSuppressExceptions(
+            acquireSegmentAction,
+            e -> log.warn(e, "Failed to close loadingSegment[%s]", acquireSegmentAction)


this looks legit, maybe we should decorate AcquireSegmentAction to have a segmentId or something and implement toString

clintropolis · 2026-01-09T21:01:30Z

multi-stage-query/src/main/java/org/apache/druid/msq/input/PhysicalInputSlice.java

+ * An {@link InputSlice} that has been prepared for reading by an {@link InputSliceReader}. Nothing contained in
+ * this class references open resources, so this class is not closeable and does not need to be closed.
+ */
+public class PhysicalInputSlice


what does 'physical' mean here? obligatory 'naming is hard' comment. Is it basically 'slice that gets processed to kick everything off by a leaf processor'? If so, should it be LeafInputSlice to match other terminology? though i guess i could see how that could be confusing too

It's definitely a "naming is hard" thing.

It's a somewhat-more-physical version of InputSlice. InputSlice is always purely logical (it just has the identifiers / paths of things that should be read, as strings and POJOs) but PhysicalInputSlice may contain the ability to actually read external data or actually load segments.

It's not just "leaf" since it includes partition inputs, which would be used by non-leaf stages.

clintropolis · 2026-01-09T21:12:15Z

multi-stage-query/src/main/java/org/apache/druid/msq/input/RegularLoadableSegment.java

+    if (cachedDataSegment != null) {
+      return Futures.immediateFuture(cachedDataSegment);
+    } else if (coordinatorClient != null) {
+      return coordinatorClient.fetchSegment(


i find myself wondering if maybe we can just do something to make SegmentManager more flexible (like say backed with a coordinator client to download datasegments on the fly) to make it so we can just do all the stuff with a segmentManager, but haven't really thought about this and maybe is too weird, either way, i think this is fine for this PR

That would be nice.

Along with the ability to manage non-segment data, like external files from cloud storage.

gianm added 2 commits December 29, 2025 07:05

github-actions bot added Area - Batch Ingestion Area - Segment Format and Ser/De Area - Ingestion Area - MSQ For multi stage queries - https://github.com/apache/druid/issues/12262 labels Dec 29, 2025

github-advanced-security bot found potential problems Dec 29, 2025

View reviewed changes

gianm added 9 commits December 29, 2025 17:11

Various fixes.

b37e7e0

Additional casting in RowBasedColumnSelectorFactory.

67c6157

Restructure.

b4f8613

Fix imports.

b6c0fac

Fix import.

8d01fd4

Fix tests.

38f4169

Include DataSegment in LoadableSegment.

2b9b7c2

Include queryInterval.

70f4a7c

Fix imports.

bf07733

gianm mentioned this pull request Jan 2, 2026

Add virtualStorageFabricEvictImmediately config. #18871

Merged

gianm added 5 commits January 3, 2026 22:00

Update method name and comments.

7888273

Merge branch 'master' into msq-read-input-channels-1x1

5b48f94

Remove DataSegmentProvider, add acquireCachedSegment to LoadableSegment.

c6e8025

Merge branch 'master' into msq-read-input-channels-1x1

a21cfe6

Fix imports.

288c3fa

github-advanced-security bot found potential problems Jan 7, 2026

View reviewed changes

gianm added 7 commits January 8, 2026 00:06

Add CoordinatorClient to Dart so it doesn't encounter missing segments.

6e070c8

Use the correct map.

ff5c981

Merge branch 'master' into msq-read-input-channels-1x1

c0284ec

Merge branch 'master' into msq-read-input-channels-1x1

94f3a26

Fix conflict.

057df40

Fix issue with context conflict.

ece92fe

Merge branch 'master' into msq-read-input-channels-1x1

2a3a866

Only close SegmentReference once, don't close inner Segment.

84964d9

clintropolis reviewed Jan 9, 2026

View reviewed changes

Comments from review.

7d76685

MSQ: Use VSF, give stages control over inputs. #18873

Are you sure you want to change the base?

MSQ: Use VSF, give stages control over inputs. #18873

Conversation

gianm commented Dec 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Check notice

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!

Check notice

Uh oh!

Check notice

Uh oh!

clintropolis left a comment

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

gianm commented Dec 29, 2025 •

edited

Loading