src/org/apache/pig/backend/hadoop/executionengine/tez/plan/operator/POFRJoinTez.java - pig - Git at Google

 /**
  * 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.pig.backend.hadoop.executionengine.tez.plan.operator;

 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.pig.backend.executionengine.ExecException;
 import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POFRJoin;
 import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLocalRearrange;
 import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POMergeJoin.TuplesToSchemaTupleList;
 import org.apache.pig.backend.hadoop.executionengine.tez.runtime.ObjectCache;
 import org.apache.pig.backend.hadoop.executionengine.tez.runtime.TezInput;
 import org.apache.pig.data.SchemaTupleBackend;
 import org.apache.pig.data.SchemaTupleClassGenerator.GenContext;
 import org.apache.pig.data.SchemaTupleFactory;
 import org.apache.pig.data.Tuple;
 import org.apache.pig.impl.io.NullableTuple;
 import org.apache.pig.impl.io.PigNullableWritable;
 import org.apache.pig.impl.logicalLayer.schema.Schema;
 import org.apache.tez.runtime.api.LogicalInput;
 import org.apache.tez.runtime.library.api.KeyValueReader;

 import com.google.common.collect.Lists;

 /**
  * POFRJoinTez is used on the backend to load replicated table from Tez
  * ShuffleUnorderedKVInput and load fragmented table from data pipeline.
  */
 public class POFRJoinTez extends POFRJoin implements TezInput {

     private static final Log log = LogFactory.getLog(POFRJoinTez.class);
     private static final long serialVersionUID = 1L;

     private List<String> inputKeys;

     // For replicated tables
     private transient List<LogicalInput> replInputs;
     private transient List<KeyValueReader> replReaders;
     private transient boolean isInputCached;
     private transient String cacheKey;

     public POFRJoinTez(POFRJoin copy, List<String> inputKeys) throws ExecException {
        super(copy);
        this.inputKeys = inputKeys;
     }

     @Override
     public String[] getTezInputs() {
         return inputKeys.toArray(new String[inputKeys.size()]);
     }

     @Override
     public void replaceInput(String oldInputKey, String newInputKey) {
         while (inputKeys.remove(oldInputKey)) {
             inputKeys.add(newInputKey);
         }
     }

     @Override
     public void addInputsToSkip(Set<String> inputsToSkip) {
         cacheKey = "replicatemap-" + inputKeys.toString();
         Object cacheValue = ObjectCache.getInstance().retrieve(cacheKey);
         if (cacheValue != null) {
             isInputCached = true;
             inputsToSkip.addAll(inputKeys);
         }
     }

     @Override
     public void attachInputs(Map<String, LogicalInput> inputs, Configuration conf)
             throws ExecException {
         if (isInputCached) {
             return;
         }
         try {
             this.replInputs = Lists.newArrayList();
             this.replReaders = Lists.newArrayList();
             for (String key : inputKeys) {
                 LogicalInput input = inputs.get(key);
                 if (!this.replInputs.contains(input)) {
                     this.replInputs.add(input);
                     KeyValueReader reader = (KeyValueReader) input.getReader();
                     this.replReaders.add(reader);
                     log.info("Attached input from vertex " + key + " : input=" + input + ", reader=" + reader);
                 }
             }
             // Do not force fetch input by reading first record. Cases like MultiQuery_Union_4 have
             // multiple POFRJoinTez loading same replicate input and will skip records
         } catch (Exception e) {
             throw new ExecException(e);
         }
     }

     /**
      * Builds the HashMaps by reading replicated inputs from broadcast edges
      *
      * @throws ExecException
      */
     @SuppressWarnings("unchecked")
     @Override
     protected void setUpHashMap() throws ExecException {

         // Re-check again in case of Split + union + replicate join
         // where same POFRJoinTez occurs in different Split sub-plans
         Object cacheValue = ObjectCache.getInstance().retrieve(cacheKey);
         if (cacheValue != null) {
             replicates =  (List<Map<? extends Object, ? extends List<Tuple>>>) cacheValue;
             log.info("Found " + (replicates.size() - 1) + " replication hash tables in Tez cache. cachekey=" + cacheKey);
             return;
         }

         log.info("Building replication hash table");

         SchemaTupleFactory[] inputSchemaTupleFactories = new SchemaTupleFactory[inputSchemas.length];
         SchemaTupleFactory[] keySchemaTupleFactories = new SchemaTupleFactory[inputSchemas.length];

         for (int i = 0; i < inputSchemas.length; i++) {
             Schema schema = inputSchemas[i];
             if (schema != null) {
                 log.debug("Using SchemaTuple for FR Join Schema: " + schema);
                 inputSchemaTupleFactories[i] =
                         SchemaTupleBackend.newSchemaTupleFactory(schema, false, GenContext.FR_JOIN);
             }
             schema = keySchemas[i];
             if (schema != null) {
                 log.debug("Using SchemaTuple for FR Join key Schema: " + schema);
                 keySchemaTupleFactories[i] =
                         SchemaTupleBackend.newSchemaTupleFactory(schema, false, GenContext.FR_JOIN);
             }
         }

         long time1 = System.currentTimeMillis();

         replicates.set(fragment, null);
         int inputIdx = 0;
         // We need to adjust the index because the number of replInputs is
         // one less than the number of inputSchemas. The inputSchemas
         // includes the fragmented table.
         int schemaIdx = 1;
         while (inputIdx < replInputs.size()) {
             SchemaTupleFactory inputSchemaTupleFactory = inputSchemaTupleFactories[schemaIdx];
             SchemaTupleFactory keySchemaTupleFactory = keySchemaTupleFactories[schemaIdx];

             Map<Object, ArrayList<Tuple>> replicate;
             if (keySchemaTupleFactory == null) {
                 replicate = new HashMap<Object, ArrayList<Tuple>>(4000);
             } else {
                 replicate = new TupleToMapKey(4000, keySchemaTupleFactory);
             }
             POLocalRearrange lr = LRs[schemaIdx];

             try {
                 while (replReaders.get(inputIdx).next()) {
                     if (getReporter() != null) {
                         getReporter().progress();
                     }

                     PigNullableWritable key = (PigNullableWritable) replReaders.get(inputIdx).getCurrentKey();
                     Object keyValue = key.getValueAsPigType();
                     if (isKeyNull(keyValue)) continue;
                     NullableTuple val = (NullableTuple) replReaders.get(inputIdx).getCurrentValue();

                     // POFRJoin#getValueTuple() is reused to construct valTuple,
                     // and it takes an indexed Tuple as parameter. So we need to
                     // construct one here.
                     Tuple retTuple = mTupleFactory.newTuple(3);
                     retTuple.set(0, key.getIndex());
                     retTuple.set(1, keyValue);
                     retTuple.set(2, val.getValueAsPigType());
                     Tuple valTuple = getValueTuple(lr, retTuple);

                     ArrayList<Tuple> values = replicate.get(keyValue);
                     if (values == null) {
                         if (inputSchemaTupleFactory == null) {
                             values = new ArrayList<Tuple>(1);
                         } else {
                             values = new TuplesToSchemaTupleList(1, inputSchemaTupleFactory);
                         }
                         replicate.put(keyValue, values);
                     }
                     values.add(valTuple);
                 }
             } catch (IOException e) {
                 throw new ExecException(e);
             }
             replicates.set(schemaIdx, replicate);
             inputIdx++;
             schemaIdx++;
         }

         long time2 = System.currentTimeMillis();
         log.info((replicates.size() - 1) + " replication hash tables built. Time taken: " + (time2 - time1));

         ObjectCache.getInstance().cache(cacheKey, replicates);
         log.info("Cached replicate hash tables in Tez ObjectRegistry with vertex scope. cachekey=" + cacheKey);
     }

     @Override
     public String name() {
         StringBuffer inputs = new StringBuffer();
         for (int i = 0; i < inputKeys.size(); i++) {
             if (i > 0) inputs.append(",");
             inputs.append(inputKeys.get(i));
         }
         return super.name() + "\t<-\t " + inputs.toString();
     }

     public List<String> getInputKeys() {
         return inputKeys;
     }

     @Override
     public POFRJoinTez clone() throws CloneNotSupportedException {
         return (POFRJoinTez) super.clone();
     }

 }
	/**
	* 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.pig.backend.hadoop.executionengine.tez.plan.operator;

	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.hadoop.conf.Configuration;
	import org.apache.pig.backend.executionengine.ExecException;
	import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POFRJoin;
	import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POLocalRearrange;
	import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POMergeJoin.TuplesToSchemaTupleList;
	import org.apache.pig.backend.hadoop.executionengine.tez.runtime.ObjectCache;
	import org.apache.pig.backend.hadoop.executionengine.tez.runtime.TezInput;
	import org.apache.pig.data.SchemaTupleBackend;
	import org.apache.pig.data.SchemaTupleClassGenerator.GenContext;
	import org.apache.pig.data.SchemaTupleFactory;
	import org.apache.pig.data.Tuple;
	import org.apache.pig.impl.io.NullableTuple;
	import org.apache.pig.impl.io.PigNullableWritable;
	import org.apache.pig.impl.logicalLayer.schema.Schema;
	import org.apache.tez.runtime.api.LogicalInput;
	import org.apache.tez.runtime.library.api.KeyValueReader;

	import com.google.common.collect.Lists;

	/**
	* POFRJoinTez is used on the backend to load replicated table from Tez
	* ShuffleUnorderedKVInput and load fragmented table from data pipeline.
	*/
	public class POFRJoinTez extends POFRJoin implements TezInput {

	private static final Log log = LogFactory.getLog(POFRJoinTez.class);
	private static final long serialVersionUID = 1L;

	private List<String> inputKeys;

	// For replicated tables
	private transient List<LogicalInput> replInputs;
	private transient List<KeyValueReader> replReaders;
	private transient boolean isInputCached;
	private transient String cacheKey;

	public POFRJoinTez(POFRJoin copy, List<String> inputKeys) throws ExecException {
	super(copy);
	this.inputKeys = inputKeys;
	}

	@Override
	public String[] getTezInputs() {
	return inputKeys.toArray(new String[inputKeys.size()]);
	}

	@Override
	public void replaceInput(String oldInputKey, String newInputKey) {
	while (inputKeys.remove(oldInputKey)) {
	inputKeys.add(newInputKey);
	}
	}

	@Override
	public void addInputsToSkip(Set<String> inputsToSkip) {
	cacheKey = "replicatemap-" + inputKeys.toString();
	Object cacheValue = ObjectCache.getInstance().retrieve(cacheKey);
	if (cacheValue != null) {
	isInputCached = true;
	inputsToSkip.addAll(inputKeys);
	}
	}

	@Override
	public void attachInputs(Map<String, LogicalInput> inputs, Configuration conf)
	throws ExecException {
	if (isInputCached) {
	return;
	}
	try {
	this.replInputs = Lists.newArrayList();
	this.replReaders = Lists.newArrayList();
	for (String key : inputKeys) {
	LogicalInput input = inputs.get(key);
	if (!this.replInputs.contains(input)) {
	this.replInputs.add(input);
	KeyValueReader reader = (KeyValueReader) input.getReader();
	this.replReaders.add(reader);
	log.info("Attached input from vertex " + key + " : input=" + input + ", reader=" + reader);
	}
	}
	// Do not force fetch input by reading first record. Cases like MultiQuery_Union_4 have
	// multiple POFRJoinTez loading same replicate input and will skip records
	} catch (Exception e) {
	throw new ExecException(e);
	}
	}

	/**
	* Builds the HashMaps by reading replicated inputs from broadcast edges
	*
	* @throws ExecException
	*/
	@SuppressWarnings("unchecked")
	@Override
	protected void setUpHashMap() throws ExecException {

	// Re-check again in case of Split + union + replicate join
	// where same POFRJoinTez occurs in different Split sub-plans
	Object cacheValue = ObjectCache.getInstance().retrieve(cacheKey);
	if (cacheValue != null) {
	replicates = (List<Map<? extends Object, ? extends List<Tuple>>>) cacheValue;
	log.info("Found " + (replicates.size() - 1) + " replication hash tables in Tez cache. cachekey=" + cacheKey);
	return;
	}

	log.info("Building replication hash table");

	SchemaTupleFactory[] inputSchemaTupleFactories = new SchemaTupleFactory[inputSchemas.length];
	SchemaTupleFactory[] keySchemaTupleFactories = new SchemaTupleFactory[inputSchemas.length];

	for (int i = 0; i < inputSchemas.length; i++) {
	Schema schema = inputSchemas[i];
	if (schema != null) {
	log.debug("Using SchemaTuple for FR Join Schema: " + schema);
	inputSchemaTupleFactories[i] =
	SchemaTupleBackend.newSchemaTupleFactory(schema, false, GenContext.FR_JOIN);
	}
	schema = keySchemas[i];
	if (schema != null) {
	log.debug("Using SchemaTuple for FR Join key Schema: " + schema);
	keySchemaTupleFactories[i] =
	SchemaTupleBackend.newSchemaTupleFactory(schema, false, GenContext.FR_JOIN);
	}
	}

	long time1 = System.currentTimeMillis();

	replicates.set(fragment, null);
	int inputIdx = 0;
	// We need to adjust the index because the number of replInputs is
	// one less than the number of inputSchemas. The inputSchemas
	// includes the fragmented table.
	int schemaIdx = 1;
	while (inputIdx < replInputs.size()) {
	SchemaTupleFactory inputSchemaTupleFactory = inputSchemaTupleFactories[schemaIdx];
	SchemaTupleFactory keySchemaTupleFactory = keySchemaTupleFactories[schemaIdx];

	Map<Object, ArrayList<Tuple>> replicate;
	if (keySchemaTupleFactory == null) {
	replicate = new HashMap<Object, ArrayList<Tuple>>(4000);
	} else {
	replicate = new TupleToMapKey(4000, keySchemaTupleFactory);
	}
	POLocalRearrange lr = LRs[schemaIdx];

	try {
	while (replReaders.get(inputIdx).next()) {
	if (getReporter() != null) {
	getReporter().progress();
	}

	PigNullableWritable key = (PigNullableWritable) replReaders.get(inputIdx).getCurrentKey();
	Object keyValue = key.getValueAsPigType();
	if (isKeyNull(keyValue)) continue;
	NullableTuple val = (NullableTuple) replReaders.get(inputIdx).getCurrentValue();

	// POFRJoin#getValueTuple() is reused to construct valTuple,
	// and it takes an indexed Tuple as parameter. So we need to
	// construct one here.
	Tuple retTuple = mTupleFactory.newTuple(3);
	retTuple.set(0, key.getIndex());
	retTuple.set(1, keyValue);
	retTuple.set(2, val.getValueAsPigType());
	Tuple valTuple = getValueTuple(lr, retTuple);

	ArrayList<Tuple> values = replicate.get(keyValue);
	if (values == null) {
	if (inputSchemaTupleFactory == null) {
	values = new ArrayList<Tuple>(1);
	} else {
	values = new TuplesToSchemaTupleList(1, inputSchemaTupleFactory);
	}
	replicate.put(keyValue, values);
	}
	values.add(valTuple);
	}
	} catch (IOException e) {
	throw new ExecException(e);
	}
	replicates.set(schemaIdx, replicate);
	inputIdx++;
	schemaIdx++;
	}

	long time2 = System.currentTimeMillis();
	log.info((replicates.size() - 1) + " replication hash tables built. Time taken: " + (time2 - time1));

	ObjectCache.getInstance().cache(cacheKey, replicates);
	log.info("Cached replicate hash tables in Tez ObjectRegistry with vertex scope. cachekey=" + cacheKey);
	}

	@Override
	public String name() {
	StringBuffer inputs = new StringBuffer();
	for (int i = 0; i < inputKeys.size(); i++) {
	if (i > 0) inputs.append(",");
	inputs.append(inputKeys.get(i));
	}
	return super.name() + "\t<-\t " + inputs.toString();
	}

	public List<String> getInputKeys() {
	return inputKeys;
	}

	@Override
	public POFRJoinTez clone() throws CloneNotSupportedException {
	return (POFRJoinTez) super.clone();
	}

	}