org.apache.calcite.schema.impl.ScalarFunctionImpl Java Examples

/** Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-3547">[CALCITE-3547]
 * SqlValidatorException because Planner cannot find UDFs added to schema</a>. */
@Test void testValidateUserDefinedFunctionInSchema() throws Exception {
  SchemaPlus rootSchema = Frameworks.createRootSchema(true);
  rootSchema.add("my_plus",
      ScalarFunctionImpl.create(Smalls.MyPlusFunction.class, "eval"));
  final FrameworkConfig config = Frameworks.newConfigBuilder()
      .defaultSchema(
          CalciteAssert.addSchema(rootSchema, CalciteAssert.SchemaSpec.HR))
      .build();
  final Planner planner = Frameworks.getPlanner(config);
  final String sql = "select \"my_plus\"(\"deptno\", 100) as \"p\"\n"
      + "from \"hr\".\"emps\"";
  SqlNode parse = planner.parse(sql);
  SqlNode validate = planner.validate(parse);
  assertThat(Util.toLinux(validate.toString()),
      equalTo("SELECT `my_plus`(`emps`.`deptno`, 100) AS `p`\n"
          + "FROM `hr`.`emps` AS `emps`"));
}
 

private void handleCreateFunction(SqlCreateFunction sqlCreateFunction) throws ClassNotFoundException {
  if(sqlCreateFunction.jarName() != null) {
    throw new UnsupportedOperationException("UDF 'USING JAR' not implemented");
  }
  Method method;
  Function function;
  if ((method=findMethod(sqlCreateFunction.className(), "evaluate")) != null) {
    function = ScalarFunctionImpl.create(method);
  } else if (findMethod(sqlCreateFunction.className(), "add") != null) {
    function = AggregateFunctionImpl.create(Class.forName(sqlCreateFunction.className()));
  } else {
    throw new RuntimeException("Invalid scalar or aggregate function");
  }
  schema.add(sqlCreateFunction.functionName().toUpperCase(), function);
  hasUdf = true;
}
 

@Override public void visit(UserFuncExpression op) throws FrontendException {
  if (op.getFuncSpec().getClassName().equals("org.apache.pig.impl.builtin.IdentityColumn")) {
    // Skip this Pig dummy function
    return;
  }
  final int numAgrs = optSize(op.getPlan().getSuccessors(op))
      + optSize(op.getPlan().getSoftLinkSuccessors(op));

  final RelDataType returnType = PigTypes.convertSchemaField(op.getFieldSchema());
  stack.push(
      PigRelUdfConverter.convertPigFunction(
          builder, op.getFuncSpec(), buildOperands(numAgrs), returnType));

  String className = op.getFuncSpec().getClassName();
  SqlOperator sqlOp = ((RexCall) stack.peek()).getOperator();
  if (sqlOp instanceof SqlUserDefinedFunction) {
    ScalarFunctionImpl sqlFunc =
        (ScalarFunctionImpl) ((SqlUserDefinedFunction) sqlOp).getFunction();
    // "Exec" method can be implemented from the parent class.
    className = sqlFunc.method.getDeclaringClass().getName();
  }
  builder.registerPigUDF(className, op.getFuncSpec());
}
 

/**
 * Gets the {@link SqlAggFunction} for the corresponding Pig aggregate
 * UDF call; returns null for invalid rex call.
 *
 * @param call Pig aggregate UDF call
 */
static SqlAggFunction getSqlAggFuncForPigUdf(RexCall call) {
  if (!(call.getOperator() instanceof PigUserDefinedFunction)) {
    return null;
  }

  final PigUserDefinedFunction pigUdf = (PigUserDefinedFunction) call.getOperator();
  if (pigUdf.funcSpec != null) {
    final String pigUdfClassName = pigUdf.funcSpec.getClassName();
    final SqlAggFunction sqlAggFunction = BUILTIN_AGG_FUNC.get(pigUdfClassName);
    if (sqlAggFunction == null) {
      final Class udfClass =
          ((ScalarFunctionImpl) pigUdf.getFunction()).method.getDeclaringClass();
      if (Accumulator.class.isAssignableFrom(udfClass)) {
        throw new UnsupportedOperationException(
            "Cannot find corresponding SqlAgg func for Pig aggegate " + pigUdfClassName);
      }
    }
    return sqlAggFunction;
  }
  return null;
}
 

public RexNode getIObject() {
    if (object == null) {
        RelDataTypeFactory typeFactory = getCluster().getTypeFactory();
        RexBuilder b = getBuilder();
        final SqlFunction GET_IOBJECT =
                new SqlUserDefinedFunction(
                        new SqlIdentifier("GET_IOBJECT", SqlParserPos.ZERO),
                        ReturnTypes.explicit(typeFactory.createTypeWithNullability(typeFactory.createJavaType(IObject.class), false)),
                        null,
                        OperandTypes.ANY_ANY,
                        ImmutableList.of(typeFactory.createTypeWithNullability(typeFactory.createJavaType(ISnapshot.class), false),
                                typeFactory.createJavaType(int.class)),
                        ScalarFunctionImpl.create(ISnapshotMethods.class, "getIObject"));
        object = b.makeCall(GET_IOBJECT, getSnapshot(), getIObjectId());
    }
    return object;
}
 

@Override
public RexNode getSnapshot() {
    if (snapshot == null) {
        RelDataTypeFactory typeFactory = getCluster().getTypeFactory();
        RexBuilder b = getBuilder();
        final SqlFunction UDF =
                new SqlUserDefinedFunction(
                        new SqlIdentifier("GET_SNAPSHOT", SqlParserPos.ZERO),
                        ReturnTypes.explicit(typeFactory.createTypeWithNullability(typeFactory.createJavaType(ISnapshot.class), false)),
                        null,
                        OperandTypes.NUMERIC,
                        ImmutableList.of(typeFactory.createJavaType(Integer.class)),
                        ScalarFunctionImpl.create(SnapshotHolder.class, "get"));
        snapshot = b.makeCall(UDF, b.makeLiteral(snapshotId, typeFactory.createSqlType(SqlTypeName.INTEGER), false));
    }
    return snapshot;
}
 

@Override
public RexNode apply(RexBuilderContext context) {
	RelOptCluster cluster = context.getCluster();
	RelDataTypeFactory typeFactory = cluster.getTypeFactory();
	final SqlFunction UDF =
			new SqlUserDefinedFunction(
					new SqlIdentifier("RESOLVE_SIMPLE", SqlParserPos.ZERO),
					ReturnTypes.explicit(typeFactory.createJavaType(Object.class)),
					null,
					OperandTypes.ANY_ANY,
					ImmutableList.of(typeFactory.createTypeWithNullability(typeFactory.createJavaType(IObject.class), false),
							typeFactory.createJavaType(int.class)),
					ScalarFunctionImpl.create(IObjectMethods.class, "resolveSimpleValue"));
	RexBuilder b = context.getBuilder();
	RexNode rexNode = b.makeCall(UDF, context.getIObject(), b.makeLiteral(name));
	return b.makeCast(dataType, rexNode);
}
 

@Override
public RexNode apply(RexBuilderContext context) {
	RelOptCluster cluster = context.getCluster();
	RelDataTypeFactory typeFactory = cluster.getTypeFactory();
	final SqlFunction UDF =
			new SqlUserDefinedFunction(
					new SqlIdentifier("RESOLVE_REFERENCE", SqlParserPos.ZERO),
					ReturnTypes.explicit(typeFactory.createJavaType(HeapReference.class)),
					null,
					OperandTypes.ANY_ANY,
					ImmutableList.of(typeFactory.createTypeWithNullability(typeFactory.createJavaType(IObject.class), false),
							typeFactory.createJavaType(String.class)),
					ScalarFunctionImpl.create(IObjectMethods.class, "resolveReferenceValue"));
	RexBuilder b = context.getBuilder();
	return b.makeCall(UDF, context.getIObject(), b.makeLiteral(name));
}
 

private SchemaPlus createRootSchema(final CalciteConnection calciteConnection) {
    final SchemaPlus rootSchema = calciteConnection.getRootSchema();
    rootSchema.add("RPATH", ScalarFunctionImpl.create(ObjectRecordPath.class, "eval"));
    rootSchema.add("RPATH_STRING", ScalarFunctionImpl.create(StringRecordPath.class, "eval"));
    rootSchema.add("RPATH_INT", ScalarFunctionImpl.create(IntegerRecordPath.class, "eval"));
    rootSchema.add("RPATH_LONG", ScalarFunctionImpl.create(LongRecordPath.class, "eval"));
    rootSchema.add("RPATH_DATE", ScalarFunctionImpl.create(DateRecordPath.class, "eval"));
    rootSchema.add("RPATH_DOUBLE", ScalarFunctionImpl.create(DoubleRecordPath.class, "eval"));
    rootSchema.add("RPATH_FLOAT", ScalarFunctionImpl.create(FloatRecordPath.class, "eval"));

    return rootSchema;
}
 

/** Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-937">[CALCITE-937]
 * User-defined function within view</a>. */
@Test void testUserDefinedFunctionInView() throws Exception {
  Class.forName("org.apache.calcite.jdbc.Driver");
  Connection connection = DriverManager.getConnection("jdbc:calcite:");
  CalciteConnection calciteConnection =
      connection.unwrap(CalciteConnection.class);
  SchemaPlus rootSchema = calciteConnection.getRootSchema();
  rootSchema.add("hr", new ReflectiveSchema(new JdbcTest.HrSchema()));

  SchemaPlus post = rootSchema.add("POST", new AbstractSchema());
  post.add("MY_INCREMENT",
      ScalarFunctionImpl.create(Smalls.MyIncrement.class, "eval"));

  final String viewSql = "select \"empid\" as EMPLOYEE_ID,\n"
      + "  \"name\" || ' ' || \"name\" as EMPLOYEE_NAME,\n"
      + "  \"salary\" as EMPLOYEE_SALARY,\n"
      + "  POST.MY_INCREMENT(\"empid\", 10) as INCREMENTED_SALARY\n"
      + "from \"hr\".\"emps\"";
  post.add("V_EMP",
      ViewTable.viewMacro(post, viewSql, ImmutableList.of(),
          ImmutableList.of("POST", "V_EMP"), null));

  final String result = ""
      + "EMPLOYEE_ID=100; EMPLOYEE_NAME=Bill Bill; EMPLOYEE_SALARY=10000.0; INCREMENTED_SALARY=110.0\n"
      + "EMPLOYEE_ID=200; EMPLOYEE_NAME=Eric Eric; EMPLOYEE_SALARY=8000.0; INCREMENTED_SALARY=220.0\n"
      + "EMPLOYEE_ID=150; EMPLOYEE_NAME=Sebastian Sebastian; EMPLOYEE_SALARY=7000.0; INCREMENTED_SALARY=165.0\n"
      + "EMPLOYEE_ID=110; EMPLOYEE_NAME=Theodore Theodore; EMPLOYEE_SALARY=11500.0; INCREMENTED_SALARY=121.0\n";

  Statement statement = connection.createStatement();
  ResultSet resultSet = statement.executeQuery(viewSql);
  assertThat(CalciteAssert.toString(resultSet), is(result));
  resultSet.close();

  ResultSet viewResultSet =
      statement.executeQuery("select * from \"POST\".\"V_EMP\"");
  assertThat(CalciteAssert.toString(viewResultSet), is(result));
  statement.close();
  connection.close();
}
 

private static SqlReturnTypeInference infer(final ScalarFunction function) {
  return opBinding -> {
    final RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
    final RelDataType type;
    if (function instanceof ScalarFunctionImpl) {
      type = ((ScalarFunctionImpl) function).getReturnType(typeFactory,
          opBinding);
    } else {
      type = function.getReturnType(typeFactory);
    }
    return toSql(typeFactory, type);
  };
}
 

public RexNode apply(RexBuilderContext context) {
	RelOptCluster cluster = context.getCluster();
	RelDataTypeFactory typeFactory = cluster.getTypeFactory();
	final SqlFunction UDF =
			new SqlUserDefinedFunction(
					new SqlIdentifier("TO_REFERENCE", SqlParserPos.ZERO),
					ReturnTypes.explicit(typeFactory.createJavaType(HeapReference.class)),
					null,
					OperandTypes.ANY,
					ImmutableList.of(typeFactory.createTypeWithNullability(typeFactory.createJavaType(IObject.class), false)),
					ScalarFunctionImpl.create(ISnapshotMethods.class, "toReference")
			);
	return context.getBuilder().makeCall(UDF, context.getIObject());
}
 

public static CalciteState sqlOverDummyTable(String sql)
        throws RelConversionException, ValidationException, SqlParseException {
    SchemaPlus schema = Frameworks.createRootSchema(true);
    JavaTypeFactory typeFactory = new JavaTypeFactoryImpl
            (RelDataTypeSystem.DEFAULT);
    StreamableTable streamableTable = new CompilerUtil.TableBuilderInfo(typeFactory)
            .field("ID", SqlTypeName.INTEGER)
            .field("NAME", typeFactory.createType(String.class))
            .field("ADDR", typeFactory.createType(String.class))
            .build();
    Table table = streamableTable.stream();
    schema.add("FOO", table);
    schema.add("BAR", table);
    schema.add("MYPLUS", ScalarFunctionImpl.create(MyPlus.class, "eval"));

    List<SqlOperatorTable> sqlOperatorTables = new ArrayList<>();
    sqlOperatorTables.add(SqlStdOperatorTable.instance());
    sqlOperatorTables.add(new CalciteCatalogReader(CalciteSchema.from(schema),
            false,
            Collections.<String>emptyList(), typeFactory));
    SqlOperatorTable chainedSqlOperatorTable = new ChainedSqlOperatorTable(sqlOperatorTables);
    FrameworkConfig config = Frameworks.newConfigBuilder().defaultSchema(
            schema).operatorTable(chainedSqlOperatorTable).build();
    Planner planner = Frameworks.getPlanner(config);
    SqlNode parse = planner.parse(sql);
    SqlNode validate = planner.validate(parse);
    RelNode tree = planner.convert(validate);
    System.out.println(RelOptUtil.toString(tree, SqlExplainLevel.ALL_ATTRIBUTES));
    return new CalciteState(schema, tree);
}
 

public SamzaSqlScalarFunctionImpl(UdfMetadata udfMetadata) {

    myIncFunction = ScalarFunctionImpl.create(udfMetadata.getUdfMethod());
    this.udfMetadata = udfMetadata;
    this.udfName = udfMetadata.getName();
    this.udfMethod = udfMetadata.getUdfMethod();
    this.getUdfMethod = Arrays.stream(SamzaSqlExecutionContext.class.getMethods())
        .filter(x -> x.getName().equals("getOrCreateUdf"))
        .findFirst()
        .get();
  }
 

@Override
protected Multimap<String, Function> getFunctionMultimap() {
	ImmutableMultimap<String, ScalarFunction> funcs = ScalarFunctionImpl.createAll(JSONFunction.class);
	Multimap<String, Function> functions = HashMultimap.create();
	for (String key : funcs.keySet()) {
		for (ScalarFunction func : funcs.get(key)) {
			functions.put(key, func);
		}
	}
	return functions;
}
 

public static PackageSchema resolveSchema(ISnapshot snapshot) {

        try {
            // Create functions for schema
            ImmutableMultimap.Builder<String, Function> builder = ImmutableMultimap.builder();
            builder.putAll(ScalarFunctionImpl.createAll(HeapFunctions.class));
            builder.putAll(CollectionsFunctions.createAll());
            builder.putAll(TableFunctions.createAll());
            builder.putAll(SnapshotFunctions.createAll(snapshot));
            builder.put("getHints", new SqlAdvisorGetHintsFunction());
            ImmutableMultimap<String, Function> functions = builder.build();

            // Create default schema
            PackageSchema defaultSchema = new PackageSchema(functions);

            // Collect all classes names
            Collection<IClass> classes = snapshot.getClasses();
            HashSet<String> classesNames = new HashSet<>();
            for (IClass iClass : classes) {
                classesNames.add(iClass.getName());
            }

            PackageSchema instanceOfPackage = defaultSchema.getPackage("instanceof");

            // Add all classes to schema
            for (String fullClassName : classesNames) {
                IClassesList classOnly = new IClassesList(snapshot, fullClassName, false);

                // Make class available via "package.name.ClassName" (full class name in a root schema)
                defaultSchema.addClass(fullClassName, classOnly);

                String simpleClassName = getClassName(fullClassName);

                // Make class available via package.name.ClassName (schema.schema.Class)
                PackageSchema packageSchema = getPackage(defaultSchema, fullClassName);
                packageSchema.addClass(simpleClassName, classOnly);

                // Add instanceof
                IClassesList withSubClasses = new IClassesList(snapshot, fullClassName, true);

                // Make class available via "instanceof.package.name.ClassName"
                defaultSchema.addClass("instanceof." + fullClassName, withSubClasses);

                // Make class available via instanceof.package.name.ClassName
                PackageSchema instanceOfSchema = getPackage(instanceOfPackage, fullClassName);
                instanceOfSchema.addClass(simpleClassName, withSubClasses);

            }

            // Add thread stacks table
            defaultSchema.getPackage("native").addClass("ThreadStackFrames", new SnapshotThreadStacksTable(snapshot));

            return defaultSchema;
        } catch (SnapshotException e) {
            throw new RuntimeException("Cannot resolve package schemes", e);
        }
    }
 

public static CalciteState sqlOverNestedTable(String sql)
        throws RelConversionException, ValidationException, SqlParseException {
    SchemaPlus schema = Frameworks.createRootSchema(true);
    JavaTypeFactory typeFactory = new JavaTypeFactoryImpl
            (RelDataTypeSystem.DEFAULT);

    StreamableTable streamableTable = new CompilerUtil.TableBuilderInfo(typeFactory)
            .field("ID", SqlTypeName.INTEGER)
            .field("MAPFIELD",
                    typeFactory.createTypeWithNullability(
                            typeFactory.createMapType(
                                    typeFactory.createTypeWithNullability(
                                            typeFactory.createSqlType(SqlTypeName.VARCHAR), true),
                                    typeFactory.createTypeWithNullability(
                                            typeFactory.createSqlType(SqlTypeName.INTEGER), true))
                            , true))
            .field("NESTEDMAPFIELD",
                    typeFactory.createTypeWithNullability(
                        typeFactory.createMapType(
                                typeFactory.createTypeWithNullability(
                                        typeFactory.createSqlType(SqlTypeName.VARCHAR), true),
                                typeFactory.createTypeWithNullability(
                                        typeFactory.createMapType(
                                                typeFactory.createTypeWithNullability(
                                                        typeFactory.createSqlType(SqlTypeName.VARCHAR), true),
                                                typeFactory.createTypeWithNullability(
                                                        typeFactory.createSqlType(SqlTypeName.INTEGER), true))
                                        , true))
                                    , true))
            .field("ARRAYFIELD", typeFactory.createTypeWithNullability(
                    typeFactory.createArrayType(
                        typeFactory.createTypeWithNullability(
                            typeFactory.createSqlType(SqlTypeName.INTEGER), true), -1L)
                    , true))
            .build();
    Table table = streamableTable.stream();
    schema.add("FOO", table);
    schema.add("BAR", table);
    schema.add("MYPLUS", ScalarFunctionImpl.create(MyPlus.class, "eval"));
    List<SqlOperatorTable> sqlOperatorTables = new ArrayList<>();
    sqlOperatorTables.add(SqlStdOperatorTable.instance());
    sqlOperatorTables.add(new CalciteCatalogReader(CalciteSchema.from(schema),
                                                   false,
                                                   Collections.<String>emptyList(), typeFactory));
    SqlOperatorTable chainedSqlOperatorTable = new ChainedSqlOperatorTable(sqlOperatorTables);
    FrameworkConfig config = Frameworks.newConfigBuilder().defaultSchema(
            schema).operatorTable(chainedSqlOperatorTable).build();
    Planner planner = Frameworks.getPlanner(config);
    SqlNode parse = planner.parse(sql);
    SqlNode validate = planner.validate(parse);
    RelNode tree = planner.convert(validate);
    System.out.println(RelOptUtil.toString(tree, SqlExplainLevel.ALL_ATTRIBUTES));
    return new CalciteState(schema, tree);
}
 

public static void registerFlinkBuiltInFunctions(SchemaPlus schema) {
    BiConsumer<String, Method> addScalarFunctionConsumer =
        (k, v) -> schema.add(k, ScalarFunctionImpl.create(v));

    addScalarFunctionConsumer.accept("LOG2", MathFunctions.LOG2);
    addScalarFunctionConsumer.accept("LOG2", MathFunctions.LOG2_DEC);
    addScalarFunctionConsumer.accept("LOG", MathFunctions.LOG);
    addScalarFunctionConsumer.accept("LOG", MathFunctions.LOG_DEC);
    addScalarFunctionConsumer.accept("LOG", MathFunctions.LOG_WITH_BASE);
    addScalarFunctionConsumer.accept("LOG", MathFunctions.LOG_WITH_BASE_DEC_DOU);
    addScalarFunctionConsumer.accept("LOG", MathFunctions.LOG_WITH_BASE_DOU_DEC);
    addScalarFunctionConsumer.accept("LOG", MathFunctions.LOG_WITH_BASE_DEC_DEC);
    addScalarFunctionConsumer.accept("SINH", MathFunctions.SINH);
    addScalarFunctionConsumer.accept("SINH", MathFunctions.SINH_DEC);
    addScalarFunctionConsumer.accept("COSH", MathFunctions.COSH);
    addScalarFunctionConsumer.accept("COSH", MathFunctions.COSH_DEC);
    addScalarFunctionConsumer.accept("TANH", MathFunctions.TANH);
    addScalarFunctionConsumer.accept("TANH", MathFunctions.TANH_DEC);
    addScalarFunctionConsumer.accept("UUID", MathFunctions.UUID);
    addScalarFunctionConsumer.accept("BIN", MathFunctions.BIN);
    addScalarFunctionConsumer.accept("HEX", MathFunctions.HEX_LONG);
    addScalarFunctionConsumer.accept("HEX", MathFunctions.HEX_STRING);

    addScalarFunctionConsumer.accept("FROM_BASE64", StringFunctions.FROMBASE64);
    addScalarFunctionConsumer.accept("TO_BASE64", StringFunctions.TOBASE64);
    addScalarFunctionConsumer.accept("LPAD", StringFunctions.LPAD);
    addScalarFunctionConsumer.accept("RPAD", StringFunctions.RPAD);
    addScalarFunctionConsumer.accept("REGEXP_REPLACE", StringFunctions.REGEXP_REPLACE);
    addScalarFunctionConsumer.accept("REGEXP_EXTRACT", StringFunctions.REGEXP_EXTRACT);

    addScalarFunctionConsumer.accept("LTRIM", BuiltInMethod.LTRIM.method);
    addScalarFunctionConsumer.accept("RTRIM", BuiltInMethod.RTRIM.method);

    addScalarFunctionConsumer.accept("MD5", StringFunctions.MD5);
    addScalarFunctionConsumer.accept("SHA1", StringFunctions.SHA1);
    addScalarFunctionConsumer.accept("SHA224", StringFunctions.SHA224);
    addScalarFunctionConsumer.accept("SHA256", StringFunctions.SHA256);
    addScalarFunctionConsumer.accept("SHA384", StringFunctions.SHA384);
    addScalarFunctionConsumer.accept("SHA512", StringFunctions.SHA512);
    addScalarFunctionConsumer.accept("SHA2", StringFunctions.SHA2);
}
 

/**
 * Creates a generic SqlUDF operator from a Pig UDF.
 *
 * @param udfName Name of the UDF
 * @param method Method "exec" for implementing the UDF
 * @param funcSpec Pig Funcspec
 * @param inputType Argument type for the input
 * @param returnType Function return data type
 */
static SqlUserDefinedFunction createGeneralPigUdf(String udfName,
    Method method, FuncSpec funcSpec, RelDataType inputType,
    RelDataType returnType) {
  return new PigUserDefinedFunction(udfName, opBinding -> returnType,
      OperandTypes.ANY, Collections.singletonList(inputType),
      ScalarFunctionImpl.createUnsafe(method), funcSpec);
}
 

/**
 * Register custom function from given static method with this {@link SQLExecEnvironment}
 *
 * @param name Name of the scalar SQL function that needs make available to SQL Statement
 * @param clazz {@link Class} which contains given static method
 * @param methodName Name of the method from given clazz
 *
 * @return Return this {@link SQLExecEnvironment}
 */
public SQLExecEnvironment registerFunction(String name, Class clazz, String methodName)
{
  Preconditions.checkNotNull(name, "Function name cannot be null");
  ScalarFunction scalarFunction = ScalarFunctionImpl.create(clazz, methodName);
  return registerFunction(name, scalarFunction);
}