org.apache.calcite.sql.validate.SqlUserDefinedFunction Java Examples

@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;
}
 

public RexCallImplementor get(final SqlOperator operator) {
  if (operator instanceof SqlUserDefinedFunction) {
    org.apache.calcite.schema.Function udf =
        ((SqlUserDefinedFunction) operator).getFunction();
    if (!(udf instanceof ImplementableFunction)) {
      throw new IllegalStateException("User defined function " + operator
          + " must implement ImplementableFunction");
    }
    CallImplementor implementor =
        ((ImplementableFunction) udf).getImplementor();
    return wrapAsRexCallImplementor(implementor);
  } else if (operator instanceof SqlTypeConstructorFunction) {
    return map.get(SqlStdOperatorTable.ROW);
  }
  return map.get(operator);
}
 

@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());
}
 

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 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));
}
 

@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);
}
 

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());
}
 

@Override
public TupleExpression visitCall(RexCall call) {
    SqlOperator op = call.getOperator();
    if (op instanceof SqlCastFunction) {
        return call.getOperands().get(0).accept(this);
    } else if (op instanceof SqlUserDefinedFunction) {
        if (op.getName().equals("QUARTER")) {
            return visitFirstRexInputRef(call);
        }
    }

    TupleExpression tupleExpression;
    switch (op.getKind()) {
    case PLUS:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.PLUS);
        break;
    case MINUS:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.MINUS);
        break;
    case TIMES:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.MULTIPLE);
        break;
    case DIVIDE:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.DIVIDE);
        break;
    case CASE:
        tupleExpression = getCaseTupleExpression(call);
        break;
    default:
        tupleExpression = getRexCallTupleExpression(call);
    }
    if (ifVerify) {
        tupleExpression.verify();
    }
    return tupleExpression;
}
 

private SqlOperator getSqlOperator(SamzaSqlScalarFunctionImpl scalarFunction, List<UdfMetadata> udfMetadataList) {
  int numArguments = scalarFunction.numberOfArguments();
  UdfMetadata udfMetadata = scalarFunction.getUdfMetadata();

  if (udfMetadata.isDisableArgCheck()) {
    return new SqlUserDefinedFunction(new SqlIdentifier(scalarFunction.getUdfName(), SqlParserPos.ZERO),
      o -> scalarFunction.getReturnType(o.getTypeFactory()), null, Checker.ANY_CHECKER, null, scalarFunction);
  } else {
    return new SqlUserDefinedFunction(new SqlIdentifier(scalarFunction.getUdfName(), SqlParserPos.ZERO),
      o -> scalarFunction.getReturnType(o.getTypeFactory()), null,
      Checker.getChecker(numArguments, numArguments, udfMetadata), null, scalarFunction);
  }
}
 

/** Converts a function to a {@link org.apache.calcite.sql.SqlOperator}.
 *
 * <p>The {@code typeFactory} argument is technical debt; see [CALCITE-2082]
 * Remove RelDataTypeFactory argument from SqlUserDefinedAggFunction
 * constructor. */
private static SqlOperator toOp(RelDataTypeFactory typeFactory,
    SqlIdentifier name, final Function function) {
  List<RelDataType> argTypes = new ArrayList<>();
  List<SqlTypeFamily> typeFamilies = new ArrayList<>();
  for (FunctionParameter o : function.getParameters()) {
    final RelDataType type = o.getType(typeFactory);
    argTypes.add(type);
    typeFamilies.add(
        Util.first(type.getSqlTypeName().getFamily(), SqlTypeFamily.ANY));
  }
  final FamilyOperandTypeChecker typeChecker =
      OperandTypes.family(typeFamilies, i ->
          function.getParameters().get(i).isOptional());
  final List<RelDataType> paramTypes = toSql(typeFactory, argTypes);
  if (function instanceof ScalarFunction) {
    return new SqlUserDefinedFunction(name, infer((ScalarFunction) function),
        InferTypes.explicit(argTypes), typeChecker, paramTypes, function);
  } else if (function instanceof AggregateFunction) {
    return new SqlUserDefinedAggFunction(name,
        infer((AggregateFunction) function), InferTypes.explicit(argTypes),
        typeChecker, (AggregateFunction) function, false, false,
        Optionality.FORBIDDEN, typeFactory);
  } else if (function instanceof TableMacro) {
    return new SqlUserDefinedTableMacro(name, ReturnTypes.CURSOR,
        InferTypes.explicit(argTypes), typeChecker, paramTypes,
        (TableMacro) function);
  } else if (function instanceof TableFunction) {
    return new SqlUserDefinedTableFunction(name, ReturnTypes.CURSOR,
        InferTypes.explicit(argTypes), typeChecker, paramTypes,
        (TableFunction) function);
  } else {
    throw new AssertionError("unknown function type " + function);
  }
}
 

/**
 * Converts a Pig UDF, given its {@link FuncSpec} and a list of relational
 * operands (function arguments). To call this function, the arguments of
 * Pig functions need to be converted into the relational types before.
 *
 * @param builder The relational builder
 * @param pigFunc Pig function description
 * @param operands Relational operands for the function
 * @param returnType Function return data type
 * @return The SQL calls equivalent to the Pig function
 */
static RexNode convertPigFunction(PigRelBuilder builder, FuncSpec pigFunc,
    ImmutableList<RexNode> operands, RelDataType returnType) throws FrontendException {
  // First, check the map for the direct mapping SQL builtin
  final SqlOperator operator = BUILTIN_FUNC.get(pigFunc.getClassName());
  if (operator != null) {
    return builder.call(operator, operands);
  }

  // If no mapping found, build the argument wrapper to convert the relation operands
  // into a Pig tuple so that the Pig function can consume it.
  try {
    // Find the implementation method for the Pig function from
    // the class defining the UDF.
    final Class clazz = Class.forName(pigFunc.getClassName());
    final Method method =
        PIG_UDF_FINDER.findPigUdfImplementationMethod(clazz);

    // Now create the argument wrapper. Depend on the type of the UDF, the
    // relational operands are converted into a Pig Tuple or Pig DataBag
    // with the appropriate wrapper.
    final SqlUserDefinedFunction convertOp =
        Accumulator.class.isAssignableFrom(clazz)
            ? PigRelSqlUdfs.createPigBagUDF(operands)
            : PigRelSqlUdfs.createPigTupleUDF(operands);
    final RexNode rexTuple = builder.call(convertOp, operands);

    // Then convert the Pig function into a @SqlUserDefinedFunction.
    SqlUserDefinedFunction userFuncOp =
        PigRelSqlUdfs.createGeneralPigUdf(clazz.getSimpleName(),
            method, pigFunc, rexTuple.getType(), returnType);

    // Ready to return SqlCall after having SqlUDF and operand
    return builder.call(userFuncOp, ImmutableList.of(rexTuple));
  } catch (ClassNotFoundException e) {
    throw new FrontendException("Cannot find the implementation for Pig UDF class: "
        + pigFunc.getClassName());
  }
}
 

/**
 * Rest of class is utility functions taken directly from CalciteCatalogReader. This is because that class consider these utilities to be private concerns.
 */
private SqlOperator toOp(SqlIdentifier name, final Function function) {
  List<RelDataType> argTypes = new ArrayList<>();
  List<SqlTypeFamily> typeFamilies = new ArrayList<>();
  for (FunctionParameter o : function.getParameters()) {
    final RelDataType type = o.getType(typeFactory);
    argTypes.add(type);
    typeFamilies.add(
        Util.first(type.getSqlTypeName().getFamily(), SqlTypeFamily.ANY));
  }
  final Predicate<Integer> optional =
      new Predicate<Integer>() {
        @Override
        public boolean apply(Integer input) {
          return function.getParameters().get(input).isOptional();
        }
      };
  final FamilyOperandTypeChecker typeChecker =
      OperandTypes.family(typeFamilies, optional);
  final List<RelDataType> paramTypes = toSql(argTypes);
  if (function instanceof ScalarFunction) {
    return new SqlUserDefinedFunction(name, infer((ScalarFunction) function),
        InferTypes.explicit(argTypes), typeChecker, paramTypes, function);
  } else if (function instanceof AggregateFunction) {
    return new SqlUserDefinedAggFunction(name,
        infer((AggregateFunction) function), InferTypes.explicit(argTypes),
        typeChecker, (AggregateFunction) function, false, false, typeFactory);
  } else if (function instanceof TableMacro) {
    return new SqlUserDefinedTableMacro(name, ReturnTypes.CURSOR,
        InferTypes.explicit(argTypes), typeChecker, paramTypes,
        (TableMacro) function);
  } else if (function instanceof TableFunction) {
    return new SqlUserDefinedTableFunction(name, ReturnTypes.CURSOR,
        InferTypes.explicit(argTypes), typeChecker, paramTypes,
        (TableFunction) function);
  } else {
    throw new AssertionError("unknown function type " + function);
  }
}
 

/**
 * Creates a Pig Tuple from a list of relational operands.
 *
 * @param operands Relational operands
 * @return Pig Tuple SqlUDF
 */
static SqlUserDefinedFunction createPigTupleUDF(ImmutableList<RexNode> operands) {
  return new PigUserDefinedFunction("PIG_TUPLE",
      infer(PigRelSqlUdfs.PIG_TUPLE_FUNC),
      OperandTypes.family(getTypeFamilies(operands)),
      getRelDataTypes(operands),
      PigRelSqlUdfs.PIG_TUPLE_FUNC);
}
 

/**
 * Creates a Pig DataBag from a list of relational operands.
 *
 * @param operands Relational operands
 * @return Pig DataBag SqlUDF
 */
static SqlUserDefinedFunction createPigBagUDF(ImmutableList<RexNode> operands) {
  return new PigUserDefinedFunction(
      "PIG_BAG",
      infer(PigRelSqlUdfs.PIG_BAG_FUNC),
      OperandTypes.family(getTypeFamilies(operands)),
      getRelDataTypes(operands),
      PigRelSqlUdfs.PIG_BAG_FUNC);
}
 

/** Converts a function to a {@link org.apache.calcite.sql.SqlOperator}.
 *
 * <p>The {@code typeFactory} argument is technical debt; see [CALCITE-2082]
 * Remove RelDataTypeFactory argument from SqlUserDefinedAggFunction
 * constructor. */
private static SqlOperator toOp(RelDataTypeFactory typeFactory,
    SqlIdentifier name, final Function function) {
  List<RelDataType> argTypes = new ArrayList<>();
  List<SqlTypeFamily> typeFamilies = new ArrayList<>();
  for (FunctionParameter o : function.getParameters()) {
    final RelDataType type = o.getType(typeFactory);
    argTypes.add(type);
    typeFamilies.add(
        Util.first(type.getSqlTypeName().getFamily(), SqlTypeFamily.ANY));
  }
  final FamilyOperandTypeChecker typeChecker =
      OperandTypes.family(typeFamilies, i ->
          function.getParameters().get(i).isOptional());
  final List<RelDataType> paramTypes = toSql(typeFactory, argTypes);
  if (function instanceof ScalarFunction) {
    return new SqlUserDefinedFunction(name, infer((ScalarFunction) function),
        InferTypes.explicit(argTypes), typeChecker, paramTypes, function);
  } else if (function instanceof AggregateFunction) {
    return new SqlUserDefinedAggFunction(name,
        infer((AggregateFunction) function), InferTypes.explicit(argTypes),
        typeChecker, (AggregateFunction) function, false, false,
        Optionality.FORBIDDEN, typeFactory);
  } else if (function instanceof TableMacro) {
    return new SqlUserDefinedTableMacro(name, ReturnTypes.CURSOR,
        InferTypes.explicit(argTypes), typeChecker, paramTypes,
        (TableMacro) function);
  } else if (function instanceof TableFunction) {
    return new SqlUserDefinedTableFunction(name, ReturnTypes.CURSOR,
        InferTypes.explicit(argTypes), typeChecker, paramTypes,
        (TableFunction) function);
  } else {
    throw new AssertionError("unknown function type " + function);
  }
}
 

@Override
public TupleExpression visitCall(RexCall call) {
    SqlOperator op = call.getOperator();
    if (op instanceof SqlCastFunction) {
        return call.getOperands().get(0).accept(this);
    } else if (op instanceof SqlUserDefinedFunction) {
        if (op.getName().equals("QUARTER")) {
            return visitFirstRexInputRef(call);
        }
    }

    TupleExpression tupleExpression;
    switch (op.getKind()) {
    case PLUS:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.PLUS);
        break;
    case MINUS:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.MINUS);
        break;
    case TIMES:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.MULTIPLE);
        break;
    case DIVIDE:
        tupleExpression = getBinaryTupleExpression(call, TupleExpression.ExpressionOperatorEnum.DIVIDE);
        break;
    case CASE:
        tupleExpression = getCaseTupleExpression(call);
        break;
    default:
        tupleExpression = getRexCallTupleExpression(call);
    }
    if (ifVerify) {
        tupleExpression.verify();
    }
    return tupleExpression;
}
 

private boolean isFlatten(RexNode rexNode) {
  return rexNode instanceof RexCall && ((RexCall) rexNode).op instanceof SqlUserDefinedFunction
      && ((RexCall) rexNode).op.getName().equalsIgnoreCase("flatten");
}
 

/**
 * 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);
}