This is an automated email from the ASF dual-hosted git repository.
capistrant 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 3075b00284b MSQ tasks dont need the entire worker context passed
through to the task context. pass only what is needed (#20650)
3075b00284b is described below
commit 3075b00284b099c152436cf147d545be3277d728
Author: Lucas Capistrant <[email protected]>
AuthorDate: Tue Oct 20 19:50:35 2026 -0500
MSQ tasks dont need the entire worker context passed through to the task
context. pass only what is needed (#20650)
---
.../msq/indexing/IndexerControllerContext.java | 47 +++++++++--
.../msq/indexing/IndexerControllerContextTest.java | 95 ++++++++++++++++++++++
2 files changed, 134 insertions(+), 8 deletions(-)
diff --git
a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerControllerContext.java
b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerControllerContext.java
index f3a7b5bd3f5..b540a7e0353 100644
---
a/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerControllerContext.java
+++
b/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerControllerContext.java
@@ -21,6 +21,7 @@ package org.apache.druid.msq.indexing;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
import com.google.inject.Injector;
import com.google.inject.Key;
import com.google.inject.TypeLiteral;
@@ -30,6 +31,7 @@ import org.apache.druid.guice.annotations.Self;
import org.apache.druid.indexing.common.TaskLockType;
import org.apache.druid.indexing.common.TaskToolbox;
import org.apache.druid.indexing.common.actions.TaskActionClient;
+import org.apache.druid.indexing.common.task.Tasks;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.io.Closer;
import org.apache.druid.java.util.common.logger.Logger;
@@ -48,11 +50,13 @@ import
org.apache.druid.msq.indexing.client.ControllerChatHandler;
import org.apache.druid.msq.indexing.client.IndexerWorkerClient;
import org.apache.druid.msq.indexing.destination.MSQDestination;
import org.apache.druid.msq.indexing.error.MSQException;
+import org.apache.druid.msq.indexing.error.MSQWarnings;
import org.apache.druid.msq.indexing.error.UnknownFault;
import org.apache.druid.msq.input.InputSpecSlicerProvider;
import org.apache.druid.msq.kernel.WorkOrder;
import org.apache.druid.msq.kernel.controller.ControllerQueryKernelConfig;
import org.apache.druid.msq.util.MultiStageQueryContext;
+import org.apache.druid.query.BaseQuery;
import org.apache.druid.query.DruidMetrics;
import org.apache.druid.query.QueryContext;
import org.apache.druid.query.context.QueryContextParameters;
@@ -83,6 +87,28 @@ public class IndexerControllerContext implements
ControllerContext
private static final Logger log = new Logger(IndexerControllerContext.class);
+ /**
+ * Worker context keys that worker tasks require in their task context.
+ */
+ private static final Set<String> WORKER_TASK_CONTEXT_KEYS = ImmutableSet.of(
+ MultiStageQueryContext.CTX_DURABLE_SHUFFLE_STORAGE,
+ MSQWarnings.CTX_MAX_PARSE_EXCEPTIONS_ALLOWED,
+ MultiStageQueryContext.CTX_IS_REINDEX,
+ MultiStageQueryContext.CTX_MAX_CONCURRENT_STAGES,
+ MultiStageQueryContext.CTX_ROW_BASED_FRAME_TYPE,
+ MultiStageQueryContext.CTX_REMOVE_NULL_BYTES,
+ MultiStageQueryContext.CTX_INCLUDE_ALL_COUNTERS,
+ MultiStageQueryContext.CTX_MAX_FRAME_SIZE,
+ MultiStageQueryContext.CTX_SELECT_DESTINATION,
+ MultiStageQueryContext.CTX_LIVE_REPORT_COUNTERS,
+ MultiStageQueryContext.CTX_MAX_THREADS,
+ MultiStageQueryContext.CTX_VIRTUAL_STORAGE_PARTIAL_DOWNLOADS,
+ Tasks.STORE_EMPTY_COLUMNS_KEY,
+ QueryContextParameters.QUERY_ID.getName(),
+ BaseQuery.SQL_QUERY_ID,
+ QueryContextParameters.DEBUG.getName()
+ );
+
private final MSQControllerTask task;
private final TaskLockType taskLockType;
private final String taskDataSource;
@@ -381,16 +407,21 @@ public class IndexerControllerContext implements
ControllerContext
final Map<String, Object> controllerTaskContext
)
{
- // Put worker context into the task context. That way, workers can get
these context keys either from
- // WorkOrder#getContext or Task#getContext.
- final Map<String, Object> taskContext = new HashMap<>(
- makeWorkerContextMap(
- querySpec,
- queryKernelConfig.isDurableStorage(),
- queryKernelConfig.getMaxConcurrentStages()
- )
+ final Map<String, Object> workerContextMap = makeWorkerContextMap(
+ querySpec,
+ queryKernelConfig.isDurableStorage(),
+ queryKernelConfig.getMaxConcurrentStages()
);
+ // Forward only the worker context keys that the task requires.
+ final Map<String, Object> taskContext = new HashMap<>();
+ for (String key : WORKER_TASK_CONTEXT_KEYS) {
+ final Object value = workerContextMap.get(key);
+ if (value != null) {
+ taskContext.put(key, value);
+ }
+ }
+
// Put the lookup loading info in the task context to facilitate selective
loading of lookups.
if (controllerTaskContext.get(LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE)
!= null) {
taskContext.put(
diff --git
a/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/IndexerControllerContextTest.java
b/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/IndexerControllerContextTest.java
new file mode 100644
index 00000000000..ac08070bc49
--- /dev/null
+++
b/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/IndexerControllerContextTest.java
@@ -0,0 +1,95 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.msq.indexing;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.druid.indexing.common.task.Tasks;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.msq.indexing.destination.TaskReportMSQDestination;
+import org.apache.druid.msq.kernel.controller.ControllerQueryKernelConfig;
+import org.apache.druid.msq.util.MultiStageQueryContext;
+import org.apache.druid.query.BaseQuery;
+import org.apache.druid.query.Druids;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.scan.ScanQuery;
+import org.apache.druid.query.spec.LegacySegmentSpec;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.sql.calcite.planner.ColumnMappings;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+public class IndexerControllerContextTest
+{
+ @Test
+ public void testMakeTaskContextCopiesOnlyWorkerStartupKeys()
+ {
+ final Map<String, Object> queryContext = ImmutableMap.of(
+ BaseQuery.QUERY_ID, "query-id",
+ BaseQuery.SQL_QUERY_ID, "sql-query-id",
+ MultiStageQueryContext.CTX_MAX_THREADS, 2,
+ Tasks.STORE_EMPTY_COLUMNS_KEY, false,
+ MultiStageQueryContext.CTX_FINALIZE_AGGREGATIONS, false,
+ MultiStageQueryContext.CTX_ARRAY_INGEST_MODE, "array"
+ );
+
+ final LegacyMSQSpec querySpec =
+ LegacyMSQSpec.builder()
+ .query(
+ Druids.newScanQueryBuilder()
+ .dataSource(new TableDataSource("foo"))
+ .intervals(new
LegacySegmentSpec(Intervals.ETERNITY))
+ .columns("dim1")
+
.resultFormat(ScanQuery.ResultFormat.RESULT_FORMAT_COMPACTED_LIST)
+ .build()
+ )
+ .queryContext(QueryContext.of(queryContext))
+
.columnMappings(ColumnMappings.identity(RowSignature.builder().add("dim1",
ColumnType.STRING).build()))
+ .destination(TaskReportMSQDestination.instance())
+ .tuningConfig(MSQTuningConfig.defaultConfig())
+ .build();
+
+ final ControllerQueryKernelConfig queryKernelConfig =
+ ControllerQueryKernelConfig.builder()
+ .maxRetainedPartitionSketchBytes(1)
+ .destination(querySpec.getDestination())
+ .build();
+
+ final Map<String, Object> taskContext =
+ IndexerControllerContext.makeTaskContext(querySpec, queryKernelConfig,
ImmutableMap.of());
+
+ // Keys that workers read from the task context at startup.
+ Assertions.assertEquals("query-id", taskContext.get(BaseQuery.QUERY_ID));
+ Assertions.assertEquals("sql-query-id",
taskContext.get(BaseQuery.SQL_QUERY_ID));
+ Assertions.assertEquals(2,
taskContext.get(MultiStageQueryContext.CTX_MAX_THREADS));
+ Assertions.assertEquals(false,
taskContext.get(Tasks.STORE_EMPTY_COLUMNS_KEY));
+ Assertions.assertEquals(
+ TaskReportMSQDestination.instance().toSelectDestination().getName(),
+ taskContext.get(MultiStageQueryContext.CTX_SELECT_DESTINATION)
+ );
+
+ // Other query context keys reach workers through the work order, not the
task context.
+
Assertions.assertFalse(taskContext.containsKey(MultiStageQueryContext.CTX_FINALIZE_AGGREGATIONS));
+
Assertions.assertFalse(taskContext.containsKey(MultiStageQueryContext.CTX_ARRAY_INGEST_MODE));
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
