This is an automated email from the ASF dual-hosted git repository.
FrankChen021 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 8c8db7b834b minor: add segmentAcquireMode query context to control MSQ
engine AcquireMode (#20562)
8c8db7b834b is described below
commit 8c8db7b834b78ad2305e6b587a390ab43251fd1f
Author: Clint Wylie <[email protected]>
AuthorDate: Tue Oct 13 23:15:50 2026 -0700
minor: add segmentAcquireMode query context to control MSQ engine
AcquireMode (#20562)
---
.../msq/EmbeddedMSQPartialDownloadsTest.java | 43 ++++++++++++++++++++++
.../druid/msq/querykit/BaseLeafStageProcessor.java | 23 +++++-------
.../groupby/GroupByPreShuffleStageProcessor.java | 10 +----
.../msq/querykit/scan/ScanQueryStageProcessor.java | 12 +-----
.../java/org/apache/druid/query/QueryContext.java | 15 ++++++++
.../java/org/apache/druid/query/QueryContexts.java | 7 ++++
.../apache/druid/segment/loading/AcquireMode.java | 5 ++-
.../org/apache/druid/query/QueryContextTest.java | 23 ++++++++++++
8 files changed, 106 insertions(+), 32 deletions(-)
diff --git
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQPartialDownloadsTest.java
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQPartialDownloadsTest.java
index af3581b6563..1c8a3c58e2b 100644
---
a/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQPartialDownloadsTest.java
+++
b/embedded-tests/src/test/java/org/apache/druid/testing/embedded/msq/EmbeddedMSQPartialDownloadsTest.java
@@ -26,6 +26,7 @@ import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.msq.indexing.report.MSQTaskReportPayload;
import org.apache.druid.msq.util.MultiStageQueryContext;
import org.apache.druid.query.DruidProcessingConfigTest;
+import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.http.SqlTaskStatus;
import org.apache.druid.server.metrics.LatchableEmitter;
import org.apache.druid.server.metrics.StorageMonitor;
@@ -59,6 +60,10 @@ import java.util.Objects;
* on-demand deep-storage range reads (so the count is positive), while the
full-download path does not (the count
* stays zero because {@code virtualStoragePartialDownloadsEnabled=false}
makes the worker mount segments as complete
* cache entries and download them in full instead of via range reads).
+ * <p>
+ * Separately, {@link QueryContexts#SEGMENT_ACQUIRE_MODE} overrides the leaf
processor's default acquire mode while
+ * partial downloads stay enabled; forcing {@code FULL} shows up as more
{@link StorageMonitor#VSF_LOAD_BYTES} than the
+ * default {@code PARTIAL} acquire for the same query.
*/
class EmbeddedMSQPartialDownloadsTest extends EmbeddedClusterTestBase
{
@@ -151,6 +156,44 @@ class EmbeddedMSQPartialDownloadsTest extends
EmbeddedClusterTestBase
);
}
+ @Test
+ void testSegmentAcquireModeQueryContextOverride()
+ {
+ final LatchableEmitter emitter = indexer.latchableEmitter();
+
+ // The GroupBy leaf processor defaults to AcquireMode.PARTIAL, so only the
columns the query touches are pulled
+ // from deep storage. Forcing FULL via the query context makes the worker
download the entire segment up front
+ // instead, even though partial downloads remain enabled.
+ final long partialLoadBytes = runMsqSelectAndGetLoadBytes(emitter,
Collections.emptyMap());
+ final long fullLoadBytes = runMsqSelectAndGetLoadBytes(
+ emitter,
+ Map.of(QueryContexts.SEGMENT_ACQUIRE_MODE, "FULL")
+ );
+
+ Assertions.assertTrue(partialLoadBytes > 0, "expected the default PARTIAL
acquire to download some bytes");
+ Assertions.assertTrue(
+ fullLoadBytes > partialLoadBytes,
+ StringUtils.format(
+ "expected the FULL acquire to download more bytes[%,d] than the
PARTIAL acquire[%,d]",
+ fullLoadBytes,
+ partialLoadBytes
+ )
+ );
+ }
+
+ /**
+ * Runs {@link #SELECT_SQL} with the given context and returns the {@link
StorageMonitor#VSF_LOAD_BYTES} the worker
+ * pulled from deep storage. Each worker task has its own ephemeral cache,
so every run starts cold.
+ */
+ private long runMsqSelectAndGetLoadBytes(LatchableEmitter emitter,
Map<String, Object> queryContext)
+ {
+ emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BYTES,
MONITOR_QUIESCE_TIMEOUT_MILLIS);
+ emitter.flush();
+ runMsqSelect(queryContext);
+ emitter.awaitMetricQuiescent(StorageMonitor.VSF_LOAD_BYTES,
MONITOR_QUIESCE_TIMEOUT_MILLIS);
+ return emitter.getMetricEventLongSum(StorageMonitor.VSF_LOAD_BYTES);
+ }
+
private void runMsqSelect(Map<String, Object> queryContext)
{
final SqlTaskStatus status = msqApis.submitTaskSql(queryContext,
SELECT_SQL, dataSource);
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java
index 608ea80fe21..42b1866fcbc 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafStageProcessor.java
@@ -266,19 +266,15 @@ public abstract class BaseLeafStageProcessor extends
BasicStageProcessor
/**
- * Create the {@link ReadableInputQueue} for this stage's base inputs.
Override to change how the queue acquires
- * segments: the default acquires them fully up front ({@link
AcquireMode#FULL}), which is safe for any processor.
- * Leaf processors that read segments through the async cursor API
- * ({@link org.apache.druid.segment.CursorFactory#makeCursorHolderAsync})
can override to {@link AcquireMode#PARTIAL}
- * so only the columns a query actually touches are downloaded.
+ * How the {@link ReadableInputQueue} for this stage's base inputs acquires
segments, unless overridden by
+ * {@link org.apache.druid.query.QueryContexts#SEGMENT_ACQUIRE_MODE}. The
default acquires them fully up front
+ * ({@link AcquireMode#FULL}), which is safe for any processor. Leaf
processors that read segments through the async
+ * cursor API ({@link
org.apache.druid.segment.CursorFactory#makeCursorHolderAsync}) can override to
+ * {@link AcquireMode#PARTIAL} so only the columns a query actually touches
are downloaded.
*/
- protected ReadableInputQueue makeReadableInputQueue(
- final StandardPartitionReader partitionReader,
- final List<PhysicalInputSlice> slices,
- final int loadahead
- )
+ protected AcquireMode getDefaultAcquireMode()
{
- return new ReadableInputQueue(partitionReader, slices, loadahead,
AcquireMode.FULL);
+ return AcquireMode.FULL;
}
/**
@@ -311,10 +307,11 @@ public abstract class BaseLeafStageProcessor extends
BasicStageProcessor
}
final List<PhysicalInputSlice> filteredSlices =
filterBaseInput(physicalInputSlices);
- return makeReadableInputQueue(
+ return new ReadableInputQueue(
new StandardPartitionReader(context),
filteredSlices,
- context.segmentLoadAheadCount()
+ context.segmentLoadAheadCount(),
+
context.workOrder().getWorkerContext().getSegmentAcquireMode(getDefaultAcquireMode())
);
}
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java
index 208c57eefe4..986084b23af 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByPreShuffleStageProcessor.java
@@ -30,13 +30,11 @@ import org.apache.druid.frame.processor.FrameProcessor;
import org.apache.druid.frame.write.FrameWriterFactory;
import org.apache.druid.msq.exec.ExecutionContext;
import org.apache.druid.msq.exec.FrameContext;
-import org.apache.druid.msq.exec.std.StandardPartitionReader;
import org.apache.druid.msq.exec.std.StandardStageRunner;
import org.apache.druid.msq.input.LoadableSegment;
import org.apache.druid.msq.input.PhysicalInputSlice;
import org.apache.druid.msq.querykit.BaseLeafStageProcessor;
import org.apache.druid.msq.querykit.ReadableInput;
-import org.apache.druid.msq.querykit.ReadableInputQueue;
import org.apache.druid.msq.util.MultiStageQueryContext;
import org.apache.druid.query.QueryToolChestWarehouse;
import org.apache.druid.query.aggregation.AggregatorFactory;
@@ -173,15 +171,11 @@ public class GroupByPreShuffleStageProcessor extends
BaseLeafStageProcessor
}
@Override
- protected ReadableInputQueue makeReadableInputQueue(
- StandardPartitionReader partitionReader,
- List<PhysicalInputSlice> slices,
- int loadahead
- )
+ protected AcquireMode getDefaultAcquireMode()
{
// GroupBy reads segments through the async cursor API, so acquire them
partially: only the columns the query
// touches are downloaded.
- return new ReadableInputQueue(partitionReader, slices, loadahead,
AcquireMode.PARTIAL);
+ return AcquireMode.PARTIAL;
}
@Override
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java
index 1e7fbe215c5..3ab7ce98b9b 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryStageProcessor.java
@@ -28,11 +28,8 @@ import org.apache.druid.frame.channel.WritableFrameChannel;
import org.apache.druid.frame.processor.FrameProcessor;
import org.apache.druid.frame.write.FrameWriterFactory;
import org.apache.druid.msq.exec.FrameContext;
-import org.apache.druid.msq.exec.std.StandardPartitionReader;
-import org.apache.druid.msq.input.PhysicalInputSlice;
import org.apache.druid.msq.querykit.BaseLeafStageProcessor;
import org.apache.druid.msq.querykit.ReadableInput;
-import org.apache.druid.msq.querykit.ReadableInputQueue;
import org.apache.druid.query.DataSource;
import org.apache.druid.query.Druids;
import org.apache.druid.query.filter.DimFilter;
@@ -45,7 +42,6 @@ import org.apache.druid.segment.loading.AcquireMode;
import org.apache.druid.sql.calcite.filtration.Filtration;
import javax.annotation.Nullable;
-import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
@JsonTypeName("scan")
@@ -129,15 +125,11 @@ public class ScanQueryStageProcessor extends
BaseLeafStageProcessor
}
@Override
- protected ReadableInputQueue makeReadableInputQueue(
- StandardPartitionReader partitionReader,
- List<PhysicalInputSlice> slices,
- int loadahead
- )
+ protected AcquireMode getDefaultAcquireMode()
{
// Scan reads segments through the async cursor API, so acquire them
partially: only the columns the query
// touches are downloaded.
- return new ReadableInputQueue(partitionReader, slices, loadahead,
AcquireMode.PARTIAL);
+ return AcquireMode.PARTIAL;
}
@Override
diff --git a/processing/src/main/java/org/apache/druid/query/QueryContext.java
b/processing/src/main/java/org/apache/druid/query/QueryContext.java
index dcb9b4f4319..cc4de1c0f6a 100644
--- a/processing/src/main/java/org/apache/druid/query/QueryContext.java
+++ b/processing/src/main/java/org/apache/druid/query/QueryContext.java
@@ -34,6 +34,7 @@ import org.apache.druid.query.context.QueryContextParameter;
import org.apache.druid.query.context.QueryContextParameters;
import org.apache.druid.query.filter.InDimFilter;
import org.apache.druid.query.filter.TypedInFilter;
+import org.apache.druid.segment.loading.AcquireMode;
import javax.annotation.Nullable;
import java.io.IOException;
@@ -997,4 +998,18 @@ public class QueryContext
}
return QueryContexts.DEFAULT_REALTIME_SEGMENTS_MODE;
}
+
+ /**
+ * Returns the {@link AcquireMode} set by {@link
QueryContexts#SEGMENT_ACQUIRE_MODE}, or {@code defaultValue} (the
+ * engine's own choice) if absent. Currently only honored by MSQ leaf stage
processors.
+ * <p>
+ * {@link AcquireMode#PARTIAL} degrades to {@link AcquireMode#FULL} where
partial downloads are unavailable, such as
+ * when they are disabled by configuration or the segment is not
range-readable. Forcing {@link AcquireMode#PARTIAL}
+ * where the engine defaults to {@link AcquireMode#FULL} fails at cursor
creation if the engine reads segments through
+ * the synchronous cursor API.
+ */
+ public AcquireMode getSegmentAcquireMode(AcquireMode defaultValue)
+ {
+ return getEnum(QueryContexts.SEGMENT_ACQUIRE_MODE, AcquireMode.class,
defaultValue);
+ }
}
diff --git a/processing/src/main/java/org/apache/druid/query/QueryContexts.java
b/processing/src/main/java/org/apache/druid/query/QueryContexts.java
index 2c2204a8683..928eb10f1d5 100644
--- a/processing/src/main/java/org/apache/druid/query/QueryContexts.java
+++ b/processing/src/main/java/org/apache/druid/query/QueryContexts.java
@@ -132,6 +132,13 @@ public class QueryContexts
public static final String FORCE_PROJECTION = "forceProjections";
public static final String USE_PROJECTION = "useProjection";
+ /**
+ * Overrides the {@link org.apache.druid.segment.loading.AcquireMode} an
engine uses to acquire segments from the
+ * segment cache; see {@link QueryContext#getSegmentAcquireMode}. Intended
for testing, debugging, and performance
+ * comparison.
+ */
+ public static final String SEGMENT_ACQUIRE_MODE = "segmentAcquireMode";
+
// Unique identifier for the query, that is used to map the global shared
resources (specifically merge buffers) to the
// query's runtime
public static final String QUERY_RESOURCE_ID = "queryResourceId";
diff --git
a/server/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
b/processing/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
similarity index 90%
rename from
server/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
rename to
processing/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
index 28b58fb0849..60c0199a671 100644
--- a/server/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
+++ b/processing/src/main/java/org/apache/druid/segment/loading/AcquireMode.java
@@ -20,9 +20,12 @@
package org.apache.druid.segment.loading;
/**
- * How {@link SegmentCacheManager#acquireSegment} and {@link
SegmentCacheManager#acquireCachedSegment} materialize a
+ * How {@code SegmentCacheManager#acquireSegment} and {@code
SegmentCacheManager#acquireCachedSegment} materialize a
* segment. Both modes return an open reference that callers must close when
finished; they differ only in how much
* of the segment is downloaded before the reference is handed back.
+ * <p>
+ * Engines that support it can be told which mode to use via
+ * {@link org.apache.druid.query.QueryContext#getSegmentAcquireMode}.
*/
public enum AcquireMode
{
diff --git
a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
index 2d2b450d375..e2f3d0c08e5 100644
--- a/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
+++ b/processing/src/test/java/org/apache/druid/query/QueryContextTest.java
@@ -38,6 +38,7 @@ import org.apache.druid.query.context.QueryContextParameters;
import org.apache.druid.query.filter.DimFilter;
import org.apache.druid.query.spec.QuerySegmentSpec;
import org.apache.druid.segment.DimensionHandlerUtils;
+import org.apache.druid.segment.loading.AcquireMode;
import org.joda.time.DateTimeZone;
import org.joda.time.Duration;
import org.joda.time.Interval;
@@ -559,6 +560,28 @@ public class QueryContextTest
);
}
+ @Test
+ public void testGetSegmentAcquireMode()
+ {
+ assertEquals(AcquireMode.FULL,
QueryContext.empty().getSegmentAcquireMode(AcquireMode.FULL));
+ assertEquals(AcquireMode.PARTIAL,
QueryContext.empty().getSegmentAcquireMode(AcquireMode.PARTIAL));
+ assertEquals(
+ AcquireMode.FULL,
+ QueryContext.of(ImmutableMap.of(QueryContexts.SEGMENT_ACQUIRE_MODE,
"full"))
+ .getSegmentAcquireMode(AcquireMode.PARTIAL)
+ );
+ assertEquals(
+ AcquireMode.PARTIAL,
+ QueryContext.of(ImmutableMap.of(QueryContexts.SEGMENT_ACQUIRE_MODE,
"PARTIAL"))
+ .getSegmentAcquireMode(AcquireMode.FULL)
+ );
+ assertThrows(
+ BadQueryContextException.class,
+ () ->
QueryContext.of(ImmutableMap.of(QueryContexts.SEGMENT_ACQUIRE_MODE, "dummy"))
+ .getSegmentAcquireMode(AcquireMode.FULL)
+ );
+ }
+
@Test
public void testGetRealtimeSegmentsModeBackwardCompat()
{
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
