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 <[email protected]> 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: [email protected] For additional commands, e-mail: [email protected]
