Disable loading lookups by default in CompactionTask

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

@@ -55,7 +55,7 @@
 import org.apache.druid.rpc.indexing.OverlordClient;
 import org.apache.druid.segment.realtime.firehose.ChatHandler;
 import org.apache.druid.server.DruidNode;
-import org.apache.druid.sql.calcite.planner.PlannerContext;
+import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
 
 import java.util.Map;
 import java.util.concurrent.ThreadLocalRandom;
@@ -271,16 +271,16 @@ public static Map<String, Object> makeTaskContext(
         .put(MultiStageQueryContext.CTX_MAX_CONCURRENT_STAGES, queryKernelConfig.getMaxConcurrentStages());
 
     // Put the lookup loading info in the task context to facilitate selective loading of lookups.
-    if (controllerTaskContext.get(PlannerContext.CTX_LOOKUP_LOADING_MODE) != null) {
+    if (controllerTaskContext.get(LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE) != null) {
       taskContextOverridesBuilder.put(
-          PlannerContext.CTX_LOOKUP_LOADING_MODE,
-          controllerTaskContext.get(PlannerContext.CTX_LOOKUP_LOADING_MODE)
+          LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE,
+          controllerTaskContext.get(LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE)
       );
     }
-    if (controllerTaskContext.get(PlannerContext.CTX_LOOKUPS_TO_LOAD) != null) {
+    if (controllerTaskContext.get(LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD) != null) {
       taskContextOverridesBuilder.put(
-          PlannerContext.CTX_LOOKUPS_TO_LOAD,
-          controllerTaskContext.get(PlannerContext.CTX_LOOKUPS_TO_LOAD)
+          LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD,
+          controllerTaskContext.get(LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD)
       );
     }
 

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/MSQWorkerTask.java

@@ -27,7 +27,6 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
 import com.google.inject.Injector;
-import org.apache.druid.error.InvalidInput;
 import org.apache.druid.indexer.TaskStatus;
 import org.apache.druid.indexing.common.TaskToolbox;
 import org.apache.druid.indexing.common.actions.TaskActionClient;
@@ -38,13 +37,9 @@
 import org.apache.druid.msq.exec.Worker;
 import org.apache.druid.msq.exec.WorkerContext;
 import org.apache.druid.msq.exec.WorkerImpl;
-import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
 import org.apache.druid.server.security.ResourceAction;
-import org.apache.druid.sql.calcite.planner.PlannerContext;
 
 import javax.annotation.Nonnull;
-import java.util.Collection;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -190,26 +185,4 @@ public int hashCode()
   {
     return Objects.hash(super.hashCode(), controllerTaskId, workerNumber, retryCount, worker);
   }
-
-  @Override
-  public LookupLoadingSpec getLookupLoadingSpec()
-  {
-    final Object lookupModeValue = getContext().get(PlannerContext.CTX_LOOKUP_LOADING_MODE);
-    if (lookupModeValue == null) {
-      return LookupLoadingSpec.ALL;
-    }
-
-    final LookupLoadingSpec.Mode lookupLoadingMode = LookupLoadingSpec.Mode.valueOf(lookupModeValue.toString());
-    if (lookupLoadingMode == LookupLoadingSpec.Mode.NONE) {
-      return LookupLoadingSpec.NONE;
-    } else if (lookupLoadingMode == LookupLoadingSpec.Mode.ONLY_REQUIRED) {
-      Collection<String> lookupsToLoad = (Collection<String>) getContext().get(PlannerContext.CTX_LOOKUPS_TO_LOAD);
-      if (lookupsToLoad == null || lookupsToLoad.isEmpty()) {
-        throw InvalidInput.exception("Set of lookups to load cannot be %s for mode[ONLY_REQUIRED].", lookupsToLoad);
-      }
-      return LookupLoadingSpec.loadOnly(new HashSet<>(lookupsToLoad));
-    } else {
-      return LookupLoadingSpec.ALL;
-    }
-  }
 }

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/sql/MSQTaskQueryMaker.java

@@ -285,9 +285,9 @@ public QueryResponse<Object[]> runQuery(final DruidQuery druidQuery)
     MSQTaskQueryMakerUtils.validateRealtimeReindex(querySpec);
 
     final Map<String, Object> context = new HashMap<>();
-    context.put(PlannerContext.CTX_LOOKUP_LOADING_MODE, plannerContext.getLookupLoadingSpec().getMode());
+    context.put(LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE, plannerContext.getLookupLoadingSpec().getMode());
     if (plannerContext.getLookupLoadingSpec().getMode() == LookupLoadingSpec.Mode.ONLY_REQUIRED) {
-      context.put(PlannerContext.CTX_LOOKUPS_TO_LOAD, plannerContext.getLookupLoadingSpec().getLookupsToLoad());
+      context.put(LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD, plannerContext.getLookupLoadingSpec().getLookupsToLoad());
     }
 
     final MSQControllerTask controllerTask = new MSQControllerTask(

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQControllerTaskTest.java

@@ -38,7 +38,6 @@
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
 import org.apache.druid.sql.calcite.planner.ColumnMapping;
 import org.apache.druid.sql.calcite.planner.ColumnMappings;
-import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.joda.time.Interval;
 import org.junit.Assert;
 import org.junit.Test;
@@ -114,8 +113,8 @@ public void testGetLookupLoadingSpecUsingLookupLoadingInfoInContext()
                    .dataSource("target")
                    .context(
                        ImmutableMap.of(
-                           PlannerContext.CTX_LOOKUPS_TO_LOAD, Arrays.asList("lookupName1", "lookupName2"),
-                           PlannerContext.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.ONLY_REQUIRED)
+                           LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD, Arrays.asList("lookupName1", "lookupName2"),
+                           LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.ONLY_REQUIRED)
                    )
                    .build()
         )

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQWorkerTaskTest.java

@@ -23,7 +23,6 @@
 import com.google.common.collect.ImmutableSet;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
-import org.apache.druid.sql.calcite.planner.PlannerContext;
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -125,7 +124,7 @@ public void testGetDefaultLookupLoadingSpec()
   @Test
   public void testGetLookupLoadingWithModeNoneInContext()
   {
-    final ImmutableMap<String, Object> context = ImmutableMap.of(PlannerContext.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.NONE);
+    final ImmutableMap<String, Object> context = ImmutableMap.of(LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.NONE);
     MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId, dataSource, workerNumber, context, retryCount);
     Assert.assertEquals(LookupLoadingSpec.NONE, msqWorkerTask.getLookupLoadingSpec());
   }
@@ -134,8 +133,8 @@ public void testGetLookupLoadingWithModeNoneInContext()
   public void testGetLookupLoadingSpecWithLookupListInContext()
   {
     final ImmutableMap<String, Object> context = ImmutableMap.of(
-        PlannerContext.CTX_LOOKUPS_TO_LOAD, Arrays.asList("lookupName1", "lookupName2"),
-        PlannerContext.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.ONLY_REQUIRED);
+        LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD, Arrays.asList("lookupName1", "lookupName2"),
+        LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.ONLY_REQUIRED);
     MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId, dataSource, workerNumber, context, retryCount);
     Assert.assertEquals(LookupLoadingSpec.Mode.ONLY_REQUIRED, msqWorkerTask.getLookupLoadingSpec().getMode());
     Assert.assertEquals(ImmutableSet.of("lookupName1", "lookupName2"), msqWorkerTask.getLookupLoadingSpec().getLookupsToLoad());
@@ -145,10 +144,10 @@ public void testGetLookupLoadingSpecWithLookupListInContext()
   public void testGetLookupLoadingSpecWithInvalidInput()
   {
     final HashMap<String, Object> context = new HashMap<>();
-    context.put(PlannerContext.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.ONLY_REQUIRED);
+    context.put(LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.ONLY_REQUIRED);
 
     // Setting CTX_LOOKUPS_TO_LOAD as null
-    context.put(PlannerContext.CTX_LOOKUPS_TO_LOAD, null);
+    context.put(LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD, null);
 
     MSQWorkerTask taskWithNullLookups = new MSQWorkerTask(controllerTaskId, dataSource, workerNumber, context, retryCount);
     DruidException exception = Assert.assertThrows(
@@ -160,7 +159,7 @@ public void testGetLookupLoadingSpecWithInvalidInput()
         exception.getMessage());
 
     // Setting CTX_LOOKUPS_TO_LOAD as empty list
-    context.put(PlannerContext.CTX_LOOKUPS_TO_LOAD, Collections.emptyList());
+    context.put(LookupLoadingSpec.CTX_LOOKUPS_TO_LOAD, Collections.emptyList());
 
     MSQWorkerTask taskWithEmptyLookups = new MSQWorkerTask(controllerTaskId, dataSource, workerNumber, context, retryCount);
     exception = Assert.assertThrows(

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java

@@ -854,7 +854,7 @@ public abstract class MSQTester<Builder extends MSQTester<Builder>>
     protected CompactionState expectedLastCompactionState = null;
     protected Set<Interval> expectedTombstoneIntervals = null;
     protected List<Object[]> expectedResultRows = null;
-    protected LookupLoadingSpec expectedLookupLoadingSpec = null;
+    protected LookupLoadingSpec expectedLookupLoadingSpec = LookupLoadingSpec.NONE;
     protected Matcher<Throwable> expectedValidationErrorMatcher = null;
     protected List<Pair<Predicate<MSQTaskReportPayload>, String>> adhocReportAssertionAndReasons = new ArrayList<>();
     protected Matcher<Throwable> expectedExecutionErrorMatcher = null;
@@ -1020,19 +1020,8 @@ public void verifyPlanningErrors()
 
     protected void verifyLookupLoadingInfoInTaskContext(Map<String, Object> context)
     {
-      String lookupLoadingMode = context.get(PlannerContext.CTX_LOOKUP_LOADING_MODE).toString();
-      List<String> lookupsToLoad = (List<String>) context.get(PlannerContext.CTX_LOOKUPS_TO_LOAD);
-      if (expectedLookupLoadingSpec != null) {
-        Assert.assertEquals(expectedLookupLoadingSpec.getMode().toString(), lookupLoadingMode);
-        if (expectedLookupLoadingSpec.getMode().equals(LookupLoadingSpec.Mode.ONLY_REQUIRED)) {
-          Assert.assertEquals(new ArrayList<>(expectedLookupLoadingSpec.getLookupsToLoad()), lookupsToLoad);
-        } else {
-          Assert.assertNull(lookupsToLoad);
-        }
-      } else {
-        Assert.assertEquals(LookupLoadingSpec.Mode.NONE.toString(), lookupLoadingMode);
-        Assert.assertNull(lookupsToLoad);
-      }
+      LookupLoadingSpec specFromContext = LookupLoadingSpec.createFromContext(context, LookupLoadingSpec.ALL);
+      Assert.assertEquals(expectedLookupLoadingSpec, specFromContext);
     }
 
     protected void verifyWorkerCount(CounterSnapshotsTree counterSnapshotsTree)

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

@@ -92,6 +92,7 @@
 import org.apache.druid.segment.transform.TransformSpec;
 import org.apache.druid.segment.writeout.SegmentWriteOutMediumFactory;
 import org.apache.druid.server.coordinator.duty.CompactSegments;
+import org.apache.druid.server.lookup.cache.LookupLoadingSpec;
 import org.apache.druid.server.security.ResourceAction;
 import org.apache.druid.timeline.DataSegment;
 import org.apache.druid.timeline.SegmentTimeline;
@@ -249,6 +250,14 @@ public CompactionTask(
     this.segmentProvider = new SegmentProvider(dataSource, this.ioConfig.getInputSpec());
     this.partitionConfigurationManager = new PartitionConfigurationManager(this.tuningConfig);
     this.segmentCacheManagerFactory = segmentCacheManagerFactory;
+
+    // Do not load any lookups in sub-tasks launched by compaction task, unless transformSpec is present.
+    // If transformSpec is present, we will not modify the context so that the sub-tasks can make the
+    // decision based on context values, loading all lookups by default.
+    // This is done to ensure backward compatibility since transformSpec can reference lookups.
+    if (transformSpec == null) {
+      addToContextIfAbsent(LookupLoadingSpec.CTX_LOOKUP_LOADING_MODE, LookupLoadingSpec.Mode.NONE.toString());
+    }
   }
 
   @VisibleForTesting

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

@@ -334,9 +334,14 @@ static TaskInfo<TaskIdentifier, TaskStatus> toTaskIdentifierInfo(TaskInfo<Task,
     );
   }
 
+  /**
+   * Specifies the list of lookups to load for this task. Tasks load ALL lookups by default.
+   * This behaviour can be overridden by passing parameters {@link LookupLoadingSpec#CTX_LOOKUP_LOADING_MODE}
+   * and {@link LookupLoadingSpec#CTX_LOOKUPS_TO_LOAD} in the task context.
+   */
   @Nullable
   default LookupLoadingSpec getLookupLoadingSpec()
   {
-    return LookupLoadingSpec.ALL;
+    return LookupLoadingSpec.createFromContext(getContext(), LookupLoadingSpec.ALL);
   }
 }