(druid) 02/04: MSQ: Fix WindowOperatorQuery planning. (#18911)

[kg...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

kgyrtkirk pushed a commit to branch 36.0.0
in repository https://gitbox.apache.org/repos/asf/druid.git

commit c2ec86e28dfe86bda614f475480b48513014332f
Author: Gian Merlino <gi...@gmail.com>
AuthorDate: Tue Jan 13 13:53:03 2026 -0800

    MSQ: Fix WindowOperatorQuery planning. (#18911)
    
    PR #18875 caused a regression in WindowOperatorQuery planning, by
    causing non-outermost WindowOperatorQuery to be planned without
    windowFunctionOperatorTransformation (see #17443). This happened because
    only the outermost query has the windowFunctionOperatorTransformation
    parameter.
    
    This patch fixes the problem by propagating the context parameter from
    the outermost queries to any inner WindowOperatorQuery instances.
    
    (cherry picked from commit 07de2ff03b1af3a03770365160589fd4549d999b)
---
 .../druid/msq/exec/QueryKitBasedMSQPlanner.java    | 58 +++++++++++++++++++++-
 1 file changed, 57 insertions(+), 1 deletion(-)

diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryKitBasedMSQPlanner.java b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryKitBasedMSQPlanner.java
index cf2ca91c8d1..7ec3f24e3de 100644
--- a/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryKitBasedMSQPlanner.java
+++ b/multi-stage-query/src/main/java/org/apache/druid/msq/exec/QueryKitBasedMSQPlanner.java
@@ -39,17 +39,25 @@ import org.apache.druid.msq.kernel.StageDefinition;
 import org.apache.druid.msq.querykit.QueryKitSpec;
 import org.apache.druid.msq.querykit.QueryKitUtils;
 import org.apache.druid.msq.querykit.ShuffleSpecFactory;
+import org.apache.druid.msq.querykit.WindowOperatorQueryKit;
 import org.apache.druid.msq.querykit.results.ExportResultsStageProcessor;
 import org.apache.druid.msq.querykit.results.QueryResultStageProcessor;
 import org.apache.druid.msq.util.MSQTaskQueryMakerUtils;
 import org.apache.druid.msq.util.MultiStageQueryContext;
+import org.apache.druid.query.DataSource;
 import org.apache.druid.query.Query;
 import org.apache.druid.query.QueryContext;
+import org.apache.druid.query.QueryDataSource;
+import org.apache.druid.query.operator.WindowOperatorQuery;
 import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.sql.calcite.planner.ColumnMappings;
 import org.apache.druid.sql.http.ResultFormat;
 import org.apache.druid.storage.ExportStorageProvider;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
 public class QueryKitBasedMSQPlanner
 {
   private final MSQSpec querySpec;
@@ -122,7 +130,7 @@ public class QueryKitBasedMSQPlanner
     try {
       queryDef = queryKitSpec.getQueryKit().makeQueryDefinition(
           queryKitSpec,
-          queryToPlan,
+          propagateWindowOperatorTransformationContext(queryToPlan),
           makeResultShuffleSpecFacory(),
           0
       );
@@ -227,4 +235,52 @@ public class QueryKitBasedMSQPlanner
       return destination.getShuffleSpecFactory(MultiStageQueryContext.getRowsPerPage(query.context()));
     }
   }
+
+  /**
+   * Propagates {@link MultiStageQueryContext#WINDOW_FUNCTION_OPERATOR_TRANSFORMATION} from the outer query to
+   * any inner {@link WindowOperatorQuery}, so {@link WindowOperatorQueryKit} can see it. The original change
+   * in https://github.com/apache/druid/pull/17443 only added it to the outer query at the Broker, and since
+   * QueryKit became injectable via Guice, {@link WindowOperatorQueryKit} can no longer "see" parameters that
+   * come from the context of the outermost query. For compatibility reasons, this propagation must be done here
+   * rather than during SQL planning.
+   */
+  private static Query<?> propagateWindowOperatorTransformationContext(final Query<?> query)
+  {
+    final Boolean val = query.context().getBoolean(MultiStageQueryContext.WINDOW_FUNCTION_OPERATOR_TRANSFORMATION);
+    if (val != null) {
+      // Update inner queries.
+      return query.withDataSource(setWindowOperatorTransformationContext(query.getDataSource(), val));
+    } else {
+      return query;
+    }
+  }
+
+  /**
+   * Sets {@link MultiStageQueryContext#WINDOW_FUNCTION_OPERATOR_TRANSFORMATION} on all {@link QueryDataSource}
+   * found within the provided {@link DataSource}.
+   */
+  private static DataSource setWindowOperatorTransformationContext(final DataSource dataSource, final boolean val)
+  {
+    final List<DataSource> children = dataSource.getChildren();
+    final List<DataSource> newChildren = new ArrayList<>();
+    for (final DataSource child : children) {
+      newChildren.add(setWindowOperatorTransformationContext(child, val));
+    }
+    if (dataSource instanceof QueryDataSource
+        && ((QueryDataSource) dataSource).getQuery() instanceof WindowOperatorQuery) {
+      return new QueryDataSource(
+          ((QueryDataSource) dataSource)
+              .getQuery()
+              .withOverriddenContext(
+                  Map.of(
+                      MultiStageQueryContext.WINDOW_FUNCTION_OPERATOR_TRANSFORMATION,
+                      val
+                  )
+              )
+              .withDataSource(Iterables.getOnlyElement(newChildren))
+      );
+    } else {
+      return dataSource.withChildren(newChildren);
+    }
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org