@Override
public ScalarFunctionImplementation specialize(
Map<String, Type> types,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
checkArgument(types.size() == 1, "Expected one type, got %s", types);
Type elementType = types.get("E");
MethodHandle methodHandle;
if (elementType.getJavaType() == void.class) {
methodHandle = METHOD_HANDLE_UNKNOWN;
} else if (elementType.getJavaType() == boolean.class) {
methodHandle = METHOD_HANDLE_BOOLEAN;
} else if (elementType.getJavaType() == long.class) {
methodHandle = METHOD_HANDLE_LONG;
} else if (elementType.getJavaType() == double.class) {
methodHandle = METHOD_HANDLE_DOUBLE;
} else if (elementType.getJavaType() == Slice.class) {
methodHandle = METHOD_HANDLE_SLICE;
} else {
methodHandle = METHOD_HANDLE_OBJECT;
}
methodHandle = methodHandle.bindTo(elementType);
requireNonNull(methodHandle, "methodHandle is null");
return new ScalarFunctionImplementation(
true, ImmutableList.of(false, false), methodHandle, isDeterministic());
}
private static Object nativeContainerToOrcValue(Type type, Object nativeValue) {
if (nativeValue == null) {
return null;
}
if (type instanceof DecimalType) {
BigInteger unscaledValue;
DecimalType decimalType = (DecimalType) type;
if (decimalType.isShort()) {
unscaledValue = BigInteger.valueOf((long) nativeValue);
} else {
unscaledValue = Decimals.decodeUnscaledValue((Slice) nativeValue);
}
return HiveDecimal.create(unscaledValue, decimalType.getScale());
}
if (type.getJavaType() == boolean.class) {
return nativeValue;
}
if (type.getJavaType() == long.class) {
return nativeValue;
}
if (type.getJavaType() == double.class) {
return nativeValue;
}
if (type.getJavaType() == Slice.class) {
Slice slice = (Slice) nativeValue;
return type instanceof VarcharType ? slice.toStringUtf8() : slice.getBytes();
}
if (isArrayType(type)) {
Block arrayBlock = (Block) nativeValue;
Type elementType = type.getTypeParameters().get(0);
List<Object> list = new ArrayList<>();
for (int i = 0; i < arrayBlock.getPositionCount(); i++) {
list.add(
nativeContainerToOrcValue(
elementType, getNativeContainerValue(elementType, arrayBlock, i)));
}
return list;
}
if (isMapType(type)) {
Block mapBlock = (Block) nativeValue;
Type keyType = type.getTypeParameters().get(0);
Type valueType = type.getTypeParameters().get(1);
Map<Object, Object> map = new HashMap<>();
for (int i = 0; i < mapBlock.getPositionCount(); i += 2) {
Object key =
nativeContainerToOrcValue(keyType, getNativeContainerValue(keyType, mapBlock, i));
Object value =
nativeContainerToOrcValue(
valueType, getNativeContainerValue(valueType, mapBlock, i + 1));
map.put(key, value);
}
return map;
}
throw new PrestoException(INTERNAL_ERROR, "Unimplemented type: " + type);
}
private static void verifyMethodSignature(
Method method,
TypeSignature returnTypeName,
List<TypeSignature> argumentTypeNames,
TypeManager typeManager) {
Type returnType = typeManager.getType(returnTypeName);
checkNotNull(returnType, "returnType is null");
List<Type> argumentTypes = resolveTypes(argumentTypeNames, typeManager);
checkArgument(
Primitives.unwrap(method.getReturnType()) == returnType.getJavaType(),
"Expected method %s return type to be %s (%s)",
method,
returnType.getJavaType().getName(),
returnType);
// skip Session argument
Class<?>[] parameterTypes = method.getParameterTypes();
Annotation[][] annotations = method.getParameterAnnotations();
if (parameterTypes.length > 0 && parameterTypes[0] == ConnectorSession.class) {
parameterTypes = Arrays.copyOfRange(parameterTypes, 1, parameterTypes.length);
annotations = Arrays.copyOfRange(annotations, 1, annotations.length);
}
for (int i = 0; i < parameterTypes.length; i++) {
Class<?> actualType = parameterTypes[i];
Type expectedType = argumentTypes.get(i);
boolean nullable =
!FluentIterable.from(Arrays.asList(annotations[i])).filter(Nullable.class).isEmpty();
// Only allow boxing for functions that need to see nulls
if (Primitives.isWrapperType(actualType)) {
checkArgument(
nullable,
"Method %s has parameter with type %s that is missing @Nullable",
method,
actualType);
}
if (nullable) {
checkArgument(
!NON_NULLABLE_ARGUMENT_TYPES.contains(actualType),
"Method %s has parameter type %s, but @Nullable is not supported on this type",
method,
actualType);
}
checkArgument(
Primitives.unwrap(actualType) == expectedType.getJavaType(),
"Expected method %s parameter %s type to be %s (%s)",
method,
i,
expectedType.getJavaType().getName(),
expectedType);
}
}
private static void appendTo(Type type, SliceOutput output, Block block) {
if (type.getJavaType() == long.class) {
output.appendLong(type.getLong(block, 0));
} else if (type.getJavaType() == double.class) {
output.appendDouble(type.getDouble(block, 0));
} else if (type.getJavaType() == Slice.class) {
output.appendBytes(type.getSlice(block, 0));
} else if (type.getJavaType() == boolean.class) {
output.appendByte(type.getBoolean(block, 0) ? 1 : 0);
} else {
throw new IllegalArgumentException("Unsupported type: " + type.getJavaType().getSimpleName());
}
}
@Override
public ScalarFunctionImplementation specialize(
BoundVariables boundVariables,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
Type fromType = boundVariables.getTypeVariable("F");
Type toType = boundVariables.getTypeVariable("T");
Class<?> returnType = Primitives.wrap(toType.getJavaType());
MethodHandle tryCastHandle;
if (fromType.equals(UNKNOWN)) {
tryCastHandle = dropArguments(constant(returnType, null), 0, Void.class);
} else {
// the resulting method needs to return a boxed type
Signature signature = functionRegistry.getCoercion(fromType, toType);
MethodHandle coercion =
functionRegistry.getScalarFunctionImplementation(signature).getMethodHandle();
coercion = coercion.asType(methodType(returnType, coercion.type()));
MethodHandle exceptionHandler =
dropArguments(constant(returnType, null), 0, RuntimeException.class);
tryCastHandle = catchException(coercion, RuntimeException.class, exceptionHandler);
}
return new ScalarFunctionImplementation(
true, ImmutableList.of(true), tryCastHandle, isDeterministic());
}
private static Block toBlock(Type type, SliceInput input, int length) {
BlockBuilder builder = type.createBlockBuilder(new BlockBuilderStatus(), 1, length);
if (type.getJavaType() == long.class) {
type.writeLong(builder, input.readLong());
} else if (type.getJavaType() == double.class) {
type.writeDouble(builder, input.readDouble());
} else if (type.getJavaType() == Slice.class) {
type.writeSlice(builder, input.readSlice(length));
} else if (type.getJavaType() == boolean.class) {
type.writeBoolean(builder, input.readByte() != 0);
} else {
throw new IllegalArgumentException("Unsupported type: " + type.getJavaType().getSimpleName());
}
return builder.build();
}
private static Object getNativeContainerValue(Type type, Block block, int position) {
if (block.isNull(position)) {
return null;
} else if (type.getJavaType() == boolean.class) {
return type.getBoolean(block, position);
} else if (type.getJavaType() == long.class) {
return type.getLong(block, position);
} else if (type.getJavaType() == double.class) {
return type.getDouble(block, position);
} else if (type.getJavaType() == Slice.class) {
return type.getSlice(block, position);
} else if (type.getJavaType() == Block.class) {
return type.getObject(block, position);
} else {
throw new AssertionError("Unimplemented type: " + type);
}
}
@Override
public ScalarFunctionImplementation specialize(
BoundVariables boundVariables,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
Type argumentType = boundVariables.getTypeVariable("T");
Type returnType = boundVariables.getTypeVariable("U");
return new ScalarFunctionImplementation(
true,
ImmutableList.of(true, false),
METHOD_HANDLE.asType(
METHOD_HANDLE
.type()
.changeReturnType(wrap(returnType.getJavaType()))
.changeParameterType(0, wrap(argumentType.getJavaType()))),
isDeterministic());
}
@Override
public ScalarFunctionImplementation specialize(
BoundVariables boundVariables,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
checkArgument(boundVariables.getTypeVariables().size() == 1, "Expected only one type");
Type type = boundVariables.getTypeVariable("T");
MethodHandle identity = MethodHandles.identity(type.getJavaType());
return new ScalarFunctionImplementation(
false, ImmutableList.of(false), identity, isDeterministic());
}
public static Row extractRow(Page page, int position, List<Type> types) {
checkArgument(page.getChannelCount() == types.size(), "channelCount does not match");
checkArgument(
position < page.getPositionCount(),
"Requested position %s from a page with positionCount %s ",
position,
page.getPositionCount());
RowBuilder rowBuilder = new RowBuilder();
for (int channel = 0; channel < page.getChannelCount(); channel++) {
Block block = page.getBlock(channel);
Type type = types.get(channel);
int size;
Object value = getNativeContainerValue(type, block, position);
if (value == null) {
size = SIZE_OF_BYTE;
} else if (type.getJavaType() == boolean.class) {
size = SIZE_OF_BYTE;
} else if (type.getJavaType() == long.class) {
size = SIZE_OF_LONG;
} else if (type.getJavaType() == double.class) {
size = SIZE_OF_DOUBLE;
} else if (type.getJavaType() == Slice.class) {
size = ((Slice) value).length();
} else if (type.getJavaType() == Block.class) {
size = ((Block) value).getSizeInBytes();
} else {
throw new AssertionError("Unimplemented type: " + type);
}
rowBuilder.add(nativeContainerToOrcValue(type, value), size);
}
Row row = rowBuilder.build();
verify(
row.getColumns().size() == types.size(),
"Column count in row: %s Expected column count: %s",
row.getColumns().size(),
types.size());
return row;
}
@Override
public FunctionInfo specialize(
Map<String, Type> types,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
checkArgument(types.size() == 1, "Can only construct arrays from exactly matching types");
ImmutableList.Builder<Class<?>> builder = ImmutableList.builder();
Type type = types.get("E");
for (int i = 0; i < arity; i++) {
if (type.getJavaType().isPrimitive()) {
builder.add(Primitives.wrap(type.getJavaType()));
} else {
builder.add(type.getJavaType());
}
}
ImmutableList<Class<?>> stackTypes = builder.build();
Class<?> clazz = generateArrayConstructor(stackTypes, type);
MethodHandle methodHandle;
try {
Method method =
clazz.getMethod("arrayConstructor", stackTypes.toArray(new Class<?>[stackTypes.size()]));
methodHandle = lookup().unreflect(method);
} catch (ReflectiveOperationException e) {
throw Throwables.propagate(e);
}
List<Boolean> nullableParameters =
ImmutableList.copyOf(Collections.nCopies(stackTypes.size(), true));
return new FunctionInfo(
arrayConstructorSignature(
parameterizedTypeName("array", type.getTypeSignature()),
Collections.nCopies(arity, type.getTypeSignature())),
"Constructs an array of the given elements",
true,
methodHandle,
true,
false,
nullableParameters);
}
@Override
public ScalarFunctionImplementation specialize(
Map<String, Type> types,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
checkArgument(types.size() == 1, "Expected one type, got %s", types);
Type elementType = types.get("E");
MethodHandle methodHandle;
if (METHOD_HANDLES.containsKey(elementType.getJavaType())) {
methodHandle = METHOD_HANDLES.get(elementType.getJavaType());
} else {
checkArgument(
!elementType.getJavaType().isPrimitive(),
"Unsupported primitive type: " + elementType.getJavaType());
methodHandle = OBJECT_METHOD_HANDLE;
}
requireNonNull(methodHandle, "methodHandle is null");
methodHandle = methodHandle.bindTo(elementType);
return new ScalarFunctionImplementation(
true, ImmutableList.of(false, false), methodHandle, isDeterministic());
}
private Block createZeroBlock(Type type, int rowsCount, Slice constantSlice) {
checkArgument(isSupportedType(type), "Unsupported type [%s]", type);
Slice slice;
// do not exceed varchar limit
if (isVarcharType(type)) {
slice =
constantSlice.slice(
0, Math.min(((VarcharType) type).getLength(), constantSlice.length()));
} else if (isLongDecimal(type)) {
slice = encodeScaledValue(ZERO);
} else {
slice = constantSlice;
}
BlockBuilder builder;
if (type instanceof FixedWidthType) {
builder = type.createBlockBuilder(new BlockBuilderStatus(), rowsCount);
} else {
builder = type.createBlockBuilder(new BlockBuilderStatus(), rowsCount, slice.length());
}
for (int i = 0; i < rowsCount; i++) {
Class<?> javaType = type.getJavaType();
if (javaType == boolean.class) {
type.writeBoolean(builder, false);
} else if (javaType == long.class) {
type.writeLong(builder, 0);
} else if (javaType == double.class) {
type.writeDouble(builder, 0.0);
} else if (javaType == Slice.class) {
requireNonNull(slice, "slice is null");
type.writeSlice(builder, slice, 0, slice.length());
} else {
throw new UnsupportedOperationException("Unknown javaType: " + javaType.getName());
}
}
return builder.build();
}
private static InternalAggregationFunction generateAggregation(Type type) {
DynamicClassLoader classLoader =
new DynamicClassLoader(ArbitraryAggregation.class.getClassLoader());
List<Type> inputTypes = ImmutableList.of(type);
MethodHandle inputFunction;
MethodHandle outputFunction;
Class<? extends AccumulatorState> stateInterface;
AccumulatorStateSerializer<?> stateSerializer;
if (type.getJavaType() == long.class) {
stateInterface = NullableLongState.class;
stateSerializer = compiler.generateStateSerializer(stateInterface, classLoader);
inputFunction = LONG_INPUT_FUNCTION;
outputFunction = LONG_OUTPUT_FUNCTION;
} else if (type.getJavaType() == double.class) {
stateInterface = NullableDoubleState.class;
stateSerializer = compiler.generateStateSerializer(stateInterface, classLoader);
inputFunction = DOUBLE_INPUT_FUNCTION;
outputFunction = DOUBLE_OUTPUT_FUNCTION;
} else if (type.getJavaType() == Slice.class) {
stateInterface = SliceState.class;
stateSerializer = compiler.generateStateSerializer(stateInterface, classLoader);
inputFunction = SLICE_INPUT_FUNCTION;
outputFunction = SLICE_OUTPUT_FUNCTION;
} else if (type.getJavaType() == boolean.class) {
stateInterface = NullableBooleanState.class;
stateSerializer = compiler.generateStateSerializer(stateInterface, classLoader);
inputFunction = BOOLEAN_INPUT_FUNCTION;
outputFunction = BOOLEAN_OUTPUT_FUNCTION;
} else {
stateInterface = BlockState.class;
stateSerializer = new BlockStateSerializer(type);
inputFunction = BLOCK_INPUT_FUNCTION;
outputFunction = BLOCK_OUTPUT_FUNCTION;
}
inputFunction = inputFunction.bindTo(type);
Type intermediateType = stateSerializer.getSerializedType();
List<ParameterMetadata> inputParameterMetadata = createInputParameterMetadata(type);
AggregationMetadata metadata =
new AggregationMetadata(
generateAggregationName(NAME, type, inputTypes),
inputParameterMetadata,
inputFunction,
inputParameterMetadata,
inputFunction,
null,
outputFunction.bindTo(type),
stateInterface,
stateSerializer,
compiler.generateStateFactory(stateInterface, classLoader),
type,
false);
GenericAccumulatorFactoryBinder factory =
new AccumulatorCompiler().generateAccumulatorFactoryBinder(metadata, classLoader);
return new InternalAggregationFunction(
NAME, inputTypes, intermediateType, type, true, false, factory);
}
private Class<?> generateProjectMethod(
CallSiteBinder callSiteBinder,
ClassDefinition classDefinition,
String methodName,
RowExpression projection,
boolean sourceIsCursor) {
MethodDefinition projectionMethod;
if (sourceIsCursor) {
projectionMethod =
classDefinition.declareMethod(
new CompilerContext(BOOTSTRAP_METHOD),
a(PUBLIC),
methodName,
type(void.class),
arg("cursor", RecordCursor.class),
arg("output", BlockBuilder.class));
} else {
ImmutableList.Builder<NamedParameterDefinition> parameters = ImmutableList.builder();
parameters.add(arg("position", int.class));
parameters.addAll(toBlockParameters(getInputChannels(projection)));
parameters.add(arg("output", BlockBuilder.class));
projectionMethod =
classDefinition.declareMethod(
new CompilerContext(BOOTSTRAP_METHOD),
a(PUBLIC),
methodName,
type(void.class),
parameters.build());
}
projectionMethod.comment("Projection: %s", projection.toString());
// generate body code
CompilerContext context = projectionMethod.getCompilerContext();
context.declareVariable(type(boolean.class), "wasNull");
Block getSessionByteCode =
new Block(context)
.pushThis()
.getField(classDefinition.getType(), "session", type(ConnectorSession.class));
ByteCodeNode body =
compileExpression(callSiteBinder, projection, sourceIsCursor, context, getSessionByteCode);
projectionMethod
.getBody()
.comment("boolean wasNull = false;")
.putVariable("wasNull", false)
.getVariable("output")
.append(body);
Type projectionType = projection.getType();
Block notNullBlock = new Block(context);
if (projectionType.getJavaType() == boolean.class) {
notNullBlock
.comment("output.append(<booleanStackValue>);")
.invokeInterface(BlockBuilder.class, "appendBoolean", BlockBuilder.class, boolean.class)
.pop();
} else if (projectionType.getJavaType() == long.class) {
notNullBlock
.comment("output.append(<longStackValue>);")
.invokeInterface(BlockBuilder.class, "appendLong", BlockBuilder.class, long.class)
.pop();
} else if (projectionType.getJavaType() == double.class) {
notNullBlock
.comment("output.append(<doubleStackValue>);")
.invokeInterface(BlockBuilder.class, "appendDouble", BlockBuilder.class, double.class)
.pop();
} else if (projectionType.getJavaType() == Slice.class) {
notNullBlock
.comment("output.append(<sliceStackValue>);")
.invokeInterface(BlockBuilder.class, "appendSlice", BlockBuilder.class, Slice.class)
.pop();
} else {
throw new UnsupportedOperationException("Type " + projectionType + " can not be output yet");
}
Block nullBlock =
new Block(context)
.comment("output.appendNull();")
.pop(projectionType.getJavaType())
.invokeInterface(BlockBuilder.class, "appendNull", BlockBuilder.class)
.pop();
projectionMethod
.getBody()
.comment("if the result was null, appendNull; otherwise append the value")
.append(
new IfStatement(
context, new Block(context).getVariable("wasNull"), nullBlock, notNullBlock))
.ret();
return projectionType.getJavaType();
}
public static Expression toExpression(Object object, Type type) {
requireNonNull(type, "type is null");
if (object instanceof Expression) {
return (Expression) object;
}
if (object == null) {
if (type.equals(UNKNOWN)) {
return new NullLiteral();
}
return new Cast(new NullLiteral(), type.getTypeSignature().toString(), false, true);
}
checkArgument(
Primitives.wrap(type.getJavaType()).isInstance(object),
"object.getClass (%s) and type.getJavaType (%s) do not agree",
object.getClass(),
type.getJavaType());
if (type.equals(BIGINT)) {
return new LongLiteral(object.toString());
}
if (type.equals(DOUBLE)) {
Double value = (Double) object;
// WARNING: the ORC predicate code depends on NaN and infinity not appearing in a tuple
// domain, so
// if you remove this, you will need to update the TupleDomainOrcPredicate
if (value.isNaN()) {
return new FunctionCall(new QualifiedName("nan"), ImmutableList.<Expression>of());
} else if (value.equals(Double.NEGATIVE_INFINITY)) {
return ArithmeticUnaryExpression.negative(
new FunctionCall(new QualifiedName("infinity"), ImmutableList.<Expression>of()));
} else if (value.equals(Double.POSITIVE_INFINITY)) {
return new FunctionCall(new QualifiedName("infinity"), ImmutableList.<Expression>of());
} else {
return new DoubleLiteral(object.toString());
}
}
if (type instanceof VarcharType) {
if (object instanceof String) {
object = Slices.utf8Slice((String) object);
}
if (object instanceof Slice) {
Slice value = (Slice) object;
int length = SliceUtf8.countCodePoints(value);
if (length == ((VarcharType) type).getLength()) {
return new StringLiteral(value.toStringUtf8());
}
return new Cast(
new StringLiteral(value.toStringUtf8()), type.getDisplayName(), false, true);
}
throw new IllegalArgumentException(
"object must be instance of Slice or String when type is VARCHAR");
}
if (type.equals(BOOLEAN)) {
return new BooleanLiteral(object.toString());
}
if (object instanceof Block) {
SliceOutput output = new DynamicSliceOutput(((Block) object).getSizeInBytes());
BlockSerdeUtil.writeBlock(output, (Block) object);
object = output.slice();
// This if condition will evaluate to true: object instanceof Slice && !type.equals(VARCHAR)
}
if (object instanceof Slice) {
// HACK: we need to serialize VARBINARY in a format that can be embedded in an expression to
// be
// able to encode it in the plan that gets sent to workers.
// We do this by transforming the in-memory varbinary into a call to
// from_base64(<base64-encoded value>)
FunctionCall fromBase64 =
new FunctionCall(
new QualifiedName("from_base64"),
ImmutableList.of(
new StringLiteral(VarbinaryFunctions.toBase64((Slice) object).toStringUtf8())));
Signature signature = FunctionRegistry.getMagicLiteralFunctionSignature(type);
return new FunctionCall(new QualifiedName(signature.getName()), ImmutableList.of(fromBase64));
}
Signature signature = FunctionRegistry.getMagicLiteralFunctionSignature(type);
Expression rawLiteral = toExpression(object, FunctionRegistry.typeForMagicLiteral(type));
return new FunctionCall(new QualifiedName(signature.getName()), ImmutableList.of(rawLiteral));
}
