HanumathRao commented on a change in pull request #1677: DRILL-7068: Support memory adjustment framework for resource manageme… URL: https://github.com/apache/drill/pull/1677#discussion_r263651133
########## File path: exec/java-exec/src/main/java/org/apache/drill/exec/planner/fragment/MemoryCalculator.java ########## @@ -0,0 +1,160 @@ +/* + * 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.drill.exec.planner.fragment; + +import org.apache.calcite.util.Pair; +import org.apache.drill.exec.ops.QueryContext; +import org.apache.drill.exec.physical.base.Exchange; +import org.apache.drill.exec.physical.base.PhysicalOperator; +import org.apache.drill.exec.physical.config.AbstractMuxExchange; +import org.apache.drill.exec.planner.AbstractOpWrapperVisitor; +import org.apache.drill.exec.planner.cost.NodeResource; +import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint; +import org.apache.drill.shaded.guava.com.google.common.base.Preconditions; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.HashMap; +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * A visitor to compute memory requirements for each operator in a minor fragment. + * This visitor will be called for each major fragment. It traverses the physical operators + * in that major fragment and computes the memory for each operator per each minor fragment. + * The minor fragment memory resources are further grouped into per Drillbit resource + * requirements. + */ +public class MemoryCalculator extends AbstractOpWrapperVisitor<Void, RuntimeException> { + + private final PlanningSet planningSet; + // List of all the buffered operators and their memory requirement per drillbit. + private final Map<DrillbitEndpoint, List<Pair<PhysicalOperator, Long>>> bufferedOperators; + private final QueryContext queryContext; + + public MemoryCalculator(PlanningSet planningSet, QueryContext context) { + this.planningSet = planningSet; + this.bufferedOperators = new HashMap<>(); + this.queryContext = context; + } + + // Helper method to compute the minor fragment count per drillbit. This method returns + // a map with key as DrillbitEndpoint and value as the width (i.e #minorFragments) + // per Drillbit. + private Map<DrillbitEndpoint, Integer> getMinorFragCountPerDrillbit(Wrapper currFragment) { + return currFragment.getAssignedEndpoints().stream() + .collect(Collectors.groupingBy(Function.identity(), + Collectors.summingInt(x -> 1))); + } + + // Helper method to merge the memory computations for each operator given memory per operator + // and the number of minor fragments per Drillbit. + private void merge(Wrapper currFrag, + Map<DrillbitEndpoint, Integer> minorFragsPerDrillBit, + Function<Entry<DrillbitEndpoint, Integer>, Long> getMemory) { + + NodeResource.merge(currFrag.getResourceMap(), + minorFragsPerDrillBit.entrySet() + .stream() + .collect(Collectors.toMap((x) -> x.getKey(), + (x) -> NodeResource.create(0, + getMemory.apply(x))))); + } + + @Override + public Void visitSendingExchange(Exchange exchange, Wrapper fragment) throws RuntimeException { + Wrapper receivingFragment = planningSet.get(fragment.getNode().getSendingExchangePair().getNode()); + merge(fragment, + getMinorFragCountPerDrillbit(fragment), + // get the memory requirements for the sender operator. + (x) -> exchange.getSenderMemory(receivingFragment.getWidth(), x.getValue())); + return visitOp(exchange, fragment); + } + + @Override + public Void visitReceivingExchange(Exchange exchange, Wrapper fragment) throws RuntimeException { + + final List<Fragment.ExchangeFragmentPair> receivingExchangePairs = fragment.getNode().getReceivingExchangePairs(); + final Map<DrillbitEndpoint, Integer> sendingFragsPerDrillBit = new HashMap<>(); + + for(Fragment.ExchangeFragmentPair pair : receivingExchangePairs) { + if (pair.getExchange() == exchange) { + Wrapper sendingFragment = planningSet.get(pair.getNode()); + Preconditions.checkArgument(sendingFragment.isEndpointsAssignmentDone()); + for (DrillbitEndpoint endpoint : sendingFragment.getAssignedEndpoints()) { + sendingFragsPerDrillBit.putIfAbsent(endpoint, 0); + sendingFragsPerDrillBit.put(endpoint, sendingFragsPerDrillBit.get(endpoint)+1); + } + } + } + final int totalSendingFrags = sendingFragsPerDrillBit.entrySet().stream() + .mapToInt((x) -> x.getValue()).reduce(0, (x, y) -> x+y); + merge(fragment, + getMinorFragCountPerDrillbit(fragment), + (x) -> exchange.getReceiverMemory(fragment.getWidth(), + // If the exchange is a MuxExchange then the sending fragments are from that particular drillbit otherwise + // sending fragments are from across the cluster. + exchange instanceof AbstractMuxExchange ? sendingFragsPerDrillBit.get(x.getKey()) : totalSendingFrags)); + return null; + } + + public List<Pair<PhysicalOperator, Long>> getBufferedOperators(DrillbitEndpoint endpoint) { + return this.bufferedOperators.getOrDefault(endpoint, new ArrayList<>()); + } + + @Override + public Void visitOp(PhysicalOperator op, Wrapper fragment) { + long memoryCost = (int)Math.ceil(op.getCost().getMemoryCost()); + if (op.isBufferedOperator(queryContext)) { + // If the operator is a buffered operator then get the memory estimates of the optimizer. + // The memory estimates of the optimizer are for the whole operator spread across all the + // minor fragments. Divide this memory estimation by fragment width to get the memory + // requirement per minor fragment. + long memoryCostPerMinorFrag = (int)Math.ceil(memoryCost/fragment.getWidth()); Review comment: Currently, both should return same value. However, I do agree if someone changes underneath in future, then it makes sense to use fragment.getAssignedEndpoints().size(). Done. ---------------------------------------------------------------- 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. For queries about this service, please contact Infrastructure at: [email protected] With regards, Apache Git Services
