src/org/apache/pig/impl/util/UDFContext.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.impl.util;

 import java.io.IOException;
 import java.io.Serializable;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Properties;

 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.MRConfiguration;

 public class UDFContext {

     private static final Log LOG = LogFactory.getLog(UDFContext.class);
     private Configuration jconf = null;
     private HashMap<UDFContextKey, Properties> udfConfs;
     private Properties clientSysProps;

     static final String CLIENT_SYS_PROPS = "pig.client.sys.props";
     static final String UDF_CONTEXT = "pig.udf.context";

     private static ThreadLocal<UDFContext> tss = new ThreadLocal<UDFContext>() {
         @Override
         public UDFContext initialValue() {
             return new UDFContext();
         }
     };

     private UDFContext() {
         udfConfs = new HashMap<UDFContextKey, Properties>();
     }

     /**
      * @return a Thread Local {@link UDFContext}
      */
     public static UDFContext getUDFContext() {
         return tss.get();
     }

     /*
      *  internal pig use only - should NOT be called from user code
      */
     public static void setUdfContext(UDFContext udfContext) {
         tss.set(udfContext);
     }

     /*
      *  internal pig use only - should NOT be called from user code
      */
     //@StaticDataCleanup
     public static void staticDataCleanup() {
         tss = new ThreadLocal<UDFContext>() {
             @Override
             public UDFContext initialValue() {
                 return new UDFContext();
             }
         };
     }

     /*
      *  internal pig use only - should NOT be called from user code
      */
     public HashMap<UDFContextKey, Properties> getUdfConfs() {
         return udfConfs;
     }


     /*
      *  internal pig use only - should NOT be called from user code
      */
     public void setClientSystemProps(Properties properties) {
         clientSysProps = properties;
     }

     /**
      * Get the System Properties (Read only) as on the client machine from where Pig
      * was launched. This will include command line properties passed at launch
      * time
      * @return client side System Properties including command line properties
      */
     public Properties getClientSystemProps() {
         return clientSysProps;
     }
     /**
      * Adds the JobConf to this singleton.  Will be
      * called on the backend by the Map and Reduce
      * functions so that UDFs can obtain the JobConf
      * on the backend.
      */
     public void addJobConf(Configuration conf) {
         jconf = conf;
     }

     /**
      * Get the JobConf.  This should only be called on
      * the backend.  It will return null on the frontend.
      * @return JobConf for this job.  This is a copy of the
      * JobConf.  Nothing written here will be kept by the system.
      * getUDFConf should be used for recording UDF specific
      * information.
      */
     public Configuration getJobConf() {
         if (jconf != null)  return new Configuration(jconf);
         else return null;
     }

     /**
      * Get a properties object that is specific to this UDF.
      * Note that if a given UDF is called multiple times in a script,
      * and each instance passes different arguments, then each will
      * be provided with different configuration object.
      * This can be used by loaders to pass their input object path
      * or URI and separate themselves from other instances of the
      * same loader.  Constructor arguments could also be used,
      * as they are available on both the front and back end.
      *
      * Note that this can only be used to share information
      * across instantiations of the same function in the front end
      * and between front end and back end.  It cannot be used to
      * share information between instantiations (that is, between
      * map and/or reduce instances) on the back end at runtime.
      * @param c of the UDF obtaining the properties object.
      * @param args String arguments that make this instance of
      * the UDF unique.
      * @return A reference to the properties object specific to
      * the calling UDF.  This is a reference, not a copy.
      * Any changes to this object will automatically be
      * propogated to other instances of the UDF calling this
      * function.
      */

     @SuppressWarnings("rawtypes")
     public Properties getUDFProperties(Class c, String[] args) {
         UDFContextKey k = generateKey(c, args);
         Properties p = udfConfs.get(k);
         if (p == null) {
             p = new Properties();
             udfConfs.put(k, p);
         }
         return p;
     }

     /**
      * Get a properties object that is specific to this UDF.
      * Note that if a given UDF is called multiple times in a script,
      * they will all be provided the same configuration object.  It
      * is up to the UDF to make sure the multiple instances do not
      * stomp on each other.
      *
      * It is guaranteed that this properties object will be separate
      * from that provided to any other UDF.
      *
      * Note that this can only be used to share information
      * across instantiations of the same function in the front end
      * and between front end and back end.  It cannot be used to
      * share information between instantiations (that is, between
      * map and/or reduce instances) on the back end at runtime.
      * @param c of the UDF obtaining the properties object.
      * @return A reference to the properties object specific to
      * the calling UDF.  This is a reference, not a copy.
      * Any changes to this object will automatically be
      * propogated to other instances of the UDF calling this
      * function.
      */
     @SuppressWarnings("rawtypes")
     public Properties getUDFProperties(Class c) {
         UDFContextKey k = generateKey(c, null);
         Properties p = udfConfs.get(k);
         if (p == null) {
             p = new Properties();
             udfConfs.put(k, p);
         }
         return p;
     }


     /**
      * Serialize the UDF specific information into an instance
      * of JobConf.  This function is intended to be called on
      * the front end in preparation for sending the data to the
      * backend.
      * @param conf JobConf to serialize into
      * @throws IOException if underlying serialization throws it
      */
     public void serialize(Configuration conf) throws IOException {
         conf.set(UDF_CONTEXT, ObjectSerializer.serialize(udfConfs));
         conf.set(CLIENT_SYS_PROPS, ObjectSerializer.serialize(clientSysProps));
     }

     /*
      * Internal pig use
      */
     public String serialize() {
         try {
             return ObjectSerializer.serialize(udfConfs);
         } catch (IOException e) {
             LOG.error("UDFContext#serialize throws error ",e);
             return null;
         }
     }

     /**
      * Populate the udfConfs field.  This function is intended to
      * be called by Map.configure or Reduce.configure on the backend.
      * It assumes that addJobConf has already been called.
      * @throws IOException if underlying deseralization throws it
      */
     @SuppressWarnings("unchecked")
     public void deserialize() throws IOException {
         udfConfs = (HashMap<UDFContextKey, Properties>)ObjectSerializer.deserialize(jconf.get(UDF_CONTEXT));
         clientSysProps = (Properties)ObjectSerializer.deserialize(
                 jconf.get(CLIENT_SYS_PROPS));
     }

     public void deserializeForSpark(String udfConfsStr, String clientSysPropsStr) throws IOException {
         if( udfConfsStr!= null && clientSysPropsStr!=null) {
             udfConfs = (HashMap<UDFContextKey, Properties>) ObjectSerializer.deserialize(udfConfsStr);
             clientSysProps = (Properties) ObjectSerializer.deserialize(
                     clientSysPropsStr);
         }
     }

     private UDFContextKey generateKey(Class<?> c, String[] args) {
         return new UDFContextKey(c.getName(), args);
     }

     public void reset() {
         udfConfs.clear();
     }

     public boolean isUDFConfEmpty() {
         return udfConfs.isEmpty();
     }

     /**
      * Convenience method for UDF code to check where it runs (see PIG-2576)
      * @return boolean type value
      */
     public boolean isFrontend() {
         // mapred.task.id is for MR1
         // mapreduce.job.application.attempt.id is for MR2
         return (this.jconf == null
                 || (jconf.get(MRConfiguration.TASK_ID) == null &&
                     jconf.get(MRConfiguration.JOB_APPLICATION_ATTEMPT_ID) == null));
     }

     /**
      * Make a shallow copy of the context.
      */
     @Override
     public UDFContext clone() {
     	UDFContext other = new UDFContext();
     	other.clientSysProps = this.clientSysProps;
     	other.jconf = this.jconf;
     	other.udfConfs = this.udfConfs;
     	return other;
     }

     /**
      * Class that acts as key for hashmap in UDFContext,
      *  it holds the class and args of the udf, and
      *  implements equals() and hashCode()
      */
     static class UDFContextKey implements Serializable{

         private static final long serialVersionUID = 1;
         private String className;
         private String[] args;

         UDFContextKey(String className, String [] args){
             this.className = className;
             this.args = args;
         }

         String getClassName() {
             return className;
         }

         String[] getArgs() {
             return args;
         }

         @Override
         public String toString() {
             return "UDFContextKey [className=" + className + ", args="
                     + Arrays.toString(args) + "]";
         }

         @Override
         public int hashCode() {
             final int prime = 31;
             int result = 1;
             result = prime * result + Arrays.hashCode(args);
             result = prime * result
                     + ((className == null) ? 0 : className.hashCode());
             return result;
         }

         @Override
         public boolean equals(Object obj) {
             if (this == obj)
                 return true;
             if (obj == null)
                 return false;
             if (getClass() != obj.getClass())
                 return false;
             UDFContextKey other = (UDFContextKey) obj;
             if (!Arrays.equals(args, other.args))
                 return false;
             if (className == null) {
                 if (other.className != null)
                     return false;
             } else if (!className.equals(other.className))
                 return false;
             return true;
         }
     }

     public Properties getClientSysProps() {
         return clientSysProps;
     }

 }
	/*
	* 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.impl.util;

	import java.io.IOException;
	import java.io.Serializable;
	import java.util.Arrays;
	import java.util.HashMap;
	import java.util.Properties;

	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.hadoop.conf.Configuration;
	import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.MRConfiguration;

	public class UDFContext {

	private static final Log LOG = LogFactory.getLog(UDFContext.class);
	private Configuration jconf = null;
	private HashMap<UDFContextKey, Properties> udfConfs;
	private Properties clientSysProps;

	static final String CLIENT_SYS_PROPS = "pig.client.sys.props";
	static final String UDF_CONTEXT = "pig.udf.context";

	private static ThreadLocal<UDFContext> tss = new ThreadLocal<UDFContext>() {
	@Override
	public UDFContext initialValue() {
	return new UDFContext();
	}
	};

	private UDFContext() {
	udfConfs = new HashMap<UDFContextKey, Properties>();
	}

	/**
	* @return a Thread Local {@link UDFContext}
	*/
	public static UDFContext getUDFContext() {
	return tss.get();
	}

	/*
	* internal pig use only - should NOT be called from user code
	*/
	public static void setUdfContext(UDFContext udfContext) {
	tss.set(udfContext);
	}

	/*
	* internal pig use only - should NOT be called from user code
	*/
	//@StaticDataCleanup
	public static void staticDataCleanup() {
	tss = new ThreadLocal<UDFContext>() {
	@Override
	public UDFContext initialValue() {
	return new UDFContext();
	}
	};
	}

	/*
	* internal pig use only - should NOT be called from user code
	*/
	public HashMap<UDFContextKey, Properties> getUdfConfs() {
	return udfConfs;
	}


	/*
	* internal pig use only - should NOT be called from user code
	*/
	public void setClientSystemProps(Properties properties) {
	clientSysProps = properties;
	}

	/**
	* Get the System Properties (Read only) as on the client machine from where Pig
	* was launched. This will include command line properties passed at launch
	* time
	* @return client side System Properties including command line properties
	*/
	public Properties getClientSystemProps() {
	return clientSysProps;
	}
	/**
	* Adds the JobConf to this singleton. Will be
	* called on the backend by the Map and Reduce
	* functions so that UDFs can obtain the JobConf
	* on the backend.
	*/
	public void addJobConf(Configuration conf) {
	jconf = conf;
	}

	/**
	* Get the JobConf. This should only be called on
	* the backend. It will return null on the frontend.
	* @return JobConf for this job. This is a copy of the
	* JobConf. Nothing written here will be kept by the system.
	* getUDFConf should be used for recording UDF specific
	* information.
	*/
	public Configuration getJobConf() {
	if (jconf != null) return new Configuration(jconf);
	else return null;
	}

	/**
	* Get a properties object that is specific to this UDF.
	* Note that if a given UDF is called multiple times in a script,
	* and each instance passes different arguments, then each will
	* be provided with different configuration object.
	* This can be used by loaders to pass their input object path
	* or URI and separate themselves from other instances of the
	* same loader. Constructor arguments could also be used,
	* as they are available on both the front and back end.
	*
	* Note that this can only be used to share information
	* across instantiations of the same function in the front end
	* and between front end and back end. It cannot be used to
	* share information between instantiations (that is, between
	* map and/or reduce instances) on the back end at runtime.
	* @param c of the UDF obtaining the properties object.
	* @param args String arguments that make this instance of
	* the UDF unique.
	* @return A reference to the properties object specific to
	* the calling UDF. This is a reference, not a copy.
	* Any changes to this object will automatically be
	* propogated to other instances of the UDF calling this
	* function.
	*/

	@SuppressWarnings("rawtypes")
	public Properties getUDFProperties(Class c, String[] args) {
	UDFContextKey k = generateKey(c, args);
	Properties p = udfConfs.get(k);
	if (p == null) {
	p = new Properties();
	udfConfs.put(k, p);
	}
	return p;
	}

	/**
	* Get a properties object that is specific to this UDF.
	* Note that if a given UDF is called multiple times in a script,
	* they will all be provided the same configuration object. It
	* is up to the UDF to make sure the multiple instances do not
	* stomp on each other.
	*
	* It is guaranteed that this properties object will be separate
	* from that provided to any other UDF.
	*
	* Note that this can only be used to share information
	* across instantiations of the same function in the front end
	* and between front end and back end. It cannot be used to
	* share information between instantiations (that is, between
	* map and/or reduce instances) on the back end at runtime.
	* @param c of the UDF obtaining the properties object.
	* @return A reference to the properties object specific to
	* the calling UDF. This is a reference, not a copy.
	* Any changes to this object will automatically be
	* propogated to other instances of the UDF calling this
	* function.
	*/
	@SuppressWarnings("rawtypes")
	public Properties getUDFProperties(Class c) {
	UDFContextKey k = generateKey(c, null);
	Properties p = udfConfs.get(k);
	if (p == null) {
	p = new Properties();
	udfConfs.put(k, p);
	}
	return p;
	}



	/**
	* Serialize the UDF specific information into an instance
	* of JobConf. This function is intended to be called on
	* the front end in preparation for sending the data to the
	* backend.
	* @param conf JobConf to serialize into
	* @throws IOException if underlying serialization throws it
	*/
	public void serialize(Configuration conf) throws IOException {
	conf.set(UDF_CONTEXT, ObjectSerializer.serialize(udfConfs));
	conf.set(CLIENT_SYS_PROPS, ObjectSerializer.serialize(clientSysProps));
	}

	/*
	* Internal pig use
	*/
	public String serialize() {
	try {
	return ObjectSerializer.serialize(udfConfs);
	} catch (IOException e) {
	LOG.error("UDFContext#serialize throws error ",e);
	return null;
	}
	}

	/**
	* Populate the udfConfs field. This function is intended to
	* be called by Map.configure or Reduce.configure on the backend.
	* It assumes that addJobConf has already been called.
	* @throws IOException if underlying deseralization throws it
	*/
	@SuppressWarnings("unchecked")
	public void deserialize() throws IOException {
	udfConfs = (HashMap<UDFContextKey, Properties>)ObjectSerializer.deserialize(jconf.get(UDF_CONTEXT));
	clientSysProps = (Properties)ObjectSerializer.deserialize(
	jconf.get(CLIENT_SYS_PROPS));
	}

	public void deserializeForSpark(String udfConfsStr, String clientSysPropsStr) throws IOException {
	if( udfConfsStr!= null && clientSysPropsStr!=null) {
	udfConfs = (HashMap<UDFContextKey, Properties>) ObjectSerializer.deserialize(udfConfsStr);
	clientSysProps = (Properties) ObjectSerializer.deserialize(
	clientSysPropsStr);
	}
	}

	private UDFContextKey generateKey(Class<?> c, String[] args) {
	return new UDFContextKey(c.getName(), args);
	}

	public void reset() {
	udfConfs.clear();
	}

	public boolean isUDFConfEmpty() {
	return udfConfs.isEmpty();
	}

	/**
	* Convenience method for UDF code to check where it runs (see PIG-2576)
	* @return boolean type value
	*/
	public boolean isFrontend() {
	// mapred.task.id is for MR1
	// mapreduce.job.application.attempt.id is for MR2
	return (this.jconf == null
	\|\| (jconf.get(MRConfiguration.TASK_ID) == null &&
	jconf.get(MRConfiguration.JOB_APPLICATION_ATTEMPT_ID) == null));
	}

	/**
	* Make a shallow copy of the context.
	*/
	@Override
	public UDFContext clone() {
	UDFContext other = new UDFContext();
	other.clientSysProps = this.clientSysProps;
	other.jconf = this.jconf;
	other.udfConfs = this.udfConfs;
	return other;
	}

	/**
	* Class that acts as key for hashmap in UDFContext,
	* it holds the class and args of the udf, and
	* implements equals() and hashCode()
	*/
	static class UDFContextKey implements Serializable{

	private static final long serialVersionUID = 1;
	private String className;
	private String[] args;

	UDFContextKey(String className, String [] args){
	this.className = className;
	this.args = args;
	}

	String getClassName() {
	return className;
	}

	String[] getArgs() {
	return args;
	}

	@Override
	public String toString() {
	return "UDFContextKey [className=" + className + ", args="
	+ Arrays.toString(args) + "]";
	}

	@Override
	public int hashCode() {
	final int prime = 31;
	int result = 1;
	result = prime * result + Arrays.hashCode(args);
	result = prime * result
	+ ((className == null) ? 0 : className.hashCode());
	return result;
	}

	@Override
	public boolean equals(Object obj) {
	if (this == obj)
	return true;
	if (obj == null)
	return false;
	if (getClass() != obj.getClass())
	return false;
	UDFContextKey other = (UDFContextKey) obj;
	if (!Arrays.equals(args, other.args))
	return false;
	if (className == null) {
	if (other.className != null)
	return false;
	} else if (!className.equals(other.className))
	return false;
	return true;
	}
	}

	public Properties getClientSysProps() {
	return clientSysProps;
	}

	}