aboutsummaryrefslogtreecommitdiff
path: root/exec/java-exec/src/main/java/org/apache/drill/exec/planner/fragment/Wrapper.java
blob: 6da8850447ad0265dab32dc6142ea3c86b1165a5 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
/*
 * 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 java.util.List;

import org.apache.drill.shaded.guava.com.google.common.collect.ImmutableList;
import org.apache.drill.exec.physical.PhysicalOperatorSetupException;
import org.apache.drill.exec.physical.base.AbstractPhysicalVisitor;
import org.apache.drill.exec.physical.base.Exchange;
import org.apache.drill.exec.physical.base.GroupScan;
import org.apache.drill.exec.physical.base.PhysicalOperator;
import org.apache.drill.exec.physical.base.Store;
import org.apache.drill.exec.physical.base.SubScan;
import org.apache.drill.exec.planner.fragment.Fragment.ExchangeFragmentPair;
import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;

import org.apache.drill.shaded.guava.com.google.common.base.Preconditions;
import org.apache.drill.shaded.guava.com.google.common.collect.Lists;

/**
 * A wrapping class that allows us to add additional information to each fragment node for planning purposes.
 */
public class Wrapper {
  static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(Wrapper.class);

  private final Fragment node;
  private final int majorFragmentId;
  private int width = -1;
  private final Stats stats;
  private boolean endpointsAssigned;
  private long initialAllocation = 0;
  private long maxAllocation = 0;

  // List of fragments this particular fragment depends on for determining its parallelization and endpoint assignments.
  private final List<Wrapper> fragmentDependencies = Lists.newArrayList();

  // a list of assigned endpoints. Technically, there could repeated endpoints in this list if we'd like to assign the
  // same fragment multiple times to the same endpoint.
  private final List<DrillbitEndpoint> endpoints = Lists.newLinkedList();

  public Wrapper(Fragment node, int majorFragmentId) {
    this.majorFragmentId = majorFragmentId;
    this.node = node;
    this.stats = new Stats();
  }

  public Stats getStats() {
    return stats;
  }

  public void resetAllocation() {
    initialAllocation = 0;
    maxAllocation = 0;
  }

  public int getMajorFragmentId() {
    return majorFragmentId;
  }

  public int getWidth() {
    return width;
  }

  public void setWidth(int width) {
    Preconditions.checkState(this.width == -1);
    this.width = width;
  }

  public Fragment getNode() {
    return node;
  }

  public long getInitialAllocation() {
    return initialAllocation;
  }

  public long getMaxAllocation() {
    return maxAllocation;
  }

  public void addAllocation(PhysicalOperator pop) {
    initialAllocation += pop.getInitialAllocation();
//    logger.debug("Incrementing initialAllocation by {} to {}. Pop: {}", pop.getInitialAllocation(), initialAllocation, pop);
    maxAllocation += pop.getMaxAllocation();
  }

  private class AssignEndpointsToScanAndStore extends AbstractPhysicalVisitor<Void, List<DrillbitEndpoint>, PhysicalOperatorSetupException>{

    @Override
    public Void visitExchange(Exchange exchange, List<DrillbitEndpoint> value) throws PhysicalOperatorSetupException {
      if(exchange == node.getSendingExchange()){
        return visitOp(exchange, value);
      }
      // stop on receiver exchange.
      return null;
    }

    @Override
    public Void visitGroupScan(GroupScan groupScan, List<DrillbitEndpoint> value) throws PhysicalOperatorSetupException {
      groupScan.applyAssignments(value);
      return super.visitGroupScan(groupScan, value);
    }

    @Override
    public Void visitSubScan(SubScan subScan, List<DrillbitEndpoint> value) throws PhysicalOperatorSetupException {
      // TODO - implement this
      return visitOp(subScan, value);
    }

    @Override
    public Void visitStore(Store store, List<DrillbitEndpoint> value) throws PhysicalOperatorSetupException {
      store.applyAssignments(value);
      return super.visitStore(store, value);
    }

    @Override
    public Void visitOp(PhysicalOperator op, List<DrillbitEndpoint> value) throws PhysicalOperatorSetupException {
      return visitChildren(op, value);
    }

  }

  public void assignEndpoints(List<DrillbitEndpoint> assignedEndpoints) throws
      PhysicalOperatorSetupException {
    Preconditions.checkState(!endpointsAssigned);
    endpointsAssigned = true;

    endpoints.addAll(assignedEndpoints);

    // Set scan and store endpoints.
    AssignEndpointsToScanAndStore visitor = new AssignEndpointsToScanAndStore();
    node.getRoot().accept(visitor, endpoints);

    // Set the endpoints for this (one at most) sending exchange.
    if (node.getSendingExchange() != null) {
      node.getSendingExchange().setupSenders(majorFragmentId, endpoints);
    }

    // Set the endpoints for each incoming exchange within this fragment.
    for (ExchangeFragmentPair e : node.getReceivingExchangePairs()) {
      e.getExchange().setupReceivers(majorFragmentId, endpoints);
    }
  }

  @Override
  public String toString() {
    return "FragmentWrapper [majorFragmentId=" + majorFragmentId + ", width=" + width + ", stats=" + stats + "]";
  }

  public List<DrillbitEndpoint> getAssignedEndpoints() {
    Preconditions.checkState(endpointsAssigned);
    return ImmutableList.copyOf(endpoints);
  }

  public DrillbitEndpoint getAssignedEndpoint(int minorFragmentId) {
    Preconditions.checkState(endpointsAssigned);
    return endpoints.get(minorFragmentId);
  }

  /**
   * Add a parallelization dependency on given fragment.
   *
   * @param dependsOn
   */
  public void addFragmentDependency(Wrapper dependsOn) {
    fragmentDependencies.add(dependsOn);
  }

  /**
   * Is the endpoints assignment done for this fragment?
   * @return True if the endpoints assignment done for this fragment. False otherwise.
   */
  public boolean isEndpointsAssignmentDone() {
    return endpointsAssigned;
  }

  /**
   * Get the list of fragements this particular fragment depends on.
   * @return The list of fragements this particular fragment depends on.
   */
  public List<Wrapper> getFragmentDependencies() {
    return ImmutableList.copyOf(fragmentDependencies);
  }
}