somu-imply commented on code in PR #15470: URL: https://github.com/apache/druid/pull/15470#discussion_r1464347067
########## extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java: ########## @@ -0,0 +1,346 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.druid.msq.querykit; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.frame.key.ClusterBy; +import org.apache.druid.frame.key.KeyColumn; +import org.apache.druid.frame.key.KeyOrder; +import org.apache.druid.java.util.common.ISE; +import org.apache.druid.java.util.common.Pair; +import org.apache.druid.msq.input.stage.StageInputSpec; +import org.apache.druid.msq.kernel.HashShuffleSpec; +import org.apache.druid.msq.kernel.QueryDefinition; +import org.apache.druid.msq.kernel.QueryDefinitionBuilder; +import org.apache.druid.msq.kernel.ShuffleSpec; +import org.apache.druid.msq.kernel.StageDefinition; +import org.apache.druid.msq.querykit.scan.ScanQueryFrameProcessorFactory; +import org.apache.druid.query.Query; +import org.apache.druid.query.operator.ColumnWithDirection; +import org.apache.druid.query.operator.NaivePartitioningOperatorFactory; +import org.apache.druid.query.operator.NaiveSortOperatorFactory; +import org.apache.druid.query.operator.OperatorFactory; +import org.apache.druid.query.operator.ScanOperatorFactory; +import org.apache.druid.query.operator.WindowOperatorQuery; +import org.apache.druid.query.operator.window.WindowOperatorFactory; +import org.apache.druid.query.scan.ScanQuery; +import org.apache.druid.segment.VirtualColumns; +import org.apache.druid.segment.column.ColumnCapabilities; +import org.apache.druid.segment.column.RowSignature; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +public class WindowOperatorQueryKit implements QueryKit<WindowOperatorQuery> +{ + private final ObjectMapper jsonMapper; + public static final String CTX_SCAN_SIGNATURE = "scanSignature"; + + public WindowOperatorQueryKit(ObjectMapper jsonMapper) + { + this.jsonMapper = jsonMapper; + } + + @Override + public QueryDefinition makeQueryDefinition( + String queryId, + WindowOperatorQuery originalQuery, + QueryKit<Query<?>> queryKit, + ShuffleSpecFactory resultShuffleSpecFactory, + int maxWorkerCount, + int minStageNumber + ) + { + // need to validate query first + // populate the group of operators to be processed as each stage + // the size of the operators is the number of serialized stages + // later we should also check if these can be parallelized + // check there is an empty over clause or not + List<List<OperatorFactory>> operatorList = new ArrayList<>(); + boolean status = validateAndReturnOperatorList(originalQuery, operatorList); + + + ShuffleSpec nextShuffleSpec = findShuffleSpecForNextWindow(operatorList.get(0), maxWorkerCount); + // add this shuffle spec to the last stage of the inner query + + final QueryDefinitionBuilder queryDefBuilder = QueryDefinition.builder().queryId(queryId); + if (nextShuffleSpec != null) { + final ClusterBy windowClusterBy = nextShuffleSpec.clusterBy(); + originalQuery = (WindowOperatorQuery) originalQuery.withOverriddenContext(ImmutableMap.of( + DataSourcePlan.NEXT_SHUFFLE_COL, + windowClusterBy + )); + } else { + nextShuffleSpec = ShuffleSpecFactories.singlePartition() + .build(ClusterBy.none(), false); + } + final DataSourcePlan dataSourcePlan = DataSourcePlan.forDataSource( + queryKit, + queryId, + originalQuery.context(), + originalQuery.getDataSource(), + originalQuery.getQuerySegmentSpec(), + originalQuery.getFilter(), + null, + maxWorkerCount, + minStageNumber, + false + ); + + // Handling for case where leaf operators are present + // leaf operators can only be a scan, + // and they execute before any other operators, + // so we introduce a separate stage for the leaf operators. + // We recreate scan query from leaf operator and use the associated + // ScanQueryFrameProcessorFactory for that particular stage. + // The next stage will work on the output of this stage + if (!originalQuery.getLeafOperators().isEmpty()) { + final RowSignature originalSign = originalQuery.getRowSignature(); + for (OperatorFactory of : originalQuery.getLeafOperators()) { Review Comment: This part of leaf operators still needs work: 1. It has only access to column names and need the column types which are not available 2. Since this runs at the controller, we cannot get the storage adapter. Also it is illogical to get the column type by going over all segments 3. Probably a better way is to include the data types in the leaf operator at the time of planning and pass through to use here -- 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]
