kfaraz commented on code in PR #16358:
URL: https://github.com/apache/druid/pull/16358#discussion_r1594881612
##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/MSQWorkerTask.java:
##########
@@ -185,4 +190,26 @@ 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) {
+ List<String> lookupsToLoad = (List<String>)
getContext().get(PlannerContext.CTX_LOOKUPS_TO_LOAD);
Review Comment:
Nit (can be addressed in follow up):
I think it would be safer to read this as a collection instead of a list. It
would also be less confusing since we actually write a set into the context but
then read back a list.
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQWorkerTaskTest.java:
##########
@@ -108,4 +114,61 @@ public void testGetInputSourceResources()
MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
Assert.assertTrue(msqWorkerTask.getInputSourceResources().isEmpty());
}
+
+ @Test
+ public void testGetDefaultLookupLoadingSpec()
+ {
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ Assert.assertEquals(LookupLoadingSpec.ALL,
msqWorkerTask.getLookupLoadingSpec());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingLookupLoadingModeNoneInContext()
+ {
+ final ImmutableMap<String, Object> context =
ImmutableMap.of(PlannerContext.CTX_LOOKUP_LOADING_MODE,
LookupLoadingSpec.Mode.NONE);
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ Assert.assertEquals(LookupLoadingSpec.NONE,
msqWorkerTask.getLookupLoadingSpec());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingNonEmptyLookupListInContext()
+ {
+ 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);
+ 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());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingInvalidInput()
+ {
+ final HashMap<String, Object> context = new HashMap<>();
+ context.put(PlannerContext.CTX_LOOKUP_LOADING_MODE,
LookupLoadingSpec.Mode.ONLY_REQUIRED);
+
+ // Setting CTX_LOOKUPS_TO_LOAD as null
+ context.put(PlannerContext.CTX_LOOKUPS_TO_LOAD, null);
+
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ DruidException exception = Assert.assertThrows(
+ DruidException.class,
+ msqWorkerTask::getLookupLoadingSpec
+ );
+ Assert.assertEquals(
+ "Set of lookups to load cannot be null for mode[ONLY_REQUIRED].",
+ exception.getMessage());
+
+ // Setting CTX_LOOKUPS_TO_LOAD as empty list
+ context.put(PlannerContext.CTX_LOOKUPS_TO_LOAD, Collections.emptyList());
+
+ msqWorkerTask = new MSQWorkerTask(controllerTaskId, dataSource,
workerNumber, context, retryCount);
Review Comment:
Nit: Might be cleaner to just use a new variable rather than the same
`msqWorkerTask`.
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQWorkerTaskTest.java:
##########
@@ -108,4 +114,61 @@ public void testGetInputSourceResources()
MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
Assert.assertTrue(msqWorkerTask.getInputSourceResources().isEmpty());
}
+
+ @Test
+ public void testGetDefaultLookupLoadingSpec()
+ {
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ Assert.assertEquals(LookupLoadingSpec.ALL,
msqWorkerTask.getLookupLoadingSpec());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingLookupLoadingModeNoneInContext()
Review Comment:
It's nice to have descriptive method names but they should still try to be
concise. It is okay to get rid of implicit/obvious words from the names.
```suggestion
public void testGetLookupLoadingWithModeInContext()
```
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQWorkerTaskTest.java:
##########
@@ -108,4 +114,61 @@ public void testGetInputSourceResources()
MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
Assert.assertTrue(msqWorkerTask.getInputSourceResources().isEmpty());
}
+
+ @Test
+ public void testGetDefaultLookupLoadingSpec()
+ {
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ Assert.assertEquals(LookupLoadingSpec.ALL,
msqWorkerTask.getLookupLoadingSpec());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingLookupLoadingModeNoneInContext()
+ {
+ final ImmutableMap<String, Object> context =
ImmutableMap.of(PlannerContext.CTX_LOOKUP_LOADING_MODE,
LookupLoadingSpec.Mode.NONE);
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ Assert.assertEquals(LookupLoadingSpec.NONE,
msqWorkerTask.getLookupLoadingSpec());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingNonEmptyLookupListInContext()
Review Comment:
```suggestion
public void testGetLookupLoadingSpecWithLookupListInContext()
```
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java:
##########
@@ -807,12 +809,34 @@ private MSQErrorReport getErrorReportOrThrow(String
controllerTaskId)
private void assertMSQSpec(MSQSpec expectedMSQSpec, MSQSpec querySpecForTask)
{
- Assert.assertEquals(expectedMSQSpec.getQuery(),
querySpecForTask.getQuery());
+ assertMSQSpecQuery(expectedMSQSpec.getQuery(),
querySpecForTask.getQuery());
Assert.assertEquals(expectedMSQSpec.getAssignmentStrategy(),
querySpecForTask.getAssignmentStrategy());
Assert.assertEquals(expectedMSQSpec.getColumnMappings(),
querySpecForTask.getColumnMappings());
Assert.assertEquals(expectedMSQSpec.getDestination(),
querySpecForTask.getDestination());
}
+ private void assertMSQSpecQuery(Query msqSpecQuery, Query taskSpecQuery)
+ {
+ Assert.assertEquals(msqSpecQuery.getId(), taskSpecQuery.getId());
Review Comment:
Why do we need to verify each field separately?
Does MSQSpec.Query.equals() fail due to different context values?
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java:
##########
@@ -1010,6 +1041,23 @@ public void verifyPlanningErrors()
assertThat(e, expectedValidationErrorMatcher);
}
+ protected void verifyLookupLoadingInfoInTaskContext(Map<String, Object>
context)
Review Comment:
Can this be private?
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java:
##########
@@ -1166,6 +1214,7 @@ public void verifyResults()
verifyCounters(reportPayload.getCounters());
MSQSpec foundSpec =
indexingServiceClient.getMSQControllerTask(controllerId).getQuerySpec();
+
verifyLookupLoadingInfoInTaskContext(indexingServiceClient.getMSQControllerTask(controllerId).getContext());
Review Comment:
The indexing client is already being invoked to get the controller task in
the previous statement. Let's just reuse that value.
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java:
##########
@@ -1010,6 +1041,23 @@ public void verifyPlanningErrors()
assertThat(e, expectedValidationErrorMatcher);
}
+ 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);
+ }
+ }
Review Comment:
Rather than this if-else chain, we can just build a `LookupLoadingSpec`
object from the values in the context and then do an equals check with the
expected lookup loading spec. You would also need to override equals and
hashCode in LookupLoadingSpec.
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java:
##########
@@ -807,12 +809,34 @@ private MSQErrorReport getErrorReportOrThrow(String
controllerTaskId)
private void assertMSQSpec(MSQSpec expectedMSQSpec, MSQSpec querySpecForTask)
{
- Assert.assertEquals(expectedMSQSpec.getQuery(),
querySpecForTask.getQuery());
+ assertMSQSpecQuery(expectedMSQSpec.getQuery(),
querySpecForTask.getQuery());
Assert.assertEquals(expectedMSQSpec.getAssignmentStrategy(),
querySpecForTask.getAssignmentStrategy());
Assert.assertEquals(expectedMSQSpec.getColumnMappings(),
querySpecForTask.getColumnMappings());
Assert.assertEquals(expectedMSQSpec.getDestination(),
querySpecForTask.getDestination());
}
+ private void assertMSQSpecQuery(Query msqSpecQuery, Query taskSpecQuery)
Review Comment:
```suggestion
private void assertMSQSpecQuery(Query expectedQuery, Query taskSpecQuery)
```
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestBase.java:
##########
@@ -1010,6 +1041,23 @@ public void verifyPlanningErrors()
assertThat(e, expectedValidationErrorMatcher);
}
+ protected void verifyLookupLoadingInfoInTaskContext(Map<String, Object>
context)
+ {
+ String lookupLoadingMode =
context.get(PlannerContext.CTX_LOOKUP_LOADING_MODE).toString();
Review Comment:
Can this never be null? Will the context always have the lookup loading mode
set?
##########
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQWorkerTaskTest.java:
##########
@@ -108,4 +114,61 @@ public void testGetInputSourceResources()
MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
Assert.assertTrue(msqWorkerTask.getInputSourceResources().isEmpty());
}
+
+ @Test
+ public void testGetDefaultLookupLoadingSpec()
+ {
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ Assert.assertEquals(LookupLoadingSpec.ALL,
msqWorkerTask.getLookupLoadingSpec());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingLookupLoadingModeNoneInContext()
+ {
+ final ImmutableMap<String, Object> context =
ImmutableMap.of(PlannerContext.CTX_LOOKUP_LOADING_MODE,
LookupLoadingSpec.Mode.NONE);
+ MSQWorkerTask msqWorkerTask = new MSQWorkerTask(controllerTaskId,
dataSource, workerNumber, context, retryCount);
+ Assert.assertEquals(LookupLoadingSpec.NONE,
msqWorkerTask.getLookupLoadingSpec());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingNonEmptyLookupListInContext()
+ {
+ 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);
+ 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());
+ }
+
+ @Test
+ public void testGetLookupLoadingSpecUsingInvalidInput()
Review Comment:
```suggestion
public void testGetLookupLoadingSpecWithInvalidInput()
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
