TEZ-873. Expose InputSplit via MRInputLegacy, and underlying splits via
[incubator-tez.git] / tez-runtime-library / src / main / java / org / apache / tez / runtime / library / input / ShuffledUnorderedKVInput.java
1 /**
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements.  See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership.  The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License.  You may obtain a copy of the License at
9 *
10 *     http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing, software
13 * distributed under the License is distributed on an "AS IS" BASIS,
14 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 * See the License for the specific language governing permissions and
16 * limitations under the License.
17 */
18
19 package org.apache.tez.runtime.library.input;
20
21 import java.util.List;
22
23 import org.apache.commons.logging.Log;
24 import org.apache.commons.logging.LogFactory;
25 import org.apache.hadoop.conf.Configuration;
26 import org.apache.tez.common.TezJobConfig;
27 import org.apache.tez.common.TezUtils;
28 import org.apache.tez.runtime.api.Event;
29 import org.apache.tez.runtime.api.LogicalInput;
30 import org.apache.tez.runtime.api.TezInputContext;
31 import org.apache.tez.runtime.library.api.KeyValueReader;
32 import org.apache.tez.runtime.library.broadcast.input.BroadcastKVReader;
33 import org.apache.tez.runtime.library.broadcast.input.BroadcastShuffleManager;
34
35 import com.google.common.base.Preconditions;
36
37 public class ShuffledUnorderedKVInput implements LogicalInput {
38
39   private static final Log LOG = LogFactory.getLog(ShuffledUnorderedKVInput.class);
40   
41   private Configuration conf;
42   private int numInputs = -1;
43   private BroadcastShuffleManager shuffleManager;
44   @SuppressWarnings("rawtypes")
45   private BroadcastKVReader kvReader;
46   
47   
48   
49   public ShuffledUnorderedKVInput() {
50   }
51
52   @SuppressWarnings("rawtypes")
53   @Override
54   public List<Event> initialize(TezInputContext inputContext) throws Exception {
55     Preconditions.checkArgument(numInputs != -1, "Number of Inputs has not been set");
56     this.conf = TezUtils.createConfFromUserPayload(inputContext.getUserPayload());
57     this.conf.setStrings(TezJobConfig.LOCAL_DIRS, inputContext.getWorkDirs());
58     
59     this.shuffleManager = new BroadcastShuffleManager(inputContext, conf, numInputs);
60     this.shuffleManager.run();
61     this.kvReader = new BroadcastKVReader(shuffleManager, conf);
62     return null;
63   }
64
65   @Override
66   public KeyValueReader getReader() throws Exception {
67     return this.kvReader;
68   }
69
70   @Override
71   public void handleEvents(List<Event> inputEvents) {
72     shuffleManager.handleEvents(inputEvents);
73   }
74
75   @Override
76   public List<Event> close() throws Exception {
77     this.shuffleManager.shutdown();
78     return null;
79   }
80
81   @Override
82   public void setNumPhysicalInputs(int numInputs) {
83     this.numInputs = numInputs;
84   }
85 }