tajo-core/src/main/java/org/apache/tajo/engine/function/FunctionLoader.java - tajo - 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.tajo.engine.function;

 import com.google.common.base.Preconditions;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import org.apache.commons.collections.Predicate;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.fs.FileStatus;
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.tajo.annotation.Nullable;
 import org.apache.tajo.catalog.CatalogUtil;
 import org.apache.tajo.catalog.FunctionDesc;
 import org.apache.tajo.common.TajoDataTypes;
 import org.apache.tajo.conf.TajoConf;
 import org.apache.tajo.engine.function.annotation.Description;
 import org.apache.tajo.engine.function.annotation.ParamOptionTypes;
 import org.apache.tajo.engine.function.annotation.ParamTypes;
 import org.apache.tajo.function.*;
 import org.apache.tajo.plan.function.python.PythonScriptEngine;
 import org.apache.tajo.util.ClassUtil;

 import java.io.IOException;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.util.*;

 import static org.apache.tajo.catalog.proto.CatalogProtos.FunctionType.GENERAL;

 public class FunctionLoader {

   private static Log LOG = LogFactory.getLog(FunctionLoader.class);
   public static final String PYTHON_FUNCTION_NAMESPACE = "python";

   /**
    * Load built-in functions
    *
    * @return
    */
   public static Map<FunctionSignature, FunctionDesc> loadBuiltinFunctions() {
     Map<FunctionSignature, FunctionDesc> map = Maps.newHashMap();

     List<FunctionDesc> dd = Lists.newArrayList();
     for (FunctionDesc f : findLegacyFunctions()) {
       map.put(f.getSignature(), f);

       if (f.getSignature().getName().equals("pow") || f.getSignature().getName().equals("pi")) {
         dd.add(f);
       }
     }

     for (FunctionDesc f : findScalarFunctions()) {
       if (map.containsKey(f.getSignature())) {
         FunctionDesc existing = map.get(f.getSignature());
         existing.getInvocation().setScalar(f.getInvocation().getScalar());
       } else {
         map.put(f.getSignature(), f);
       }
     }

     return map;
   }

   /**
    * Load functions defined by users.
    *
    * @param conf
    * @return
    * @throws IOException
    */
   public static Optional<List<FunctionDesc>> loadUserDefinedFunctions(TajoConf conf)
       throws IOException {
     List<FunctionDesc> functionList = new LinkedList<>();

     String[] codePaths = conf.getStrings(TajoConf.ConfVars.PYTHON_CODE_DIR.varname);
     if (codePaths != null) {
       FileSystem localFS = FileSystem.getLocal(conf);
       for (String codePathStr : codePaths) {

         Path codePath;
         try {
           codePath = new Path(codePathStr);
         } catch (IllegalArgumentException e) {
           LOG.warn(e);
           continue;
         }

         List<Path> filePaths = new ArrayList<>();
         if (localFS.isDirectory(codePath)) {
           for (FileStatus file : localFS.listStatus(codePath,
               (Path path) -> path.getName().endsWith(PythonScriptEngine.FILE_EXTENSION))) {
             filePaths.add(file.getPath());
           }
         } else {
           filePaths.add(codePath);
         }
         for (Path filePath : filePaths) {
           PythonScriptEngine.registerFunctions(filePath.toUri(), FunctionLoader.PYTHON_FUNCTION_NAMESPACE).
               forEach(functionList::add);
         }
       }
     }

     return Optional.of(functionList);
   }

   public static Set<FunctionDesc> findScalarFunctions() {
     Set<FunctionDesc> functions = Sets.newHashSet();

     Set<Method> scalarFunctions = findPublicStaticMethods("org.apache.tajo.engine.function",
         (Object object) -> ((Method) object).getAnnotation(ScalarFunction.class) != null);

     for (Method method : scalarFunctions) {
       ScalarFunction annotation = method.getAnnotation(ScalarFunction.class);
       functions.addAll(buildFunctionDescs(annotation, method));
     }

     return functions;
   }

   private static Set<Method> findPublicStaticMethods(String packageName, Predicate predicate) {
     Set<Class> found = findFunctionCollections(packageName);
     Set<Method> filtered = Sets.newHashSet();

     for (Class clazz : found) {
       for (Method method : clazz.getMethods()) {
         if (isPublicStaticMethod(method) && (predicate == null || predicate.evaluate(method))) {
           filtered.add(method);
         }
       }
     }

     return filtered;
   }

   private static boolean isPublicStaticMethod(Method method) {
     return Modifier.isPublic(method.getModifiers()) && Modifier.isStatic(method.getModifiers());
   }

   private static Set<Class> findFunctionCollections(String packageName) {
     return ClassUtil.findClasses(null, packageName, (Object object) ->
         ((Class)object).getAnnotation(FunctionCollection.class) != null);
   }

   private static Collection<FunctionDesc> buildFunctionDescs(ScalarFunction annotation, Method method) {
     List<FunctionDesc> functionDescs = Lists.newArrayList();

     FunctionInvocation invocation = new FunctionInvocation();
     invocation.setScalar(extractStaticMethodInvocation(method));
     FunctionSupplement supplement = extractSupplement(annotation);

     // primary name
     functionDescs.add(new FunctionDesc(extractSignature(annotation, null), invocation, supplement));

     // for multiple aliases
     for (String alias : annotation.synonyms()) {
       functionDescs.add(new FunctionDesc(extractSignature(annotation, alias), invocation, supplement));
     }

     return functionDescs;
   }

   private static FunctionSignature extractSignature(ScalarFunction annotation, @Nullable String alias) {
     return new FunctionSignature(
         GENERAL,
         alias != null ? alias : annotation.name(),
         CatalogUtil.newSimpleDataType(annotation.returnType()),
         CatalogUtil.newSimpleDataTypeArray(annotation.paramTypes())
     );
   }

   private static FunctionSupplement extractSupplement(ScalarFunction function) {
     return new FunctionSupplement(
             function.shortDescription(),
             function.detail(),
             function.example());
   }

   private static StaticMethodInvocationDesc extractStaticMethodInvocation(Method method) {
     Preconditions.checkArgument(Modifier.isPublic(method.getModifiers()));
     Preconditions.checkArgument(Modifier.isStatic(method.getModifiers()));

     String methodName = method.getName();
     Class returnClass = method.getReturnType();
     Class [] paramClasses = method.getParameterTypes();
     return new StaticMethodInvocationDesc(method.getDeclaringClass(), methodName, returnClass, paramClasses);
   }

   /**
    * This method finds and build FunctionDesc for the legacy function and UD(A)F system.
    *
    * @return A list of FunctionDescs
    */
   public static List<FunctionDesc> findLegacyFunctions() {
     List<FunctionDesc> sqlFuncs = new ArrayList<>();

     Set<Class> functionClasses = ClassUtil.findClasses(Function.class, "org.apache.tajo.engine.function");

     for (Class eachClass : functionClasses) {
       if(eachClass.isInterface() || Modifier.isAbstract(eachClass.getModifiers())) {
         continue;
       }
       Function function = null;
       try {
         function = (Function)eachClass.newInstance();
       } catch (Exception e) {
         LOG.warn(eachClass + " cannot instantiate Function class because of " + e.getMessage(), e);
         continue;
       }
       String functionName = function.getClass().getAnnotation(Description.class).functionName();
       String[] synonyms = function.getClass().getAnnotation(Description.class).synonyms();
       String description = function.getClass().getAnnotation(Description.class).description();
       String detail = function.getClass().getAnnotation(Description.class).detail();
       String example = function.getClass().getAnnotation(Description.class).example();
       TajoDataTypes.Type returnType = function.getClass().getAnnotation(Description.class).returnType();
       ParamTypes[] paramArray = function.getClass().getAnnotation(Description.class).paramTypes();

       String[] allFunctionNames = null;
       if(synonyms != null && synonyms.length > 0) {
         allFunctionNames = new String[1 + synonyms.length];
         allFunctionNames[0] = functionName;
         System.arraycopy(synonyms, 0, allFunctionNames, 1, synonyms.length);
       } else {
         allFunctionNames = new String[]{functionName};
       }

       for(String eachFunctionName: allFunctionNames) {
         for (ParamTypes params : paramArray) {
           ParamOptionTypes[] paramOptionArray;
           if(params.paramOptionTypes() == null ||
               params.paramOptionTypes().getClass().getAnnotation(ParamTypes.class) == null) {
             paramOptionArray = new ParamOptionTypes[0];
           } else {
             paramOptionArray = params.paramOptionTypes().getClass().getAnnotation(ParamTypes.class).paramOptionTypes();
           }

           TajoDataTypes.Type[] paramTypes = params.paramTypes();
           if (paramOptionArray.length > 0)
             paramTypes = params.paramTypes().clone();

           for (int i=0; i < paramOptionArray.length + 1; i++) {
             FunctionDesc functionDesc = new FunctionDesc(eachFunctionName,
                 function.getClass(), function.getFunctionType(),
                 CatalogUtil.newSimpleDataType(returnType),
                 paramTypes.length == 0 ? CatalogUtil.newSimpleDataTypeArray() : CatalogUtil.newSimpleDataTypeArray(paramTypes));

             functionDesc.setDescription(description);
             functionDesc.setExample(example);
             functionDesc.setDetail(detail);
             sqlFuncs.add(functionDesc);

             if (i != paramOptionArray.length) {
               paramTypes = new TajoDataTypes.Type[paramTypes.length +
                   paramOptionArray[i].paramOptionTypes().length];
               System.arraycopy(params.paramTypes(), 0, paramTypes, 0, paramTypes.length);
               System.arraycopy(paramOptionArray[i].paramOptionTypes(), 0, paramTypes, paramTypes.length,
                   paramOptionArray[i].paramOptionTypes().length);
             }
           }
         }
       }
     }

     return sqlFuncs;
   }
 }
	/***
	* 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.tajo.engine.function;

	import com.google.common.base.Preconditions;
	import com.google.common.collect.Lists;
	import com.google.common.collect.Maps;
	import com.google.common.collect.Sets;
	import org.apache.commons.collections.Predicate;
	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.hadoop.fs.FileStatus;
	import org.apache.hadoop.fs.FileSystem;
	import org.apache.hadoop.fs.Path;
	import org.apache.tajo.annotation.Nullable;
	import org.apache.tajo.catalog.CatalogUtil;
	import org.apache.tajo.catalog.FunctionDesc;
	import org.apache.tajo.common.TajoDataTypes;
	import org.apache.tajo.conf.TajoConf;
	import org.apache.tajo.engine.function.annotation.Description;
	import org.apache.tajo.engine.function.annotation.ParamOptionTypes;
	import org.apache.tajo.engine.function.annotation.ParamTypes;
	import org.apache.tajo.function.*;
	import org.apache.tajo.plan.function.python.PythonScriptEngine;
	import org.apache.tajo.util.ClassUtil;

	import java.io.IOException;
	import java.lang.reflect.Method;
	import java.lang.reflect.Modifier;
	import java.util.*;

	import static org.apache.tajo.catalog.proto.CatalogProtos.FunctionType.GENERAL;

	public class FunctionLoader {

	private static Log LOG = LogFactory.getLog(FunctionLoader.class);
	public static final String PYTHON_FUNCTION_NAMESPACE = "python";

	/**
	* Load built-in functions
	*
	* @return
	*/
	public static Map<FunctionSignature, FunctionDesc> loadBuiltinFunctions() {
	Map<FunctionSignature, FunctionDesc> map = Maps.newHashMap();

	List<FunctionDesc> dd = Lists.newArrayList();
	for (FunctionDesc f : findLegacyFunctions()) {
	map.put(f.getSignature(), f);

	if (f.getSignature().getName().equals("pow") \|\| f.getSignature().getName().equals("pi")) {
	dd.add(f);
	}
	}

	for (FunctionDesc f : findScalarFunctions()) {
	if (map.containsKey(f.getSignature())) {
	FunctionDesc existing = map.get(f.getSignature());
	existing.getInvocation().setScalar(f.getInvocation().getScalar());
	} else {
	map.put(f.getSignature(), f);
	}
	}

	return map;
	}

	/**
	* Load functions defined by users.
	*
	* @param conf
	* @return
	* @throws IOException
	*/
	public static Optional<List<FunctionDesc>> loadUserDefinedFunctions(TajoConf conf)
	throws IOException {
	List<FunctionDesc> functionList = new LinkedList<>();

	String[] codePaths = conf.getStrings(TajoConf.ConfVars.PYTHON_CODE_DIR.varname);
	if (codePaths != null) {
	FileSystem localFS = FileSystem.getLocal(conf);
	for (String codePathStr : codePaths) {

	Path codePath;
	try {
	codePath = new Path(codePathStr);
	} catch (IllegalArgumentException e) {
	LOG.warn(e);
	continue;
	}

	List<Path> filePaths = new ArrayList<>();
	if (localFS.isDirectory(codePath)) {
	for (FileStatus file : localFS.listStatus(codePath,
	(Path path) -> path.getName().endsWith(PythonScriptEngine.FILE_EXTENSION))) {
	filePaths.add(file.getPath());
	}
	} else {
	filePaths.add(codePath);
	}
	for (Path filePath : filePaths) {
	PythonScriptEngine.registerFunctions(filePath.toUri(), FunctionLoader.PYTHON_FUNCTION_NAMESPACE).
	forEach(functionList::add);
	}
	}
	}

	return Optional.of(functionList);
	}

	public static Set<FunctionDesc> findScalarFunctions() {
	Set<FunctionDesc> functions = Sets.newHashSet();

	Set<Method> scalarFunctions = findPublicStaticMethods("org.apache.tajo.engine.function",
	(Object object) -> ((Method) object).getAnnotation(ScalarFunction.class) != null);

	for (Method method : scalarFunctions) {
	ScalarFunction annotation = method.getAnnotation(ScalarFunction.class);
	functions.addAll(buildFunctionDescs(annotation, method));
	}

	return functions;
	}

	private static Set<Method> findPublicStaticMethods(String packageName, Predicate predicate) {
	Set<Class> found = findFunctionCollections(packageName);
	Set<Method> filtered = Sets.newHashSet();

	for (Class clazz : found) {
	for (Method method : clazz.getMethods()) {
	if (isPublicStaticMethod(method) && (predicate == null \|\| predicate.evaluate(method))) {
	filtered.add(method);
	}
	}
	}

	return filtered;
	}

	private static boolean isPublicStaticMethod(Method method) {
	return Modifier.isPublic(method.getModifiers()) && Modifier.isStatic(method.getModifiers());
	}

	private static Set<Class> findFunctionCollections(String packageName) {
	return ClassUtil.findClasses(null, packageName, (Object object) ->
	((Class)object).getAnnotation(FunctionCollection.class) != null);
	}

	private static Collection<FunctionDesc> buildFunctionDescs(ScalarFunction annotation, Method method) {
	List<FunctionDesc> functionDescs = Lists.newArrayList();

	FunctionInvocation invocation = new FunctionInvocation();
	invocation.setScalar(extractStaticMethodInvocation(method));
	FunctionSupplement supplement = extractSupplement(annotation);

	// primary name
	functionDescs.add(new FunctionDesc(extractSignature(annotation, null), invocation, supplement));

	// for multiple aliases
	for (String alias : annotation.synonyms()) {
	functionDescs.add(new FunctionDesc(extractSignature(annotation, alias), invocation, supplement));
	}

	return functionDescs;
	}

	private static FunctionSignature extractSignature(ScalarFunction annotation, @Nullable String alias) {
	return new FunctionSignature(
	GENERAL,
	alias != null ? alias : annotation.name(),
	CatalogUtil.newSimpleDataType(annotation.returnType()),
	CatalogUtil.newSimpleDataTypeArray(annotation.paramTypes())
	);
	}

	private static FunctionSupplement extractSupplement(ScalarFunction function) {
	return new FunctionSupplement(
	function.shortDescription(),
	function.detail(),
	function.example());
	}

	private static StaticMethodInvocationDesc extractStaticMethodInvocation(Method method) {
	Preconditions.checkArgument(Modifier.isPublic(method.getModifiers()));
	Preconditions.checkArgument(Modifier.isStatic(method.getModifiers()));

	String methodName = method.getName();
	Class returnClass = method.getReturnType();
	Class [] paramClasses = method.getParameterTypes();
	return new StaticMethodInvocationDesc(method.getDeclaringClass(), methodName, returnClass, paramClasses);
	}

	/**
	* This method finds and build FunctionDesc for the legacy function and UD(A)F system.
	*
	* @return A list of FunctionDescs
	*/
	public static List<FunctionDesc> findLegacyFunctions() {
	List<FunctionDesc> sqlFuncs = new ArrayList<>();

	Set<Class> functionClasses = ClassUtil.findClasses(Function.class, "org.apache.tajo.engine.function");

	for (Class eachClass : functionClasses) {
	if(eachClass.isInterface() \|\| Modifier.isAbstract(eachClass.getModifiers())) {
	continue;
	}
	Function function = null;
	try {
	function = (Function)eachClass.newInstance();
	} catch (Exception e) {
	LOG.warn(eachClass + " cannot instantiate Function class because of " + e.getMessage(), e);
	continue;
	}
	String functionName = function.getClass().getAnnotation(Description.class).functionName();
	String[] synonyms = function.getClass().getAnnotation(Description.class).synonyms();
	String description = function.getClass().getAnnotation(Description.class).description();
	String detail = function.getClass().getAnnotation(Description.class).detail();
	String example = function.getClass().getAnnotation(Description.class).example();
	TajoDataTypes.Type returnType = function.getClass().getAnnotation(Description.class).returnType();
	ParamTypes[] paramArray = function.getClass().getAnnotation(Description.class).paramTypes();

	String[] allFunctionNames = null;
	if(synonyms != null && synonyms.length > 0) {
	allFunctionNames = new String[1 + synonyms.length];
	allFunctionNames[0] = functionName;
	System.arraycopy(synonyms, 0, allFunctionNames, 1, synonyms.length);
	} else {
	allFunctionNames = new String[]{functionName};
	}

	for(String eachFunctionName: allFunctionNames) {
	for (ParamTypes params : paramArray) {
	ParamOptionTypes[] paramOptionArray;
	if(params.paramOptionTypes() == null \|\|
	params.paramOptionTypes().getClass().getAnnotation(ParamTypes.class) == null) {
	paramOptionArray = new ParamOptionTypes[0];
	} else {
	paramOptionArray = params.paramOptionTypes().getClass().getAnnotation(ParamTypes.class).paramOptionTypes();
	}

	TajoDataTypes.Type[] paramTypes = params.paramTypes();
	if (paramOptionArray.length > 0)
	paramTypes = params.paramTypes().clone();

	for (int i=0; i < paramOptionArray.length + 1; i++) {
	FunctionDesc functionDesc = new FunctionDesc(eachFunctionName,
	function.getClass(), function.getFunctionType(),
	CatalogUtil.newSimpleDataType(returnType),
	paramTypes.length == 0 ? CatalogUtil.newSimpleDataTypeArray() : CatalogUtil.newSimpleDataTypeArray(paramTypes));

	functionDesc.setDescription(description);
	functionDesc.setExample(example);
	functionDesc.setDetail(detail);
	sqlFuncs.add(functionDesc);

	if (i != paramOptionArray.length) {
	paramTypes = new TajoDataTypes.Type[paramTypes.length +
	paramOptionArray[i].paramOptionTypes().length];
	System.arraycopy(params.paramTypes(), 0, paramTypes, 0, paramTypes.length);
	System.arraycopy(paramOptionArray[i].paramOptionTypes(), 0, paramTypes, paramTypes.length,
	paramOptionArray[i].paramOptionTypes().length);
	}
	}
	}
	}
	}

	return sqlFuncs;
	}
	}