Merge remote-tracking branch 'upstream/master' into fix-metric-emission

.github/scripts/setup_generate_license.sh

@@ -18,6 +18,9 @@
 set -e
 
 sudo apt-get update && sudo apt-get install python3 -y
-curl https://bootstrap.pypa.io/pip/3.5/get-pip.py | sudo -H python3
+# creating python virtual env
+python3 -m venv ~/.python3venv
+source ~/.python3venv/bin/activate
+sudo apt install python3-pip
 pip3 install wheel  # install wheel first explicitly
 pip3 install --upgrade pyyaml

.github/workflows/unit-and-integration-tests-unified.yml

@@ -19,6 +19,7 @@ on:
     paths-ignore:
       - '**/*.md'
       - 'dev/**'
+      - 'distribution/bin/**'
       - 'docs/**'
       - 'examples/**/jupyter-notebooks/**'
       - 'web-console/**'
@@ -31,6 +32,7 @@ on:
     paths-ignore:
       - '**/*.md'
       - 'dev/**'
+      - 'distribution/bin/**'
       - 'docs/**'
       - 'examples/**/jupyter-notebooks/**'
       - 'web-console/**'

benchmarks/src/test/java/org/apache/druid/benchmark/DataSourcesSnapshotBenchmark.java

@@ -41,9 +41,7 @@ import org.openjdk.jmh.infra.Blackhole;
 
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 @State(Scope.Benchmark)
@@ -54,7 +52,7 @@ import java.util.concurrent.TimeUnit;
 @Measurement(iterations = 50)
 public class DataSourcesSnapshotBenchmark
 {
-  private static Interval TEST_SEGMENT_INTERVAL = Intervals.of("2012-03-15T00:00:00.000/2012-03-16T00:00:00.000");
+  private static final Interval TEST_SEGMENT_INTERVAL = Intervals.of("2012-03-15T00:00:00.000/2012-03-16T00:00:00.000");
 
   @Param({"500", "1000"})
   private int numDataSources;
@@ -69,11 +67,10 @@ public class DataSourcesSnapshotBenchmark
   {
     long start = System.currentTimeMillis();
 
-    Map<String, ImmutableDruidDataSource> dataSources = new HashMap<>();
+    final List<DataSegment> segments = new ArrayList<>();
 
     for (int i = 0; i < numDataSources; i++) {
       String dataSource = StringUtils.format("ds-%d", i);
-      List<DataSegment> segments = new ArrayList<>();
 
       for (int j = 0; j < numSegmentPerDataSource; j++) {
         segments.add(
@@ -90,11 +87,9 @@ public class DataSourcesSnapshotBenchmark
             )
         );
       }
-
-      dataSources.put(dataSource, new ImmutableDruidDataSource(dataSource, Collections.emptyMap(), segments));
     }
 
-    snapshot = new DataSourcesSnapshot(dataSources);
+    snapshot = DataSourcesSnapshot.fromUsedSegments(segments);
 
     System.out.println("Setup Time " + (System.currentTimeMillis() - start) + " ms");
   }

benchmarks/src/test/java/org/apache/druid/server/coordinator/NewestSegmentFirstPolicyBenchmark.java

@@ -20,7 +20,6 @@
 package org.apache.druid.server.coordinator;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import org.apache.druid.client.DataSourcesSnapshot;
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.apache.druid.java.util.common.DateTimes;
@@ -128,7 +127,7 @@ public class NewestSegmentFirstPolicyBenchmark
         }
       }
     }
-    dataSources = DataSourcesSnapshot.fromUsedSegments(segments, ImmutableMap.of()).getUsedSegmentsTimelinesPerDataSource();
+    dataSources = DataSourcesSnapshot.fromUsedSegments(segments).getUsedSegmentsTimelinesPerDataSource();
   }
 
   @Benchmark

distribution/bin/find-missing-backports.py

@@ -64,17 +64,18 @@ def find_next_url(links):
     return None
 
 
-if len(sys.argv) != 5:
-  sys.stderr.write('usage: program <github-username> <previous-release-branch> <current-release-branch> <milestone-number>\n')
-  sys.stderr.write("  e.g., program myusername 0.17.0 0.18.0 30")
-  sys.stderr.write("  e.g., The milestone number for Druid 30 is 56, since the milestone has the url https://github.com/apache/druid/milestone/56\n")
-  sys.stderr.write("  It is also necessary to set a GIT_TOKEN environment variable containing a personal access token.")
+if len(sys.argv) != 4:
+  sys.stderr.write('Incorrect program arguments.\n')
+  sys.stderr.write('Usage: program <github-username> <previous-major-release-branch> <current-major-release-branch>\n')
+  sys.stderr.write("  e.g., program myusername 29.0.0 30.0.0\n")
+  sys.stderr.write("  Ensure that the title of the milestone is the same as the release branch.\n")
+  sys.stderr.write("  Ensure that a GIT_TOKEN environment variable containing a personal access token has been set.\n")
   sys.exit(1)
 
 github_username = sys.argv[1]
 previous_branch = sys.argv[2]
 release_branch = sys.argv[3]
-milestone_number = int(sys.argv[4])
+milestone_title = release_branch
 
 master_branch = "master"
 command = "git log {}..{} --oneline | tail -1".format(master_branch, previous_branch)
@@ -102,9 +103,6 @@ for commit_msg in all_release_commits.splitlines():
 
 print("Number of release PR subjects: {}".format(len(release_pr_subjects)))
 # Get all closed PRs and filter out with milestone
-milestone_url = "https://api.github.com/repos/apache/druid/milestones/{}".format(milestone_number)
-resp = requests.get(milestone_url, auth=(github_username, os.environ["GIT_TOKEN"])).json()
-milestone_title = resp['title']
 pr_items = []
 page = 0
 while True:

distribution/docker/environment

@@ -37,8 +37,6 @@ druid_metadata_storage_connector_connectURI=jdbc:postgresql://postgres:5432/drui
 druid_metadata_storage_connector_user=druid
 druid_metadata_storage_connector_password=FoolishPassword
 
-druid_coordinator_balancer_strategy=cachingCost
-
 druid_indexer_runner_javaOptsArray=["-server", "-Xmx1g", "-Xms1g", "-XX:MaxDirectMemorySize=3g", "-Duser.timezone=UTC", "-Dfile.encoding=UTF-8", "-Djava.util.logging.manager=org.apache.logging.log4j.jul.LogManager"]
 druid_indexer_fork_property_druid_processing_buffer_sizeBytes=256MiB
 

docs/api-reference/automatic-compaction-api.md

@@ -27,7 +27,13 @@ import TabItem from '@theme/TabItem';
   ~ under the License.
   -->
 
-This topic describes the status and configuration API endpoints for [automatic compaction](../data-management/automatic-compaction.md) in Apache Druid. You can configure automatic compaction in the Druid web console or API.
+This topic describes the status and configuration API endpoints for [automatic compaction using Coordinator duties](../data-management/automatic-compaction.md#auto-compaction-using-coordinator-duties) in Apache Druid. You can configure automatic compaction in the Druid web console or API.
+
+:::info Experimental
+
+Instead of the automatic compaction API, you can use the supervisor API to submit auto-compaction jobs using compaction supervisors. For more information, see [Auto-compaction using compaction supervisors](../data-management/automatic-compaction.md#auto-compaction-using-compaction-supervisors).
+
+:::
 
 In this topic, `http://ROUTER_IP:ROUTER_PORT` is a placeholder for your Router service address and port. Replace it with the information for your deployment. For example, use `http://localhost:8888` for quickstart deployments.
 

docs/configuration/index.md

@@ -1050,7 +1050,7 @@ The following table shows the supported configurations for auto-compaction.
 
 |Property|Description|Required|
 |--------|-----------|--------|
-|type|The task type, this should always be `index_parallel`.|yes|
+|type|The task type. If you're using Coordinator duties for auto-compaction, set it to `index_parallel`. If you're using compaction supervisors, set it to `autocompact`. |yes|
 |`maxRowsInMemory`|Used in determining when intermediate persists to disk should occur. Normally user does not need to set this, but depending on the nature of data, if rows are short in terms of bytes, user may not want to store a million rows in memory and this value should be set.|no (default = 1000000)|
 |`maxBytesInMemory`|Used in determining when intermediate persists to disk should occur. Normally this is computed internally and user does not need to set it. This value represents number of bytes to aggregate in heap memory before persisting. This is based on a rough estimate of memory usage and not actual usage. The maximum heap memory usage for indexing is `maxBytesInMemory` * (2 + `maxPendingPersists`)|no (default = 1/6 of max JVM memory)|
 |`splitHintSpec`|Used to give a hint to control the amount of data that each first phase task reads. This hint could be ignored depending on the implementation of the input source. See [Split hint spec](../ingestion/native-batch.md#split-hint-spec) for more details.|no (default = size-based split hint spec)|
@@ -1067,6 +1067,7 @@ The following table shows the supported configurations for auto-compaction.
 |`taskStatusCheckPeriodMs`|Polling period in milliseconds to check running task statuses.|no (default = 1000)|
 |`chatHandlerTimeout`|Timeout for reporting the pushed segments in worker tasks.|no (default = PT10S)|
 |`chatHandlerNumRetries`|Retries for reporting the pushed segments in worker tasks.|no (default = 5)|
+|`engine` | Engine for compaction. Can be either `native` or `msq`. `msq`  uses the MSQ task engine and is only supported with [compaction supervisors](../data-management/automatic-compaction.md#auto-compaction-using-compaction-supervisors). | no (default = native)|
 
 ###### Automatic compaction granularitySpec
 
@@ -2218,7 +2219,6 @@ Supported query contexts:
 |Key|Description|Default|
 |---|-----------|-------|
 |`druid.expressions.useStrictBooleans`|Controls the behavior of Druid boolean operators and functions, if set to `true` all boolean values are either `1` or `0`. This configuration has been deprecated and will be removed in a future release, taking on the `true` behavior. See [expression documentation](../querying/math-expr.md#logical-operator-modes) for more information.|true|
-|`druid.expressions.allowNestedArrays`|If enabled, Druid array expressions can create nested arrays. This configuration has been deprecated and will be removed in a future release, taking on the `true` behavior.|true|
 
 ### Router
 

docs/data-management/automatic-compaction.md

@@ -22,19 +22,7 @@ title: "Automatic compaction"
   ~ under the License.
   -->
 
-In Apache Druid, compaction is a special type of ingestion task that reads data from a Druid datasource and writes it back into the same datasource. A common use case for this is to [optimally size segments](../operations/segment-optimization.md) after ingestion to improve query performance. Automatic compaction, or auto-compaction, refers to the system for automatic execution of compaction tasks managed by the [Druid Coordinator](../design/coordinator.md).
-This topic guides you through setting up automatic compaction for your Druid cluster. See the [examples](#examples) for common use cases for automatic compaction.
-
-## How Druid manages automatic compaction
-
-The Coordinator [indexing period](../configuration/index.md#coordinator-operation), `druid.coordinator.period.indexingPeriod`, controls the frequency of compaction tasks.
-The default indexing period is 30 minutes, meaning that the Coordinator first checks for segments to compact at most 30 minutes from when auto-compaction is enabled.
-This time period affects other Coordinator duties including merge and conversion tasks.
-To configure the auto-compaction time period without interfering with `indexingPeriod`, see [Set frequency of compaction runs](#set-frequency-of-compaction-runs).
-
-At every invocation of auto-compaction, the Coordinator initiates a [segment search](../design/coordinator.md#segment-search-policy-in-automatic-compaction) to determine eligible segments to compact.
-When there are eligible segments to compact, the Coordinator issues compaction tasks based on available worker capacity.
-If a compaction task takes longer than the indexing period, the Coordinator waits for it to finish before resuming the period for segment search.
+In Apache Druid, compaction is a special type of ingestion task that reads data from a Druid datasource and writes it back into the same datasource. A common use case for this is to [optimally size segments](../operations/segment-optimization.md) after ingestion to improve query performance. Automatic compaction, or auto-compaction, refers to the system for automatic execution of compaction tasks issued by Druid itself. In addition to auto-compaction, you can perform [manual compaction](./manual-compaction.md) using the Overlord APIs.
 
 :::info
  Auto-compaction skips datasources that have a segment granularity of `ALL`.
@@ -42,53 +30,9 @@ If a compaction task takes longer than the indexing period, the Coordinator wait
 
 As a best practice, you should set up auto-compaction for all Druid datasources. You can run compaction tasks manually for cases where you want to allocate more system resources. For example, you may choose to run multiple compaction tasks in parallel to compact an existing datasource for the first time. See [Compaction](compaction.md) for additional details and use cases.
 
+This topic guides you through setting up automatic compaction for your Druid cluster. See the [examples](#examples) for common use cases for automatic compaction.
 
-## Enable automatic compaction
-
-You can enable automatic compaction for a datasource using the web console or programmatically via an API.
-This process differs for manual compaction tasks, which can be submitted from the [Tasks view of the web console](../operations/web-console.md) or the [Tasks API](../api-reference/tasks-api.md).
-
-### Web console
-
-Use the web console to enable automatic compaction for a datasource as follows.
-
-1. Click **Datasources** in the top-level navigation.
-2. In the **Compaction** column, click the edit icon for the datasource to compact.
-3. In the **Compaction config** dialog, configure the auto-compaction settings. The dialog offers a form view as well as a JSON view. Editing the form updates the JSON specification, and editing the JSON updates the form field, if present. Form fields not present in the JSON indicate default values. You may add additional properties to the JSON for auto-compaction settings not displayed in the form. See [Configure automatic compaction](#configure-automatic-compaction) for supported settings for auto-compaction.
-4. Click **Submit**.
-5. Refresh the **Datasources** view. The **Compaction** column for the datasource changes from “Not enabled” to “Awaiting first run.”
-
-The following screenshot shows the compaction config dialog for a datasource with auto-compaction enabled.
-![Compaction config in web console](../assets/compaction-config.png)
-
-To disable auto-compaction for a datasource, click **Delete** from the **Compaction config** dialog. Druid does not retain your auto-compaction configuration.
-
-### Compaction configuration API
-
-Use the [Automatic compaction API](../api-reference/automatic-compaction-api.md#manage-automatic-compaction) to configure automatic compaction.
-To enable auto-compaction for a datasource, create a JSON object with the desired auto-compaction settings.
-See [Configure automatic compaction](#configure-automatic-compaction) for the syntax of an auto-compaction spec.
-Send the JSON object as a payload in a [`POST` request](../api-reference/automatic-compaction-api.md#create-or-update-automatic-compaction-configuration) to `/druid/coordinator/v1/config/compaction`.
-The following example configures auto-compaction for the `wikipedia` datasource:
-
-```sh
-curl --location --request POST 'http://localhost:8081/druid/coordinator/v1/config/compaction' \
---header 'Content-Type: application/json' \
---data-raw '{
-    "dataSource": "wikipedia",
-    "granularitySpec": {
-        "segmentGranularity": "DAY"
-    }
-}'
-```
-
-To disable auto-compaction for a datasource, send a [`DELETE` request](../api-reference/automatic-compaction-api.md#remove-automatic-compaction-configuration) to `/druid/coordinator/v1/config/compaction/{dataSource}`. Replace `{dataSource}` with the name of the datasource for which to disable auto-compaction. For example:
-
-```sh
-curl --location --request DELETE 'http://localhost:8081/druid/coordinator/v1/config/compaction/wikipedia'
-```
-
-## Configure automatic compaction
+## Auto-compaction syntax
 
 You can configure automatic compaction dynamically without restarting Druid.
 The automatic compaction system uses the following syntax:
@@ -108,6 +52,14 @@ The automatic compaction system uses the following syntax:
 }
 ```
 
+:::info Experimental
+
+The MSQ task engine is available as a compaction engine when you run automatic compaction as a compaction supervisor. For more information, see [Auto-compaction using compaction supervisors](#auto-compaction-using-compaction-supervisors).
+
+:::
+
+For automatic compaction using Coordinator duties, you submit the spec to the [Compaction config UI](#manage-auto-compaction-using-the-web-console) or the [Compaction configuration API](#manage-auto-compaction-using-coordinator-apis).
+
 Most fields in the auto-compaction configuration correlate to a typical [Druid ingestion spec](../ingestion/ingestion-spec.md).
 The following properties only apply to auto-compaction:
 * `skipOffsetFromLatest`
@@ -131,7 +83,62 @@ maximize performance and minimize disk usage of the `compact` tasks launched by
 
 For more details on each of the specs in an auto-compaction configuration, see [Automatic compaction dynamic configuration](../configuration/index.md#automatic-compaction-dynamic-configuration).
 
-### Set frequency of compaction runs
+## Auto-compaction using Coordinator duties
+
+You can control how often the Coordinator checks to see if auto-compaction is needed. The Coordinator [indexing period](../configuration/index.md#coordinator-operation), `druid.coordinator.period.indexingPeriod`, controls the frequency of compaction tasks.
+The default indexing period is 30 minutes, meaning that the Coordinator first checks for segments to compact at most 30 minutes from when auto-compaction is enabled.
+This time period also affects other Coordinator duties such as cleanup of unused segments and stale pending segments.
+To configure the auto-compaction time period without interfering with `indexingPeriod`, see [Set frequency of compaction runs](#change-compaction-frequency).
+
+At every invocation of auto-compaction, the Coordinator initiates a [segment search](../design/coordinator.md#segment-search-policy-in-automatic-compaction) to determine eligible segments to compact.
+When there are eligible segments to compact, the Coordinator issues compaction tasks based on available worker capacity.
+If a compaction task takes longer than the indexing period, the Coordinator waits for it to finish before resuming the period for segment search.
+
+No additional configuration is needed to run automatic compaction tasks using the Coordinator and native engine. This is the default behavior for Druid.
+You can configure it for a datasource through the web console or programmatically via an API.
+This process differs for manual compaction tasks, which can be submitted from the [Tasks view of the web console](../operations/web-console.md) or the [Tasks API](../api-reference/tasks-api.md).
+
+### Manage auto-compaction using the web console
+
+Use the web console to enable automatic compaction for a datasource as follows:
+
+1. Click **Datasources** in the top-level navigation.
+2. In the **Compaction** column, click the edit icon for the datasource to compact.
+3. In the **Compaction config** dialog, configure the auto-compaction settings. The dialog offers a form view as well as a JSON view. Editing the form updates the JSON specification, and editing the JSON updates the form field, if present. Form fields not present in the JSON indicate default values. You may add additional properties to the JSON for auto-compaction settings not displayed in the form. See [Configure automatic compaction](#auto-compaction-syntax) for supported settings for auto-compaction.
+4. Click **Submit**.
+5. Refresh the **Datasources** view. The **Compaction** column for the datasource changes from “Not enabled” to “Awaiting first run.”
+
+The following screenshot shows the compaction config dialog for a datasource with auto-compaction enabled.
+![Compaction config in web console](../assets/compaction-config.png)
+
+To disable auto-compaction for a datasource, click **Delete** from the **Compaction config** dialog. Druid does not retain your auto-compaction configuration.
+
+### Manage auto-compaction using Coordinator APIs  
+
+Use the [Automatic compaction API](../api-reference/automatic-compaction-api.md#manage-automatic-compaction) to configure automatic compaction.
+To enable auto-compaction for a datasource, create a JSON object with the desired auto-compaction settings.
+See [Configure automatic compaction](#auto-compaction-syntax) for the syntax of an auto-compaction spec.
+Send the JSON object as a payload in a [`POST` request](../api-reference/automatic-compaction-api.md#create-or-update-automatic-compaction-configuration) to `/druid/coordinator/v1/config/compaction`.
+The following example configures auto-compaction for the `wikipedia` datasource:
+
+```sh
+curl --location --request POST 'http://localhost:8081/druid/coordinator/v1/config/compaction' \
+--header 'Content-Type: application/json' \
+--data-raw '{
+    "dataSource": "wikipedia",
+    "granularitySpec": {
+        "segmentGranularity": "DAY"
+    }
+}'
+```
+
+To disable auto-compaction for a datasource, send a [`DELETE` request](../api-reference/automatic-compaction-api.md#remove-automatic-compaction-configuration) to `/druid/coordinator/v1/config/compaction/{dataSource}`. Replace `{dataSource}` with the name of the datasource for which to disable auto-compaction. For example:
+
+```sh
+curl --location --request DELETE 'http://localhost:8081/druid/coordinator/v1/config/compaction/wikipedia'
+```
+
+### Change compaction frequency
 
 If you want the Coordinator to check for compaction more frequently than its indexing period, create a separate group to handle compaction duties.
 Set the time period of the duty group in the `coordinator/runtime.properties` file.
@@ -142,6 +149,15 @@ druid.coordinator.compaction.duties=["compactSegments"]
 druid.coordinator.compaction.period=PT60S
 ```
 
+### View Coordinator duty auto-compaction stats
+
+After the Coordinator has initiated auto-compaction, you can view compaction statistics for the datasource, including the number of bytes, segments, and intervals already compacted and those awaiting compaction. The Coordinator also reports the total bytes, segments, and intervals not eligible for compaction in accordance with its [segment search policy](../design/coordinator.md#segment-search-policy-in-automatic-compaction).
+
+In the web console, the Datasources view displays auto-compaction statistics. The Tasks view shows the task information for compaction tasks that were triggered by the automatic compaction system.
+
+To get statistics by API, send a [`GET` request](../api-reference/automatic-compaction-api.md#view-automatic-compaction-status) to `/druid/coordinator/v1/compaction/status`. To filter the results to a particular datasource, pass the datasource name as a query parameter to the request—for example, `/druid/coordinator/v1/compaction/status?dataSource=wikipedia`.
+
+
 ## Avoid conflicts with ingestion
 
 Compaction tasks may be interrupted when they interfere with ingestion. For example, this occurs when an ingestion task needs to write data to a segment for a time interval locked for compaction. If there are continuous failures that prevent compaction from making progress, consider one of the following strategies:
@@ -169,15 +185,6 @@ The Coordinator compacts segments from newest to oldest. In the auto-compaction
 
 To set `skipOffsetFromLatest`, consider how frequently you expect the stream to receive late arriving data. If your stream only occasionally receives late arriving data, the auto-compaction system robustly compacts your data even though data is ingested outside the `skipOffsetFromLatest` window. For most realtime streaming ingestion use cases, it is reasonable to set `skipOffsetFromLatest` to a few hours or a day.
 
-
-## View automatic compaction statistics
-
-After the Coordinator has initiated auto-compaction, you can view compaction statistics for the datasource, including the number of bytes, segments, and intervals already compacted and those awaiting compaction. The Coordinator also reports the total bytes, segments, and intervals not eligible for compaction in accordance with its [segment search policy](../design/coordinator.md#segment-search-policy-in-automatic-compaction).
-
-In the web console, the Datasources view displays auto-compaction statistics. The Tasks view shows the task information for compaction tasks that were triggered by the automatic compaction system.
-
-To get statistics by API, send a [`GET` request](../api-reference/automatic-compaction-api.md#view-automatic-compaction-status) to `/druid/coordinator/v1/compaction/status`. To filter the results to a particular datasource, pass the datasource name as a query parameter to the request—for example, `/druid/coordinator/v1/compaction/status?dataSource=wikipedia`.
-
 ## Examples
 
 The following examples demonstrate potential use cases in which auto-compaction may improve your Druid performance. See more details in [Compaction strategies](../data-management/compaction.md#compaction-guidelines). The examples in this section do not change the underlying data.
@@ -221,6 +228,137 @@ The following auto-compaction configuration compacts updates the `wikipedia` seg
 }
 ```
 
+## Auto-compaction using compaction supervisors  
+
+:::info Experimental
+Compaction supervisors are experimental. For production use, we recommend [auto-compaction using Coordinator duties](#auto-compaction-using-coordinator-duties).
+:::
+
+You can run automatic compaction using compaction supervisors on the Overlord rather than Coordinator duties. Compaction supervisors provide the following benefits over Coordinator duties:
+
+* Can use the supervisor framework to get information about the auto-compaction, such as status or state
+* More easily suspend or resume compaction for a datasource
+* Can use either the native compaction engine or the [MSQ task engine](#use-msq-for-auto-compaction)
+* More reactive and submits tasks as soon as a compaction slot is available
+* Tracked compaction task status to avoid re-compacting an interval repeatedly
+
+
+To use compaction supervisors, set the following properties in your Overlord runtime properties:
+  *  `druid.supervisor.compaction.enabled` to `true` so that compaction tasks can be run as supervisor tasks
+  *  `druid.supervisor.compaction.engine` to  `msq` to specify the MSQ task engine as the compaction engine or to `native` to use the native engine. This is the default engine if the `engine` field is omitted from your compaction config
+
+Compaction supervisors use the same syntax as auto-compaction using  Coordinator duties with one key difference: you submit the auto-compaction as a a supervisor spec. In the spec, set the `type` to `autocompact` and include the auto-compaction config in the `spec`.
+
+To submit an automatic compaction task, you can submit a supervisor spec through the [web console](#manage-compaction-supervisors-with-the-web-console) or the [supervisor API](#manage-compaction-supervisors-with-supervisor-apis).
+
+
+### Manage compaction supervisors with the web console
+
+To submit a supervisor spec for MSQ task engine automatic compaction, perform the following steps:
+
+1. In the web console, go to the **Supervisors** tab.
+1. Click **...** > **Submit JSON supervisor**.
+1. In the dialog, include the following:
+     - The type of supervisor spec by setting `"type": "autocompact"`
+     - The compaction configuration by adding it to the `spec` field
+    ```json
+    {
+     "type": "autocompact",
+     "spec": {
+       "dataSource": YOUR_DATASOURCE,
+       "tuningConfig": {...},
+       "granularitySpec": {...},
+       "engine": <native|msq>,
+       ...
+    }
+    ```
+1. Submit the supervisor.
+
+To stop the automatic compaction task, suspend or terminate the supervisor through the UI or API.
+
+### Manage compaction supervisors with supervisor APIs
+
+Submitting an automatic compaction as a supervisor task uses the same endpoint as supervisor tasks for streaming ingestion.
+
+The following example configures auto-compaction for the `wikipedia` datasource:
+
+```sh
+curl --location --request POST 'http://localhost:8081/druid/indexer/v1/supervisor' \
+--header 'Content-Type: application/json' \
+--data-raw '{
+   "type": "autocompact",                     // required
+   "suspended": false,                        // optional 
+   "spec": {                                  // required
+       "dataSource": "wikipedia",             // required
+       "tuningConfig": {...},                 // optional
+       "granularitySpec": {...},              // optional
+       "engine": <native|msq>,                // optional
+       ...
+   }
+}'
+```
+
+Note that if you omit `spec.engine`, Druid uses the default compaction engine. You can control the default compaction engine with the `druid.supervisor.compaction.engine` Overlord runtime property. If `spec.engine` and `druid.supervisor.compaction.engine` are omitted, Druid defaults to the native engine.
+
+To stop the automatic compaction task, suspend or terminate the supervisor through the UI or API.
+
+### Use MSQ for auto-compaction
+
+The MSQ task engine is available as a compaction engine if you configure auto-compaction to use compaction supervisors. To use the MSQ task engine for automatic compaction, make sure the following requirements are met:
+
+* [Load the MSQ task engine extension](../multi-stage-query/index.md#load-the-extension).
+* In your Overlord runtime properties, set the following properties:
+  *  `druid.supervisor.compaction.enabled` to `true` so that compaction tasks can be run as a supervisor task.
+  *  Optionally, set `druid.supervisor.compaction.engine` to `msq` to specify the MSQ task engine as the default compaction engine. If you don't do this, you'll need to set `spec.engine` to `msq` for each compaction supervisor spec where you want to use the MSQ task engine.
+* Have at least two compaction task slots available or set `compactionConfig.taskContext.maxNumTasks` to two or more. The MSQ task engine requires at least two tasks to run, one controller task and one worker task.
+
+You can use [MSQ task engine context parameters](../multi-stage-query/reference.md#context-parameters) in `spec.taskContext` when configuring your datasource for automatic compaction, such as setting the maximum number of tasks using the `spec.taskContext.maxNumTasks` parameter. Some of the MSQ task engine context parameters overlap with automatic compaction parameters. When these settings overlap, set one or the other.
+
+
+#### MSQ task engine limitations
+
+<!--This list also exists in multi-stage-query/known-issues-->
+
+When using the MSQ task engine for auto-compaction, keep the following limitations in mind:
+
+- The `metricSpec` field is only supported for certain aggregators. For more information, see [Supported aggregators](#supported-aggregators).
+- Only dynamic and range-based partitioning are supported.
+- Set `rollup`  to `true` if and only if `metricSpec` is not empty or null.
+- You can only partition on string dimensions. However, multi-valued string dimensions are not supported.
+- The `maxTotalRows` config is not supported in `DynamicPartitionsSpec`. Use `maxRowsPerSegment` instead.
+- Segments can only be sorted on `__time` as the first column.
+
+#### Supported aggregators
+
+Auto-compaction using the MSQ task engine supports only aggregators that satisfy the following properties: 
+* __Mergeability__: can combine partial aggregates
+* __Idempotency__: produces the same results on repeated runs of the aggregator on previously aggregated values in a column
+
+This is exemplified by the following `longSum` aggregator:
+
+```
+{"name": "added", "type": "longSum", "fieldName": "added"}
+```
+
+where `longSum` being capable of combining partial results satisfies mergeability, while input and output column being the same (`added`) ensures idempotency.
+
+The following are some examples of aggregators that aren't supported since at least one of the required conditions aren't satisfied:
+
+*  `longSum` aggregator where the `added` column rolls up into `sum_added` column discarding the input `added` column, violating idempotency, as subsequent runs would no longer find the `added` column:
+    ```
+    {"name": "sum_added", "type": "longSum", "fieldName": "added"}
+    ```
+* Partial sketches which cannot themselves be used to combine partial aggregates and need merging aggregators -- such as `HLLSketchMerge` required for `HLLSketchBuild` aggregator below -- violating mergeability:
+    ```
+    {"name": "added", "type": "HLLSketchBuild", "fieldName": "added"}
+    ```
+* Count aggregator since it cannot be used to combine partial aggregates and it rolls up into a different `count` column discarding the input column(s), violating both mergeability and idempotency.
+    ```
+    {"type": "count", "name": "count"}
+    ```
+
+
+
 ## Learn more
 
 See the following topics for more information:

docs/design/architecture.md

@@ -41,7 +41,7 @@ Druid has several types of services:
 * [Router](../design/router.md) routes requests to Brokers, Coordinators, and Overlords.
 * [Historical](../design/historical.md) stores queryable data.
 * [Middle Manager](../design/middlemanager.md) and [Peon](../design/peons.md) ingest data.
-* [Indexer](../design/indexer.md) serves an alternative to the Middle Manager + Peon task execution system.
+* [Indexer](../design/indexer.md) serves as an alternative to the Middle Manager + Peon task execution system.
 
 You can view services in the **Services** tab in the web console: 
 
@@ -105,10 +105,9 @@ for reading from external data sources and publishing new Druid segments.
 [**Indexer**](../design/indexer.md) services are an alternative to Middle Managers and Peons. Instead of
 forking separate JVM processes per-task, the Indexer runs tasks as individual threads within a single JVM process.
 
-The Indexer is designed to be easier to configure and deploy compared to the Middle Manager + Peon system and to better enable resource sharing across tasks. The Indexer is a newer feature and is currently designated [experimental](../development/experimental.md) due to the fact that its memory management system is still under
-development. It will continue to mature in future versions of Druid.
+The Indexer is designed to be easier to configure and deploy compared to the MiddleManager + Peon system and to better enable resource sharing across tasks, which can help streaming ingestion. The Indexer is currently designated [experimental](../development/experimental.md).
 
-Typically, you would deploy either Middle Managers or Indexers, but not both.
+Typically, you would deploy one of the following: MiddleManagers, [MiddleManager-less ingestion using Kubernetes](../development/extensions-contrib/k8s-jobs.md), or Indexers. You wouldn't deploy more than one of these options.
 
 ## Colocation of services
 

docs/design/indexer.md

@@ -24,8 +24,7 @@ sidebar_label: "Indexer"
   -->
 
 :::info
- The Indexer is an optional and [experimental](../development/experimental.md) feature.
- Its memory management system is still under development and will be significantly enhanced in later releases.
+ The Indexer is an optional and experimental feature. If you're primarily performing batch ingestion, we recommend you use either the MiddleManager and Peon task execution system or [MiddleManager-less ingestion using Kubernetes](../development/extensions-contrib/k8s-jobs.md). If you're primarily doing streaming ingestion, you may want to try either [MiddleManager-less ingestion using Kubernetes](../development/extensions-contrib/k8s-jobs.md) or the Indexer service.
 :::
 
 The Apache Druid Indexer service is an alternative to the Middle Manager + Peon task execution system. Instead of forking a separate JVM process per-task, the Indexer runs tasks as separate threads within a single JVM process.

docs/development/extensions-core/kubernetes.md

@@ -54,7 +54,7 @@ Additionally, this extension has following configuration.
 
 ### Gotchas
 
-- Label/Annotation path in each pod spec MUST EXIST, which is easily satisfied if there is at least one label/annotation in the pod spec already. This limitation may be removed in future.
+- Label/Annotation path in each pod spec MUST EXIST, which is easily satisfied if there is at least one label/annotation in the pod spec already. 
 - All Druid Pods belonging to one Druid cluster must be inside same kubernetes namespace.
 - All Druid Pods need permissions to be able to add labels to self-pod, List and Watch other Pods, create and read ConfigMap for leader election. Assuming, "default" service account is used by Druid pods, you might need to add following or something similar Kubernetes Role and Role Binding.
 

docs/ingestion/concurrent-append-replace.md

@@ -22,10 +22,6 @@ title: Concurrent append and replace
   ~ under the License.
   -->
 
-:::info
-Concurrent append and replace is an [experimental feature](../development/experimental.md) available for JSON-based batch, streaming, and SQL-based ingestion.
-:::
-
 Concurrent append and replace safely replaces the existing data in an interval of a datasource while new data is being appended to that interval. One of the most common applications of this feature is appending new data (such as with streaming ingestion) to an interval while compaction of that interval is already in progress. Druid segments the data ingested during this time dynamically. The subsequent compaction run segments the data into the  granularity you specified.
 
 To set up concurrent append and replace, use the context flag `useConcurrentLocks`. Druid will then determine the correct lock type for you, either append or replace. Although you can set the type of lock manually, we don't recommend it. 
@@ -38,7 +34,7 @@ If you want to append data to a datasource while compaction is running, you need
 
 In the **Compaction config** for a datasource, enable  **Use concurrent locks (experimental)**.
 
-For details on accessing the compaction config in the UI, see [Enable automatic compaction with the web console](../data-management/automatic-compaction.md#web-console).
+For details on accessing the compaction config in the UI, see [Enable automatic compaction with the web console](../data-management/automatic-compaction.md#manage-auto-compaction-using-the-web-console).
 
 ### Update the compaction settings with the API
  

docs/ingestion/supervisor.md

@@ -23,22 +23,22 @@ sidebar_label: Supervisor
   ~ under the License.
   -->
 
-A supervisor manages streaming ingestion from external streaming sources into Apache Druid.
-Supervisors oversee the state of indexing tasks to coordinate handoffs, manage failures, and ensure that the scalability and replication requirements are maintained.
+Apache Druid uses supervisors to manage streaming ingestion from external streaming sources into Druid.
+Supervisors oversee the state of indexing tasks to coordinate handoffs, manage failures, and ensure that the scalability and replication requirements are maintained. They can also be used to perform [automatic compaction](../data-management/automatic-compaction.md) after data has been ingested.
 
 This topic uses the Apache Kafka term offset to refer to the identifier for records in a partition. If you are using Amazon Kinesis, the equivalent is sequence number.
 
 ## Supervisor spec
 
-Druid uses a JSON specification, often referred to as the supervisor spec, to define streaming ingestion tasks.
-The supervisor spec specifies how Druid should consume, process, and index streaming data.
+Druid uses a JSON specification, often referred to as the supervisor spec, to define tasks used for streaming ingestion or auto-compaction.
+The supervisor spec specifies how Druid should consume, process, and index data from an external stream or Druid itself.
 
 The following table outlines the high-level configuration options for a supervisor spec:
 
 |Property|Type|Description|Required|
 |--------|----|-----------|--------|
-|`type`|String|The supervisor type. One of `kafka`or `kinesis`.|Yes|
-|`spec`|Object|The container object for the supervisor configuration.|Yes|
+|`type`|String|The supervisor type. For streaming ingestion, this can be either `kafka`, `kinesis`, or `rabbit`. For automatic compaction, set the type to `autocompact`. |Yes|
+|`spec`|Object|The container object for the supervisor configuration. For automatic compaction, this is the same as the compaction configuration. |Yes|
 |`spec.dataSchema`|Object|The schema for the indexing task to use during ingestion. See [`dataSchema`](../ingestion/ingestion-spec.md#dataschema) for more information.|Yes|
 |`spec.ioConfig`|Object|The I/O configuration object to define the connection and I/O-related settings for the supervisor and indexing tasks.|Yes|
 |`spec.tuningConfig`|Object|The tuning configuration object to define performance-related settings for the supervisor and indexing tasks.|No|

docs/multi-stage-query/known-issues.md

@@ -68,3 +68,16 @@ properties, and the `indexSpec` [`tuningConfig`](../ingestion/ingestion-spec.md#
 - The maximum number of elements in a window cannot exceed a value of 100,000. 
 - To avoid `leafOperators` in MSQ engine, window functions have an extra scan stage after the window stage for cases 
 where native engine has a non-empty `leafOperator`.
+
+## Automatic compaction
+
+<!-- If you update this list, also update data-management/automatic-compaction.md -->
+
+The following known issues and limitations affect automatic compaction with the MSQ task engine:
+
+- The `metricSpec` field is only supported for certain aggregators. For more information, see [Supported aggregators](../data-management/automatic-compaction.md#supported-aggregators).
+- Only dynamic and range-based partitioning are supported.
+- Set `rollup`  to `true` if and only if `metricSpec` is not empty or null.
+- You can only partition on string dimensions. However, multi-valued string dimensions are not supported.
+- The `maxTotalRows` config is not supported in `DynamicPartitionsSpec`. Use `maxRowsPerSegment` instead.
+- Segments can only be sorted on `__time` as the first column.

docs/querying/datasource.md

@@ -431,25 +431,21 @@ and how to detect it.
 3. One common reason for implicit subquery generation is if the types of the two halves of an equality do not match.
 For example, since lookup keys are always strings, the condition `druid.d JOIN lookup.l ON d.field = l.field` will
 perform best if `d.field` is a string.
-4. The join operator must evaluate the condition for each row. In the future, we expect
-to implement both early and deferred condition evaluation, which we expect to improve performance considerably for
-common use cases.
+4. The join operator must evaluate the condition for each row. 
 5. Currently, Druid does not support pushing down predicates (condition and filter) past a Join (i.e. into
 Join's children). Druid only supports pushing predicates into the join if they originated from
 above the join. Hence, the location of predicates and filters in your Druid SQL is very important.
 Also, as a result of this, comma joins should be avoided.
 
-#### Future work for joins
+#### Limitations for joins
 
-Joins are an area of active development in Druid. The following features are missing today but may appear in
-future versions:
+Joins in Druid have the following limitations:
 
-- Reordering of join operations to get the most performant plan.
-- Preloaded dimension tables that are wider than lookups (i.e. supporting more than a single key and single value).
-- RIGHT OUTER and FULL OUTER joins in the native query engine. Currently, they are partially implemented. Queries run
+- The order of joins is not entirely optimized. Join operations are not reordered to get the most performant plan.
+- Preloaded dimension tables that are wider than lookups (i.e. supporting more than a single key and single value) are not supported.
+- RIGHT OUTER and FULL OUTER joins in the native query engine are not fully implemented. Queries run
   but results are not always correct.
-- Performance-related optimizations as mentioned in the [previous section](#join-performance).
-- Join conditions on a column containing a multi-value dimension.
+- Join conditions on a column can't contain a multi-value dimension.
 
 ### `unnest`
 

docs/querying/sql-scalar.md

@@ -166,13 +166,13 @@ overhead.
 |`TIME_CEIL(timestamp_expr, period[, origin[, timezone]])`|Rounds up a timestamp, returning it as a new timestamp. Period can be any ISO 8601 period, like P3M (quarters) or PT12H (half-days). Specify `origin` as a timestamp to set the reference time for rounding. For example, `TIME_CEIL(__time, 'PT1H', TIMESTAMP '2016-06-27 00:30:00')` measures an hourly period from 00:30-01:30 instead of 00:00-01:00. See [Period granularities](granularities.md) for details on the default starting boundaries. The time zone, if provided, should be a time zone name like `America/Los_Angeles` or an offset like `-08:00`. This function is similar to `CEIL` but is more flexible.|
 |`TIME_FLOOR(timestamp_expr, period[, origin[, timezone]])`|Rounds down a timestamp, returning it as a new timestamp. Period can be any ISO 8601 period, like P3M (quarters) or PT12H (half-days). Specify `origin` as a timestamp to set the reference time for rounding. For example, `TIME_FLOOR(__time, 'PT1H', TIMESTAMP '2016-06-27 00:30:00')` measures an hourly period from 00:30-01:30 instead of 00:00-01:00. See [Period granularities](granularities.md) for details on the default starting boundaries. The time zone, if provided, should be a time zone name like `America/Los_Angeles` or an offset like `-08:00`. This function is similar to `FLOOR` but is more flexible.|
 |`TIME_SHIFT(timestamp_expr, period, step[, timezone])`|Shifts a timestamp by a period (step times), returning it as a new timestamp. The `period` parameter can be any ISO 8601 period. The `step` parameter can be negative. The time zone, if provided, should be a time zone name like `America/Los_Angeles` or an offset like `-08:00`.|
-|`TIME_EXTRACT(timestamp_expr, unit[, timezone])`|Extracts a time part from `expr`, returning it as a number. Unit can be EPOCH, SECOND, MINUTE, HOUR, DAY (day of month), DOW (day of week), DOY (day of year), WEEK (week of [week year](https://en.wikipedia.org/wiki/ISO_week_date)), MONTH (1 through 12), QUARTER (1 through 4), or YEAR. The time zone, if provided, should be a time zone name like `America/Los_Angeles` or an offset like `-08:00`. The `unit` and `timezone` parameters must be provided as quoted literals, such as `TIME_EXTRACT(__time, 'HOUR')` or `TIME_EXTRACT(__time, 'HOUR', 'America/Los_Angeles')`. This function is similar to `EXTRACT` but is more flexible. |
+|`TIME_EXTRACT(timestamp_expr, unit[, timezone])`| Extracts a time part from `expr`, returning it as a number. Unit can be EPOCH, MILLISECOND, SECOND, MINUTE, HOUR, DAY (day of month), DOW (day of week), DOY (day of year), WEEK (week of [week year](https://en.wikipedia.org/wiki/ISO_week_date)), MONTH (1 through 12), QUARTER (1 through 4), or YEAR. The time zone, if provided, should be a time zone name like `America/Los_Angeles` or an offset like `-08:00`. The `unit` and `timezone` parameters must be provided as quoted literals, such as `TIME_EXTRACT(__time, 'HOUR')` or `TIME_EXTRACT(__time, 'HOUR', 'America/Los_Angeles')`. This function is similar to `EXTRACT` but is more flexible. |
 |`TIME_PARSE(string_expr[, pattern[, timezone]])`|Parses a string into a timestamp using a given [Joda DateTimeFormat pattern](http://www.joda.org/joda-time/apidocs/org/joda/time/format/DateTimeFormat.html), or ISO 8601 (e.g. `2000-01-02T03:04:05Z`) if the pattern is not provided. The `timezone` parameter is used as the time zone for strings that do not already include a time zone offset. If provided, `timezone` should be a time zone name like `America/Los_Angeles` or an offset like `-08:00`. The `pattern` and `timezone` parameters must be literals. Strings that cannot be parsed as timestamps return NULL.|
 |`TIME_FORMAT(timestamp_expr[, pattern[, timezone]])`|Formats a timestamp as a string with a given [Joda DateTimeFormat pattern](http://www.joda.org/joda-time/apidocs/org/joda/time/format/DateTimeFormat.html), or ISO 8601 (e.g. `2000-01-02T03:04:05Z`) if the pattern is not provided. If provided, the `timezone` parameter should be a time zone name like `America/Los_Angeles` or an offset like `-08:00`. The `pattern` and `timezone` parameters must be literals.|
 |`TIME_IN_INTERVAL(timestamp_expr, interval)`|Returns whether a timestamp is contained within a particular interval. The interval must be a literal string containing any ISO 8601 interval, such as `'2001-01-01/P1D'` or `'2001-01-01T01:00:00/2001-01-02T01:00:00'`. The start instant of the interval is inclusive and the end instant is exclusive.|
 |`MILLIS_TO_TIMESTAMP(millis_expr)`|Converts a number of milliseconds since the epoch (1970-01-01 00:00:00 UTC) into a timestamp.|
 |`TIMESTAMP_TO_MILLIS(timestamp_expr)`|Converts a timestamp into a number of milliseconds since the epoch.|
-|`EXTRACT(unit FROM timestamp_expr)`|Extracts a time part from `expr`, returning it as a number. Unit can be EPOCH, MICROSECOND, MILLISECOND, SECOND, MINUTE, HOUR, DAY (day of month), DOW (day of week), ISODOW (ISO day of week), DOY (day of year), WEEK (week of year), MONTH, QUARTER, YEAR, ISOYEAR, DECADE, CENTURY or MILLENNIUM. Units must be provided unquoted, like `EXTRACT(HOUR FROM __time)`.|
+|`EXTRACT(unit FROM timestamp_expr)`| Extracts a time part from `expr`, returning it as a number. Unit can be EPOCH, MILLISECOND, SECOND, MINUTE, HOUR, DAY (day of month), DOW (day of week), ISODOW (ISO day of week), DOY (day of year), WEEK (week of year), MONTH, QUARTER, YEAR, ISOYEAR, DECADE, CENTURY or MILLENNIUM. Units must be provided unquoted, like `EXTRACT(HOUR FROM __time)`.                                                                                                                                                                                                                                                                                                     |
 |`FLOOR(timestamp_expr TO unit)`|Rounds down a timestamp, returning it as a new timestamp. The `unit` parameter must be unquoted and can be SECOND, MINUTE, HOUR, DAY, WEEK, MONTH, QUARTER, or YEAR.|
 |`CEIL(timestamp_expr TO unit)`|Rounds up a timestamp, returning it as a new timestamp. The `unit` parameter must be unquoted and can be SECOND, MINUTE, HOUR, DAY, WEEK, MONTH, QUARTER, or YEAR.|
 |`TIMESTAMPADD(unit, count, timestamp)`|Adds a `count` number of time `unit` to timestamp, equivalent to `timestamp + count * unit`. The `unit` parameter must be unquoted and can be SECOND, MINUTE, HOUR, DAY, WEEK, MONTH, QUARTER, or YEAR.|

docs/querying/sql-window-functions.md

@@ -28,10 +28,6 @@ description: Reference for window functions
 Apache Druid supports two query languages: [Druid SQL](sql.md) and [native queries](querying.md).
 This document describes the SQL language.
 
-Window functions are an [experimental](../development/experimental.md) feature.
-Development and testing are still at early stage. Feel free to try window functions and provide your feedback.
-Windows functions are not currently supported by multi-stage-query engine so you cannot use them in SQL-based ingestion.
-
 :::
 
 Window functions in Apache Druid produce values based upon the relationship of one row within a window of rows to the other rows within the same window. A window is a group of related rows within a result set. For example, rows with the same value for a specific dimension.
@@ -424,8 +420,4 @@ The number of rows considered for the `moving5` window for the `count5` column:
 
 The following are known issues with window functions:
 
--  Aggregates with ORDER BY specified are processed in the window: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW  
-     This behavior differs from other databases that use the default of RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW.  
-     In cases where the order column is unique there is no difference between RANGE / ROWS; windows with RANGE specifications are handled as ROWS.
-- LEAD/LAG ignores the default value
-- LAST_VALUE returns the last value of the window even when you include an ORDER BY clause
+- SELECT * queries without a WHERE clause are not supported. If you want to retrieve all columns in this case, specify the column names.

docs/release-info/migr-front-coded-dict.md

@@ -47,7 +47,7 @@ You can specify the following optional properties:
 
 For example:
 
-```
+```json
 "tuningConfig": {
   "indexSpec": {
     "stringDictionaryEncoding": {
@@ -59,6 +59,39 @@ For example:
 }
 ```
 
+For SQL based ingestion, you can add the `indexSpec` to your query context.
+In the Web Console, select *Edit context* from the context from the *Engine:* menu and enter the `indexSpec`. For example:
+
+```json
+{
+...
+"indexSpec": {
+  "stringDictionaryEncoding": {
+  "type": "frontCoded",
+  "bucketSize": 4,
+  "formatVersion": 1
+  }
+}
+}
+```
+
+For API calls to the SQL-based ingestion API, include the `indexSpec` in the context in the request payload. For example:
+
+```json
+{
+"query": ...
+"context": {
+  "maxNumTasks": 3
+  "indexSpec": {
+  "stringDictionaryEncoding": {
+    "type": "frontCoded",
+    "bucketSize": 4,
+    "formatVersion": 1}
+    }
+  }
+}
+```
+
 ## Upgrade from Druid 25.0.0
 
 Druid 26.0.0 introduced a new version of the front-coded dictionary, version 1, offering typically faster read speeds and smaller storage sizes.

docs/tutorials/tutorial-sketches-theta.md

@@ -36,7 +36,7 @@ Imagine you are interested in the number of visitors that watched episodes of a
 
 There is no way to answer these questions by just looking at the aggregated numbers. You would have to go back to the detail data and scan every single row. If the data volume is high enough, this may take a very long time, meaning that an interactive data exploration is not possible.
 
-An additional nuisance is that unique counts don't work well with rollups. For this example, it would be great if you could have just one row of data per 15 minute interval[^1], show, and episode. After all, you are not interested in the individual user IDs, just the unique counts.
+An additional nuisance is that unique counts don't work well with rollups. For this example, it would be great if you could have just one row of data per 15 minute interval<sup>[^1]</sup>, show, and episode. After all, you are not interested in the individual user IDs, just the unique counts.
 
 [^1]: Why 15 minutes and not just 1 hour? Intervals of 15 minutes work better with international timezones because those are not always aligned by hour. India, for instance, is 30 minutes off, and Nepal is even 45 minutes off. With 15 minute aggregates, you can get hourly sums for any of those timezones, too!
 
@@ -60,9 +60,11 @@ In this tutorial, you will learn how to do the following:
 
 ## Prerequisites
 
-For this tutorial, you should have already downloaded Druid as described in
-the [single-machine quickstart](index.md) and have it running on your local machine.
-It will also be helpful to have finished [Tutorial: Loading a file](../tutorials/tutorial-batch.md) and [Tutorial: Querying data](../tutorials/tutorial-query.md).
+Before proceeding, download Druid as described in the [single-machine quickstart](index.md) and have it running on your local machine. You don't need to load any data into the Druid cluster.
+
+It's helpful to have finished [Tutorial: Loading a file](../tutorials/tutorial-batch.md) and [Tutorial: Querying data](../tutorials/tutorial-query.md).
+
+## Sample data
 
 This tutorial works with the following data:
 - **date**: a timestamp. In this case it's just dates but as mentioned earlier, a finer granularity makes sense in real life.
@@ -95,103 +97,35 @@ date,uid,show,episode
 
 ## Ingest data using Theta sketches
 
-1. Navigate to the **Load data** wizard in the web console.
-2. Select `Paste data` as the data source and paste the given data:
-
-![Load data view with pasted data](../assets/tutorial-theta-01.png)
-
-3. Leave the source type as `inline` and click **Apply** and **Next: Parse data**.
-4. Parse the data as CSV, with included headers:
-
-![Parse raw data](../assets/tutorial-theta-02.png)
-
-5. Accept the default values in the **Parse time**, **Transform**, and **Filter** stages.
-6. In the **Configure schema** stage, enable rollup and confirm your choice in the dialog. Then set the query granularity to `day`.
-
-![Configure schema for rollup and query granularity](../assets/tutorial-theta-03.png)
-
-7. Add the Theta sketch during this stage. Select **Add metric**.
-8. Define the new metric as a Theta sketch with the following details:
-   * **Name**: `theta_uid`
-   * **Type**: `thetaSketch`
-   * **Field name**: `uid`
-   * **Size**: Accept the default value, `16384`.
-   * **Is input theta sketch**: Accept the default value, `False`.
-
-![Create Theta sketch metric](../assets/tutorial-theta-04.png)
-
-9. Click **Apply** to add the new metric to the data model.
+Load the sample dataset using the [`INSERT INTO`](../multi-stage-query/reference.md/#insert) statement and the [`EXTERN`](../multi-stage-query/reference.md/#extern-function) function to ingest the sample data inline. In the [Druid web console](../operations/web-console.md), go to the **Query** view and run the following query:
 
 
-10. You are not interested in individual user ID's, only the unique counts. Right now, `uid` is still in the data model. To remove it, click on the `uid` column in the data model and delete it using the trashcan icon on the right:
-
-![Delete uid column](../assets/tutorial-theta-05.png)
-
-11. For the remaining stages of the **Load data** wizard, set the following options:
-    * **Partition**: Set **Segment granularity** to `day`.
-    * **Tune**: Leave the default options.
-    * **Publish**: Set the datasource name to `ts_tutorial`.
-
-On the **Edit spec** page, your final input spec should match the following:
-
-```json
-{
-  "type": "index_parallel",
-  "spec": {
-    "ioConfig": {
-      "type": "index_parallel",
-      "inputSource": {
-        "type": "inline",
-        "data": "date,uid,show,episode\n2022-05-19,alice,Game of Thrones,S1E1\n2022-05-19,alice,Game of Thrones,S1E2\n2022-05-19,alice,Game of Thrones,S1E1\n2022-05-19,bob,Bridgerton,S1E1\n2022-05-20,alice,Game of Thrones,S1E1\n2022-05-20,carol,Bridgerton,S1E2\n2022-05-20,dan,Bridgerton,S1E1\n2022-05-21,alice,Game of Thrones,S1E1\n2022-05-21,carol,Bridgerton,S1E1\n2022-05-21,erin,Game of Thrones,S1E1\n2022-05-21,alice,Bridgerton,S1E1\n2022-05-22,bob,Game of Thrones,S1E1\n2022-05-22,bob,Bridgerton,S1E1\n2022-05-22,carol,Bridgerton,S1E2\n2022-05-22,bob,Bridgerton,S1E1\n2022-05-22,erin,Game of Thrones,S1E1\n2022-05-22,erin,Bridgerton,S1E2\n2022-05-23,erin,Game of Thrones,S1E1\n2022-05-23,alice,Game of Thrones,S1E1"
-      },
-      "inputFormat": {
-        "type": "csv",
-        "findColumnsFromHeader": true
-      }
-    },
-    "tuningConfig": {
-      "type": "index_parallel",
-      "partitionsSpec": {
-        "type": "hashed"
-      },
-      "forceGuaranteedRollup": true
-    },
-    "dataSchema": {
-      "dataSource": "ts_tutorial",
-      "timestampSpec": {
-        "column": "date",
-        "format": "auto"
-      },
-      "dimensionsSpec": {
-        "dimensions": [
-          "show",
-          "episode"
-        ]
-      },
-      "granularitySpec": {
-        "queryGranularity": "day",
-        "rollup": true,
-        "segmentGranularity": "day"
-      },
-      "metricsSpec": [
-        {
-          "name": "count",
-          "type": "count"
-        },
-        {
-          "type": "thetaSketch",
-          "name": "theta_uid",
-          "fieldName": "uid"
-        }
-      ]
-    }
-  }
-}
+```sql
+INSERT INTO "ts_tutorial"
+WITH "source" AS (SELECT * FROM TABLE(
+  EXTERN(
+    '{"type":"inline","data":"date,uid,show,episode\n2022-05-19,alice,Game of Thrones,S1E1\n2022-05-19,alice,Game of Thrones,S1E2\n2022-05-19,alice,Game of Thrones,S1E1\n2022-05-19,bob,Bridgerton,S1E1\n2022-05-20,alice,Game of Thrones,S1E1\n2022-05-20,carol,Bridgerton,S1E2\n2022-05-20,dan,Bridgerton,S1E1\n2022-05-21,alice,Game of Thrones,S1E1\n2022-05-21,carol,Bridgerton,S1E1\n2022-05-21,erin,Game of Thrones,S1E1\n2022-05-21,alice,Bridgerton,S1E1\n2022-05-22,bob,Game of Thrones,S1E1\n2022-05-22,bob,Bridgerton,S1E1\n2022-05-22,carol,Bridgerton,S1E2\n2022-05-22,bob,Bridgerton,S1E1\n2022-05-22,erin,Game of Thrones,S1E1\n2022-05-22,erin,Bridgerton,S1E2\n2022-05-23,erin,Game of Thrones,S1E1\n2022-05-23,alice,Game of Thrones,S1E1"}',
+    '{"type":"csv","findColumnsFromHeader":true}'
+  )
+) EXTEND ("date" VARCHAR, "show" VARCHAR, "episode" VARCHAR, "uid" VARCHAR))
+SELECT
+  TIME_FLOOR(TIME_PARSE("date"), 'P1D') AS "__time",
+  "show",
+  "episode",
+  COUNT(*) AS "count",
+  DS_THETA("uid") AS "theta_uid"
+FROM "source"
+GROUP BY 1, 2, 3
+PARTITIONED BY DAY
 ```
 
-Notice the `theta_uid` object in the `metricsSpec` list, that defines the `thetaSketch` aggregator on the `uid` column during ingestion.
+Notice the `theta_uid` column in the `SELECT` statement. It defines the `thetaSketch` aggregator on the `uid` column during ingestion.
+In this scenario you are not interested in individual user IDs, only the unique counts.
+Instead you create Theta sketches on the values of `uid` using the `DS_THETA` function.
 
-Click **Submit** to start the ingestion.
+[`DS_THETA`](../development/extensions-core/datasketches-theta.md#aggregator) has an optional second parameter that controls the accuracy and size of the sketches.
+
+The `GROUP BY` statement groups the entries for each episode of a show watched on the same day.
 
 ## Query the Theta sketch column
 
@@ -209,36 +143,22 @@ Let's first see what the data looks like in Druid. Run the following SQL stateme
 SELECT * FROM ts_tutorial
 ```
 
-![View data with SELECT all query](../assets/tutorial-theta-06.png)
+![View data with SELECT all query](../assets/tutorial-theta-03.png)
 
 The Theta sketch column `theta_uid` appears as a Base64-encoded string; behind it is a bitmap.
 
-The following query to compute the distinct counts of user IDs uses `APPROX_COUNT_DISTINCT_DS_THETA` and groups by the other dimensions:
-```sql
-SELECT __time,
-       "show",
-       "episode",
-       APPROX_COUNT_DISTINCT_DS_THETA(theta_uid) AS users
-FROM   ts_tutorial
-GROUP  BY 1, 2, 3
-```
-
-![Count distinct with Theta sketches](../assets/tutorial-theta-07.png)
-
-In the preceding query, `APPROX_COUNT_DISTINCT_DS_THETA` is equivalent to calling `DS_THETA` and `THETA_SKETCH_ESIMATE` as follows:
+The following query uses `THETA_SKETCH_ESTIMATE` to compute the distinct counts of user IDs and groups by the other dimensions:
 
 ```sql
-SELECT __time,
-       "show", 
-       "episode",
-       THETA_SKETCH_ESTIMATE(DS_THETA(theta_uid)) AS users
-FROM   ts_tutorial
-GROUP  BY 1, 2, 3
+SELECT
+  __time,
+  "show",
+  "episode",
+  THETA_SKETCH_ESTIMATE(theta_uid) AS users
+FROM ts_tutorial
 ```
 
-That is, `APPROX_COUNT_DISTINCT_DS_THETA` applies the following:
-* `DS_THETA`: Creates a new Theta sketch from the column of Theta sketches
-* `THETA_SKETCH_ESTIMATE`: Calculates the distinct count estimate from the output of `DS_THETA`
+![Count distinct with Theta sketches](../assets/tutorial-theta-04.png)
 
 ### Filtered metrics
 
@@ -247,7 +167,17 @@ Druid has the capability to use [filtered metrics](../querying/sql-aggregations.
  In the case of Theta sketches, the filter clause has to be inserted between the aggregator and the estimator.
 :::
 
-As an example, query the total unique users that watched _Bridgerton:_
+As an example, query the total unique users that watched _Bridgerton_:
+
+```sql
+SELECT APPROX_COUNT_DISTINCT_DS_THETA(theta_uid) FILTER(WHERE "show" = 'Bridgerton') AS users
+FROM ts_tutorial
+```
+
+![Count distinct with Theta sketches and filters](../assets/tutorial-theta-05.png)
+
+
+In the preceding query, `APPROX_COUNT_DISTINCT_DS_THETA` is equivalent to calling `DS_THETA` and `THETA_SKETCH_ESIMATE` as follows:
 
 ```sql
 SELECT THETA_SKETCH_ESTIMATE(
@@ -256,7 +186,12 @@ SELECT THETA_SKETCH_ESTIMATE(
 FROM ts_tutorial
 ```
 
-![Count distinct with Theta sketches and filters](../assets/tutorial-theta-08.png)
+The `APPROX_COUNT_DISTINCT_DS_THETA` function applies the following:
+
+* `DS_THETA`: Creates a new Theta sketch from the column of Theta sketches.
+* `THETA_SKETCH_ESTIMATE`: Calculates the distinct count estimate from the output of `DS_THETA`.
+
+Note that the filter clause limits an aggregation query to only the rows that match the filter.
 
 ### Set operations
 
@@ -274,7 +209,7 @@ SELECT THETA_SKETCH_ESTIMATE(
 FROM ts_tutorial
 ```
 
-![Count distinct with Theta sketches, filters, and set operations](../assets/tutorial-theta-09.png)
+![Count distinct with Theta sketches, filters, and set operations](../assets/tutorial-theta-06.png)
 
 Again, the set function is spliced in between the aggregator and the estimator.
 
@@ -290,7 +225,7 @@ SELECT THETA_SKETCH_ESTIMATE(
 FROM ts_tutorial
 ```
 
-![Count distinct with Theta sketches, filters, and set operations](../assets/tutorial-theta-10.png)
+![Count distinct with Theta sketches, filters, and set operations](../assets/tutorial-theta-07.png)
 
 And finally, there is `THETA_SKETCH_NOT` which computes the set difference of two or more segments.
 The result describes how many visitors watched episode 1 of Bridgerton but not episode 2.
@@ -306,7 +241,7 @@ SELECT THETA_SKETCH_ESTIMATE(
 FROM ts_tutorial
 ```
 
-![Count distinct with Theta sketches, filters, and set operations](../assets/tutorial-theta-11.png)
+![Count distinct with Theta sketches, filters, and set operations](../assets/tutorial-theta-08.png)
 
 ## Conclusions
 
@@ -314,7 +249,7 @@ FROM ts_tutorial
 - This allows us to use rollup and discard the individual values, just retaining statistical approximations in the sketches.
 - With Theta sketch set operations, affinity analysis is easier, for example, to answer questions such as which segments correlate or overlap by how much.
 
-## Further reading
+## Learn more
 
 See the following topics for more information:
 * [Theta sketch](../development/extensions-core/datasketches-theta.md) for reference on ingestion and native queries on Theta sketches in Druid.
@@ -326,4 +261,3 @@ See the following topics for more information:
 ## Acknowledgments
 
 This tutorial is adapted from a [blog post](https://blog.hellmar-becker.de/2022/06/05/druid-data-cookbook-counting-unique-visitors-for-overlapping-segments/) by community member Hellmar Becker.
-

extensions-contrib/aliyun-oss-extensions/pom.xml

@@ -100,11 +100,6 @@
             <artifactId>jsr305</artifactId>
             <scope>provided</scope>
         </dependency>
-        <dependency>
-            <groupId>commons-lang</groupId>
-            <artifactId>commons-lang</artifactId>
-            <scope>provided</scope>
-        </dependency>
             
         <!-- Tests -->
         <dependency>

extensions-contrib/druid-deltalake-extensions/pom.xml

@@ -35,7 +35,7 @@
   <modelVersion>4.0.0</modelVersion>
 
   <properties>
-    <delta-kernel.version>3.2.0</delta-kernel.version>
+    <delta-kernel.version>3.2.1</delta-kernel.version>
   </properties>
 
   <dependencies>
@@ -49,12 +49,6 @@
       <artifactId>delta-kernel-defaults</artifactId>
       <version>${delta-kernel.version}</version>
     </dependency>
-    <dependency>
-      <groupId>io.delta</groupId>
-      <artifactId>delta-storage</artifactId>
-      <version>${delta-kernel.version}</version>
-    </dependency>
-
     <dependency>
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-client-api</artifactId>

extensions-contrib/druid-deltalake-extensions/src/main/java/org/apache/druid/delta/input/DeltaInputSource.java

@@ -42,7 +42,6 @@ import io.delta.kernel.types.StructField;
 import io.delta.kernel.types.StructType;
 import io.delta.kernel.utils.CloseableIterator;
 import io.delta.kernel.utils.FileStatus;
-import io.delta.storage.LogStore;
 import org.apache.druid.data.input.ColumnsFilter;
 import org.apache.druid.data.input.InputFormat;
 import org.apache.druid.data.input.InputRowSchema;
@@ -340,20 +339,10 @@ public class DeltaInputSource implements SplittableInputSource<DeltaSplit>
 
   private Snapshot getSnapshotForTable(final Table table, final Engine engine)
   {
-    // Setting the LogStore class loader before calling the Delta Kernel snapshot API is required as a workaround with
-    // the 3.2.0 Delta Kernel because the Kernel library cannot instantiate the LogStore class otherwise. Please see
-    // https://github.com/delta-io/delta/issues/3299 for details. This workaround can be removed once the issue is fixed.
-    final ClassLoader currCtxCl = Thread.currentThread().getContextClassLoader();
-    try {
-      Thread.currentThread().setContextClassLoader(LogStore.class.getClassLoader());
-      if (snapshotVersion != null) {
-        return table.getSnapshotAsOfVersion(engine, snapshotVersion);
-      } else {
-        return table.getLatestSnapshot(engine);
-      }
-    }
-    finally {
-      Thread.currentThread().setContextClassLoader(currCtxCl);
+    if (snapshotVersion != null) {
+      return table.getSnapshotAsOfVersion(engine, snapshotVersion);
+    } else {
+      return table.getLatestSnapshot(engine);
     }
   }
 

extensions-contrib/kubernetes-overlord-extensions/pom.xml

@@ -76,11 +76,6 @@
       <version>${project.parent.version}</version>
       <scope>provided</scope>
     </dependency>
-    <dependency>
-      <groupId>commons-lang</groupId>
-      <artifactId>commons-lang</artifactId>
-      <scope>provided</scope>
-    </dependency>
     <dependency>
       <groupId>commons-io</groupId>
       <artifactId>commons-io</artifactId>

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesAndWorkerTaskRunner.java

@@ -268,6 +268,12 @@ public class KubernetesAndWorkerTaskRunner implements TaskLogStreamer, WorkerTas
     return Math.max(0, k8sCapacity) + Math.max(0, workerCapacity);
   }
 
+  @Override
+  public int getMaximumCapacityWithAutoscale()
+  {
+    return workerTaskRunner.getMaximumCapacityWithAutoscale() + kubernetesTaskRunner.getMaximumCapacityWithAutoscale();
+  }
+
   @Override
   public int getUsedCapacity()
   {

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesOverlordModule.java

@@ -19,6 +19,8 @@
 
 package org.apache.druid.k8s.overlord;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Supplier;
 import com.google.inject.Binder;
 import com.google.inject.Inject;
 import com.google.inject.Injector;
@@ -39,7 +41,10 @@ import org.apache.druid.guice.JsonConfigurator;
 import org.apache.druid.guice.LazySingleton;
 import org.apache.druid.guice.PolyBind;
 import org.apache.druid.guice.annotations.LoadScope;
+import org.apache.druid.guice.annotations.Self;
+import org.apache.druid.guice.annotations.Smile;
 import org.apache.druid.indexing.common.config.FileTaskLogsConfig;
+import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.tasklogs.FileTaskLogs;
 import org.apache.druid.indexing.overlord.RemoteTaskRunnerFactory;
 import org.apache.druid.indexing.overlord.TaskRunnerFactory;
@@ -47,6 +52,7 @@ import org.apache.druid.indexing.overlord.WorkerTaskRunner;
 import org.apache.druid.indexing.overlord.config.TaskQueueConfig;
 import org.apache.druid.indexing.overlord.hrtr.HttpRemoteTaskRunnerFactory;
 import org.apache.druid.initialization.DruidModule;
+import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.lifecycle.Lifecycle;
 import org.apache.druid.java.util.common.logger.Logger;
@@ -54,11 +60,19 @@ import org.apache.druid.k8s.overlord.common.DruidKubernetesClient;
 import org.apache.druid.k8s.overlord.execution.KubernetesTaskExecutionConfigResource;
 import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
 import org.apache.druid.k8s.overlord.runnerstrategy.RunnerStrategy;
+import org.apache.druid.k8s.overlord.taskadapter.DynamicConfigPodTemplateSelector;
+import org.apache.druid.k8s.overlord.taskadapter.MultiContainerTaskAdapter;
+import org.apache.druid.k8s.overlord.taskadapter.PodTemplateTaskAdapter;
+import org.apache.druid.k8s.overlord.taskadapter.SingleContainerTaskAdapter;
+import org.apache.druid.k8s.overlord.taskadapter.TaskAdapter;
+import org.apache.druid.server.DruidNode;
+import org.apache.druid.server.log.StartupLoggingConfig;
 import org.apache.druid.tasklogs.NoopTaskLogs;
 import org.apache.druid.tasklogs.TaskLogKiller;
 import org.apache.druid.tasklogs.TaskLogPusher;
 import org.apache.druid.tasklogs.TaskLogs;
 
+import java.util.Locale;
 import java.util.Properties;
 
 
@@ -162,6 +176,70 @@ public class KubernetesOverlordModule implements DruidModule
            : injector.getInstance(RemoteTaskRunnerFactory.class);
   }
 
+  /**
+   * Provides a TaskAdapter instance for the KubernetesTaskRunner.
+   */
+  @Provides
+  @LazySingleton
+  TaskAdapter provideTaskAdapter(
+      DruidKubernetesClient client,
+      Properties properties,
+      KubernetesTaskRunnerConfig kubernetesTaskRunnerConfig,
+      TaskConfig taskConfig,
+      StartupLoggingConfig startupLoggingConfig,
+      @Self DruidNode druidNode,
+      @Smile ObjectMapper smileMapper,
+      TaskLogs taskLogs,
+      Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef
+  )
+  {
+    String adapter = properties.getProperty(String.format(
+        Locale.ROOT,
+        "%s.%s.adapter.type",
+        IndexingServiceModuleHelper.INDEXER_RUNNER_PROPERTY_PREFIX,
+        "k8s"
+    ));
+
+    if (adapter != null && !MultiContainerTaskAdapter.TYPE.equals(adapter) && kubernetesTaskRunnerConfig.isSidecarSupport()) {
+      throw new IAE(
+          "Invalid pod adapter [%s], only pod adapter [%s] can be specified when sidecarSupport is enabled",
+          adapter,
+          MultiContainerTaskAdapter.TYPE
+      );
+    }
+
+    if (MultiContainerTaskAdapter.TYPE.equals(adapter) || kubernetesTaskRunnerConfig.isSidecarSupport()) {
+      return new MultiContainerTaskAdapter(
+          client,
+          kubernetesTaskRunnerConfig,
+          taskConfig,
+          startupLoggingConfig,
+          druidNode,
+          smileMapper,
+          taskLogs
+      );
+    } else if (PodTemplateTaskAdapter.TYPE.equals(adapter)) {
+      return new PodTemplateTaskAdapter(
+          kubernetesTaskRunnerConfig,
+          taskConfig,
+          druidNode,
+          smileMapper,
+          taskLogs,
+          new DynamicConfigPodTemplateSelector(properties, dynamicConfigRef)
+      );
+    } else {
+      return new SingleContainerTaskAdapter(
+          client,
+          kubernetesTaskRunnerConfig,
+          taskConfig,
+          startupLoggingConfig,
+          druidNode,
+          smileMapper,
+          taskLogs
+      );
+    }
+  }
+
   private static class RunnerStrategyProvider implements Provider<RunnerStrategy>
   {
     private KubernetesAndWorkerTaskRunnerConfig runnerConfig;

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycle.java

@@ -20,6 +20,7 @@
 package org.apache.druid.k8s.overlord;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Optional;
 import com.google.common.base.Preconditions;
 import io.fabric8.kubernetes.api.model.Pod;
@@ -31,6 +32,7 @@ import org.apache.commons.io.IOUtils;
 import org.apache.druid.indexer.TaskLocation;
 import org.apache.druid.indexer.TaskStatus;
 import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.emitter.EmittingLogger;
 import org.apache.druid.k8s.overlord.common.DruidK8sConstants;
@@ -177,6 +179,11 @@ public class KubernetesPeonLifecycle
   protected synchronized TaskStatus join(long timeout) throws IllegalStateException
   {
     try {
+      /* It's okay to store taskLocation because podIP only changes on pod restart, and we have to set restartPolicy to Never
+        since Druid doesn't support retrying tasks from a external system (K8s). We can explore adding a fabric8 watcher
+        if we decide we need to change this later.
+      **/
+      taskLocation = getTaskLocationFromK8s();
       updateState(new State[]{State.NOT_STARTED, State.PENDING}, State.RUNNING);
 
       JobResponse jobResponse = kubernetesClient.waitForPeonJobCompletion(
@@ -254,24 +261,8 @@ public class KubernetesPeonLifecycle
     if we decide we need to change this later.
     **/
     if (taskLocation == null) {
-      Optional<Pod> maybePod = kubernetesClient.getPeonPod(taskId.getK8sJobName());
-      if (!maybePod.isPresent()) {
-        return TaskLocation.unknown();
-      }
-
-      Pod pod = maybePod.get();
-      PodStatus podStatus = pod.getStatus();
-
-      if (podStatus == null || podStatus.getPodIP() == null) {
-        return TaskLocation.unknown();
-      }
-      taskLocation = TaskLocation.create(
-          podStatus.getPodIP(),
-          DruidK8sConstants.PORT,
-          DruidK8sConstants.TLS_PORT,
-          Boolean.parseBoolean(pod.getMetadata().getAnnotations().getOrDefault(DruidK8sConstants.TLS_ENABLED, "false")),
-          pod.getMetadata() != null ? pod.getMetadata().getName() : ""
-      );
+      log.warn("Unknown task location for [%s]", taskId);
+      return TaskLocation.unknown();
     }
 
     return taskLocation;
@@ -378,4 +369,28 @@ public class KubernetesPeonLifecycle
     );
     stateListener.stateChanged(state.get(), taskId.getOriginalTaskId());
   }
+
+  @VisibleForTesting
+  protected TaskLocation getTaskLocationFromK8s()
+  {
+    Pod pod = kubernetesClient.getPeonPodWithRetries(taskId.getK8sJobName());
+    PodStatus podStatus = pod.getStatus();
+
+    if (podStatus == null || podStatus.getPodIP() == null) {
+      throw new ISE("Could not find location of running task [%s]", taskId);
+    }
+
+    return TaskLocation.create(
+        podStatus.getPodIP(),
+        DruidK8sConstants.PORT,
+        DruidK8sConstants.TLS_PORT,
+        Boolean.parseBoolean(
+            pod.getMetadata() != null && pod.getMetadata().getAnnotations() != null ?
+                pod.getMetadata().getAnnotations().getOrDefault(DruidK8sConstants.TLS_ENABLED, "false") :
+                "false"
+        ),
+        pod.getMetadata() != null ? pod.getMetadata().getName() : ""
+    );
+
+  }
 }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerFactory.java

@@ -20,73 +20,47 @@
 package org.apache.druid.k8s.overlord;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.base.Supplier;
 import com.google.inject.Inject;
-import org.apache.druid.guice.IndexingServiceModuleHelper;
 import org.apache.druid.guice.annotations.EscalatedGlobal;
-import org.apache.druid.guice.annotations.Self;
 import org.apache.druid.guice.annotations.Smile;
-import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.overlord.TaskRunnerFactory;
-import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 import org.apache.druid.java.util.http.client.HttpClient;
 import org.apache.druid.k8s.overlord.common.DruidKubernetesClient;
 import org.apache.druid.k8s.overlord.common.KubernetesPeonClient;
-import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
-import org.apache.druid.k8s.overlord.taskadapter.MultiContainerTaskAdapter;
-import org.apache.druid.k8s.overlord.taskadapter.PodTemplateTaskAdapter;
-import org.apache.druid.k8s.overlord.taskadapter.SingleContainerTaskAdapter;
 import org.apache.druid.k8s.overlord.taskadapter.TaskAdapter;
-import org.apache.druid.server.DruidNode;
-import org.apache.druid.server.log.StartupLoggingConfig;
 import org.apache.druid.tasklogs.TaskLogs;
 
-import java.util.Locale;
-import java.util.Properties;
-
 public class KubernetesTaskRunnerFactory implements TaskRunnerFactory<KubernetesTaskRunner>
 {
   public static final String TYPE_NAME = "k8s";
   private final ObjectMapper smileMapper;
   private final HttpClient httpClient;
   private final KubernetesTaskRunnerConfig kubernetesTaskRunnerConfig;
-  private final StartupLoggingConfig startupLoggingConfig;
   private final TaskLogs taskLogs;
-  private final DruidNode druidNode;
-  private final TaskConfig taskConfig;
-  private final Properties properties;
   private final DruidKubernetesClient druidKubernetesClient;
   private final ServiceEmitter emitter;
-  private final Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef;
   private KubernetesTaskRunner runner;
+  private final TaskAdapter taskAdapter;
 
   @Inject
   public KubernetesTaskRunnerFactory(
       @Smile ObjectMapper smileMapper,
       @EscalatedGlobal final HttpClient httpClient,
       KubernetesTaskRunnerConfig kubernetesTaskRunnerConfig,
-      StartupLoggingConfig startupLoggingConfig,
       TaskLogs taskLogs,
-      @Self DruidNode druidNode,
-      TaskConfig taskConfig,
-      Properties properties,
       DruidKubernetesClient druidKubernetesClient,
       ServiceEmitter emitter,
-      Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef
+      TaskAdapter taskAdapter
   )
   {
     this.smileMapper = smileMapper;
     this.httpClient = httpClient;
     this.kubernetesTaskRunnerConfig = kubernetesTaskRunnerConfig;
-    this.startupLoggingConfig = startupLoggingConfig;
     this.taskLogs = taskLogs;
-    this.druidNode = druidNode;
-    this.taskConfig = taskConfig;
-    this.properties = properties;
     this.druidKubernetesClient = druidKubernetesClient;
     this.emitter = emitter;
-    this.dynamicConfigRef = dynamicConfigRef;
+    this.taskAdapter = taskAdapter;
   }
 
   @Override
@@ -101,7 +75,7 @@ public class KubernetesTaskRunnerFactory implements TaskRunnerFactory<Kubernetes
     );
 
     runner = new KubernetesTaskRunner(
-        buildTaskAdapter(druidKubernetesClient),
+        taskAdapter,
         kubernetesTaskRunnerConfig,
         peonClient,
         httpClient,
@@ -117,53 +91,4 @@ public class KubernetesTaskRunnerFactory implements TaskRunnerFactory<Kubernetes
     return runner;
   }
 
-  private TaskAdapter buildTaskAdapter(DruidKubernetesClient client)
-  {
-    String adapter = properties.getProperty(String.format(
-        Locale.ROOT,
-        "%s.%s.adapter.type",
-        IndexingServiceModuleHelper.INDEXER_RUNNER_PROPERTY_PREFIX,
-        TYPE_NAME
-    ));
-
-    if (adapter != null && !MultiContainerTaskAdapter.TYPE.equals(adapter) && kubernetesTaskRunnerConfig.isSidecarSupport()) {
-      throw new IAE(
-          "Invalid pod adapter [%s], only pod adapter [%s] can be specified when sidecarSupport is enabled",
-          adapter,
-          MultiContainerTaskAdapter.TYPE
-      );
-    }
-
-    if (MultiContainerTaskAdapter.TYPE.equals(adapter) || kubernetesTaskRunnerConfig.isSidecarSupport()) {
-      return new MultiContainerTaskAdapter(
-          client,
-          kubernetesTaskRunnerConfig,
-          taskConfig,
-          startupLoggingConfig,
-          druidNode,
-          smileMapper,
-          taskLogs
-      );
-    } else if (PodTemplateTaskAdapter.TYPE.equals(adapter)) {
-      return new PodTemplateTaskAdapter(
-          kubernetesTaskRunnerConfig,
-          taskConfig,
-          druidNode,
-          smileMapper,
-          properties,
-          taskLogs,
-          dynamicConfigRef
-      );
-    } else {
-      return new SingleContainerTaskAdapter(
-          client,
-          kubernetesTaskRunnerConfig,
-          taskConfig,
-          startupLoggingConfig,
-          druidNode,
-          smileMapper,
-          taskLogs
-      );
-    }
-  }
 }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/taskadapter/DynamicConfigPodTemplateSelector.java

@@ -0,0 +1,123 @@
+/*
+ * 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.k8s.overlord.taskadapter;
+
+import com.google.common.base.Optional;
+import com.google.common.base.Supplier;
+import io.fabric8.kubernetes.api.model.PodTemplate;
+import io.fabric8.kubernetes.client.utils.Serialization;
+import org.apache.druid.guice.IndexingServiceModuleHelper;
+import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
+import org.apache.druid.k8s.overlord.execution.PodTemplateSelectStrategy;
+
+import java.io.File;
+import java.nio.file.Files;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Properties;
+import java.util.Set;
+
+public class DynamicConfigPodTemplateSelector implements PodTemplateSelector
+{
+
+  private static final String TASK_PROPERTY = IndexingServiceModuleHelper.INDEXER_RUNNER_PROPERTY_PREFIX
+                                              + ".k8s.podTemplate.";
+
+  private final Properties properties;
+  private HashMap<String, PodTemplate> podTemplates;
+  private Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef;
+
+  public DynamicConfigPodTemplateSelector(
+      Properties properties,
+      Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef
+  )
+  {
+    this.properties = properties;
+    this.dynamicConfigRef = dynamicConfigRef;
+    initializeTemplatesFromFileSystem();
+  }
+
+  private void initializeTemplatesFromFileSystem()
+  {
+    Set<String> taskAdapterTemplateKeys = getTaskAdapterTemplates(properties);
+    if (!taskAdapterTemplateKeys.contains("base")) {
+      throw new IAE(
+          "Pod template task adapter requires a base pod template to be specified under druid.indexer.runner.k8s.podTemplate.base");
+    }
+
+    HashMap<String, PodTemplate> podTemplateMap = new HashMap<>();
+    for (String taskAdapterTemplateKey : taskAdapterTemplateKeys) {
+      Optional<PodTemplate> template = loadPodTemplate(taskAdapterTemplateKey, properties);
+      if (template.isPresent()) {
+        podTemplateMap.put(taskAdapterTemplateKey, template.get());
+      }
+    }
+    podTemplates = podTemplateMap;
+  }
+
+  private Set<String> getTaskAdapterTemplates(Properties properties)
+  {
+    Set<String> taskAdapterTemplates = new HashSet<>();
+
+    for (String runtimeProperty : properties.stringPropertyNames()) {
+      if (runtimeProperty.startsWith(TASK_PROPERTY)) {
+        String[] taskAdapterPropertyPaths = runtimeProperty.split("\\.");
+        taskAdapterTemplates.add(taskAdapterPropertyPaths[taskAdapterPropertyPaths.length - 1]);
+      }
+    }
+
+    return taskAdapterTemplates;
+  }
+
+  private Optional<PodTemplate> loadPodTemplate(String key, Properties properties)
+  {
+    String property = TASK_PROPERTY + key;
+    String podTemplateFile = properties.getProperty(property);
+    if (podTemplateFile == null) {
+      throw new IAE("Pod template file not specified for [%s]", property);
+
+    }
+    try {
+      return Optional.of(Serialization.unmarshal(
+          Files.newInputStream(new File(podTemplateFile).toPath()),
+          PodTemplate.class
+      ));
+    }
+    catch (Exception e) {
+      throw new IAE(e, "Failed to load pod template file for [%s] at [%s]", property, podTemplateFile);
+    }
+  }
+
+  @Override
+  public Optional<PodTemplateWithName> getPodTemplateForTask(Task task)
+  {
+    PodTemplateSelectStrategy podTemplateSelectStrategy;
+    KubernetesTaskRunnerDynamicConfig dynamicConfig = dynamicConfigRef.get();
+    if (dynamicConfig == null || dynamicConfig.getPodTemplateSelectStrategy() == null) {
+      podTemplateSelectStrategy = KubernetesTaskRunnerDynamicConfig.DEFAULT_STRATEGY;
+    } else {
+      podTemplateSelectStrategy = dynamicConfig.getPodTemplateSelectStrategy();
+    }
+
+    return Optional.of(podTemplateSelectStrategy.getPodTemplateForTask(task, podTemplates));
+  }
+}

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/taskadapter/MultiContainerTaskAdapter.java

@@ -33,7 +33,7 @@ import io.fabric8.kubernetes.api.model.VolumeBuilder;
 import io.fabric8.kubernetes.api.model.VolumeMount;
 import io.fabric8.kubernetes.api.model.VolumeMountBuilder;
 import io.fabric8.kubernetes.api.model.batch.v1.Job;
-import org.apache.commons.lang.StringEscapeUtils;
+import org.apache.commons.text.StringEscapeUtils;
 import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.k8s.overlord.KubernetesTaskRunnerConfig;

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/taskadapter/PodTemplateSelector.java

@@ -0,0 +1,40 @@
+/*
+ * 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.k8s.overlord.taskadapter;
+
+import com.google.common.base.Optional;
+import org.apache.druid.indexing.common.task.Task;
+
+/**
+ * Interface for selecting a Pod template based on a given task.
+ * Implementations of this interface are responsible for determining the appropriate
+ * Pod template to use for a specific task based on task characteristics.
+ */
+public interface PodTemplateSelector
+{
+  /**
+   * Selects a Pod template for the given task.
+   * @param task The task for which to select a Pod template..
+   * @return An Optional containing the selected PodTemplateWithName if a suitable
+   *         template is found, or an empty Optional if no appropriate template
+   *         is available for the given task.
+   */
+  Optional<PodTemplateWithName> getPodTemplateForTask(Task task);
+}

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/taskadapter/PodTemplateTaskAdapter.java

@@ -20,7 +20,7 @@
 package org.apache.druid.k8s.overlord.taskadapter;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.base.Supplier;
+import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import io.fabric8.kubernetes.api.model.EnvVar;
@@ -28,17 +28,13 @@ import io.fabric8.kubernetes.api.model.EnvVarBuilder;
 import io.fabric8.kubernetes.api.model.EnvVarSourceBuilder;
 import io.fabric8.kubernetes.api.model.ObjectFieldSelector;
 import io.fabric8.kubernetes.api.model.Pod;
-import io.fabric8.kubernetes.api.model.PodTemplate;
 import io.fabric8.kubernetes.api.model.batch.v1.Job;
 import io.fabric8.kubernetes.api.model.batch.v1.JobBuilder;
-import io.fabric8.kubernetes.client.utils.Serialization;
 import org.apache.commons.io.IOUtils;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.error.InternalServerError;
-import org.apache.druid.guice.IndexingServiceModuleHelper;
 import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.task.Task;
-import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.k8s.overlord.KubernetesTaskRunnerConfig;
@@ -46,24 +42,15 @@ import org.apache.druid.k8s.overlord.common.Base64Compression;
 import org.apache.druid.k8s.overlord.common.DruidK8sConstants;
 import org.apache.druid.k8s.overlord.common.K8sTaskId;
 import org.apache.druid.k8s.overlord.common.KubernetesOverlordUtils;
-import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
-import org.apache.druid.k8s.overlord.execution.PodTemplateSelectStrategy;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.tasklogs.TaskLogs;
 
-import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.Charset;
-import java.nio.file.Files;
 import java.util.Collection;
-import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
-import java.util.Properties;
-import java.util.Set;
 
 /**
  * A PodTemplate {@link TaskAdapter} to transform tasks to kubernetes jobs and kubernetes pods to tasks
@@ -85,33 +72,28 @@ public class PodTemplateTaskAdapter implements TaskAdapter
 
   private static final Logger log = new Logger(PodTemplateTaskAdapter.class);
 
-  private static final String TASK_PROPERTY = IndexingServiceModuleHelper.INDEXER_RUNNER_PROPERTY_PREFIX + ".k8s.podTemplate.";
-
   private final KubernetesTaskRunnerConfig taskRunnerConfig;
   private final TaskConfig taskConfig;
   private final DruidNode node;
   private final ObjectMapper mapper;
-  private final HashMap<String, PodTemplate> templates;
   private final TaskLogs taskLogs;
-  private final Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef;
+  private final PodTemplateSelector podTemplateSelector;
 
   public PodTemplateTaskAdapter(
       KubernetesTaskRunnerConfig taskRunnerConfig,
       TaskConfig taskConfig,
       DruidNode node,
       ObjectMapper mapper,
-      Properties properties,
       TaskLogs taskLogs,
-      Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef
+      PodTemplateSelector podTemplateSelector
   )
   {
     this.taskRunnerConfig = taskRunnerConfig;
     this.taskConfig = taskConfig;
     this.node = node;
     this.mapper = mapper;
-    this.templates = initializePodTemplates(properties);
     this.taskLogs = taskLogs;
-    this.dynamicConfigRef = dynamicConfigRef;
+    this.podTemplateSelector = podTemplateSelector;
   }
 
   /**
@@ -130,15 +112,15 @@ public class PodTemplateTaskAdapter implements TaskAdapter
   @Override
   public Job fromTask(Task task) throws IOException
   {
-    PodTemplateSelectStrategy podTemplateSelectStrategy;
-    KubernetesTaskRunnerDynamicConfig dynamicConfig = dynamicConfigRef.get();
-    if (dynamicConfig == null || dynamicConfig.getPodTemplateSelectStrategy() == null) {
-      podTemplateSelectStrategy = KubernetesTaskRunnerDynamicConfig.DEFAULT_STRATEGY;
-    } else {
-      podTemplateSelectStrategy = dynamicConfig.getPodTemplateSelectStrategy();
+    Optional<PodTemplateWithName> selectedPodTemplate = podTemplateSelector.getPodTemplateForTask(task);
+    if (selectedPodTemplate == null || !selectedPodTemplate.isPresent()) {
+      throw InternalServerError.exception(
+          "Could not find pod template for task [%s]."
+              + " Check the overlord logs for errors selecting the pod template",
+          task.getId()
+      );
     }
-
-    PodTemplateWithName podTemplateWithName = podTemplateSelectStrategy.getPodTemplateForTask(task, templates);
+    PodTemplateWithName podTemplateWithName = podTemplateSelector.getPodTemplateForTask(task).get();
 
     return new JobBuilder()
         .withNewMetadata()
@@ -198,7 +180,7 @@ public class PodTemplateTaskAdapter implements TaskAdapter
 
   private Task toTaskUsingDeepStorage(Job from) throws IOException
   {
-    com.google.common.base.Optional<InputStream> taskBody = taskLogs.streamTaskPayload(getTaskId(from).getOriginalTaskId());
+    Optional<InputStream> taskBody = taskLogs.streamTaskPayload(getTaskId(from).getOriginalTaskId());
     if (!taskBody.isPresent()) {
       throw InternalServerError.exception(
           "Could not load task payload from deep storage for job [%s]."
@@ -224,51 +206,6 @@ public class PodTemplateTaskAdapter implements TaskAdapter
     return new K8sTaskId(taskId);
   }
 
-  private HashMap<String, PodTemplate> initializePodTemplates(Properties properties)
-  {
-    Set<String> taskAdapterTemplateKeys = getTaskAdapterTemplates(properties);
-    if (!taskAdapterTemplateKeys.contains("base")) {
-      throw new IAE("Pod template task adapter requires a base pod template to be specified under druid.indexer.runner.k8s.podTemplate.base");
-    }
-
-    HashMap<String, PodTemplate> podTemplateMap = new HashMap<>();
-    for (String taskAdapterTemplateKey : taskAdapterTemplateKeys) {
-      Optional<PodTemplate> template = loadPodTemplate(taskAdapterTemplateKey, properties);
-      template.ifPresent(podTemplate -> podTemplateMap.put(taskAdapterTemplateKey, podTemplate));
-    }
-    return podTemplateMap;
-  }
-
-  private static Set<String> getTaskAdapterTemplates(Properties properties)
-  {
-    Set<String> taskAdapterTemplates = new HashSet<>();
-
-    for (String runtimeProperty : properties.stringPropertyNames()) {
-      if (runtimeProperty.startsWith(TASK_PROPERTY)) {
-        String[] taskAdapterPropertyPaths = runtimeProperty.split("\\.");
-        taskAdapterTemplates.add(taskAdapterPropertyPaths[taskAdapterPropertyPaths.length - 1]);
-      }
-    }
-
-    return taskAdapterTemplates;
-  }
-
-  private Optional<PodTemplate> loadPodTemplate(String key, Properties properties)
-  {
-    String property = TASK_PROPERTY + key;
-    String podTemplateFile = properties.getProperty(property);
-    if (podTemplateFile == null) {
-      throw new IAE("Pod template file not specified for [%s]", property);
-
-    }
-    try {
-      return Optional.of(Serialization.unmarshal(Files.newInputStream(new File(podTemplateFile).toPath()), PodTemplate.class));
-    }
-    catch (Exception e) {
-      throw new IAE(e, "Failed to load pod template file for [%s] at [%s]", property, podTemplateFile);
-    }
-  }
-
   private Collection<EnvVar> getEnv(Task task) throws IOException
   {
     List<EnvVar> envVars = Lists.newArrayList(

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesAndWorkerTaskRunnerTest.java

@@ -371,4 +371,14 @@ public class KubernetesAndWorkerTaskRunnerTest extends EasyMockSupport
     runner.updateLocation(task, TaskLocation.unknown());
     verifyAll();
   }
+
+  @Test
+  public void test_getMaximumCapacity()
+  {
+    EasyMock.expect(kubernetesTaskRunner.getMaximumCapacityWithAutoscale()).andReturn(1);
+    EasyMock.expect(workerTaskRunner.getMaximumCapacityWithAutoscale()).andReturn(1);
+    replayAll();
+    Assert.assertEquals(2, runner.getMaximumCapacityWithAutoscale());
+    verifyAll();
+  }
 }

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesOverlordModuleTest.java

@@ -38,6 +38,10 @@ import org.apache.druid.indexing.overlord.hrtr.HttpRemoteTaskRunnerFactory;
 import org.apache.druid.jackson.JacksonModule;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 import org.apache.druid.java.util.http.client.HttpClient;
+import org.apache.druid.k8s.overlord.taskadapter.MultiContainerTaskAdapter;
+import org.apache.druid.k8s.overlord.taskadapter.PodTemplateTaskAdapter;
+import org.apache.druid.k8s.overlord.taskadapter.SingleContainerTaskAdapter;
+import org.apache.druid.k8s.overlord.taskadapter.TaskAdapter;
 import org.apache.druid.metadata.MetadataStorageConnector;
 import org.apache.druid.metadata.MetadataStorageTablesConfig;
 import org.apache.druid.server.DruidNode;
@@ -47,6 +51,7 @@ import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.net.URL;
 import java.util.Properties;
 
 @RunWith(EasyMockRunner.class)
@@ -101,6 +106,99 @@ public class KubernetesOverlordModuleTest
     injector.getInstance(KubernetesAndWorkerTaskRunnerFactory.class);
   }
 
+  @Test
+  public void test_build_withMultiContainerAdapterType_returnsWithMultiContainerTaskAdapter()
+  {
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.adapter.type", "overlordMultiContainer");
+    props.setProperty("druid.indexer.runner.namespace", "NAMESPACE");
+
+    injector = makeInjectorWithProperties(props, false, true);
+    TaskAdapter taskAdapter = injector.getInstance(
+        TaskAdapter.class);
+
+    Assert.assertNotNull(taskAdapter);
+    Assert.assertTrue(taskAdapter instanceof MultiContainerTaskAdapter);
+  }
+
+  @Test
+  public void test_build_withSingleContainerAdapterType_returnsKubernetesTaskRunnerWithSingleContainerTaskAdapter()
+  {
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.adapter.type", "overlordSingleContainer");
+    props.setProperty("druid.indexer.runner.namespace", "NAMESPACE");
+    injector = makeInjectorWithProperties(props, false, true);
+    TaskAdapter taskAdapter = injector.getInstance(
+        TaskAdapter.class);
+
+    Assert.assertNotNull(taskAdapter);
+    Assert.assertTrue(taskAdapter instanceof SingleContainerTaskAdapter);
+  }
+
+  @Test
+  public void test_build_withSingleContainerAdapterTypeAndSidecarSupport_throwsProvisionException()
+  {
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.adapter.type", "overlordSingleContainer");
+    props.setProperty("druid.indexer.runner.sidecarSupport", "true");
+    props.setProperty("druid.indexer.runner.namespace", "NAMESPACE");
+    injector = makeInjectorWithProperties(props, false, true);
+
+    Assert.assertThrows(
+        "Invalid pod adapter [overlordSingleContainer], only pod adapter [overlordMultiContainer] can be specified when sidecarSupport is enabled",
+        ProvisionException.class,
+        () -> injector.getInstance(TaskAdapter.class)
+    );
+  }
+
+  @Test
+  public void test_build_withSidecarSupport_returnsKubernetesTaskRunnerWithMultiContainerTaskAdapter()
+  {
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.sidecarSupport", "true");
+    props.setProperty("druid.indexer.runner.namespace", "NAMESPACE");
+    injector = makeInjectorWithProperties(props, false, true);
+
+
+    TaskAdapter adapter = injector.getInstance(TaskAdapter.class);
+
+    Assert.assertNotNull(adapter);
+    Assert.assertTrue(adapter instanceof MultiContainerTaskAdapter);
+  }
+
+  @Test
+  public void test_build_withoutSidecarSupport_returnsKubernetesTaskRunnerWithSingleContainerTaskAdapter()
+  {
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.sidecarSupport", "false");
+    props.setProperty("druid.indexer.runner.namespace", "NAMESPACE");
+    injector = makeInjectorWithProperties(props, false, true);
+
+
+    TaskAdapter adapter = injector.getInstance(TaskAdapter.class);
+
+    Assert.assertNotNull(adapter);
+    Assert.assertTrue(adapter instanceof SingleContainerTaskAdapter);
+  }
+
+  @Test
+  public void test_build_withPodTemplateAdapterType_returnsKubernetesTaskRunnerWithPodTemplateTaskAdapter()
+  {
+    URL url = this.getClass().getClassLoader().getResource("basePodTemplate.yaml");
+
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.adapter.type", "customTemplateAdapter");
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", url.getPath());
+    props.setProperty("druid.indexer.runner.namespace", "NAMESPACE");
+    injector = makeInjectorWithProperties(props, false, true);
+
+
+    TaskAdapter adapter = injector.getInstance(TaskAdapter.class);
+
+    Assert.assertNotNull(adapter);
+    Assert.assertTrue(adapter instanceof PodTemplateTaskAdapter);
+  }
+
   private Injector makeInjectorWithProperties(
       final Properties props,
       boolean isWorkerTypeRemote,

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesPeonLifecycleTest.java

@@ -23,6 +23,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Optional;
 import io.fabric8.kubernetes.api.model.Pod;
 import io.fabric8.kubernetes.api.model.PodBuilder;
+import io.fabric8.kubernetes.api.model.PodStatus;
 import io.fabric8.kubernetes.api.model.batch.v1.Job;
 import io.fabric8.kubernetes.api.model.batch.v1.JobBuilder;
 import io.fabric8.kubernetes.client.dsl.LogWatch;
@@ -57,6 +58,7 @@ import java.util.concurrent.atomic.AtomicReference;
 public class KubernetesPeonLifecycleTest extends EasyMockSupport
 {
   private static final String ID = "id";
+  private static final String IP = "ip";
   private static final TaskStatus SUCCESS = TaskStatus.success(ID);
 
   @Mock KubernetesPeonClient kubernetesClient;
@@ -286,6 +288,9 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     stateListener.stateChanged(KubernetesPeonLifecycle.State.STOPPED, ID);
     EasyMock.expectLastCall().once();
 
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(
+        new PodBuilder().editOrNewStatusLike(getPodStatusWithIP()).endStatus().build()
+    );
     replayAll();
 
     TaskStatus taskStatus = peonLifecycle.join(0L);
@@ -337,7 +342,9 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     EasyMock.expectLastCall().once();
     logWatch.close();
     EasyMock.expectLastCall();
-
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(
+        new PodBuilder().editOrNewStatusLike(getPodStatusWithIP()).endStatus().build()
+    );
     Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());
 
     replayAll();
@@ -393,7 +400,9 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     EasyMock.expectLastCall();
 
     Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());
-
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(
+        new PodBuilder().editOrNewStatusLike(getPodStatusWithIP()).endStatus().build()
+    ).anyTimes();
     replayAll();
 
     TaskStatus taskStatus = peonLifecycle.join(0L);
@@ -445,7 +454,9 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     EasyMock.expectLastCall().once();
 
     Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());
-
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(
+        new PodBuilder().editOrNewStatusLike(getPodStatusWithIP()).endStatus().build()
+    );
     replayAll();
 
     TaskStatus taskStatus = peonLifecycle.join(0L);
@@ -493,7 +504,9 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     EasyMock.expectLastCall().once();
     logWatch.close();
     EasyMock.expectLastCall();
-
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(
+        new PodBuilder().editOrNewStatusLike(getPodStatusWithIP()).endStatus().build()
+    );
     Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());
 
     replayAll();
@@ -545,7 +558,9 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     EasyMock.expectLastCall().once();
     logWatch.close();
     EasyMock.expectLastCall();
-
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(
+        new PodBuilder().editOrNewStatusLike(getPodStatusWithIP()).endStatus().build()
+    );
     Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());
 
     replayAll();
@@ -585,7 +600,9 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     EasyMock.expectLastCall().once();
     logWatch.close();
     EasyMock.expectLastCall();
-
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(
+        new PodBuilder().editOrNewStatusLike(getPodStatusWithIP()).endStatus().build()
+    );
     Assert.assertEquals(KubernetesPeonLifecycle.State.NOT_STARTED, peonLifecycle.getState());
 
     replayAll();
@@ -768,7 +785,7 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
   }
 
   @Test
-  public void test_getTaskLocation_withRunningTaskState_withoutPeonPod_returnsUnknown()
+  public void test_getTaskLocation_withRunningTaskState_taskLocationUnset_returnsUnknown()
       throws NoSuchFieldException, IllegalAccessException
   {
     KubernetesPeonLifecycle peonLifecycle = new KubernetesPeonLifecycle(
@@ -780,8 +797,6 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     );
     setPeonLifecycleState(peonLifecycle, KubernetesPeonLifecycle.State.RUNNING);
 
-    EasyMock.expect(kubernetesClient.getPeonPod(k8sTaskId.getK8sJobName())).andReturn(Optional.absent());
-
     replayAll();
 
     Assert.assertEquals(TaskLocation.unknown(), peonLifecycle.getTaskLocation());
@@ -790,35 +805,7 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
   }
 
   @Test
-  public void test_getTaskLocation_withRunningTaskState_withPeonPodWithoutStatus_returnsUnknown()
-      throws NoSuchFieldException, IllegalAccessException
-  {
-    KubernetesPeonLifecycle peonLifecycle = new KubernetesPeonLifecycle(
-        task,
-        kubernetesClient,
-        taskLogs,
-        mapper,
-        stateListener
-    );
-    setPeonLifecycleState(peonLifecycle, KubernetesPeonLifecycle.State.RUNNING);
-
-    Pod pod = new PodBuilder()
-        .withNewMetadata()
-        .withName(ID)
-        .endMetadata()
-        .build();
-
-    EasyMock.expect(kubernetesClient.getPeonPod(k8sTaskId.getK8sJobName())).andReturn(Optional.of(pod));
-
-    replayAll();
-
-    Assert.assertEquals(TaskLocation.unknown(), peonLifecycle.getTaskLocation());
-
-    verifyAll();
-  }
-
-  @Test
-  public void test_getTaskLocation_withRunningTaskState_withPeonPodWithStatus_returnsLocation()
+  public void test_getTaskLocationFromK8s()
       throws NoSuchFieldException, IllegalAccessException
   {
     KubernetesPeonLifecycle peonLifecycle = new KubernetesPeonLifecycle(
@@ -839,12 +826,11 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
         .endStatus()
         .build();
 
-    EasyMock.expect(kubernetesClient.getPeonPod(k8sTaskId.getK8sJobName())).andReturn(Optional.of(pod));
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(pod).once();
 
     replayAll();
 
-    TaskLocation location = peonLifecycle.getTaskLocation();
-
+    TaskLocation location = peonLifecycle.getTaskLocationFromK8s();
     Assert.assertEquals("ip", location.getHost());
     Assert.assertEquals(8100, location.getPort());
     Assert.assertEquals(-1, location.getTlsPort());
@@ -854,43 +840,7 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
   }
 
   @Test
-  public void test_getTaskLocation_saveTaskLocation()
-      throws NoSuchFieldException, IllegalAccessException
-  {
-    KubernetesPeonLifecycle peonLifecycle = new KubernetesPeonLifecycle(
-        task,
-        kubernetesClient,
-        taskLogs,
-        mapper,
-        stateListener
-    );
-    setPeonLifecycleState(peonLifecycle, KubernetesPeonLifecycle.State.RUNNING);
-
-    Pod pod = new PodBuilder()
-        .withNewMetadata()
-        .withName(ID)
-        .endMetadata()
-        .withNewStatus()
-        .withPodIP("ip")
-        .endStatus()
-        .build();
-
-    EasyMock.expect(kubernetesClient.getPeonPod(k8sTaskId.getK8sJobName())).andReturn(Optional.of(pod)).once();
-
-    replayAll();
-
-    TaskLocation location = peonLifecycle.getTaskLocation();
-    peonLifecycle.getTaskLocation();
-    Assert.assertEquals("ip", location.getHost());
-    Assert.assertEquals(8100, location.getPort());
-    Assert.assertEquals(-1, location.getTlsPort());
-    Assert.assertEquals(ID, location.getK8sPodName());
-
-    verifyAll();
-  }
-
-  @Test
-  public void test_getTaskLocation_withRunningTaskState_withPeonPodWithStatusWithTLSAnnotation_returnsLocation()
+  public void test_getTaskLocationFromK8s_withPeonPodWithStatusWithTLSAnnotation()
       throws NoSuchFieldException, IllegalAccessException
   {
     KubernetesPeonLifecycle peonLifecycle = new KubernetesPeonLifecycle(
@@ -912,11 +862,11 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
         .endStatus()
         .build();
 
-    EasyMock.expect(kubernetesClient.getPeonPod(k8sTaskId.getK8sJobName())).andReturn(Optional.of(pod));
+    EasyMock.expect(kubernetesClient.getPeonPodWithRetries(k8sTaskId.getK8sJobName())).andReturn(pod).once();
 
     replayAll();
 
-    TaskLocation location = peonLifecycle.getTaskLocation();
+    TaskLocation location = peonLifecycle.getTaskLocationFromK8s();
 
     Assert.assertEquals("ip", location.getHost());
     Assert.assertEquals(-1, location.getPort());
@@ -938,7 +888,6 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
         stateListener
     );
     setPeonLifecycleState(peonLifecycle, KubernetesPeonLifecycle.State.STOPPED);
-    EasyMock.expect(kubernetesClient.getPeonPod(k8sTaskId.getK8sJobName())).andReturn(Optional.absent()).once();
 
     replayAll();
     Assert.assertEquals(TaskLocation.unknown(), peonLifecycle.getTaskLocation());
@@ -952,4 +901,11 @@ public class KubernetesPeonLifecycleTest extends EasyMockSupport
     stateField.setAccessible(true);
     stateField.set(peonLifecycle, new AtomicReference<>(state));
   }
+
+  private PodStatus getPodStatusWithIP()
+  {
+    PodStatus podStatus = new PodStatus();
+    podStatus.setPodIP(IP);
+    return podStatus;
+  }
 }

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerFactoryTest.java

@@ -24,13 +24,10 @@ import com.google.common.base.Supplier;
 import org.apache.druid.indexing.common.TestUtils;
 import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.config.TaskConfigBuilder;
-import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 import org.apache.druid.k8s.overlord.common.DruidKubernetesClient;
 import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
-import org.apache.druid.k8s.overlord.taskadapter.MultiContainerTaskAdapter;
-import org.apache.druid.k8s.overlord.taskadapter.PodTemplateTaskAdapter;
-import org.apache.druid.k8s.overlord.taskadapter.SingleContainerTaskAdapter;
+import org.apache.druid.k8s.overlord.taskadapter.TaskAdapter;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.log.StartupLoggingConfig;
 import org.apache.druid.tasklogs.NoopTaskLogs;
@@ -40,7 +37,6 @@ import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 
-import java.net.URL;
 import java.util.Properties;
 
 public class KubernetesTaskRunnerFactoryTest
@@ -56,6 +52,7 @@ public class KubernetesTaskRunnerFactoryTest
   private DruidKubernetesClient druidKubernetesClient;
   @Mock private ServiceEmitter emitter;
   @Mock private Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef;
+  @Mock private TaskAdapter taskAdapter;
 
   @Before
   public void setup()
@@ -87,14 +84,10 @@ public class KubernetesTaskRunnerFactoryTest
         objectMapper,
         null,
         kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
         taskLogs,
-        druidNode,
-        taskConfig,
-        properties,
         druidKubernetesClient,
         emitter,
-        dynamicConfigRef
+        taskAdapter
     );
 
     KubernetesTaskRunner expectedRunner = factory.build();
@@ -102,199 +95,4 @@ public class KubernetesTaskRunnerFactoryTest
 
     Assert.assertEquals(expectedRunner, actualRunner);
   }
-
-  @Test
-  public void test_build_withoutSidecarSupport_returnsKubernetesTaskRunnerWithSingleContainerTaskAdapter()
-  {
-    KubernetesTaskRunnerFactory factory = new KubernetesTaskRunnerFactory(
-        objectMapper,
-        null,
-        kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
-        taskLogs,
-        druidNode,
-        taskConfig,
-        properties,
-        druidKubernetesClient,
-        emitter,
-        dynamicConfigRef
-    );
-
-    KubernetesTaskRunner runner = factory.build();
-
-    Assert.assertNotNull(runner);
-    Assert.assertTrue(runner.adapter instanceof SingleContainerTaskAdapter);
-  }
-
-  @Test
-  public void test_build_withSidecarSupport_returnsKubernetesTaskRunnerWithMultiContainerTaskAdapter()
-  {
-    kubernetesTaskRunnerConfig = KubernetesTaskRunnerConfig.builder()
-        .withCapacity(1)
-        .withSidecarSupport(true)
-        .build();
-
-    KubernetesTaskRunnerFactory factory = new KubernetesTaskRunnerFactory(
-        objectMapper,
-        null,
-        kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
-        taskLogs,
-        druidNode,
-        taskConfig,
-        properties,
-        druidKubernetesClient,
-        emitter,
-        dynamicConfigRef
-    );
-
-    KubernetesTaskRunner runner = factory.build();
-
-    Assert.assertNotNull(runner);
-    Assert.assertTrue(runner.adapter instanceof MultiContainerTaskAdapter);
-  }
-
-  @Test
-  public void test_build_withSingleContainerAdapterType_returnsKubernetesTaskRunnerWithSingleContainerTaskAdapter()
-  {
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.adapter.type", "overlordSingleContainer");
-
-    KubernetesTaskRunnerFactory factory = new KubernetesTaskRunnerFactory(
-        objectMapper,
-        null,
-        kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
-        taskLogs,
-        druidNode,
-        taskConfig,
-        props,
-        druidKubernetesClient,
-        emitter,
-        dynamicConfigRef
-    );
-
-    KubernetesTaskRunner runner = factory.build();
-
-    Assert.assertNotNull(runner);
-    Assert.assertTrue(runner.adapter instanceof SingleContainerTaskAdapter);
-  }
-
-  @Test
-  public void test_build_withSingleContainerAdapterTypeAndSidecarSupport_throwsIAE()
-  {
-    kubernetesTaskRunnerConfig = KubernetesTaskRunnerConfig.builder()
-        .withCapacity(1)
-        .withSidecarSupport(true)
-        .build();
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.adapter.type", "overlordSingleContainer");
-
-    KubernetesTaskRunnerFactory factory = new KubernetesTaskRunnerFactory(
-        objectMapper,
-        null,
-        kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
-        taskLogs,
-        druidNode,
-        taskConfig,
-        props,
-        druidKubernetesClient,
-        emitter,
-        dynamicConfigRef
-    );
-
-    Assert.assertThrows(
-        "Invalid pod adapter [overlordSingleContainer], only pod adapter [overlordMultiContainer] can be specified when sidecarSupport is enabled",
-        IAE.class,
-        factory::build
-    );
-  }
-
-  @Test
-  public void test_build_withMultiContainerAdapterType_returnsKubernetesTaskRunnerWithMultiContainerTaskAdapter()
-  {
-    kubernetesTaskRunnerConfig = KubernetesTaskRunnerConfig.builder()
-        .withCapacity(1)
-        .withSidecarSupport(true)
-        .build();
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.adapter.type", "overlordMultiContainer");
-
-    KubernetesTaskRunnerFactory factory = new KubernetesTaskRunnerFactory(
-        objectMapper,
-        null,
-        kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
-        taskLogs,
-        druidNode,
-        taskConfig,
-        props,
-        druidKubernetesClient,
-        emitter,
-        dynamicConfigRef
-    );
-
-    KubernetesTaskRunner runner = factory.build();
-
-    Assert.assertNotNull(runner);
-    Assert.assertTrue(runner.adapter instanceof MultiContainerTaskAdapter);
-  }
-
-  @Test
-  public void test_build_withMultiContainerAdapterTypeAndSidecarSupport_returnsKubernetesTaskRunnerWithMultiContainerTaskAdapter()
-  {
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.adapter.type", "overlordMultiContainer");
-
-    KubernetesTaskRunnerFactory factory = new KubernetesTaskRunnerFactory(
-        objectMapper,
-        null,
-        kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
-        taskLogs,
-        druidNode,
-        taskConfig,
-        props,
-        druidKubernetesClient,
-        emitter,
-        dynamicConfigRef
-    );
-
-    KubernetesTaskRunner runner = factory.build();
-
-    Assert.assertNotNull(runner);
-    Assert.assertTrue(runner.adapter instanceof MultiContainerTaskAdapter);
-  }
-
-  @Test
-  public void test_build_withPodTemplateAdapterType_returnsKubernetesTaskRunnerWithPodTemplateTaskAdapter()
-  {
-    URL url = this.getClass().getClassLoader().getResource("basePodTemplate.yaml");
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.adapter.type", "customTemplateAdapter");
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", url.getPath());
-
-    KubernetesTaskRunnerFactory factory = new KubernetesTaskRunnerFactory(
-        objectMapper,
-        null,
-        kubernetesTaskRunnerConfig,
-        startupLoggingConfig,
-        taskLogs,
-        druidNode,
-        taskConfig,
-        props,
-        druidKubernetesClient,
-        emitter,
-        dynamicConfigRef
-    );
-
-    KubernetesTaskRunner runner = factory.build();
-
-    Assert.assertNotNull(runner);
-    Assert.assertTrue(runner.adapter instanceof PodTemplateTaskAdapter);
-  }
 }

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/taskadapter/DruidPeonClientIntegrationTest.java

@@ -50,8 +50,8 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 
 import java.io.File;
-import java.io.IOException;
 import java.io.InputStream;
+import java.io.UncheckedIOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -147,7 +147,7 @@ public class DruidPeonClientIntegrationTest
                                  .map(Integer::parseInt)
                                  .collect(Collectors.toList()));
       }
-      catch (IOException e) {
+      catch (UncheckedIOException e) {
         throw new RuntimeException(e);
       }
     });

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/taskadapter/DynamicConfigPodTemplateSelectorTest.java

@@ -0,0 +1,253 @@
+/*
+ * 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.k8s.overlord.taskadapter;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Optional;
+import com.google.common.base.Supplier;
+import io.fabric8.kubernetes.api.model.PodTemplate;
+import io.fabric8.kubernetes.api.model.PodTemplateBuilder;
+import io.fabric8.kubernetes.api.model.VolumeBuilder;
+import org.apache.druid.indexing.common.TestUtils;
+import org.apache.druid.indexing.common.task.NoopTask;
+import org.apache.druid.indexing.common.task.Task;
+import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.k8s.overlord.common.K8sTestUtils;
+import org.apache.druid.k8s.overlord.execution.DefaultKubernetesTaskRunnerDynamicConfig;
+import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
+import org.apache.druid.k8s.overlord.execution.Selector;
+import org.apache.druid.k8s.overlord.execution.SelectorBasedPodTemplateSelectStrategy;
+import org.junit.Assert;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+import org.mockito.internal.util.collections.Sets;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.Properties;
+
+public class DynamicConfigPodTemplateSelectorTest
+{
+  @TempDir
+  private Path tempDir;
+  private ObjectMapper mapper;
+  private PodTemplate podTemplateSpec;
+  private Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef;
+
+  @BeforeEach
+  public void setup()
+  {
+    mapper = new TestUtils().getTestObjectMapper();
+    podTemplateSpec = K8sTestUtils.fileToResource("basePodTemplate.yaml", PodTemplate.class);
+    dynamicConfigRef = () -> new DefaultKubernetesTaskRunnerDynamicConfig(KubernetesTaskRunnerDynamicConfig.DEFAULT_STRATEGY);
+  }
+
+  @Test
+  public void test_fromTask_withoutBasePodTemplateInRuntimeProperites_raisesIAE()
+  {
+    Exception exception = Assert.assertThrows(
+        "No base prop should throw an IAE",
+        IAE.class,
+        () -> new DynamicConfigPodTemplateSelector(
+        new Properties(),
+        dynamicConfigRef
+    ));
+    Assertions.assertEquals(exception.getMessage(), "Pod template task adapter requires a base pod template to be specified under druid.indexer.runner.k8s.podTemplate.base");
+  }
+
+  @Test
+  public void test_fromTask_withBasePodTemplateInRuntimeProperites_withEmptyFile_raisesIAE() throws IOException
+  {
+    Path templatePath = Files.createFile(tempDir.resolve("empty.yaml"));
+
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+
+    Exception exception = Assert.assertThrows(
+        "Empty base pod template should throw a exception",
+        IAE.class,
+        () -> new DynamicConfigPodTemplateSelector(
+            props,
+            dynamicConfigRef
+    ));
+
+    Assertions.assertTrue(exception.getMessage().contains("Failed to load pod template file for"));
+  }
+
+  @Test
+  public void test_fromTask_withBasePodTemplateInRuntimeProperites() throws IOException
+  {
+    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
+    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
+
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+
+    DynamicConfigPodTemplateSelector adapter = new DynamicConfigPodTemplateSelector(
+        props,
+        dynamicConfigRef
+    );
+
+    Task task = new NoopTask("id", "id", "datasource", 0, 0, null);
+    Optional<PodTemplateWithName> actual = adapter.getPodTemplateForTask(task);
+    PodTemplate expected = K8sTestUtils.fileToResource("expectedNoopPodTemplate.yaml", PodTemplate.class);
+
+    Assertions.assertTrue(actual.isPresent());
+    Assertions.assertEquals("base", actual.get().getName());
+    Assertions.assertEquals(expected, actual.get().getPodTemplate());
+  }
+
+  @Test
+  public void test_fromTask_withIndexKafkaPodTemplateInRuntimeProperties() throws IOException
+  {
+    Path baseTemplatePath = Files.createFile(tempDir.resolve("base.yaml"));
+    mapper.writeValue(baseTemplatePath.toFile(), podTemplateSpec);
+
+    Path kafkaTemplatePath = Files.createFile(tempDir.resolve("kafka.yaml"));
+    PodTemplate kafkaPodTemplate = new PodTemplateBuilder(podTemplateSpec)
+        .editTemplate()
+        .editSpec()
+        .setNewVolumeLike(0, new VolumeBuilder().withName("volume").build())
+        .endVolume()
+        .endSpec()
+        .endTemplate()
+        .build();
+    mapper.writeValue(kafkaTemplatePath.toFile(), kafkaPodTemplate);
+
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", baseTemplatePath.toString());
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.index_kafka", kafkaTemplatePath.toString());
+
+    DynamicConfigPodTemplateSelector selector = new DynamicConfigPodTemplateSelector(
+        props,
+        dynamicConfigRef
+    );
+
+    Task kafkaTask = new NoopTask("id", "id", "datasource", 0, 0, null) {
+      @Override
+      public String getType()
+      {
+        return "index_kafka";
+      }
+    };
+
+    Task noopTask = new NoopTask("id", "id", "datasource", 0, 0, null);
+    Optional<PodTemplateWithName> podTemplateWithName = selector.getPodTemplateForTask(kafkaTask);
+    Assertions.assertTrue(podTemplateWithName.isPresent());
+    Assertions.assertEquals(1, podTemplateWithName.get().getPodTemplate().getTemplate().getSpec().getVolumes().size(), 1);
+    Assertions.assertEquals("index_kafka", podTemplateWithName.get().getName());
+
+
+    podTemplateWithName = selector.getPodTemplateForTask(noopTask);
+    Assertions.assertTrue(podTemplateWithName.isPresent());
+    Assertions.assertEquals(0, podTemplateWithName.get().getPodTemplate().getTemplate().getSpec().getVolumes().size(), 1);
+    Assertions.assertEquals("base", podTemplateWithName.get().getName());
+  }
+
+  @Test
+  public void test_fromTask_withNoopPodTemplateInRuntimeProperties_withEmptyFile_raisesIAE() throws IOException
+  {
+    Path baseTemplatePath = Files.createFile(tempDir.resolve("base.yaml"));
+    Path noopTemplatePath = Files.createFile(tempDir.resolve("noop.yaml"));
+    mapper.writeValue(baseTemplatePath.toFile(), podTemplateSpec);
+
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", baseTemplatePath.toString());
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.noop", noopTemplatePath.toString());
+
+    Assert.assertThrows(IAE.class, () -> new DynamicConfigPodTemplateSelector(
+        props,
+        dynamicConfigRef
+    ));
+  }
+
+  @Test
+  public void test_fromTask_withNoopPodTemplateInRuntimeProperites() throws IOException
+  {
+    Path templatePath = Files.createFile(tempDir.resolve("noop.yaml"));
+    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
+
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.noop", templatePath.toString());
+
+    DynamicConfigPodTemplateSelector podTemplateSelector = new DynamicConfigPodTemplateSelector(
+        props,
+        dynamicConfigRef
+    );
+
+    Task task = new NoopTask("id", "id", "datasource", 0, 0, null);
+    Optional<PodTemplateWithName> podTemplateWithName = podTemplateSelector.getPodTemplateForTask(task);
+    Assertions.assertTrue(podTemplateWithName.isPresent());
+    PodTemplate expected = K8sTestUtils.fileToResource("expectedNoopPodTemplate.yaml", PodTemplate.class);
+
+    Assertions.assertTrue(podTemplateWithName.isPresent());
+    Assertions.assertEquals("noop", podTemplateWithName.get().getName());
+    Assertions.assertEquals(expected, podTemplateWithName.get().getPodTemplate());
+  }
+
+  @Test
+  public void test_fromTask_matchPodTemplateBasedOnStrategy() throws IOException
+  {
+    String dataSource = "my_table";
+    Path baseTemplatePath = Files.createFile(tempDir.resolve("base.yaml"));
+    mapper.writeValue(baseTemplatePath.toFile(), podTemplateSpec);
+
+    Path lowThroughputTemplatePath = Files.createFile(tempDir.resolve("low-throughput.yaml"));
+    PodTemplate lowThroughputPodTemplate = new PodTemplateBuilder(podTemplateSpec)
+        .editTemplate()
+        .editSpec()
+        .setNewVolumeLike(0, new VolumeBuilder().withName("volume").build())
+        .endVolume()
+        .endSpec()
+        .endTemplate()
+        .build();
+    mapper.writeValue(lowThroughputTemplatePath.toFile(), lowThroughputPodTemplate);
+
+    Properties props = new Properties();
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", baseTemplatePath.toString());
+    props.setProperty("druid.indexer.runner.k8s.podTemplate.lowThroughput", lowThroughputTemplatePath.toString());
+    dynamicConfigRef = () -> new DefaultKubernetesTaskRunnerDynamicConfig(new SelectorBasedPodTemplateSelectStrategy(
+        Collections.singletonList(
+            new Selector("lowThrougput", null, null, Sets.newSet(dataSource)
+            )
+        )
+    ));
+
+    DynamicConfigPodTemplateSelector podTemplateSelector = new DynamicConfigPodTemplateSelector(
+        props,
+        dynamicConfigRef
+    );
+
+    Task taskWithMatchedDatasource = new NoopTask("id", "id", dataSource, 0, 0, null);
+    Task noopTask = new NoopTask("id", "id", "datasource", 0, 0, null);
+    Optional<PodTemplateWithName> actual = podTemplateSelector.getPodTemplateForTask(taskWithMatchedDatasource);
+    Assertions.assertTrue(actual.isPresent());
+    Assertions.assertEquals(1, actual.get().getPodTemplate().getTemplate().getSpec().getVolumes().size(), 1);
+
+    actual = podTemplateSelector.getPodTemplateForTask(noopTask);
+    Assertions.assertTrue(actual.isPresent());
+    Assertions.assertEquals(0, actual.get().getPodTemplate().getTemplate().getSpec().getVolumes().size(), 1);
+  }
+}

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/taskadapter/K8sTaskAdapterTest.java

@@ -40,8 +40,8 @@ import io.fabric8.kubernetes.api.model.batch.v1.JobBuilder;
 import io.fabric8.kubernetes.api.model.batch.v1.JobList;
 import io.fabric8.kubernetes.client.KubernetesClient;
 import io.fabric8.kubernetes.client.server.mock.EnableKubernetesMockClient;
-import org.apache.commons.lang.RandomStringUtils;
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.RandomStringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.indexing.common.TestUtils;
 import org.apache.druid.indexing.common.config.TaskConfig;

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/taskadapter/PodTemplateTaskAdapterTest.java

@@ -21,30 +21,22 @@ package org.apache.druid.k8s.overlord.taskadapter;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Optional;
-import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
 import io.fabric8.kubernetes.api.model.PodTemplate;
-import io.fabric8.kubernetes.api.model.PodTemplateBuilder;
-import io.fabric8.kubernetes.api.model.VolumeBuilder;
 import io.fabric8.kubernetes.api.model.batch.v1.Job;
 import io.fabric8.kubernetes.api.model.batch.v1.JobBuilder;
-import org.apache.commons.lang.RandomStringUtils;
+import org.apache.commons.lang3.RandomStringUtils;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.indexing.common.TestUtils;
 import org.apache.druid.indexing.common.config.TaskConfig;
 import org.apache.druid.indexing.common.config.TaskConfigBuilder;
 import org.apache.druid.indexing.common.task.NoopTask;
 import org.apache.druid.indexing.common.task.Task;
-import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.k8s.overlord.KubernetesTaskRunnerConfig;
 import org.apache.druid.k8s.overlord.common.Base64Compression;
 import org.apache.druid.k8s.overlord.common.DruidK8sConstants;
 import org.apache.druid.k8s.overlord.common.K8sTaskId;
 import org.apache.druid.k8s.overlord.common.K8sTestUtils;
-import org.apache.druid.k8s.overlord.execution.DefaultKubernetesTaskRunnerDynamicConfig;
-import org.apache.druid.k8s.overlord.execution.KubernetesTaskRunnerDynamicConfig;
-import org.apache.druid.k8s.overlord.execution.Selector;
-import org.apache.druid.k8s.overlord.execution.SelectorBasedPodTemplateSelectStrategy;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.coordination.BroadcastDatasourceLoadingSpec;
 import org.apache.druid.tasklogs.TaskLogs;
@@ -53,32 +45,25 @@ import org.junit.Assert;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-import org.junit.jupiter.api.io.TempDir;
 import org.mockito.Mockito;
-import org.mockito.internal.util.collections.Sets;
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.nio.charset.Charset;
-import java.nio.file.Files;
-import java.nio.file.Path;
 import java.util.Collections;
 import java.util.Map;
-import java.util.Properties;
 import java.util.stream.Collectors;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class PodTemplateTaskAdapterTest
 {
-  @TempDir private Path tempDir;
   private KubernetesTaskRunnerConfig taskRunnerConfig;
   private PodTemplate podTemplateSpec;
   private TaskConfig taskConfig;
   private DruidNode node;
   private ObjectMapper mapper;
   private TaskLogs taskLogs;
-  private Supplier<KubernetesTaskRunnerDynamicConfig> dynamicConfigRef;
 
   @BeforeEach
   public void setup()
@@ -98,85 +83,12 @@ public class PodTemplateTaskAdapterTest
     podTemplateSpec = K8sTestUtils.fileToResource("basePodTemplate.yaml", PodTemplate.class);
 
     taskLogs = EasyMock.createMock(TaskLogs.class);
-    dynamicConfigRef = () -> new DefaultKubernetesTaskRunnerDynamicConfig(KubernetesTaskRunnerDynamicConfig.DEFAULT_STRATEGY);
-  }
-
-  @Test
-  public void test_fromTask_withoutBasePodTemplateInRuntimeProperites_raisesIAE()
-  {
-    Exception exception = Assert.assertThrows(
-        "No base prop should throw an IAE",
-        IAE.class,
-        () -> new PodTemplateTaskAdapter(
-        taskRunnerConfig,
-        taskConfig,
-        node,
-        mapper,
-        new Properties(),
-        taskLogs,
-        dynamicConfigRef
-    ));
-    Assert.assertEquals(exception.getMessage(), "Pod template task adapter requires a base pod template to be specified under druid.indexer.runner.k8s.podTemplate.base");
-  }
-
-  @Test
-  public void test_fromTask_withBasePodTemplateInRuntimeProperites_withEmptyFile_raisesIAE() throws IOException
-  {
-    Path templatePath = Files.createFile(tempDir.resolve("empty.yaml"));
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
-
-    Exception exception = Assert.assertThrows(
-        "Empty base pod template should throw a exception",
-        IAE.class,
-        () -> new PodTemplateTaskAdapter(
-        taskRunnerConfig,
-        taskConfig,
-        node,
-        mapper,
-        props,
-        taskLogs,
-        dynamicConfigRef
-    ));
-
-    Assert.assertTrue(exception.getMessage().contains("Failed to load pod template file for"));
-  }
-
-  @Test
-  public void test_fromTask_withBasePodTemplateInRuntimeProperites() throws IOException
-  {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
-
-    PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
-        taskRunnerConfig,
-        taskConfig,
-        node,
-        mapper,
-        props,
-        taskLogs,
-        dynamicConfigRef
-    );
-
-    Task task = new NoopTask("id", "id", "datasource", 0, 0, null);
-    Job actual = adapter.fromTask(task);
-    Job expected = K8sTestUtils.fileToResource("expectedNoopJobBase.yaml", Job.class);
-
-    assertJobSpecsEqual(actual, expected);
   }
 
   @Test
   public void test_fromTask_withBasePodTemplateInRuntimeProperites_andTlsEnabled() throws IOException
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
@@ -191,9 +103,8 @@ public class PodTemplateTaskAdapterTest
             true
         ),
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Task task = new NoopTask("id", "id", "datasource", 0, 0, null);
@@ -203,73 +114,18 @@ public class PodTemplateTaskAdapterTest
     assertJobSpecsEqual(actual, expected);
   }
 
-  @Test
-  public void test_fromTask_withNoopPodTemplateInRuntimeProperties_withEmptyFile_raisesIAE() throws IOException
-  {
-    Path baseTemplatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    Path noopTemplatePath = Files.createFile(tempDir.resolve("noop.yaml"));
-    mapper.writeValue(baseTemplatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", baseTemplatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.noop", noopTemplatePath.toString());
-
-    Assert.assertThrows(IAE.class, () -> new PodTemplateTaskAdapter(
-        taskRunnerConfig,
-        taskConfig,
-        node,
-        mapper,
-        props,
-        taskLogs,
-        dynamicConfigRef
-    ));
-  }
-
-  @Test
-  public void test_fromTask_withNoopPodTemplateInRuntimeProperites() throws IOException
-  {
-    Path templatePath = Files.createFile(tempDir.resolve("noop.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.noop", templatePath.toString());
-
-    PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
-        taskRunnerConfig,
-        taskConfig,
-        node,
-        mapper,
-        props,
-        taskLogs,
-        dynamicConfigRef
-    );
-
-    Task task = new NoopTask("id", "id", "datasource", 0, 0, null);
-    Job actual = adapter.fromTask(task);
-    Job expected = K8sTestUtils.fileToResource("expectedNoopJob.yaml", Job.class);
-
-    assertJobSpecsEqual(actual, expected);
-  }
-
   @Test
   public void test_fromTask_withNoopPodTemplateInRuntimeProperites_dontSetTaskJSON() throws IOException
   {
-    Path templatePath = Files.createFile(tempDir.resolve("noop.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.noop", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Task task = new NoopTask(
@@ -288,22 +144,18 @@ public class PodTemplateTaskAdapterTest
   }
 
   @Test
-  public void test_fromTask_withoutAnnotations_throwsDruidException() throws IOException
+  public void test_fromTask_withoutAnnotations_throwsDruidException()
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Job job = K8sTestUtils.fileToResource("baseJobWithoutAnnotations.yaml", Job.class);
@@ -313,21 +165,17 @@ public class PodTemplateTaskAdapterTest
   }
 
   @Test
-  public void test_getTaskId() throws IOException
+  public void test_getTaskId()
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
     Job job = new JobBuilder()
         .editSpec().editTemplate().editMetadata()
@@ -338,21 +186,16 @@ public class PodTemplateTaskAdapterTest
   }
 
   @Test
-  public void test_getTaskId_noAnnotations() throws IOException
+  public void test_getTaskId_noAnnotations()
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
     Job job = new JobBuilder()
         .editSpec().editTemplate().editMetadata()
@@ -363,21 +206,16 @@ public class PodTemplateTaskAdapterTest
   }
 
   @Test
-  public void test_getTaskId_missingTaskIdAnnotation() throws IOException
+  public void test_getTaskId_missingTaskIdAnnotation()
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
     Job job = new JobBuilder()
         .editSpec().editTemplate().editMetadata()
@@ -389,22 +227,17 @@ public class PodTemplateTaskAdapterTest
   }
 
   @Test
-  public void test_toTask_withoutTaskAnnotation_throwsIOE() throws IOException
+  public void test_toTask_withoutTaskAnnotation_throwsIOE()
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.put("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Job baseJob = K8sTestUtils.fileToResource("baseJobWithoutAnnotations.yaml", Job.class);
@@ -424,20 +257,15 @@ public class PodTemplateTaskAdapterTest
   @Test
   public void test_toTask() throws IOException
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.put("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Job job = K8sTestUtils.fileToResource("baseJob.yaml", Job.class);
@@ -450,11 +278,7 @@ public class PodTemplateTaskAdapterTest
   @Test
   public void test_toTask_useTaskPayloadManager() throws IOException
   {
-    Path templatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.put("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
     Task expected = new NoopTask("id", null, "datasource", 0, 0, ImmutableMap.of());
     TaskLogs mockTestLogs = Mockito.mock(TaskLogs.class);
@@ -467,9 +291,8 @@ public class PodTemplateTaskAdapterTest
         taskConfig,
         node,
         mapper,
-        props,
         mockTestLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Job job = K8sTestUtils.fileToResource("expectedNoopJob.yaml", Job.class);
@@ -480,21 +303,15 @@ public class PodTemplateTaskAdapterTest
   @Test
   public void test_fromTask_withRealIds() throws IOException
   {
-    Path templatePath = Files.createFile(tempDir.resolve("noop.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.noop", templatePath.toString());
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Task task = new NoopTask(
@@ -513,23 +330,75 @@ public class PodTemplateTaskAdapterTest
   }
 
   @Test
-  public void test_fromTask_taskSupportsQueries() throws IOException
+  public void test_fromTask_noTemplate()
   {
-    Path templatePath = Files.createFile(tempDir.resolve("noop.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
 
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.queryable", templatePath.toString());
+    PodTemplateSelector podTemplateSelector = EasyMock.createMock(PodTemplateSelector.class);
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
+    );
+
+    Task task = new NoopTask(
+        "id",
+        "groupId",
+        "data_source",
+        0,
+        0,
+        null
+    );
+    EasyMock.expect(podTemplateSelector.getPodTemplateForTask(EasyMock.anyObject()))
+        .andReturn(Optional.absent());
+
+    Assert.assertThrows(DruidException.class, () -> adapter.fromTask(task));
+  }
+
+  @Test
+  public void test_fromTask_null()
+  {
+
+    PodTemplateSelector podTemplateSelector = EasyMock.createMock(PodTemplateSelector.class);
+
+    PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
+        taskRunnerConfig,
+        taskConfig,
+        node,
+        mapper,
+        taskLogs,
+        podTemplateSelector
+    );
+
+    Task task = new NoopTask(
+        "id",
+        "groupId",
+        "data_source",
+        0,
+        0,
+        null
+    );
+    EasyMock.expect(podTemplateSelector.getPodTemplateForTask(EasyMock.anyObject()))
+        .andReturn(null);
+
+    Assert.assertThrows(DruidException.class, () -> adapter.fromTask(task));
+  }
+
+  @Test
+  public void test_fromTask_taskSupportsQueries() throws IOException
+  {
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
+
+    PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
+        taskRunnerConfig,
+        taskConfig,
+        node,
+        mapper,
+        taskLogs,
+        podTemplateSelector
     );
 
     Task task = EasyMock.mock(Task.class);
@@ -554,21 +423,16 @@ public class PodTemplateTaskAdapterTest
   @Test
   public void test_fromTask_withBroadcastDatasourceLoadingModeAll() throws IOException
   {
-    Path templatePath = Files.createFile(tempDir.resolve("noop.yaml"));
-    mapper.writeValue(templatePath.toFile(), podTemplateSpec);
+    TestPodTemplateSelector podTemplateSelector = new TestPodTemplateSelector(podTemplateSpec);
 
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", templatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.queryable", templatePath.toString());
 
     PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
         taskRunnerConfig,
         taskConfig,
         node,
         mapper,
-        props,
         taskLogs,
-        dynamicConfigRef
+        podTemplateSelector
     );
 
     Task task = EasyMock.mock(Task.class);
@@ -590,98 +454,6 @@ public class PodTemplateTaskAdapterTest
                                           .collect(Collectors.toList()).get(0).getValue());
   }
 
-  @Test
-  public void test_fromTask_withIndexKafkaPodTemplateInRuntimeProperties() throws IOException
-  {
-    Path baseTemplatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(baseTemplatePath.toFile(), podTemplateSpec);
-
-    Path kafkaTemplatePath = Files.createFile(tempDir.resolve("kafka.yaml"));
-    PodTemplate kafkaPodTemplate = new PodTemplateBuilder(podTemplateSpec)
-          .editTemplate()
-          .editSpec()
-          .setNewVolumeLike(0, new VolumeBuilder().withName("volume").build())
-          .endVolume()
-          .endSpec()
-          .endTemplate()
-          .build();
-    mapper.writeValue(kafkaTemplatePath.toFile(), kafkaPodTemplate);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", baseTemplatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.index_kafka", kafkaTemplatePath.toString());
-
-    PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
-          taskRunnerConfig,
-          taskConfig,
-          node,
-          mapper,
-          props,
-          taskLogs,
-          dynamicConfigRef
-    );
-
-    Task kafkaTask = new NoopTask("id", "id", "datasource", 0, 0, null) {
-      @Override
-      public String getType()
-      {
-        return "index_kafka";
-      }
-    };
-
-    Task noopTask = new NoopTask("id", "id", "datasource", 0, 0, null);
-    Job actual = adapter.fromTask(kafkaTask);
-    Assert.assertEquals(1, actual.getSpec().getTemplate().getSpec().getVolumes().size(), 1);
-
-    actual = adapter.fromTask(noopTask);
-    Assert.assertEquals(0, actual.getSpec().getTemplate().getSpec().getVolumes().size(), 1);
-  }
-
-  @Test
-  public void test_fromTask_matchPodTemplateBasedOnStrategy() throws IOException
-  {
-    String dataSource = "my_table";
-    Path baseTemplatePath = Files.createFile(tempDir.resolve("base.yaml"));
-    mapper.writeValue(baseTemplatePath.toFile(), podTemplateSpec);
-
-    Path lowThroughputTemplatePath = Files.createFile(tempDir.resolve("low-throughput.yaml"));
-    PodTemplate lowThroughputPodTemplate = new PodTemplateBuilder(podTemplateSpec)
-        .editTemplate()
-        .editSpec()
-        .setNewVolumeLike(0, new VolumeBuilder().withName("volume").build())
-        .endVolume()
-        .endSpec()
-        .endTemplate()
-        .build();
-    mapper.writeValue(lowThroughputTemplatePath.toFile(), lowThroughputPodTemplate);
-
-    Properties props = new Properties();
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.base", baseTemplatePath.toString());
-    props.setProperty("druid.indexer.runner.k8s.podTemplate.lowThroughput", lowThroughputTemplatePath.toString());
-    dynamicConfigRef = () -> new DefaultKubernetesTaskRunnerDynamicConfig(new SelectorBasedPodTemplateSelectStrategy(
-        Collections.singletonList(
-            new Selector("lowThrougput", null, null, Sets.newSet(dataSource)
-            ))));
-
-    PodTemplateTaskAdapter adapter = new PodTemplateTaskAdapter(
-        taskRunnerConfig,
-        taskConfig,
-        node,
-        mapper,
-        props,
-        taskLogs,
-        dynamicConfigRef
-    );
-
-    Task taskWithMatchedDatasource = new NoopTask("id", "id", dataSource, 0, 0, null);
-    Task noopTask = new NoopTask("id", "id", "datasource", 0, 0, null);
-    Job actual = adapter.fromTask(taskWithMatchedDatasource);
-    Assert.assertEquals(1, actual.getSpec().getTemplate().getSpec().getVolumes().size(), 1);
-
-    actual = adapter.fromTask(noopTask);
-    Assert.assertEquals(0, actual.getSpec().getTemplate().getSpec().getVolumes().size(), 1);
-  }
-
   private void assertJobSpecsEqual(Job actual, Job expected) throws IOException
   {
     Map<String, String> actualAnnotations = actual.getSpec().getTemplate().getMetadata().getAnnotations();

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/taskadapter/TestPodTemplateSelector.java

@@ -0,0 +1,41 @@
+/*
+ * 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.k8s.overlord.taskadapter;
+
+import com.google.common.base.Optional;
+import io.fabric8.kubernetes.api.model.PodTemplate;
+import org.apache.druid.indexing.common.task.Task;
+
+public class TestPodTemplateSelector implements PodTemplateSelector
+{
+
+  private final PodTemplate basePodTemplate;
+
+  public TestPodTemplateSelector(PodTemplate basePodTemplate)
+  {
+    this.basePodTemplate = basePodTemplate;
+  }
+
+  @Override
+  public Optional<PodTemplateWithName> getPodTemplateForTask(Task task)
+  {
+    return Optional.of(new PodTemplateWithName("base", basePodTemplate));
+  }
+}

extensions-contrib/kubernetes-overlord-extensions/src/test/resources/expectedNoopJobBase.yaml

@@ -1,57 +0,0 @@
-apiVersion: batch/v1
-kind: Job
-metadata:
-  name: "id-3e70afe5cd823dfc7dd308eea616426b"
-  labels:
-    druid.k8s.peons: "true"
-    druid.task.id: "id"
-    druid.task.type: "noop"
-    druid.task.group.id: "id"
-    druid.task.datasource: "datasource"
-
-  annotations:
-    task.id: "id"
-    task.type: "noop"
-    task.group.id: "id"
-    task.datasource: "datasource"
-    task.jobTemplate: base
-spec:
-  activeDeadlineSeconds: 14400
-  backoffLimit: 0
-  ttlSecondsAfterFinished: 172800
-  template:
-    metadata:
-      labels:
-        druid.k8s.peons: "true"
-        druid.task.id: "id"
-        druid.task.type: "noop"
-        druid.task.group.id: "id"
-        druid.task.datasource: "datasource"
-      annotations:
-        task: "H4sIAAAAAAAAAD2MvQ4CIRCE32VqijsTG1qLi7W+wArEbHICrmC8EN7dJf40k/lmJtNQthxgEVPKMGCvXsXgKqnm4x89FTqlKm6MBzw+YCA1nvmm8W4/TQYuxRJeBbZ17cJ3ZhvoSbzShVcu2zLOf9cS7pUl+ANlclrCzr2/AQUK0FqZAAAA"
-        tls.enabled: "false"
-        task.id: "id"
-        task.type: "noop"
-        task.group.id: "id"
-        task.datasource: "datasource"
-        task.jobTemplate: base
-    spec:
-      containers:
-        - command:
-            - sleep
-            - "3600"
-          env:
-            - name: "TASK_DIR"
-              value: "/tmp"
-            - name: "TASK_ID"
-              value: "id"
-            - name: "LOAD_BROADCAST_DATASOURCE_MODE"
-              value: "ALL"
-            - name: "LOAD_BROADCAST_SEGMENTS"
-              value: "false"
-            - name: "TASK_JSON"
-              valueFrom:
-                fieldRef:
-                  fieldPath: "metadata.annotations['task']"
-          image: one
-          name: primary

extensions-contrib/kubernetes-overlord-extensions/src/test/resources/expectedNoopJobLongIds.yaml

@@ -13,7 +13,7 @@ metadata:
     task.type: "noop"
     task.group.id: "api-issued_kill_wikipedia3_omjobnbc_1000-01-01T00:00:00.000Z_2023-05-14T00:00:00.000Z_2023-05-15T17:03:01.220Z"
     task.datasource: "data_source"
-    task.jobTemplate: noop
+    task.jobTemplate: base
 spec:
   activeDeadlineSeconds: 14400
   backoffLimit: 0
@@ -33,7 +33,7 @@ spec:
         task.type: "noop"
         task.group.id: "api-issued_kill_wikipedia3_omjobnbc_1000-01-01T00:00:00.000Z_2023-05-14T00:00:00.000Z_2023-05-15T17:03:01.220Z"
         task.datasource: "data_source"
-        task.jobTemplate: noop
+        task.jobTemplate: base
     spec:
       containers:
         - command:

extensions-contrib/kubernetes-overlord-extensions/src/test/resources/expectedNoopJobNoTaskJson.yaml

@@ -13,7 +13,7 @@ metadata:
     task.type: "noop"
     task.group.id: "id"
     task.datasource: "datasource"
-    task.jobTemplate: noop
+    task.jobTemplate: base
 spec:
   activeDeadlineSeconds: 14400
   backoffLimit: 0
@@ -32,7 +32,7 @@ spec:
         task.type: "noop"
         task.group.id: "id"
         task.datasource: "datasource"
-        task.jobTemplate: noop
+        task.jobTemplate: base
     spec:
       containers:
         - command:

extensions-contrib/kubernetes-overlord-extensions/src/test/resources/expectedNoopPodTemplate.yaml

@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: PodTemplate
+metadata:
+  name: test
+template:
+  spec:
+    containers:
+      - command:
+          - sleep
+          - "3600"
+        image: one
+        name: primary

extensions-contrib/prometheus-emitter/src/main/java/org/apache/druid/emitter/prometheus/Metrics.java

@@ -64,7 +64,7 @@ public class Metrics
 
   public Metrics(String namespace, String path, boolean isAddHostAsLabel, boolean isAddServiceAsLabel, Map<String, String> extraLabels)
   {
-    Map<String, DimensionsAndCollector> registeredMetrics = new HashMap<>();
+    Map<String, DimensionsAndCollector> parsedRegisteredMetrics = new HashMap<>();
     Map<String, Metric> metrics = readConfig(path);
 
     if (extraLabels == null) {
@@ -116,10 +116,10 @@ public class Metrics
       }
 
       if (collector != null) {
-        registeredMetrics.put(name, new DimensionsAndCollector(dimensions, collector, metric.conversionFactor));
+        parsedRegisteredMetrics.put(name, new DimensionsAndCollector(dimensions, collector, metric.conversionFactor));
       }
     }
-    this.registeredMetrics = Collections.unmodifiableMap(registeredMetrics);
+    this.registeredMetrics = Collections.unmodifiableMap(parsedRegisteredMetrics);
   }
 
   private Map<String, Metric> readConfig(String path)

extensions-contrib/prometheus-emitter/src/test/java/org/apache/druid/emitter/prometheus/MetricsTest.java

@@ -20,6 +20,7 @@
 package org.apache.druid.emitter.prometheus;
 
 import io.prometheus.client.Histogram;
+import org.apache.druid.java.util.common.ISE;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -93,4 +94,20 @@ public class MetricsTest
 
     Assert.assertTrue(actualMessage.contains(expectedMessage));
   }
+
+  @Test
+  public void testMetricsConfigurationWithNonExistentMetric()
+  {
+    Metrics metrics = new Metrics("test_4", null, true, true, null);
+    DimensionsAndCollector nonExistentDimsCollector = metrics.getByName("non/existent", "historical");
+    Assert.assertNull(nonExistentDimsCollector);
+  }
+
+  @Test
+  public void testMetricsConfigurationWithUnSupportedType()
+  {
+    Assert.assertThrows(ISE.class, () -> {
+      new Metrics("test_5", "src/test/resources/defaultMetricsTest.json", true, true, null);
+    });
+  }
 }

extensions-contrib/prometheus-emitter/src/test/resources/defaultMetricsTest.json

@@ -0,0 +1,3 @@
+{
+  "query/nonExistent" : { "dimensions" : ["dataSource"], "type" : "nonExistent", "help":  "Non supported type."}
+}

extensions-contrib/redis-cache/pom.xml

@@ -77,8 +77,8 @@
             <scope>provided</scope>
         </dependency>
         <dependency>
-            <groupId>commons-lang</groupId>
-            <artifactId>commons-lang</artifactId>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
             <scope>provided</scope>
         </dependency>
         <dependency>

extensions-contrib/redis-cache/src/main/java/org/apache/druid/client/cache/RedisCacheFactory.java

@@ -19,7 +19,7 @@
 
 package org.apache.druid.client.cache;
 
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.druid.java.util.common.IAE;
 import redis.clients.jedis.ConnectionPoolConfig;
 import redis.clients.jedis.HostAndPort;

extensions-contrib/statsd-emitter/src/main/resources/defaultMetricDimensions.json

@@ -159,7 +159,7 @@
   "jetty/threadPool/queueSize": { "dimensions" : [], "type" : "gauge" },
 
   "coordinator/time" : { "dimensions" : [], "type" : "timer"},
-  "coordinator/global/time" : { "dimensions" : [], "type" : "timer"},
+  "coordinator/global/time" : { "dimensions" : ["dutyGroup"], "type" : "timer"},
 
   "tier/required/capacity" : { "dimensions" : ["tier"], "type" : "gauge" },
   "tier/total/capacity" : { "dimensions" : ["tier"], "type" : "gauge" },

extensions-core/avro-extensions/pom.xml

@@ -217,8 +217,8 @@
       <scope>provided</scope>
     </dependency>
     <dependency>
-      <groupId>commons-lang</groupId>
-      <artifactId>commons-lang</artifactId>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-lang3</artifactId>
       <scope>provided</scope>
     </dependency>
     <dependency>

extensions-core/avro-extensions/src/main/java/org/apache/druid/data/input/avro/AvroValueInputFormat.java

@@ -22,7 +22,7 @@ package org.apache.druid.data.input.avro;
 import org.apache.avro.Schema;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.avro.mapreduce.AvroJob;
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.hadoop.fs.FSDataInputStream;
 import org.apache.hadoop.fs.FileSystem;

extensions-core/azure-extensions/pom.xml

@@ -131,11 +131,6 @@
             <artifactId>jsr305</artifactId>
             <scope>provided</scope>
         </dependency>
-        <dependency>
-            <groupId>commons-lang</groupId>
-            <artifactId>commons-lang</artifactId>
-            <scope>provided</scope>
-        </dependency>
         <dependency>
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-lang3</artifactId>

extensions-core/azure-extensions/src/main/java/org/apache/druid/storage/azure/AzureDataSegmentPusher.java

@@ -73,7 +73,7 @@ public class AzureDataSegmentPusher implements DataSegmentPusher
   public String getPathForHadoop()
   {
     String prefix = segmentConfig.getPrefix();
-    boolean prefixIsNullOrEmpty = org.apache.commons.lang.StringUtils.isEmpty(prefix);
+    boolean prefixIsNullOrEmpty = org.apache.commons.lang3.StringUtils.isEmpty(prefix);
     String hadoopPath = StringUtils.format(
         "%s://%s@%s.%s/%s",
         AzureUtils.AZURE_STORAGE_HADOOP_PROTOCOL,
@@ -129,7 +129,7 @@ public class AzureDataSegmentPusher implements DataSegmentPusher
   public DataSegment pushToPath(File indexFilesDir, DataSegment segment, String storageDirSuffix) throws IOException
   {
     String prefix = segmentConfig.getPrefix();
-    boolean prefixIsNullOrEmpty = org.apache.commons.lang.StringUtils.isEmpty(prefix);
+    boolean prefixIsNullOrEmpty = org.apache.commons.lang3.StringUtils.isEmpty(prefix);
     final String azurePath = JOINER.join(
         prefixIsNullOrEmpty ? null : StringUtils.maybeRemoveTrailingSlash(prefix),
         storageDirSuffix

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactory.java

@@ -278,6 +278,26 @@ public abstract class HllSketchAggregatorFactory extends AggregatorFactory
            && stringEncoding == that.stringEncoding;
   }
 
+  @Nullable
+  @Override
+  public AggregatorFactory substituteCombiningFactory(AggregatorFactory preAggregated)
+  {
+    if (this == preAggregated) {
+      return getCombiningFactory();
+    }
+    if (getClass() != preAggregated.getClass()) {
+      return null;
+    }
+    HllSketchAggregatorFactory that = (HllSketchAggregatorFactory) preAggregated;
+    if (lgK <= that.lgK &&
+        stringEncoding == that.stringEncoding &&
+        Objects.equals(fieldName, that.fieldName)
+    ) {
+      return getCombiningFactory();
+    }
+    return null;
+  }
+
   @Override
   public int hashCode()
   {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/kll/KllSketchAggregatorFactory.java

@@ -226,6 +226,23 @@ abstract class KllSketchAggregatorFactory<SketchType extends KllSketch, ValueTyp
     return new CacheKeyBuilder(cacheTypeId).appendString(name).appendString(fieldName).appendInt(k).build();
   }
 
+  @Nullable
+  @Override
+  public AggregatorFactory substituteCombiningFactory(AggregatorFactory preAggregated)
+  {
+    if (this == preAggregated) {
+      return getCombiningFactory();
+    }
+    if (getClass() != preAggregated.getClass()) {
+      return null;
+    }
+    KllSketchAggregatorFactory<?, ?> that = (KllSketchAggregatorFactory<?, ?>) preAggregated;
+    if (Objects.equals(fieldName, that.fieldName) && k == that.k && maxStreamLength <= that.maxStreamLength) {
+      return getCombiningFactory();
+    }
+    return null;
+  }
+
   @Override
   public boolean equals(final Object o)
   {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/quantiles/DoublesSketchAggregatorFactory.java

@@ -424,6 +424,25 @@ public class DoublesSketchAggregatorFactory extends AggregatorFactory
     return new CacheKeyBuilder(cacheTypeId).appendString(name).appendString(fieldName).appendInt(k).build();
   }
 
+  @Nullable
+  @Override
+  public AggregatorFactory substituteCombiningFactory(AggregatorFactory preAggregated)
+  {
+    if (this == preAggregated) {
+      return getCombiningFactory();
+    }
+
+    if (getClass() != preAggregated.getClass()) {
+      return null;
+    }
+
+    DoublesSketchAggregatorFactory that = (DoublesSketchAggregatorFactory) preAggregated;
+    if (k <= that.k && maxStreamLength <= that.getMaxStreamLength() && Objects.equals(fieldName, that.fieldName)) {
+      return getCombiningFactory();
+    }
+    return null;
+  }
+
   @Override
   public boolean equals(Object o)
   {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchAggregatorFactory.java

@@ -49,6 +49,7 @@ import java.nio.ByteBuffer;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Objects;
 
 public abstract class SketchAggregatorFactory extends AggregatorFactory
 {
@@ -266,6 +267,22 @@ public abstract class SketchAggregatorFactory extends AggregatorFactory
                      .array();
   }
 
+  @Override
+  public AggregatorFactory substituteCombiningFactory(AggregatorFactory preAggregated)
+  {
+    if (this == preAggregated) {
+      return getCombiningFactory();
+    }
+    if (getClass() != preAggregated.getClass()) {
+      return null;
+    }
+    SketchMergeAggregatorFactory that = (SketchMergeAggregatorFactory) preAggregated;
+    if (Objects.equals(fieldName, that.fieldName) && size <= that.size) {
+      return getCombiningFactory();
+    }
+    return null;
+  }
+
   @Override
   public String toString()
   {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchMergeAggregatorFactory.java

@@ -29,6 +29,7 @@ import org.apache.druid.query.aggregation.AggregatorUtil;
 import org.apache.druid.segment.column.ColumnType;
 
 import javax.annotation.Nullable;
+import java.util.Objects;
 
 public class SketchMergeAggregatorFactory extends SketchAggregatorFactory
 {
@@ -165,6 +166,25 @@ public class SketchMergeAggregatorFactory extends SketchAggregatorFactory
     );
   }
 
+  @Override
+  public AggregatorFactory substituteCombiningFactory(AggregatorFactory preAggregated)
+  {
+    if (this == preAggregated) {
+      return getCombiningFactory();
+    }
+    if (getClass() != preAggregated.getClass()) {
+      return null;
+    }
+    SketchMergeAggregatorFactory that = (SketchMergeAggregatorFactory) preAggregated;
+    if (Objects.equals(fieldName, that.fieldName) &&
+        size <= that.size &&
+        isInputThetaSketch == that.isInputThetaSketch
+    ) {
+      return getCombiningFactory();
+    }
+    return null;
+  }
+
   @Override
   public boolean equals(Object o)
   {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchAggregatorFactory.java

@@ -307,6 +307,29 @@ public class ArrayOfDoublesSketchAggregatorFactory extends AggregatorFactory
     return ColumnType.DOUBLE;
   }
 
+  @Nullable
+  @Override
+  public AggregatorFactory substituteCombiningFactory(AggregatorFactory preAggregated)
+  {
+    if (this == preAggregated) {
+      return getCombiningFactory();
+    }
+
+    if (getClass() != preAggregated.getClass()) {
+      return null;
+    }
+
+    ArrayOfDoublesSketchAggregatorFactory that = (ArrayOfDoublesSketchAggregatorFactory) preAggregated;
+    if (nominalEntries <= that.nominalEntries &&
+        numberOfValues == that.numberOfValues &&
+        Objects.equals(fieldName, that.fieldName) &&
+        Objects.equals(metricColumns, that.metricColumns)
+    ) {
+      return getCombiningFactory();
+    }
+    return null;
+  }
+
   @Override
   public String toString()
   {

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchAggregatorFactoryTest.java

@@ -231,6 +231,42 @@ public class HllSketchAggregatorFactoryTest
     Assert.assertArrayEquals(target.getCacheKey(), other.getCacheKey());
   }
 
+  @Test
+  public void testCanSubstitute()
+  {
+    HllSketchBuildAggregatorFactory factory = new HllSketchBuildAggregatorFactory(
+        NAME,
+        FIELD_NAME,
+        LG_K,
+        TGT_HLL_TYPE,
+        STRING_ENCODING,
+        true,
+        true
+    );
+    HllSketchBuildAggregatorFactory other = new HllSketchBuildAggregatorFactory(
+        "other name",
+        FIELD_NAME,
+        LG_K,
+        TGT_HLL_TYPE,
+        STRING_ENCODING,
+        false,
+        false
+    );
+
+    HllSketchBuildAggregatorFactory incompatible = new HllSketchBuildAggregatorFactory(
+        NAME,
+        "different field",
+        LG_K,
+        TGT_HLL_TYPE,
+        STRING_ENCODING,
+        false,
+        false
+    );
+    Assert.assertNotNull(other.substituteCombiningFactory(factory));
+    Assert.assertNotNull(factory.substituteCombiningFactory(other));
+    Assert.assertNull(factory.substituteCombiningFactory(incompatible));
+  }
+
   @Test
   public void testToString()
   {

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/hll/sql/HllSketchSqlAggregatorTest.java

@@ -322,6 +322,7 @@ public class HllSketchSqlAggregatorTest extends BaseCalciteQueryTest
   @Test
   public void testApproxCountDistinctHllSketch()
   {
+    cannotVectorizeUnlessFallback();
     final String sql = "SELECT\n"
                        + "  SUM(cnt),\n"
                        + "  APPROX_COUNT_DISTINCT_DS_HLL(dim2),\n" // uppercase
@@ -1138,6 +1139,7 @@ public class HllSketchSqlAggregatorTest extends BaseCalciteQueryTest
   @Test
   public void testHllEstimateAsVirtualColumnWithGroupByOrderBy()
   {
+    cannotVectorizeUnlessFallback();
     testQuery(
         "SELECT"
         + " HLL_SKETCH_ESTIMATE(hllsketch_dim1), count(*)"

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/kll/KllDoublesSketchAggregatorFactoryTest.java

@@ -153,4 +153,22 @@ public class KllDoublesSketchAggregatorFactoryTest
         new TimeseriesQueryQueryToolChest().resultArraySignature(query)
     );
   }
+
+  @Test
+  public void testCanSubstitute()
+  {
+    AggregatorFactory sketch = new KllDoublesSketchAggregatorFactory("sketch", "x", 200, null);
+    AggregatorFactory sketch2 = new KllDoublesSketchAggregatorFactory("other", "x", 200, null);
+    AggregatorFactory sketch3 = new KllDoublesSketchAggregatorFactory("sketch", "x", 200, 1_000L);
+    AggregatorFactory sketch4 = new KllDoublesSketchAggregatorFactory("sketch", "y", 200, null);
+    AggregatorFactory sketch5 = new KllDoublesSketchAggregatorFactory("sketch", "x", 300, null);
+
+    Assert.assertNotNull(sketch.substituteCombiningFactory(sketch2));
+    Assert.assertNotNull(sketch3.substituteCombiningFactory(sketch2));
+    Assert.assertNotNull(sketch3.substituteCombiningFactory(sketch));
+    Assert.assertNotNull(sketch2.substituteCombiningFactory(sketch));
+    Assert.assertNull(sketch.substituteCombiningFactory(sketch3));
+    Assert.assertNull(sketch.substituteCombiningFactory(sketch4));
+    Assert.assertNull(sketch.substituteCombiningFactory(sketch5));
+  }
 }

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/quantiles/DoublesSketchAggregatorFactoryTest.java

@@ -201,4 +201,19 @@ public class DoublesSketchAggregatorFactoryTest
     ac.fold(new TestDoublesSketchColumnValueSelector());
     Assert.assertNotNull(ac.getObject());
   }
+
+  @Test
+  public void testCanSubstitute()
+  {
+    final DoublesSketchAggregatorFactory sketch = new DoublesSketchAggregatorFactory("sketch", "x", 1024, 1000L, null);
+    final DoublesSketchAggregatorFactory sketch2 = new DoublesSketchAggregatorFactory("other", "x", 1024, 2000L, null);
+    final DoublesSketchAggregatorFactory sketch3 = new DoublesSketchAggregatorFactory("another", "x", 2048, 1000L, null);
+    final DoublesSketchAggregatorFactory incompatible = new DoublesSketchAggregatorFactory("incompatible", "y", 1024, 1000L, null);
+
+    Assert.assertNotNull(sketch.substituteCombiningFactory(sketch2));
+    Assert.assertNotNull(sketch.substituteCombiningFactory(sketch3));
+    Assert.assertNull(sketch2.substituteCombiningFactory(sketch3));
+    Assert.assertNull(sketch.substituteCombiningFactory(incompatible));
+    Assert.assertNull(sketch3.substituteCombiningFactory(sketch));
+  }
 }

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/theta/SketchAggregatorFactoryTest.java

@@ -24,6 +24,7 @@ import org.apache.druid.error.DruidException;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.query.Druids;
 import org.apache.druid.query.aggregation.AggregatorAndSize;
+import org.apache.druid.query.aggregation.AggregatorFactory;
 import org.apache.druid.query.aggregation.CountAggregatorFactory;
 import org.apache.druid.query.aggregation.datasketches.theta.oldapi.OldSketchBuildAggregatorFactory;
 import org.apache.druid.query.aggregation.datasketches.theta.oldapi.OldSketchMergeAggregatorFactory;
@@ -213,4 +214,18 @@ public class SketchAggregatorFactoryTest
     Throwable exception = Assert.assertThrows(DruidException.class, () -> AGGREGATOR_16384.factorizeVector(vectorFactory));
     Assert.assertEquals("Unsupported input [x] of type [COMPLEX<json>] for aggregator [COMPLEX<thetaSketchBuild>].", exception.getMessage());
   }
+
+  @Test
+  public void testCanSubstitute()
+  {
+    AggregatorFactory sketch1 = new SketchMergeAggregatorFactory("sketch", "x", 16, true, false, 2);
+    AggregatorFactory sketch2 = new SketchMergeAggregatorFactory("other", "x", null, false, false, null);
+    AggregatorFactory sketch3 = new SketchMergeAggregatorFactory("sketch", "x", null, false, false, 3);
+    AggregatorFactory sketch4 = new SketchMergeAggregatorFactory("sketch", "y", null, false, false, null);
+
+    Assert.assertNotNull(sketch1.substituteCombiningFactory(sketch2));
+    Assert.assertNotNull(sketch1.substituteCombiningFactory(sketch3));
+    Assert.assertNull(sketch1.substituteCombiningFactory(sketch4));
+    Assert.assertNull(sketch2.substituteCombiningFactory(sketch1));
+  }
 }

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/theta/sql/ThetaSketchSqlAggregatorTest.java

@@ -177,6 +177,7 @@ public class ThetaSketchSqlAggregatorTest extends BaseCalciteQueryTest
   @Test
   public void testApproxCountDistinctThetaSketch()
   {
+    cannotVectorizeUnlessFallback();
     final String sql = "SELECT\n"
                        + "  SUM(cnt),\n"
                        + "  APPROX_COUNT_DISTINCT_DS_THETA(dim2),\n"
@@ -1157,6 +1158,7 @@ public class ThetaSketchSqlAggregatorTest extends BaseCalciteQueryTest
   @Test
   public void testThetaEstimateAsVirtualColumnWithGroupByOrderBy()
   {
+    cannotVectorizeUnlessFallback();
     testQuery(
         "SELECT"
         + " THETA_SKETCH_ESTIMATE(thetasketch_dim1), count(*)"

extensions-core/datasketches/src/test/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchAggregatorFactoryTest.java

@@ -118,4 +118,21 @@ public class ArrayOfDoublesSketchAggregatorFactoryTest
     Assert.assertEquals(factory, factory.withName("name"));
     Assert.assertEquals("newTest", factory.withName("newTest").getName());
   }
+
+  @Test
+  public void testCanSubstitute()
+  {
+    AggregatorFactory sketch = new ArrayOfDoublesSketchAggregatorFactory("sketch", "x", null, null, null);
+    AggregatorFactory sketch2 = new ArrayOfDoublesSketchAggregatorFactory("sketch2", "x", null, null, null);
+    AggregatorFactory other = new ArrayOfDoublesSketchAggregatorFactory("other", "x", 8192, null, null);
+    AggregatorFactory incompatible = new ArrayOfDoublesSketchAggregatorFactory("incompatible", "x", 2048, null, null);
+    AggregatorFactory incompatible2 = new ArrayOfDoublesSketchAggregatorFactory("sketch", "y", null, null, null);
+    Assert.assertNotNull(sketch.substituteCombiningFactory(other));
+    Assert.assertNotNull(sketch.substituteCombiningFactory(sketch2));
+    Assert.assertNull(sketch.substituteCombiningFactory(incompatible));
+    Assert.assertNotNull(sketch.substituteCombiningFactory(sketch));
+    Assert.assertNull(other.substituteCombiningFactory(sketch));
+    Assert.assertNull(sketch.substituteCombiningFactory(incompatible2));
+    Assert.assertNull(other.substituteCombiningFactory(incompatible2));
+  }
 }

extensions-core/datasketches/src/test/java/org/apache/druid/segment/DatasketchesProjectionTest.java

@@ -0,0 +1,489 @@
+/*
+ * 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.segment;
+
+import com.google.common.collect.ImmutableMap;
+import org.apache.datasketches.common.Family;
+import org.apache.datasketches.hll.HllSketch;
+import org.apache.datasketches.kll.KllDoublesSketch;
+import org.apache.datasketches.quantiles.DoublesSketch;
+import org.apache.datasketches.theta.SetOperation;
+import org.apache.datasketches.theta.Union;
+import org.apache.datasketches.thetacommon.ThetaUtil;
+import org.apache.datasketches.tuple.arrayofdoubles.ArrayOfDoublesSketch;
+import org.apache.datasketches.tuple.arrayofdoubles.ArrayOfDoublesUpdatableSketch;
+import org.apache.datasketches.tuple.arrayofdoubles.ArrayOfDoublesUpdatableSketchBuilder;
+import org.apache.druid.collections.CloseableDefaultBlockingPool;
+import org.apache.druid.collections.CloseableStupidPool;
+import org.apache.druid.collections.NonBlockingPool;
+import org.apache.druid.data.input.impl.AggregateProjectionSpec;
+import org.apache.druid.data.input.impl.DimensionSchema;
+import org.apache.druid.data.input.impl.DimensionsSpec;
+import org.apache.druid.data.input.impl.DoubleDimensionSchema;
+import org.apache.druid.data.input.impl.FloatDimensionSchema;
+import org.apache.druid.data.input.impl.LongDimensionSchema;
+import org.apache.druid.data.input.impl.StringDimensionSchema;
+import org.apache.druid.java.util.common.FileUtils;
+import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.granularity.Granularities;
+import org.apache.druid.java.util.common.guava.Sequence;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.query.DruidProcessingConfig;
+import org.apache.druid.query.QueryContexts;
+import org.apache.druid.query.aggregation.AggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.hll.HllSketchBuildAggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.hll.HllSketchHolder;
+import org.apache.druid.query.aggregation.datasketches.hll.HllSketchModule;
+import org.apache.druid.query.aggregation.datasketches.kll.KllDoublesSketchAggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.kll.KllSketchModule;
+import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchAggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchModule;
+import org.apache.druid.query.aggregation.datasketches.theta.SketchHolder;
+import org.apache.druid.query.aggregation.datasketches.theta.SketchMergeAggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.theta.SketchModule;
+import org.apache.druid.query.aggregation.datasketches.tuple.ArrayOfDoublesSketchAggregatorFactory;
+import org.apache.druid.query.aggregation.datasketches.tuple.ArrayOfDoublesSketchModule;
+import org.apache.druid.query.groupby.GroupByQuery;
+import org.apache.druid.query.groupby.GroupByQueryConfig;
+import org.apache.druid.query.groupby.GroupByResourcesReservationPool;
+import org.apache.druid.query.groupby.GroupingEngine;
+import org.apache.druid.query.groupby.ResultRow;
+import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.RowSignature;
+import org.apache.druid.segment.incremental.IncrementalIndex;
+import org.apache.druid.segment.incremental.IncrementalIndexCursorFactory;
+import org.apache.druid.segment.incremental.IncrementalIndexSchema;
+import org.apache.druid.testing.InitializedNullHandlingTest;
+import org.junit.AfterClass;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * like {@link CursorFactoryProjectionTest} but for sketch aggs
+ */
+@RunWith(Parameterized.class)
+public class DatasketchesProjectionTest extends InitializedNullHandlingTest
+{
+  private static final Closer CLOSER = Closer.create();
+
+  private static final List<AggregateProjectionSpec> PROJECTIONS = Collections.singletonList(
+      new AggregateProjectionSpec(
+          "a_projection",
+          VirtualColumns.create(
+              Granularities.toVirtualColumn(Granularities.HOUR, "__gran")
+          ),
+          Arrays.asList(
+              new LongDimensionSchema("__gran"),
+              new StringDimensionSchema("a")
+          ),
+          new AggregatorFactory[]{
+              new HllSketchBuildAggregatorFactory("_b_hll", "b", null, null, null, null, false),
+              new SketchMergeAggregatorFactory("_b_theta", "b", null, null, false, null),
+              new DoublesSketchAggregatorFactory("_d_doubles", "d", null),
+              new ArrayOfDoublesSketchAggregatorFactory("_bcd_aod", "b", null, Arrays.asList("c", "d"), null),
+              new KllDoublesSketchAggregatorFactory("_d_kll", "d", null, null)
+          }
+      )
+  );
+
+  private static final List<AggregateProjectionSpec> AUTO_PROJECTIONS = PROJECTIONS.stream().map(projection -> {
+    return new AggregateProjectionSpec(
+        projection.getName(),
+        projection.getVirtualColumns(),
+        projection.getGroupingColumns()
+                  .stream()
+                  .map(x -> new AutoTypeColumnSchema(x.getName(), null))
+                  .collect(Collectors.toList()),
+        projection.getAggregators()
+    );
+  }).collect(Collectors.toList());
+
+  @Parameterized.Parameters(name = "name: {0}, sortByDim: {3}, autoSchema: {4}")
+  public static Collection<?> constructorFeeder()
+  {
+    HllSketchModule.registerSerde();
+    TestHelper.JSON_MAPPER.registerModules(new HllSketchModule().getJacksonModules());
+    SketchModule.registerSerde();
+    TestHelper.JSON_MAPPER.registerModules(new SketchModule().getJacksonModules());
+    KllSketchModule.registerSerde();
+    TestHelper.JSON_MAPPER.registerModules(new KllSketchModule().getJacksonModules());
+    DoublesSketchModule.registerSerde();
+    TestHelper.JSON_MAPPER.registerModules(new DoublesSketchModule().getJacksonModules());
+    ArrayOfDoublesSketchModule.registerSerde();
+    TestHelper.JSON_MAPPER.registerModules(new ArrayOfDoublesSketchModule().getJacksonModules());
+
+    final List<Object[]> constructors = new ArrayList<>();
+    final DimensionsSpec.Builder dimensionsBuilder =
+        DimensionsSpec.builder()
+                      .setDimensions(
+                          Arrays.asList(
+                              new StringDimensionSchema("a"),
+                              new StringDimensionSchema("b"),
+                              new LongDimensionSchema("c"),
+                              new DoubleDimensionSchema("d"),
+                              new FloatDimensionSchema("e")
+                          )
+                      );
+    DimensionsSpec dimsTimeOrdered = dimensionsBuilder.build();
+    DimensionsSpec dimsOrdered = dimensionsBuilder.setForceSegmentSortByTime(false).build();
+
+
+    List<DimensionSchema> autoDims = dimsOrdered.getDimensions()
+                                                .stream()
+                                                .map(x -> new AutoTypeColumnSchema(x.getName(), null))
+                                                .collect(Collectors.toList());
+    for (boolean incremental : new boolean[]{true, false}) {
+      for (boolean sortByDim : new boolean[]{true, false}) {
+        for (boolean autoSchema : new boolean[]{true, false}) {
+          final DimensionsSpec dims;
+          if (sortByDim) {
+            if (autoSchema) {
+              dims = dimsOrdered.withDimensions(autoDims);
+            } else {
+              dims = dimsOrdered;
+            }
+          } else {
+            if (autoSchema) {
+              dims = dimsTimeOrdered.withDimensions(autoDims);
+            } else {
+              dims = dimsTimeOrdered;
+            }
+          }
+          if (incremental) {
+            IncrementalIndex index = CLOSER.register(makeBuilder(dims, autoSchema).buildIncrementalIndex());
+            constructors.add(new Object[]{
+                "incrementalIndex",
+                new IncrementalIndexCursorFactory(index),
+                new IncrementalIndexTimeBoundaryInspector(index),
+                sortByDim,
+                autoSchema
+            });
+          } else {
+            QueryableIndex index = CLOSER.register(makeBuilder(dims, autoSchema).buildMMappedIndex());
+            constructors.add(new Object[]{
+                "queryableIndex",
+                new QueryableIndexCursorFactory(index),
+                QueryableIndexTimeBoundaryInspector.create(index),
+                sortByDim,
+                autoSchema
+            });
+          }
+        }
+      }
+    }
+    return constructors;
+  }
+
+  @AfterClass
+  public static void cleanup() throws IOException
+  {
+    CLOSER.close();
+  }
+
+  private static IndexBuilder makeBuilder(DimensionsSpec dimensionsSpec, boolean autoSchema)
+  {
+    File tmp = FileUtils.createTempDir();
+    CLOSER.register(tmp::delete);
+    return IndexBuilder.create()
+                       .tmpDir(tmp)
+                       .schema(
+                           IncrementalIndexSchema.builder()
+                                                 .withDimensionsSpec(dimensionsSpec)
+                                                 .withRollup(false)
+                                                 .withMinTimestamp(CursorFactoryProjectionTest.TIMESTAMP.getMillis())
+                                                 .withProjections(autoSchema ? AUTO_PROJECTIONS : PROJECTIONS)
+                                                 .build()
+                       )
+                       .rows(CursorFactoryProjectionTest.ROWS);
+  }
+
+  public final CursorFactory projectionsCursorFactory;
+  public final TimeBoundaryInspector projectionsTimeBoundaryInspector;
+
+  private final GroupingEngine groupingEngine;
+
+  private final NonBlockingPool<ByteBuffer> nonBlockingPool;
+  public final boolean sortByDim;
+  public final boolean autoSchema;
+
+  @Rule
+  public final CloserRule closer = new CloserRule(false);
+
+  public DatasketchesProjectionTest(
+      String name,
+      CursorFactory projectionsCursorFactory,
+      TimeBoundaryInspector projectionsTimeBoundaryInspector,
+      boolean sortByDim,
+      boolean autoSchema
+  )
+  {
+    this.projectionsCursorFactory = projectionsCursorFactory;
+    this.projectionsTimeBoundaryInspector = projectionsTimeBoundaryInspector;
+    this.sortByDim = sortByDim;
+    this.autoSchema = autoSchema;
+    this.nonBlockingPool = closer.closeLater(
+        new CloseableStupidPool<>(
+            "GroupByQueryEngine-bufferPool",
+            () -> ByteBuffer.allocate(1 << 24)
+        )
+    );
+    this.groupingEngine = new GroupingEngine(
+        new DruidProcessingConfig(),
+        GroupByQueryConfig::new,
+        new GroupByResourcesReservationPool(
+            closer.closeLater(
+                new CloseableDefaultBlockingPool<>(
+                    () -> ByteBuffer.allocate(1 << 24),
+                    5
+                )
+            ),
+            new GroupByQueryConfig()
+        ),
+        TestHelper.makeJsonMapper(),
+        TestHelper.makeSmileMapper(),
+        (query, future) -> {
+        }
+    );
+  }
+
+  @Test
+  public void testProjectionSingleDim()
+  {
+    // test can use the single dimension projection
+    final GroupByQuery query =
+        GroupByQuery.builder()
+                    .setDataSource("test")
+                    .setGranularity(Granularities.ALL)
+                    .setInterval(Intervals.ETERNITY)
+                    .addDimension("a")
+                    .setAggregatorSpecs(
+                        new HllSketchBuildAggregatorFactory("b_distinct", "b", null, null, null, true, true),
+                        new SketchMergeAggregatorFactory("b_distinct_theta", "b", null, null, null, null),
+                        new DoublesSketchAggregatorFactory("d_doubles", "d", null, null, null),
+                        new ArrayOfDoublesSketchAggregatorFactory("b_doubles", "b", null, Arrays.asList("c", "d"), null),
+                        new KllDoublesSketchAggregatorFactory("d", "d", null, null)
+                    )
+                    .build();
+    final CursorBuildSpec buildSpec = GroupingEngine.makeCursorBuildSpec(query, null);
+    try (final CursorHolder cursorHolder = projectionsCursorFactory.makeCursorHolder(buildSpec)) {
+      final Cursor cursor = cursorHolder.asCursor();
+      int rowCount = 0;
+      while (!cursor.isDone()) {
+        rowCount++;
+        cursor.advance();
+      }
+      Assert.assertEquals(3, rowCount);
+    }
+    final Sequence<ResultRow> resultRows = groupingEngine.process(
+        query,
+        projectionsCursorFactory,
+        projectionsTimeBoundaryInspector,
+        nonBlockingPool,
+        null
+    );
+    final List<ResultRow> results = resultRows.toList();
+    Assert.assertEquals(2, results.size());
+    List<Object[]> expectedResults = getSingleDimExpected();
+    final RowSignature querySignature = query.getResultRowSignature(RowSignature.Finalization.NO);
+    for (int i = 0; i < expectedResults.size(); i++) {
+      assertResults(
+          expectedResults.get(i),
+          results.get(i).getArray(),
+          querySignature
+      );
+    }
+  }
+
+  @Test
+  public void testProjectionSingleDimNoProjections()
+  {
+    // test can use the single dimension projection
+    final GroupByQuery query =
+        GroupByQuery.builder()
+                    .setDataSource("test")
+                    .setGranularity(Granularities.ALL)
+                    .setInterval(Intervals.ETERNITY)
+                    .addDimension("a")
+                    .setAggregatorSpecs(
+                        new HllSketchBuildAggregatorFactory("b_distinct", "b", null, null, null, true, true),
+                        new SketchMergeAggregatorFactory("b_distinct_theta", "b", null, null, null, null),
+                        new DoublesSketchAggregatorFactory("d_doubles", "d", null, null, null),
+                        new ArrayOfDoublesSketchAggregatorFactory("b_doubles", "b", null, Arrays.asList("c", "d"), null),
+                        new KllDoublesSketchAggregatorFactory("d", "d", null, null)
+                    )
+                    .setContext(ImmutableMap.of(QueryContexts.NO_PROJECTIONS, true))
+                    .build();
+    final CursorBuildSpec buildSpec = GroupingEngine.makeCursorBuildSpec(query, null);
+    try (final CursorHolder cursorHolder = projectionsCursorFactory.makeCursorHolder(buildSpec)) {
+      final Cursor cursor = cursorHolder.asCursor();
+      int rowCount = 0;
+      while (!cursor.isDone()) {
+        rowCount++;
+        cursor.advance();
+      }
+      Assert.assertEquals(8, rowCount);
+    }
+    final Sequence<ResultRow> resultRows = groupingEngine.process(
+        query,
+        projectionsCursorFactory,
+        projectionsTimeBoundaryInspector,
+        nonBlockingPool,
+        null
+    );
+    final List<ResultRow> results = resultRows.toList();
+    Assert.assertEquals(2, results.size());
+    List<Object[]> expectedResults = getSingleDimExpected();
+    final RowSignature querySignature = query.getResultRowSignature(RowSignature.Finalization.NO);
+    for (int i = 0; i < expectedResults.size(); i++) {
+      assertResults(
+          expectedResults.get(i),
+          results.get(i).getArray(),
+          querySignature
+      );
+    }
+  }
+
+  private List<Object[]> getSingleDimExpected()
+  {
+    HllSketch hll1 = new HllSketch(HllSketch.DEFAULT_LG_K);
+    Union theta1 = (Union) SetOperation.builder().build(Family.UNION);
+    DoublesSketch d1 = DoublesSketch.builder().setK(DoublesSketchAggregatorFactory.DEFAULT_K).build();
+    ArrayOfDoublesUpdatableSketch ad1 = new ArrayOfDoublesUpdatableSketchBuilder().setNominalEntries(ThetaUtil.DEFAULT_NOMINAL_ENTRIES)
+                                                                                  .setNumberOfValues(2)
+                                                                                  .build();
+    KllDoublesSketch kll1 = KllDoublesSketch.newHeapInstance();
+    hll1.update("aa");
+    hll1.update("bb");
+    hll1.update("cc");
+    hll1.update("dd");
+    theta1.update("aa");
+    theta1.update("bb");
+    theta1.update("cc");
+    theta1.update("dd");
+    d1.update(1.0);
+    d1.update(1.1);
+    d1.update(2.2);
+    d1.update(1.1);
+    d1.update(2.2);
+    ad1.update("aa", new double[]{1.0, 1.0});
+    ad1.update("bb", new double[]{1.0, 1.1});
+    ad1.update("cc", new double[]{2.0, 2.2});
+    ad1.update("aa", new double[]{1.0, 1.1});
+    ad1.update("dd", new double[]{2.0, 2.2});
+    kll1.update(1.0);
+    kll1.update(1.1);
+    kll1.update(2.2);
+    kll1.update(1.1);
+    kll1.update(2.2);
+    HllSketch hll2 = new HllSketch(HllSketch.DEFAULT_LG_K);
+    Union theta2 = (Union) SetOperation.builder().build(Family.UNION);
+    DoublesSketch d2 = DoublesSketch.builder().setK(DoublesSketchAggregatorFactory.DEFAULT_K).build();
+    ArrayOfDoublesUpdatableSketch ad2 = new ArrayOfDoublesUpdatableSketchBuilder().setNominalEntries(ThetaUtil.DEFAULT_NOMINAL_ENTRIES)
+                                                                                  .setNumberOfValues(2)
+                                                                                  .build();
+    KllDoublesSketch kll2 = KllDoublesSketch.newHeapInstance();
+    hll2.update("aa");
+    hll2.update("bb");
+    theta2.update("aa");
+    theta2.update("bb");
+    d2.update(3.3);
+    d2.update(4.4);
+    d2.update(5.5);
+    ad2.update("aa", new double[]{3.0, 3.3});
+    ad2.update("aa", new double[]{4.0, 4.4});
+    ad2.update("bb", new double[]{5.0, 5.5});
+    kll2.update(3.3);
+    kll2.update(4.4);
+    kll2.update(5.5);
+
+    return Arrays.asList(
+        new Object[]{"a", HllSketchHolder.of(hll1), SketchHolder.of(theta1), d1, ad1, kll1},
+        new Object[]{"b", HllSketchHolder.of(hll2), SketchHolder.of(theta2), d2, ad2, kll2}
+    );
+
+  }
+
+  private void assertResults(Object[] expected, Object[] actual, RowSignature signature)
+  {
+    Assert.assertEquals(expected.length, actual.length);
+    for (int i = 0; i < expected.length; i++) {
+      if (signature.getColumnType(i).get().equals(ColumnType.ofComplex(HllSketchModule.BUILD_TYPE_NAME))) {
+        Assert.assertEquals(
+            ((HllSketchHolder) expected[i]).getEstimate(),
+            ((HllSketchHolder) actual[i]).getEstimate(),
+            0.01
+        );
+      } else if (signature.getColumnType(i).get().equals(DoublesSketchModule.TYPE)) {
+        Assert.assertEquals(
+            ((DoublesSketch) expected[i]).getMinItem(),
+            ((DoublesSketch) actual[i]).getMinItem(),
+            0.01
+        );
+        Assert.assertEquals(
+            ((DoublesSketch) expected[i]).getMaxItem(),
+            ((DoublesSketch) actual[i]).getMaxItem(),
+            0.01
+        );
+      } else if (signature.getColumnType(i).get().equals(ArrayOfDoublesSketchModule.BUILD_TYPE)) {
+        Assert.assertEquals(
+            ((ArrayOfDoublesSketch) expected[i]).getEstimate(),
+            ((ArrayOfDoublesSketch) actual[i]).getEstimate(),
+            0.01
+        );
+        Assert.assertEquals(
+            ((ArrayOfDoublesSketch) expected[i]).getLowerBound(0),
+            ((ArrayOfDoublesSketch) actual[i]).getLowerBound(0),
+            0.01
+        );
+        Assert.assertEquals(
+            ((ArrayOfDoublesSketch) expected[i]).getUpperBound(0),
+            ((ArrayOfDoublesSketch) actual[i]).getUpperBound(0),
+            0.01
+        );
+      } else if (signature.getColumnType(i).get().equals(KllSketchModule.DOUBLES_TYPE)) {
+        Assert.assertEquals(
+            ((KllDoublesSketch) expected[i]).getMinItem(),
+            ((KllDoublesSketch) actual[i]).getMinItem(),
+            0.01
+        );
+        Assert.assertEquals(
+            ((KllDoublesSketch) expected[i]).getMaxItem(),
+            ((KllDoublesSketch) actual[i]).getMaxItem(),
+            0.01
+        );
+      } else {
+        Assert.assertEquals(expected[i], actual[i]);
+      }
+    }
+  }
+}

extensions-core/google-extensions/pom.xml

@@ -57,12 +57,6 @@
             <artifactId>commons-io</artifactId>
             <scope>provided</scope>
         </dependency>
-        <dependency>
-            <groupId>commons-lang</groupId>
-            <artifactId>commons-lang</artifactId>
-            <version>2.6</version>
-            <scope>provided</scope>
-        </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-annotations</artifactId>

extensions-core/hdfs-storage/pom.xml

@@ -94,8 +94,8 @@
             <scope>provided</scope>
         </dependency>
         <dependency>
-            <groupId>commons-lang</groupId>
-            <artifactId>commons-lang</artifactId>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
             <scope>provided</scope>
         </dependency>
         <dependency>

extensions-core/hdfs-storage/src/main/java/org/apache/druid/storage/hdfs/HdfsDataSegmentKiller.java

@@ -22,7 +22,7 @@ package org.apache.druid.storage.hdfs;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
 import com.google.inject.Inject;
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.druid.guice.Hdfs;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.emitter.EmittingLogger;

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/data/input/kafkainput/KafkaInputFormat.java

@@ -21,11 +21,13 @@ package org.apache.druid.data.input.kafkainput;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
+import org.apache.druid.data.input.ColumnsFilter;
 import org.apache.druid.data.input.InputEntity;
 import org.apache.druid.data.input.InputEntityReader;
 import org.apache.druid.data.input.InputFormat;
 import org.apache.druid.data.input.InputRowSchema;
 import org.apache.druid.data.input.impl.ByteEntity;
+import org.apache.druid.data.input.impl.DimensionsSpec;
 import org.apache.druid.data.input.impl.JsonInputFormat;
 import org.apache.druid.data.input.impl.TimestampSpec;
 import org.apache.druid.data.input.kafka.KafkaRecordEntity;
@@ -33,7 +35,6 @@ import org.apache.druid.indexing.seekablestream.SettableByteEntity;
 import org.apache.druid.java.util.common.DateTimes;
 
 import javax.annotation.Nullable;
-
 import java.io.File;
 import java.util.Objects;
 
@@ -111,7 +112,12 @@ public class KafkaInputFormat implements InputFormat
                 (record.getRecord().key() == null) ?
                     null :
                     JsonInputFormat.withLineSplittable(keyFormat, false).createReader(
-                        newInputRowSchema,
+                        // for keys, discover all fields; in KafkaInputReader we will pick the first one.
+                        new InputRowSchema(
+                            dummyTimestampSpec,
+                            DimensionsSpec.builder().useSchemaDiscovery(true).build(),
+                            ColumnsFilter.all()
+                        ),
                         new ByteEntity(record.getRecord().key()),
                         temporaryDirectory
                     ),

extensions-core/kafka-indexing-service/src/main/java/org/apache/druid/data/input/kafkainput/KafkaInputReader.java

@@ -58,14 +58,13 @@ public class KafkaInputReader implements InputEntityReader
   private final String topicColumnName;
 
   /**
-   *
-   * @param inputRowSchema Actual schema from the ingestion spec
-   * @param source kafka record containing header, key & value that is wrapped inside SettableByteEntity
+   * @param inputRowSchema       Actual schema from the ingestion spec
+   * @param source               kafka record containing header, key & value that is wrapped inside SettableByteEntity
    * @param headerParserSupplier Function to get Header parser for parsing the header section, kafkaInputFormat allows users to skip header parsing section and hence an be null
-   * @param keyParserSupplier Function to get Key parser for key section, can be null as well. Key parser supplier can also return a null key parser.
-   * @param valueParser Value parser is a required section in kafkaInputFormat. It cannot be null.
-   * @param keyColumnName Default key column name
-   * @param timestampColumnName Default kafka record's timestamp column name
+   * @param keyParserSupplier    Function to get Key parser for key section, can be null as well. Key parser supplier can also return a null key parser.
+   * @param valueParser          Value parser is a required section in kafkaInputFormat. It cannot be null.
+   * @param keyColumnName        Default key column name
+   * @param timestampColumnName  Default kafka record's timestamp column name
    */
   public KafkaInputReader(
       InputRowSchema inputRowSchema,
@@ -144,14 +143,9 @@ public class KafkaInputReader implements InputEntityReader
       try (CloseableIterator<InputRow> keyIterator = keyParser.read()) {
         // Key currently only takes the first row and ignores the rest.
         if (keyIterator.hasNext()) {
-          // Return type for the key parser should be of type MapBasedInputRow
-          // Parsers returning other types are not compatible currently.
-          MapBasedInputRow keyRow = (MapBasedInputRow) keyIterator.next();
+          final InputRow keyRow = keyIterator.next();
           // Add the key to the mergeList only if the key string is not already present
-          mergedHeaderMap.putIfAbsent(
-              keyColumnName,
-              keyRow.getEvent().entrySet().stream().findFirst().get().getValue()
-          );
+          mergedHeaderMap.computeIfAbsent(keyColumnName, ignored -> getFirstValue(keyRow));
         }
       }
       catch (ClassCastException e) {
@@ -344,4 +338,15 @@ public class KafkaInputReader implements InputEntityReader
       }
     };
   }
+
+  /**
+   * Get the first value from an {@link InputRow}. This is the first element from {@link InputRow#getDimensions()}
+   * if there are any. If there are not any, returns null. This method is used to extract keys.
+   */
+  @Nullable
+  static Object getFirstValue(final InputRow row)
+  {
+    final List<String> dimensions = row.getDimensions();
+    return !dimensions.isEmpty() ? row.getRaw(dimensions.get(0)) : null;
+  }
 }

extensions-core/kafka-indexing-service/src/test/java/org/apache/druid/data/input/kafkainput/KafkaInputFormatTest.java

@@ -693,6 +693,105 @@ public class KafkaInputFormatTest
     }
   }
 
+  @Test
+  public void testKeyInCsvFormat() throws IOException
+  {
+    format = new KafkaInputFormat(
+        new KafkaStringHeaderFormat(null),
+        // Key Format
+        new CsvInputFormat(
+            // name of the field doesn't matter, it just has to be something
+            Collections.singletonList("foo"),
+            null,
+            false,
+            false,
+            0,
+            null
+        ),
+        // Value Format
+        new JsonInputFormat(
+            new JSONPathSpec(true, ImmutableList.of()),
+            null,
+            null,
+            false,
+            false
+        ),
+        "kafka.newheader.",
+        "kafka.newkey.key",
+        "kafka.newts.timestamp",
+        "kafka.newtopic.topic"
+    );
+
+    Headers headers = new RecordHeaders(SAMPLE_HEADERS);
+    KafkaRecordEntity inputEntity =
+        makeInputEntity(
+            // x,y,z are ignored; key will be "sampleKey"
+            StringUtils.toUtf8("sampleKey,x,y,z"),
+            SIMPLE_JSON_VALUE_BYTES,
+            headers
+        );
+
+    final InputEntityReader reader = format.createReader(
+        new InputRowSchema(
+            new TimestampSpec("timestamp", "iso", null),
+            new DimensionsSpec(
+                DimensionsSpec.getDefaultSchemas(
+                    ImmutableList.of(
+                        "bar",
+                        "foo",
+                        "kafka.newkey.key",
+                        "kafka.newheader.encoding",
+                        "kafka.newheader.kafkapkc",
+                        "kafka.newts.timestamp",
+                        "kafka.newtopic.topic"
+                    )
+                )
+            ),
+            ColumnsFilter.all()
+        ),
+        newSettableByteEntity(inputEntity),
+        null
+    );
+
+    final int numExpectedIterations = 1;
+    try (CloseableIterator<InputRow> iterator = reader.read()) {
+      int numActualIterations = 0;
+      while (iterator.hasNext()) {
+
+        final InputRow row = iterator.next();
+        Assert.assertEquals(
+            Arrays.asList(
+                "bar",
+                "foo",
+                "kafka.newkey.key",
+                "kafka.newheader.encoding",
+                "kafka.newheader.kafkapkc",
+                "kafka.newts.timestamp",
+                "kafka.newtopic.topic"
+            ),
+            row.getDimensions()
+        );
+        // Payload verifications
+        // this isn't super realistic, since most of these columns are not actually defined in the dimensionSpec
+        // but test reading them anyway since it isn't technically illegal
+
+        Assert.assertEquals(DateTimes.of("2021-06-25"), row.getTimestamp());
+        Assert.assertEquals("x", Iterables.getOnlyElement(row.getDimension("foo")));
+        Assert.assertEquals("4", Iterables.getOnlyElement(row.getDimension("baz")));
+        Assert.assertTrue(row.getDimension("bar").isEmpty());
+
+        verifyHeader(row);
+
+        // Key verification
+        Assert.assertEquals("sampleKey", Iterables.getOnlyElement(row.getDimension("kafka.newkey.key")));
+
+        numActualIterations++;
+      }
+
+      Assert.assertEquals(numExpectedIterations, numActualIterations);
+    }
+  }
+
   @Test
   public void testValueInCsvFormat() throws IOException
   {

extensions-core/kinesis-indexing-service/src/test/java/org/apache/druid/indexing/kinesis/supervisor/KinesisSupervisorTest.java

@@ -358,6 +358,7 @@ public class KinesisSupervisorTest extends EasyMockSupport
     Thread.sleep(1 * 1000);
     int taskCountAfterScale = supervisor.getIoConfig().getTaskCount();
     Assert.assertEquals(2, taskCountAfterScale);
+    autoscaler.stop();
   }
 
   @Test
@@ -435,6 +436,7 @@ public class KinesisSupervisorTest extends EasyMockSupport
     Thread.sleep(1 * 1000);
     int taskCountAfterScale = supervisor.getIoConfig().getTaskCount();
     Assert.assertEquals(1, taskCountAfterScale);
+    autoscaler.stop();
   }
 
   @Test

extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java

@@ -28,8 +28,10 @@ import org.apache.druid.server.lookup.cache.loading.LoadingCache;
 
 import javax.annotation.Nullable;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -73,15 +75,19 @@ public class LoadingLookup extends LookupExtractor
       return null;
     }
 
-    final String presentVal;
-    try {
-      presentVal = loadingCache.get(keyEquivalent, new ApplyCallable(keyEquivalent));
+    final String presentVal = this.loadingCache.getIfPresent(keyEquivalent);
+    if (presentVal != null) {
       return NullHandling.emptyToNullIfNeeded(presentVal);
     }
-    catch (ExecutionException e) {
-      LOGGER.debug("value not found for key [%s]", key);
+
+    final String val = this.dataFetcher.fetch(keyEquivalent);
+    if (val == null) {
       return null;
     }
+
+    this.loadingCache.putAll(Collections.singletonMap(keyEquivalent, val));
+
+    return NullHandling.emptyToNullIfNeeded(val);
   }
 
   @Override
@@ -108,13 +114,16 @@ public class LoadingLookup extends LookupExtractor
   @Override
   public boolean supportsAsMap()
   {
-    return false;
+    return true;
   }
 
   @Override
   public Map<String, String> asMap()
   {
-    throw new UnsupportedOperationException("Cannot get map view");
+    final Map<String, String> map = new HashMap<>();
+    Optional.ofNullable(this.dataFetcher.fetchAll())
+            .ifPresent(data -> data.forEach(entry -> map.put(entry.getKey(), entry.getValue())));
+    return map;
   }
 
   @Override
@@ -123,24 +132,6 @@ public class LoadingLookup extends LookupExtractor
     return LookupExtractionModule.getRandomCacheKey();
   }
 
-  private class ApplyCallable implements Callable<String>
-  {
-    private final String key;
-
-    public ApplyCallable(String key)
-    {
-      this.key = key;
-    }
-
-    @Override
-    public String call()
-    {
-      // When SQL compatible null handling is disabled,
-      // avoid returning null and return an empty string to cache it.
-      return NullHandling.nullToEmptyIfNeeded(dataFetcher.fetch(key));
-    }
-  }
-
   public synchronized void close()
   {
     if (isOpen.getAndSet(false)) {

extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java

@@ -33,13 +33,15 @@ import org.junit.rules.ExpectedException;
 
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 
 public class LoadingLookupTest extends InitializedNullHandlingTest
 {
   DataFetcher dataFetcher = EasyMock.createMock(DataFetcher.class);
-  LoadingCache lookupCache = EasyMock.createStrictMock(LoadingCache.class);
+  LoadingCache lookupCache = EasyMock.createMock(LoadingCache.class);
   LoadingCache reverseLookupCache = EasyMock.createStrictMock(LoadingCache.class);
   LoadingLookup loadingLookup = new LoadingLookup(dataFetcher, lookupCache, reverseLookupCache);
 
@@ -47,9 +49,9 @@ public class LoadingLookupTest extends InitializedNullHandlingTest
   public ExpectedException expectedException = ExpectedException.none();
 
   @Test
-  public void testApplyEmptyOrNull() throws ExecutionException
+  public void testApplyEmptyOrNull()
   {
-    EasyMock.expect(lookupCache.get(EasyMock.eq(""), EasyMock.anyObject(Callable.class)))
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("")))
             .andReturn("empty").atLeastOnce();
     EasyMock.replay(lookupCache);
     Assert.assertEquals("empty", loadingLookup.apply(""));
@@ -73,14 +75,40 @@ public class LoadingLookupTest extends InitializedNullHandlingTest
   }
 
   @Test
-  public void testApply() throws ExecutionException
+  public void testApply()
   {
-    EasyMock.expect(lookupCache.get(EasyMock.eq("key"), EasyMock.anyObject(Callable.class))).andReturn("value").once();
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn("value").once();
     EasyMock.replay(lookupCache);
     Assert.assertEquals(ImmutableMap.of("key", "value"), loadingLookup.applyAll(ImmutableSet.of("key")));
     EasyMock.verify(lookupCache);
   }
 
+  @Test
+  public void testApplyWithNullValue()
+  {
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn(null).once();
+    EasyMock.expect(dataFetcher.fetch("key")).andReturn(null).once();
+    EasyMock.replay(lookupCache, dataFetcher);
+    Assert.assertNull(loadingLookup.apply("key"));
+    EasyMock.verify(lookupCache, dataFetcher);
+  }
+
+  @Test
+  public void testApplyTriggersCacheMissAndSubsequentCacheHit()
+  {
+    Map<String, String> map = new HashMap<>();
+    map.put("key", "value");
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn(null).once();
+    EasyMock.expect(dataFetcher.fetch("key")).andReturn("value").once();
+    lookupCache.putAll(map);
+    EasyMock.expectLastCall().andVoid();
+    EasyMock.expect(lookupCache.getIfPresent("key")).andReturn("value").once();
+    EasyMock.replay(lookupCache, dataFetcher);
+    Assert.assertEquals(loadingLookup.apply("key"), "value");
+    Assert.assertEquals(loadingLookup.apply("key"), "value");
+    EasyMock.verify(lookupCache, dataFetcher);
+  }
+
   @Test
   public void testUnapplyAll() throws ExecutionException
   {
@@ -105,17 +133,6 @@ public class LoadingLookupTest extends InitializedNullHandlingTest
     EasyMock.verify(lookupCache, reverseLookupCache);
   }
 
-  @Test
-  public void testApplyWithExecutionError() throws ExecutionException
-  {
-    EasyMock.expect(lookupCache.get(EasyMock.eq("key"), EasyMock.anyObject(Callable.class)))
-            .andThrow(new ExecutionException(null))
-            .once();
-    EasyMock.replay(lookupCache);
-    Assert.assertNull(loadingLookup.apply("key"));
-    EasyMock.verify(lookupCache);
-  }
-
   @Test
   public void testUnApplyWithExecutionError() throws ExecutionException
   {
@@ -136,13 +153,19 @@ public class LoadingLookupTest extends InitializedNullHandlingTest
   @Test
   public void testSupportsAsMap()
   {
-    Assert.assertFalse(loadingLookup.supportsAsMap());
+    Assert.assertTrue(loadingLookup.supportsAsMap());
   }
 
   @Test
   public void testAsMap()
   {
-    expectedException.expect(UnsupportedOperationException.class);
-    loadingLookup.asMap();
+    final Map<String, String> fetchedData = new HashMap<>();
+    fetchedData.put("dummy", "test");
+    fetchedData.put("key", null);
+    fetchedData.put(null, "value");
+    EasyMock.expect(dataFetcher.fetchAll()).andReturn(fetchedData.entrySet());
+    EasyMock.replay(dataFetcher);
+    Assert.assertEquals(loadingLookup.asMap(), fetchedData);
+    EasyMock.verify(dataFetcher);
   }
 }

extensions-core/multi-stage-query/pom.xml

@@ -196,11 +196,6 @@
             <artifactId>fastutil-core</artifactId>
             <scope>provided</scope>
         </dependency>
-        <dependency>
-            <groupId>commons-lang</groupId>
-            <artifactId>commons-lang</artifactId>
-            <scope>provided</scope>
-        </dependency>
         <dependency>
             <groupId>commons-io</groupId>
             <artifactId>commons-io</artifactId>

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/Dart.java

@@ -0,0 +1,37 @@
+/*
+ * 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.dart;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Binding annotation for implements of interfaces that are Dart (MSQ-on-Broker-and-Historicals) focused.
+ */
+@Target({ElementType.FIELD, ElementType.PARAMETER, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@BindingAnnotation
+public @interface Dart
+{
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/DartResourcePermissionMapper.java

@@ -0,0 +1,57 @@
+/*
+ * 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.dart;
+
+import com.google.common.collect.ImmutableList;
+import org.apache.druid.msq.dart.controller.http.DartSqlResource;
+import org.apache.druid.msq.dart.worker.http.DartWorkerResource;
+import org.apache.druid.msq.rpc.ResourcePermissionMapper;
+import org.apache.druid.msq.rpc.WorkerResource;
+import org.apache.druid.server.security.Action;
+import org.apache.druid.server.security.Resource;
+import org.apache.druid.server.security.ResourceAction;
+
+import java.util.List;
+
+public class DartResourcePermissionMapper implements ResourcePermissionMapper
+{
+  /**
+   * Permissions for admin APIs in {@link DartWorkerResource} and {@link WorkerResource}. Note that queries from
+   * end users go through {@link DartSqlResource}, which wouldn't use these mappings.
+   */
+  @Override
+  public List<ResourceAction> getAdminPermissions()
+  {
+    return ImmutableList.of(
+        new ResourceAction(Resource.STATE_RESOURCE, Action.READ),
+        new ResourceAction(Resource.STATE_RESOURCE, Action.WRITE)
+    );
+  }
+
+  /**
+   * Permissions for per-query APIs in {@link DartWorkerResource} and {@link WorkerResource}. Note that queries from
+   * end users go through {@link DartSqlResource}, which wouldn't use these mappings.
+   */
+  @Override
+  public List<ResourceAction> getQueryPermissions(String queryId)
+  {
+    return getAdminPermissions();
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/ControllerHolder.java

@@ -0,0 +1,174 @@
+/*
+ * 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.dart.controller;
+
+import com.google.common.base.Preconditions;
+import org.apache.druid.msq.dart.worker.DartWorkerClient;
+import org.apache.druid.msq.dart.worker.WorkerId;
+import org.apache.druid.msq.exec.Controller;
+import org.apache.druid.msq.exec.ControllerContext;
+import org.apache.druid.msq.exec.QueryListener;
+import org.apache.druid.msq.indexing.error.MSQErrorReport;
+import org.apache.druid.msq.indexing.error.WorkerFailedFault;
+import org.apache.druid.server.security.AuthenticationResult;
+import org.joda.time.DateTime;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Holder for {@link Controller}, stored in {@link DartControllerRegistry}.
+ */
+public class ControllerHolder
+{
+  public enum State
+  {
+    /**
+     * Query has been accepted, but not yet {@link Controller#run(QueryListener)}.
+     */
+    ACCEPTED,
+
+    /**
+     * Query has had {@link Controller#run(QueryListener)} called.
+     */
+    RUNNING,
+
+    /**
+     * Query has been canceled.
+     */
+    CANCELED
+  }
+
+  private final Controller controller;
+  private final ControllerContext controllerContext;
+  private final String sqlQueryId;
+  private final String sql;
+  private final String controllerHost;
+  private final AuthenticationResult authenticationResult;
+  private final DateTime startTime;
+  private final AtomicReference<State> state = new AtomicReference<>(State.ACCEPTED);
+
+  public ControllerHolder(
+      final Controller controller,
+      final ControllerContext controllerContext,
+      final String sqlQueryId,
+      final String sql,
+      final String controllerHost,
+      final AuthenticationResult authenticationResult,
+      final DateTime startTime
+  )
+  {
+    this.controller = Preconditions.checkNotNull(controller, "controller");
+    this.controllerContext = controllerContext;
+    this.sqlQueryId = Preconditions.checkNotNull(sqlQueryId, "sqlQueryId");
+    this.sql = sql;
+    this.controllerHost = controllerHost;
+    this.authenticationResult = authenticationResult;
+    this.startTime = Preconditions.checkNotNull(startTime, "startTime");
+  }
+
+  public Controller getController()
+  {
+    return controller;
+  }
+
+  public String getSqlQueryId()
+  {
+    return sqlQueryId;
+  }
+
+  public String getSql()
+  {
+    return sql;
+  }
+
+  public String getControllerHost()
+  {
+    return controllerHost;
+  }
+
+  public AuthenticationResult getAuthenticationResult()
+  {
+    return authenticationResult;
+  }
+
+  public DateTime getStartTime()
+  {
+    return startTime;
+  }
+
+  public State getState()
+  {
+    return state.get();
+  }
+
+  /**
+   * Call when a worker has gone offline. Closes its client and sends a {@link Controller#workerError}
+   * to the controller.
+   */
+  public void workerOffline(final WorkerId workerId)
+  {
+    final String workerIdString = workerId.toString();
+
+    if (controllerContext instanceof DartControllerContext) {
+      // For DartControllerContext, newWorkerClient() returns the same instance every time.
+      // This will always be DartControllerContext in production; the instanceof check is here because certain
+      // tests use a different context class.
+      ((DartWorkerClient) controllerContext.newWorkerClient()).closeClient(workerId.getHostAndPort());
+    }
+
+    if (controller.hasWorker(workerIdString)) {
+      controller.workerError(
+          MSQErrorReport.fromFault(
+              workerIdString,
+              workerId.getHostAndPort(),
+              null,
+              new WorkerFailedFault(workerIdString, "Worker went offline")
+          )
+      );
+    }
+  }
+
+  /**
+   * Places this holder into {@link State#CANCELED}. Calls {@link Controller#stop()} if it was previously in
+   * state {@link State#RUNNING}.
+   */
+  public void cancel()
+  {
+    if (state.getAndSet(State.CANCELED) == State.RUNNING) {
+      controller.stop();
+    }
+  }
+
+  /**
+   * Calls {@link Controller#run(QueryListener)}, and returns true, if this holder was previously in state
+   * {@link State#ACCEPTED}. Otherwise returns false.
+   *
+   * @return whether {@link Controller#run(QueryListener)} was called.
+   */
+  public boolean run(final QueryListener listener) throws Exception
+  {
+    if (state.compareAndSet(State.ACCEPTED, State.RUNNING)) {
+      controller.run(listener);
+      return true;
+    } else {
+      return false;
+    }
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/ControllerMessageListener.java

@@ -0,0 +1,68 @@
+/*
+ * 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.dart.controller;
+
+import com.google.inject.Inject;
+import org.apache.druid.messages.client.MessageListener;
+import org.apache.druid.msq.dart.controller.messages.ControllerMessage;
+import org.apache.druid.msq.dart.worker.WorkerId;
+import org.apache.druid.msq.exec.Controller;
+import org.apache.druid.msq.indexing.error.MSQErrorReport;
+import org.apache.druid.server.DruidNode;
+
+/**
+ * Listener for worker-to-controller messages.
+ * Also responsible for calling {@link Controller#workerError(MSQErrorReport)} when a worker server goes away.
+ */
+public class ControllerMessageListener implements MessageListener<ControllerMessage>
+{
+  private final DartControllerRegistry controllerRegistry;
+
+  @Inject
+  public ControllerMessageListener(final DartControllerRegistry controllerRegistry)
+  {
+    this.controllerRegistry = controllerRegistry;
+  }
+
+  @Override
+  public void messageReceived(ControllerMessage message)
+  {
+    final ControllerHolder holder = controllerRegistry.get(message.getQueryId());
+    if (holder != null) {
+      message.handle(holder.getController());
+    }
+  }
+
+  @Override
+  public void serverAdded(DruidNode node)
+  {
+    // Nothing to do.
+  }
+
+  @Override
+  public void serverRemoved(DruidNode node)
+  {
+    for (final ControllerHolder holder : controllerRegistry.getAllHolders()) {
+      final Controller controller = holder.getController();
+      final WorkerId workerId = WorkerId.fromDruidNode(node, controller.queryId());
+      holder.workerOffline(workerId);
+    }
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java

@@ -0,0 +1,249 @@
+/*
+ * 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.dart.controller;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Injector;
+import org.apache.druid.client.BrokerServerView;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.indexing.common.TaskLockType;
+import org.apache.druid.indexing.common.actions.TaskActionClient;
+import org.apache.druid.java.util.common.io.Closer;
+import org.apache.druid.java.util.emitter.service.ServiceEmitter;
+import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
+import org.apache.druid.msq.dart.worker.DartWorkerClient;
+import org.apache.druid.msq.dart.worker.WorkerId;
+import org.apache.druid.msq.exec.Controller;
+import org.apache.druid.msq.exec.ControllerContext;
+import org.apache.druid.msq.exec.ControllerMemoryParameters;
+import org.apache.druid.msq.exec.MemoryIntrospector;
+import org.apache.druid.msq.exec.WorkerFailureListener;
+import org.apache.druid.msq.exec.WorkerManager;
+import org.apache.druid.msq.indexing.IndexerControllerContext;
+import org.apache.druid.msq.indexing.MSQSpec;
+import org.apache.druid.msq.indexing.destination.TaskReportMSQDestination;
+import org.apache.druid.msq.input.InputSpecSlicer;
+import org.apache.druid.msq.kernel.controller.ControllerQueryKernelConfig;
+import org.apache.druid.msq.querykit.QueryKit;
+import org.apache.druid.msq.querykit.QueryKitSpec;
+import org.apache.druid.msq.util.MultiStageQueryContext;
+import org.apache.druid.query.Query;
+import org.apache.druid.query.QueryContext;
+import org.apache.druid.server.DruidNode;
+import org.apache.druid.server.coordination.DruidServerMetadata;
+import org.apache.druid.server.coordination.ServerType;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Dart implementation of {@link ControllerContext}.
+ * Each instance is scoped to a query.
+ */
+public class DartControllerContext implements ControllerContext
+{
+  /**
+   * Default for {@link ControllerQueryKernelConfig#getMaxConcurrentStages()}.
+   */
+  public static final int DEFAULT_MAX_CONCURRENT_STAGES = 2;
+
+  /**
+   * Default for {@link MultiStageQueryContext#getTargetPartitionsPerWorkerWithDefault(QueryContext, int)}.
+   */
+  public static final int DEFAULT_TARGET_PARTITIONS_PER_WORKER = 1;
+
+  /**
+   * Context parameter for maximum number of nonleaf workers.
+   */
+  public static final String CTX_MAX_NON_LEAF_WORKER_COUNT = "maxNonLeafWorkers";
+
+  /**
+   * Default to scatter/gather style: fan in to a single worker after the leaf stage(s).
+   */
+  public static final int DEFAULT_MAX_NON_LEAF_WORKER_COUNT = 1;
+
+  private final Injector injector;
+  private final ObjectMapper jsonMapper;
+  private final DruidNode selfNode;
+  private final DartWorkerClient workerClient;
+  private final BrokerServerView serverView;
+  private final MemoryIntrospector memoryIntrospector;
+  private final ServiceMetricEvent.Builder metricBuilder;
+  private final ServiceEmitter emitter;
+
+  public DartControllerContext(
+      final Injector injector,
+      final ObjectMapper jsonMapper,
+      final DruidNode selfNode,
+      final DartWorkerClient workerClient,
+      final MemoryIntrospector memoryIntrospector,
+      final BrokerServerView serverView,
+      final ServiceEmitter emitter
+  )
+  {
+    this.injector = injector;
+    this.jsonMapper = jsonMapper;
+    this.selfNode = selfNode;
+    this.workerClient = workerClient;
+    this.serverView = serverView;
+    this.memoryIntrospector = memoryIntrospector;
+    this.metricBuilder = new ServiceMetricEvent.Builder();
+    this.emitter = emitter;
+  }
+
+  @Override
+  public ControllerQueryKernelConfig queryKernelConfig(
+      final String queryId,
+      final MSQSpec querySpec
+  )
+  {
+    final List<DruidServerMetadata> servers = serverView.getDruidServerMetadatas();
+
+    // Lock in the list of workers when creating the kernel config. There is a race here: the serverView itself is
+    // allowed to float. If a segment moves to a new server that isn't part of our list after the WorkerManager is
+    // created, we won't be able to find a valid server for certain segments. This isn't expected to be a problem,
+    // since the serverView is referenced shortly after the worker list is created.
+    final List<String> workerIds = new ArrayList<>(servers.size());
+    for (final DruidServerMetadata server : servers) {
+      if (server.getType() == ServerType.HISTORICAL) {
+        workerIds.add(WorkerId.fromDruidServerMetadata(server, queryId).toString());
+      }
+    }
+
+    // Shuffle workerIds, so we don't bias towards specific servers when running multiple queries concurrently. For any
+    // given query, lower-numbered workers tend to do more work, because the controller prefers using lower-numbered
+    // workers when maxWorkerCount for a stage is less than the total number of workers.
+    Collections.shuffle(workerIds);
+
+    final ControllerMemoryParameters memoryParameters =
+        ControllerMemoryParameters.createProductionInstance(
+            memoryIntrospector,
+            workerIds.size()
+        );
+
+    final int maxConcurrentStages = MultiStageQueryContext.getMaxConcurrentStagesWithDefault(
+        querySpec.getQuery().context(),
+        DEFAULT_MAX_CONCURRENT_STAGES
+    );
+
+    return ControllerQueryKernelConfig
+        .builder()
+        .controllerHost(selfNode.getHostAndPortToUse())
+        .workerIds(workerIds)
+        .pipeline(maxConcurrentStages > 1)
+        .destination(TaskReportMSQDestination.instance())
+        .maxConcurrentStages(maxConcurrentStages)
+        .maxRetainedPartitionSketchBytes(memoryParameters.getPartitionStatisticsMaxRetainedBytes())
+        .workerContextMap(IndexerControllerContext.makeWorkerContextMap(querySpec, false, maxConcurrentStages))
+        .build();
+  }
+
+  @Override
+  public ObjectMapper jsonMapper()
+  {
+    return jsonMapper;
+  }
+
+  @Override
+  public Injector injector()
+  {
+    return injector;
+  }
+
+  @Override
+  public void emitMetric(final String metric, final Number value)
+  {
+    emitter.emit(metricBuilder.setMetric(metric, value));
+  }
+
+  @Override
+  public DruidNode selfNode()
+  {
+    return selfNode;
+  }
+
+  @Override
+  public InputSpecSlicer newTableInputSpecSlicer(WorkerManager workerManager)
+  {
+    return DartTableInputSpecSlicer.createFromWorkerIds(workerManager.getWorkerIds(), serverView);
+  }
+
+  @Override
+  public TaskActionClient taskActionClient()
+  {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public WorkerManager newWorkerManager(
+      String queryId,
+      MSQSpec querySpec,
+      ControllerQueryKernelConfig queryKernelConfig,
+      WorkerFailureListener workerFailureListener
+  )
+  {
+    // We're ignoring WorkerFailureListener. Dart worker failures are routed into the controller by
+    // ControllerMessageListener, which receives a notification when a worker goes offline.
+    return new DartWorkerManager(queryKernelConfig.getWorkerIds(), workerClient);
+  }
+
+  @Override
+  public DartWorkerClient newWorkerClient()
+  {
+    return workerClient;
+  }
+
+  @Override
+  public void registerController(Controller controller, Closer closer)
+  {
+    closer.register(workerClient);
+  }
+
+  @Override
+  public QueryKitSpec makeQueryKitSpec(
+      final QueryKit<Query<?>> queryKit,
+      final String queryId,
+      final MSQSpec querySpec,
+      final ControllerQueryKernelConfig queryKernelConfig
+  )
+  {
+    final QueryContext queryContext = querySpec.getQuery().context();
+    return new QueryKitSpec(
+        queryKit,
+        queryId,
+        queryKernelConfig.getWorkerIds().size(),
+        queryContext.getInt(
+            CTX_MAX_NON_LEAF_WORKER_COUNT,
+            DEFAULT_MAX_NON_LEAF_WORKER_COUNT
+        ),
+        MultiStageQueryContext.getTargetPartitionsPerWorkerWithDefault(
+            queryContext,
+            DEFAULT_TARGET_PARTITIONS_PER_WORKER
+        )
+    );
+  }
+
+  @Override
+  public TaskLockType taskLockType()
+  {
+    throw DruidException.defensive("TaskLockType is not used with class[%s]", getClass().getName());
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactory.java

@@ -0,0 +1,31 @@
+/*
+ * 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.dart.controller;
+
+import org.apache.druid.msq.dart.controller.sql.DartQueryMaker;
+import org.apache.druid.msq.exec.ControllerContext;
+
+/**
+ * Class for creating {@link ControllerContext} in {@link DartQueryMaker}.
+ */
+public interface DartControllerContextFactory
+{
+  ControllerContext newContext(String queryId);
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java

@@ -0,0 +1,83 @@
+/*
+ * 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.dart.controller;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Inject;
+import com.google.inject.Injector;
+import org.apache.druid.client.BrokerServerView;
+import org.apache.druid.guice.annotations.EscalatedGlobal;
+import org.apache.druid.guice.annotations.Json;
+import org.apache.druid.guice.annotations.Self;
+import org.apache.druid.guice.annotations.Smile;
+import org.apache.druid.java.util.emitter.service.ServiceEmitter;
+import org.apache.druid.msq.dart.worker.DartWorkerClient;
+import org.apache.druid.msq.exec.ControllerContext;
+import org.apache.druid.msq.exec.MemoryIntrospector;
+import org.apache.druid.rpc.ServiceClientFactory;
+import org.apache.druid.server.DruidNode;
+
+public class DartControllerContextFactoryImpl implements DartControllerContextFactory
+{
+  private final Injector injector;
+  private final ObjectMapper jsonMapper;
+  private final ObjectMapper smileMapper;
+  private final DruidNode selfNode;
+  private final ServiceClientFactory serviceClientFactory;
+  private final BrokerServerView serverView;
+  private final MemoryIntrospector memoryIntrospector;
+  private final ServiceEmitter emitter;
+
+  @Inject
+  public DartControllerContextFactoryImpl(
+      final Injector injector,
+      @Json final ObjectMapper jsonMapper,
+      @Smile final ObjectMapper smileMapper,
+      @Self final DruidNode selfNode,
+      @EscalatedGlobal final ServiceClientFactory serviceClientFactory,
+      final MemoryIntrospector memoryIntrospector,
+      final BrokerServerView serverView,
+      final ServiceEmitter emitter
+  )
+  {
+    this.injector = injector;
+    this.jsonMapper = jsonMapper;
+    this.smileMapper = smileMapper;
+    this.selfNode = selfNode;
+    this.serviceClientFactory = serviceClientFactory;
+    this.serverView = serverView;
+    this.memoryIntrospector = memoryIntrospector;
+    this.emitter = emitter;
+  }
+
+  @Override
+  public ControllerContext newContext(final String queryId)
+  {
+    return new DartControllerContext(
+        injector,
+        jsonMapper,
+        selfNode,
+        new DartWorkerClient(queryId, serviceClientFactory, smileMapper, selfNode.getHostAndPortToUse()),
+        memoryIntrospector,
+        serverView,
+        emitter
+    );
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerRegistry.java

@@ -0,0 +1,72 @@
+/*
+ * 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.dart.controller;
+
+import org.apache.druid.error.DruidException;
+import org.apache.druid.msq.exec.Controller;
+
+import javax.annotation.Nullable;
+import java.util.Collection;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Registry for actively-running {@link Controller}.
+ */
+public class DartControllerRegistry
+{
+  private final ConcurrentHashMap<String, ControllerHolder> controllerMap = new ConcurrentHashMap<>();
+
+  /**
+   * Add a controller. Throws {@link DruidException} if a controller with the same {@link Controller#queryId()} is
+   * already registered.
+   */
+  public void register(ControllerHolder holder)
+  {
+    if (controllerMap.putIfAbsent(holder.getController().queryId(), holder) != null) {
+      throw DruidException.defensive("Controller[%s] already registered", holder.getController().queryId());
+    }
+  }
+
+  /**
+   * Remove a controller from the registry.
+   */
+  public void deregister(ControllerHolder holder)
+  {
+    // Remove only if the current mapping for the queryId is this specific controller.
+    controllerMap.remove(holder.getController().queryId(), holder);
+  }
+
+  /**
+   * Return a specific controller holder, or null if it doesn't exist.
+   */
+  @Nullable
+  public ControllerHolder get(final String queryId)
+  {
+    return controllerMap.get(queryId);
+  }
+
+  /**
+   * Returns all actively-running {@link Controller}.
+   */
+  public Collection<ControllerHolder> getAllHolders()
+  {
+    return controllerMap.values();
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartMessageRelayFactoryImpl.java

@@ -0,0 +1,82 @@
+/*
+ * 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.dart.controller;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Inject;
+import org.apache.druid.guice.annotations.EscalatedGlobal;
+import org.apache.druid.guice.annotations.Self;
+import org.apache.druid.guice.annotations.Smile;
+import org.apache.druid.messages.client.MessageRelay;
+import org.apache.druid.messages.client.MessageRelayClientImpl;
+import org.apache.druid.messages.client.MessageRelayFactory;
+import org.apache.druid.msq.dart.controller.messages.ControllerMessage;
+import org.apache.druid.msq.dart.worker.http.DartWorkerResource;
+import org.apache.druid.rpc.FixedServiceLocator;
+import org.apache.druid.rpc.ServiceClient;
+import org.apache.druid.rpc.ServiceClientFactory;
+import org.apache.druid.rpc.ServiceLocation;
+import org.apache.druid.rpc.StandardRetryPolicy;
+import org.apache.druid.server.DruidNode;
+
+/**
+ * Production implementation of {@link MessageRelayFactory}.
+ */
+public class DartMessageRelayFactoryImpl implements MessageRelayFactory<ControllerMessage>
+{
+  private final String clientHost;
+  private final ControllerMessageListener messageListener;
+  private final ServiceClientFactory clientFactory;
+  private final String basePath;
+  private final ObjectMapper smileMapper;
+
+  @Inject
+  public DartMessageRelayFactoryImpl(
+      @Self DruidNode selfNode,
+      @EscalatedGlobal ServiceClientFactory clientFactory,
+      @Smile ObjectMapper smileMapper,
+      ControllerMessageListener messageListener
+  )
+  {
+    this.clientHost = selfNode.getHostAndPortToUse();
+    this.messageListener = messageListener;
+    this.clientFactory = clientFactory;
+    this.smileMapper = smileMapper;
+    this.basePath = DartWorkerResource.PATH + "/relay";
+  }
+
+  @Override
+  public MessageRelay<ControllerMessage> newRelay(DruidNode clientNode)
+  {
+    final ServiceLocation location = ServiceLocation.fromDruidNode(clientNode).withBasePath(basePath);
+    final ServiceClient client = clientFactory.makeClient(
+        clientNode.getHostAndPortToUse(),
+        new FixedServiceLocator(location),
+        StandardRetryPolicy.unlimited()
+    );
+
+    return new MessageRelay<>(
+        clientHost,
+        clientNode,
+        new MessageRelayClientImpl<>(client, smileMapper, ControllerMessage.class),
+        messageListener
+    );
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartMessageRelays.java

@@ -0,0 +1,40 @@
+/*
+ * 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.dart.controller;
+
+import org.apache.druid.discovery.DruidNodeDiscoveryProvider;
+import org.apache.druid.discovery.NodeRole;
+import org.apache.druid.messages.client.MessageRelayFactory;
+import org.apache.druid.messages.client.MessageRelays;
+import org.apache.druid.msq.dart.controller.messages.ControllerMessage;
+
+/**
+ * Specialized {@link MessageRelays} for Dart controllers.
+ */
+public class DartMessageRelays extends MessageRelays<ControllerMessage>
+{
+  public DartMessageRelays(
+      final DruidNodeDiscoveryProvider discoveryProvider,
+      final MessageRelayFactory<ControllerMessage> messageRelayFactory
+  )
+  {
+    super(() -> discoveryProvider.getForNodeRole(NodeRole.HISTORICAL), messageRelayFactory);
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartTableInputSpecSlicer.java

@@ -0,0 +1,292 @@
+/*
+ * 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.dart.controller;
+
+import com.google.common.collect.FluentIterable;
+import com.google.common.collect.ImmutableList;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import org.apache.druid.client.TimelineServerView;
+import org.apache.druid.client.selector.QueryableDruidServer;
+import org.apache.druid.client.selector.ServerSelector;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.JodaUtils;
+import org.apache.druid.msq.dart.worker.DartQueryableSegment;
+import org.apache.druid.msq.dart.worker.WorkerId;
+import org.apache.druid.msq.exec.SegmentSource;
+import org.apache.druid.msq.exec.WorkerManager;
+import org.apache.druid.msq.input.InputSlice;
+import org.apache.druid.msq.input.InputSpec;
+import org.apache.druid.msq.input.InputSpecSlicer;
+import org.apache.druid.msq.input.NilInputSlice;
+import org.apache.druid.msq.input.table.RichSegmentDescriptor;
+import org.apache.druid.msq.input.table.SegmentsInputSlice;
+import org.apache.druid.msq.input.table.TableInputSpec;
+import org.apache.druid.query.TableDataSource;
+import org.apache.druid.query.filter.DimFilterUtils;
+import org.apache.druid.server.coordination.DruidServerMetadata;
+import org.apache.druid.timeline.DataSegment;
+import org.apache.druid.timeline.TimelineLookup;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Set;
+import java.util.function.ToIntFunction;
+
+/**
+ * Slices {@link TableInputSpec} into {@link SegmentsInputSlice} for persistent servers using
+ * {@link TimelineServerView}.
+ */
+public class DartTableInputSpecSlicer implements InputSpecSlicer
+{
+  private static final int UNKNOWN = -1;
+
+  /**
+   * Worker host:port -> worker number. This is the reverse of the mapping from {@link WorkerManager#getWorkerIds()}.
+   */
+  private final Object2IntMap<String> workerIdToNumber;
+
+  /**
+   * Server view for identifying which segments exist and which servers (workers) have which segments.
+   */
+  private final TimelineServerView serverView;
+
+  DartTableInputSpecSlicer(final Object2IntMap<String> workerIdToNumber, final TimelineServerView serverView)
+  {
+    this.workerIdToNumber = workerIdToNumber;
+    this.serverView = serverView;
+  }
+
+  public static DartTableInputSpecSlicer createFromWorkerIds(
+      final List<String> workerIds,
+      final TimelineServerView serverView
+  )
+  {
+    final Object2IntMap<String> reverseWorkers = new Object2IntOpenHashMap<>();
+    reverseWorkers.defaultReturnValue(UNKNOWN);
+
+    for (int i = 0; i < workerIds.size(); i++) {
+      reverseWorkers.put(WorkerId.fromString(workerIds.get(i)).getHostAndPort(), i);
+    }
+
+    return new DartTableInputSpecSlicer(reverseWorkers, serverView);
+  }
+
+  @Override
+  public boolean canSliceDynamic(final InputSpec inputSpec)
+  {
+    return false;
+  }
+
+  @Override
+  public List<InputSlice> sliceStatic(final InputSpec inputSpec, final int maxNumSlices)
+  {
+    final TableInputSpec tableInputSpec = (TableInputSpec) inputSpec;
+    final TimelineLookup<String, ServerSelector> timeline =
+        serverView.getTimeline(new TableDataSource(tableInputSpec.getDataSource()).getAnalysis()).orElse(null);
+
+    if (timeline == null) {
+      return Collections.emptyList();
+    }
+
+    final Set<DartQueryableSegment> prunedSegments =
+        findQueryableDataSegments(
+            tableInputSpec,
+            timeline,
+            serverSelector -> findWorkerForServerSelector(serverSelector, maxNumSlices)
+        );
+
+    final List<List<DartQueryableSegment>> assignments = new ArrayList<>(maxNumSlices);
+    while (assignments.size() < maxNumSlices) {
+      assignments.add(null);
+    }
+
+    int nextRoundRobinWorker = 0;
+    for (final DartQueryableSegment segment : prunedSegments) {
+      final int worker;
+      if (segment.getWorkerNumber() == UNKNOWN) {
+        // Segment is not available on any worker. Assign to some worker, round-robin. Today, that server will throw
+        // an error about the segment not being findable, but perhaps one day, it will be able to load the segment
+        // on demand.
+        worker = nextRoundRobinWorker;
+        nextRoundRobinWorker = (nextRoundRobinWorker + 1) % maxNumSlices;
+      } else {
+        worker = segment.getWorkerNumber();
+      }
+
+      if (assignments.get(worker) == null) {
+        assignments.set(worker, new ArrayList<>());
+      }
+
+      assignments.get(worker).add(segment);
+    }
+
+    return makeSegmentSlices(tableInputSpec.getDataSource(), assignments);
+  }
+
+  @Override
+  public List<InputSlice> sliceDynamic(
+      final InputSpec inputSpec,
+      final int maxNumSlices,
+      final int maxFilesPerSlice,
+      final long maxBytesPerSlice
+  )
+  {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Return the worker ID that corresponds to a particular {@link ServerSelector}, or {@link #UNKNOWN} if none does.
+   *
+   * @param serverSelector the server selector
+   * @param maxNumSlices   maximum number of worker IDs to use
+   */
+  int findWorkerForServerSelector(final ServerSelector serverSelector, final int maxNumSlices)
+  {
+    final QueryableDruidServer<?> server = serverSelector.pick(null);
+
+    if (server == null) {
+      return UNKNOWN;
+    }
+
+    final String serverHostAndPort = server.getServer().getHost();
+    final int workerNumber = workerIdToNumber.getInt(serverHostAndPort);
+
+    // The worker number may be UNKNOWN in a race condition, such as the set of Historicals changing while
+    // the query is being planned. I don't think it can be >= maxNumSlices, but if it is, treat it like UNKNOWN.
+    if (workerNumber != UNKNOWN && workerNumber < maxNumSlices) {
+      return workerNumber;
+    } else {
+      return UNKNOWN;
+    }
+  }
+
+  /**
+   * Pull the list of {@link DataSegment} that we should query, along with a clipping interval for each one, and
+   * a worker to get it from.
+   */
+  static Set<DartQueryableSegment> findQueryableDataSegments(
+      final TableInputSpec tableInputSpec,
+      final TimelineLookup<?, ServerSelector> timeline,
+      final ToIntFunction<ServerSelector> toWorkersFunction
+  )
+  {
+    final FluentIterable<DartQueryableSegment> allSegments =
+        FluentIterable.from(JodaUtils.condenseIntervals(tableInputSpec.getIntervals()))
+                      .transformAndConcat(timeline::lookup)
+                      .transformAndConcat(
+                          holder ->
+                              FluentIterable
+                                  .from(holder.getObject())
+                                  .filter(chunk -> shouldIncludeSegment(chunk.getObject()))
+                                  .transform(chunk -> {
+                                    final ServerSelector serverSelector = chunk.getObject();
+                                    final DataSegment dataSegment = serverSelector.getSegment();
+                                    final int worker = toWorkersFunction.applyAsInt(serverSelector);
+                                    return new DartQueryableSegment(dataSegment, holder.getInterval(), worker);
+                                  })
+                                  .filter(segment -> !segment.getSegment().isTombstone())
+                      );
+
+    return DimFilterUtils.filterShards(
+        tableInputSpec.getFilter(),
+        tableInputSpec.getFilterFields(),
+        allSegments,
+        segment -> segment.getSegment().getShardSpec(),
+        new HashMap<>()
+    );
+  }
+
+  /**
+   * Create a list of {@link SegmentsInputSlice} and {@link NilInputSlice} assignments.
+   *
+   * @param dataSource  datasource to read
+   * @param assignments list of assignment lists, one per slice
+   *
+   * @return a list of the same length as "assignments"
+   *
+   * @throws IllegalStateException if any provided segments do not match the provided datasource
+   */
+  static List<InputSlice> makeSegmentSlices(
+      final String dataSource,
+      final List<List<DartQueryableSegment>> assignments
+  )
+  {
+    final List<InputSlice> retVal = new ArrayList<>(assignments.size());
+
+    for (final List<DartQueryableSegment> assignment : assignments) {
+      if (assignment == null || assignment.isEmpty()) {
+        retVal.add(NilInputSlice.INSTANCE);
+      } else {
+        final List<RichSegmentDescriptor> descriptors = new ArrayList<>();
+        for (final DartQueryableSegment segment : assignment) {
+          if (!dataSource.equals(segment.getSegment().getDataSource())) {
+            throw new ISE("Expected dataSource[%s] but got[%s]", dataSource, segment.getSegment().getDataSource());
+          }
+
+          descriptors.add(toRichSegmentDescriptor(segment));
+        }
+
+        retVal.add(new SegmentsInputSlice(dataSource, descriptors, ImmutableList.of()));
+      }
+    }
+
+    return retVal;
+  }
+
+  /**
+   * Returns a {@link RichSegmentDescriptor}, which is used by {@link SegmentsInputSlice}.
+   */
+  static RichSegmentDescriptor toRichSegmentDescriptor(final DartQueryableSegment segment)
+  {
+    return new RichSegmentDescriptor(
+        segment.getSegment().getInterval(),
+        segment.getInterval(),
+        segment.getSegment().getVersion(),
+        segment.getSegment().getShardSpec().getPartitionNum()
+    );
+  }
+
+  /**
+   * Whether to include a segment from the timeline. Segments are included if they are not tombstones, and are also not
+   * purely realtime segments.
+   */
+  static boolean shouldIncludeSegment(final ServerSelector serverSelector)
+  {
+    if (serverSelector.getSegment().isTombstone()) {
+      return false;
+    }
+
+    int numRealtimeServers = 0;
+    int numOtherServers = 0;
+
+    for (final DruidServerMetadata server : serverSelector.getAllServers()) {
+      if (SegmentSource.REALTIME.getUsedServerTypes().contains(server.getType())) {
+        numRealtimeServers++;
+      } else {
+        numOtherServers++;
+      }
+    }
+
+    return numOtherServers > 0 || (numOtherServers + numRealtimeServers == 0);
+  }
+}

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartWorkerManager.java

@@ -0,0 +1,202 @@
+/*
+ * 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.dart.controller;
+
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.SettableFuture;
+import it.unimi.dsi.fastutil.ints.Int2ObjectAVLTreeMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.error.DruidException;
+import org.apache.druid.indexer.TaskState;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.msq.dart.worker.DartWorkerClient;
+import org.apache.druid.msq.exec.ControllerContext;
+import org.apache.druid.msq.exec.WorkerClient;
+import org.apache.druid.msq.exec.WorkerManager;
+import org.apache.druid.msq.exec.WorkerStats;
+import org.apache.druid.msq.indexing.WorkerCount;
+import org.apache.druid.utils.CloseableUtils;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * Dart implementation of the {@link WorkerManager} returned by {@link ControllerContext#newWorkerManager}.
+ *
+ * This manager does not actually launch workers. The workers are housed on long-lived servers outside of this
+ * manager's control. This manager merely reports on their existence.
+ */
+public class DartWorkerManager implements WorkerManager
+{
+  private static final Logger log = new Logger(DartWorkerManager.class);
+
+  private final List<String> workerIds;
+  private final DartWorkerClient workerClient;
+  private final Object2IntMap<String> workerIdToNumber;
+  private final AtomicReference<State> state = new AtomicReference<>(State.NEW);
+  private final SettableFuture<?> stopFuture = SettableFuture.create();
+
+  enum State
+  {
+    NEW,
+    STARTED,
+    STOPPED
+  }
+
+  public DartWorkerManager(
+      final List<String> workerIds,
+      final DartWorkerClient workerClient
+  )
+  {
+    this.workerIds = workerIds;
+    this.workerClient = workerClient;
+    this.workerIdToNumber = new Object2IntOpenHashMap<>();
+    this.workerIdToNumber.defaultReturnValue(UNKNOWN_WORKER_NUMBER);
+
+    for (int i = 0; i < workerIds.size(); i++) {
+      workerIdToNumber.put(workerIds.get(i), i);
+    }
+  }
+
+  @Override
+  public ListenableFuture<?> start()
+  {
+    if (!state.compareAndSet(State.NEW, State.STARTED)) {
+      throw new ISE("Cannot start from state[%s]", state.get());
+    }
+
+    return stopFuture;
+  }
+
+  @Override
+  public void launchWorkersIfNeeded(int workerCount)
+  {
+    // Nothing to do, just validate the count.
+    if (workerCount > workerIds.size()) {
+      throw DruidException.defensive(
+          "Desired workerCount[%s] must be less than or equal to actual workerCount[%s]",
+          workerCount,
+          workerIds.size()
+      );
+    }
+  }
+
+  @Override
+  public void waitForWorkers(Set<Integer> workerNumbers)
+  {
+    // Nothing to wait for, just validate the numbers.
+    for (final int workerNumber : workerNumbers) {
+      if (workerNumber >= workerIds.size()) {
+        throw DruidException.defensive(
+            "Desired workerNumber[%s] must be less than workerCount[%s]",
+            workerNumber,
+            workerIds.size()
+        );
+      }
+    }
+  }
+
+  @Override
+  public List<String> getWorkerIds()
+  {
+    return workerIds;
+  }
+
+  @Override
+  public WorkerCount getWorkerCount()
+  {
+    return new WorkerCount(workerIds.size(), 0);
+  }
+
+  @Override
+  public int getWorkerNumber(String workerId)
+  {
+    return workerIdToNumber.getInt(workerId);
+  }
+
+  @Override
+  public boolean isWorkerActive(String workerId)
+  {
+    return workerIdToNumber.containsKey(workerId);
+  }
+
+  @Override
+  public Map<Integer, List<WorkerStats>> getWorkerStats()
+  {
+    final Int2ObjectMap<List<WorkerStats>> retVal = new Int2ObjectAVLTreeMap<>();
+
+    for (int i = 0; i < workerIds.size(); i++) {
+      retVal.put(i, Collections.singletonList(new WorkerStats(workerIds.get(i), TaskState.RUNNING, -1, -1)));
+    }
+
+    return retVal;
+  }
+
+  /**
+   * Stop method. Possibly signals workers to stop, but does not actually wait for them to exit.
+   *
+   * If "interrupt" is false, does nothing special (other than setting {@link #stopFuture}). The assumption is that
+   * a previous call to {@link WorkerClient#postFinish} would have caused the worker to exit.
+   *
+   * If "interrupt" is true, sends {@link DartWorkerClient#stopWorker(String)} to workers to stop the current query ID.
+   *
+   * @param interrupt whether to interrupt currently-running work
+   */
+  @Override
+  public void stop(boolean interrupt)
+  {
+    if (state.compareAndSet(State.STARTED, State.STOPPED)) {
+      if (interrupt) {
+        final List<ListenableFuture<?>> futures = new ArrayList<>();
+
+        // Send stop commands to all workers. This ensures they exit promptly, and do not get left in a zombie state.
+        // For this reason, the workerClient uses an unlimited retry policy. If a stop command is lost, a worker
+        // could get stuck in a zombie state without its controller. This state would persist until the server that
+        // ran the controller shuts down or restarts. At that time, the listener in DartWorkerRunner.BrokerListener
+        // calls "controllerFailed()" on the Worker, and the zombie worker would exit.
+
+        for (final String workerId : workerIds) {
+          futures.add(workerClient.stopWorker(workerId));
+        }
+
+        // Block until messages are acknowledged, or until the worker we're communicating with has failed.
+
+        try {
+          FutureUtils.getUnchecked(Futures.successfulAsList(futures), false);
+        }
+        catch (Throwable ignored) {
+          // Suppress errors.
+        }
+      }
+
+      CloseableUtils.closeAndSuppressExceptions(workerClient, e -> log.warn(e, "Failed to close workerClient"));
+      stopFuture.set(null);
+    }
+  }
+}