[Backport] minor: MSQ tasks dont need the entire worker context passed through to the task context. pass only what is needed

multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerControllerContext.java

@@ -21,6 +21,7 @@
 
 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.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;
@@ -54,6 +56,7 @@
 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;
@@ -84,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;
@@ -382,16 +407,21 @@ public static Map<String, Object> makeTaskContext(
       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(

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));
+  }
+}