Address review comment: Pass boolean flag instead of QueryContext fro…

…m ControllerImpl
apache · Nov 11, 2024 · 1244cd2 · 1244cd2
1 parent 46a18a4
commit 1244cd2
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/...nsions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java b/...nsions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/ControllerImpl.java
@@ -1217,7 +1217,13 @@ private QueryKit<Query<?>> makeQueryControllerToolKit(QueryContext queryContext)
         ImmutableMap.<Class<? extends Query>, QueryKit>builder()
                     .put(ScanQuery.class, new ScanQueryKit(context.jsonMapper()))
                     .put(GroupByQuery.class, new GroupByQueryKit(context.jsonMapper()))
-                    .put(WindowOperatorQuery.class, new WindowOperatorQueryKit(context.jsonMapper(), queryContext))
+                    .put(
+                        WindowOperatorQuery.class,
+                        new WindowOperatorQueryKit(
+                            context.jsonMapper(),
+                            MultiStageQueryContext.isWindowFunctionOperatorTransformationEnabled(queryContext)
+                        )
+                    )
                     .build();
 
     return new MultiQueryKit(kitMap);

diff --git a/...multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java b/...multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java
@@ -34,7 +34,6 @@
 import org.apache.druid.msq.kernel.ShuffleSpec;
 import org.apache.druid.msq.kernel.StageDefinition;
 import org.apache.druid.msq.util.MultiStageQueryContext;
-import org.apache.druid.query.QueryContext;
 import org.apache.druid.query.operator.AbstractPartitioningOperatorFactory;
 import org.apache.druid.query.operator.AbstractSortOperatorFactory;
 import org.apache.druid.query.operator.ColumnWithDirection;
@@ -57,12 +56,12 @@ public class WindowOperatorQueryKit implements QueryKit<WindowOperatorQuery>
 {
   private static final Logger log = new Logger(WindowOperatorQueryKit.class);
   private final ObjectMapper jsonMapper;
-  private final QueryContext queryContext;
+  private final boolean isOperatorTransformationEnabled;
 
-  public WindowOperatorQueryKit(ObjectMapper jsonMapper, QueryContext queryContext)
+  public WindowOperatorQueryKit(ObjectMapper jsonMapper, boolean isOperatorTransformationEnabled)
   {
     this.jsonMapper = jsonMapper;
-    this.queryContext = queryContext;
+    this.isOperatorTransformationEnabled = isOperatorTransformationEnabled;
   }
 
   @Override
@@ -177,7 +176,7 @@ public QueryDefinition makeQueryDefinition(
                                                             .flatMap(of -> of.getPartitionColumns().stream())
                                                             .collect(Collectors.toList());
 
-      final List<OperatorFactory> operatorFactories = MultiStageQueryContext.isWindowFunctionOperatorTransformationEnabled(queryContext)
+      final List<OperatorFactory> operatorFactories = isOperatorTransformationEnabled
                                                       ? getTransformedOperatorFactoryListForStageDefinition(operatorList.get(i), maxRowsMaterialized)
                                                       : operatorList.get(i);