@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 void generateGetSerializedType(
ClassDefinition definition, List<StateField> fields, CallSiteBinder callSiteBinder) {
CompilerContext compilerContext = new CompilerContext();
Block body =
definition
.declareMethod(compilerContext, a(PUBLIC), "getSerializedType", type(Type.class))
.getBody();
Type type;
if (fields.size() > 1) {
type = VARCHAR;
} else {
Class<?> stackType = fields.get(0).getType();
if (stackType == long.class) {
type = BIGINT;
} else if (stackType == double.class) {
type = DOUBLE;
} else if (stackType == boolean.class) {
type = BOOLEAN;
} else if (stackType == byte.class) {
type = BIGINT;
} else if (stackType == Slice.class) {
type = VARCHAR;
} else {
throw new IllegalArgumentException("Unsupported type: " + stackType);
}
}
body.comment("return %s", type.getName())
.append(constantType(new CompilerContext(BOOTSTRAP_METHOD), callSiteBinder, type))
.retObject();
}
private static Block serializeList(
Type type, BlockBuilder builder, Object object, ListObjectInspector inspector) {
List<?> list = inspector.getList(object);
if (list == null) {
requireNonNull(builder, "parent builder is null").appendNull();
return null;
}
List<Type> typeParameters = type.getTypeParameters();
checkArgument(typeParameters.size() == 1, "list must have exactly 1 type parameter");
Type elementType = typeParameters.get(0);
ObjectInspector elementInspector = inspector.getListElementObjectInspector();
BlockBuilder currentBuilder;
if (builder != null) {
currentBuilder = builder.beginBlockEntry();
} else {
currentBuilder = elementType.createBlockBuilder(new BlockBuilderStatus(), list.size());
}
for (Object element : list) {
serializeObject(elementType, currentBuilder, element, elementInspector);
}
if (builder != null) {
builder.closeEntry();
return null;
} else {
Block resultBlock = currentBuilder.build();
return resultBlock;
}
}
@Override
public boolean equals(Object obj) {
boolean retval = true;
if (obj instanceof Field) {
Field field = (Field) obj;
if (type.equals(field.getType())) {
if (this.isNull() && field.isNull()) {
retval = true;
} else if (this.isNull() ^ field.isNull()) {
retval = false;
} else if (type.equals(VARBINARY)) {
// special case for byte arrays
// aren't they so fancy
retval = Arrays.equals((byte[]) value, (byte[]) field.getObject());
} else if (type.equals(DATE) || type.equals(TIME) || type.equals(TIMESTAMP)) {
retval = value.toString().equals(field.getObject().toString());
} else {
if (value instanceof Block) {
retval = equals((Block) value, (Block) field.getObject());
} else {
retval = value.equals(field.getObject());
}
}
}
}
return retval;
}
public static Signature nullIfSignature(Type returnType, Type firstType, Type secondType) {
return new Signature(
NULL_IF,
returnType.getTypeSignature(),
firstType.getTypeSignature(),
secondType.getTypeSignature());
}
public static Signature subscriptSignature(Type returnType, Type leftType, Type rightType) {
return internalOperator(
SUBSCRIPT.name(),
returnType.getTypeSignature(),
leftType.getTypeSignature(),
rightType.getTypeSignature());
}
public ListenableFuture<?> partitionPage(Page page) {
requireNonNull(page, "page is null");
Page partitionFunctionArgs = getPartitionFunctionArguments(page);
for (int position = 0; position < page.getPositionCount(); position++) {
if (nullChannel.isPresent() && page.getBlock(nullChannel.getAsInt()).isNull(position)) {
for (PageBuilder pageBuilder : pageBuilders) {
pageBuilder.declarePosition();
for (int channel = 0; channel < sourceTypes.size(); channel++) {
Type type = sourceTypes.get(channel);
type.appendTo(page.getBlock(channel), position, pageBuilder.getBlockBuilder(channel));
}
}
} else {
int partition = partitionFunction.getPartition(partitionFunctionArgs, position);
PageBuilder pageBuilder = pageBuilders.get(partition);
pageBuilder.declarePosition();
for (int channel = 0; channel < sourceTypes.size(); channel++) {
Type type = sourceTypes.get(channel);
type.appendTo(page.getBlock(channel), position, pageBuilder.getBlockBuilder(channel));
}
}
}
return flush(false);
}
public ParquetListConverter(Type prestoType, String columnName, GroupType listType) {
checkArgument(
listType.getFieldCount() == 1,
"Expected LIST column '%s' to only have one field, but has %s fields",
columnName,
listType.getFieldCount());
checkArgument(ARRAY.equals(prestoType.getTypeSignature().getBase()));
this.arrayType = prestoType;
// The Parquet specification requires that the element value of a
// LIST type be wrapped in an inner repeated group, like so:
//
// optional group listField (LIST) {
// repeated group list {
// optional int element
// }
// }
//
// However, some parquet libraries don't follow this spec. The
// compatibility rules used here are specified in the Parquet
// documentation at http://git.io/vOpNz.
parquet.schema.Type elementType = listType.getType(0);
if (isElementType(elementType, listType.getName())) {
elementConverter =
createConverter(
prestoType.getTypeParameters().get(0), columnName + ".element", elementType);
} else {
elementConverter =
new ParquetListEntryConverter(
prestoType.getTypeParameters().get(0), columnName, elementType.asGroupType());
}
}
private static void assertBlockEquals(Type type, Block actual, Block expected) {
for (int position = 0; position < actual.getPositionCount(); position++) {
assertEquals(
type.getObjectValue(SESSION, actual, position),
type.getObjectValue(SESSION, expected, position));
}
}
private Type coerceToSingleType(
AnalysisContext context, String message, List<Expression> expressions) {
// determine super type
Type superType = UNKNOWN;
for (Expression expression : expressions) {
Optional<Type> newSuperType = getCommonSuperType(superType, process(expression, context));
if (!newSuperType.isPresent()) {
throw new SemanticException(TYPE_MISMATCH, expression, message, superType);
}
superType = newSuperType.get();
}
// verify all expressions can be coerced to the superType
for (Expression expression : expressions) {
Type type = process(expression, context);
if (!type.equals(superType)) {
if (!canCoerce(type, superType)) {
throw new SemanticException(TYPE_MISMATCH, expression, message, superType);
}
expressionCoercions.put(expression, superType);
}
}
return superType;
}
private Type coerceToSingleType(
AnalysisContext context, Node node, String message, Expression first, Expression second) {
Type firstType = null;
if (first != null) {
firstType = process(first, context);
}
Type secondType = null;
if (second != null) {
secondType = process(second, context);
}
if (firstType == null) {
return secondType;
}
if (secondType == null) {
return firstType;
}
if (firstType.equals(secondType)) {
return firstType;
}
// coerce types if possible
if (canCoerce(firstType, secondType)) {
expressionCoercions.put(first, secondType);
return secondType;
}
if (canCoerce(secondType, firstType)) {
expressionCoercions.put(second, firstType);
return firstType;
}
throw new SemanticException(TYPE_MISMATCH, node, message, firstType, secondType);
}
@Test
public void testRowTypeLookup() throws Exception {
functionAssertions.getMetadata().getType(parseTypeSignature("row<bigint>('a')"));
Type type = functionAssertions.getMetadata().getType(parseTypeSignature("row<bigint>('b')"));
assertEquals(type.getTypeSignature().getParameters().size(), 1);
assertEquals(
type.getTypeSignature().getParameters().get(0).getNamedTypeSignature().getName(), "b");
}
public static Signature betweenSignature(Type valueType, Type minType, Type maxType) {
return internalOperator(
"BETWEEN",
parseTypeSignature(StandardTypes.BOOLEAN),
valueType.getTypeSignature(),
minType.getTypeSignature(),
maxType.getTypeSignature());
}
private void checkFieldType(int field, Type expected) {
Type actual = getType(field);
checkArgument(
actual.equals(expected),
"Expected field %s to be type %s but is %s",
field,
expected,
actual);
}
private static Block createVarWidthValueBlock(int positionCount, int mapSize) {
Type valueType = createUnboundedVarcharType();
BlockBuilder valueBlockBuilder =
valueType.createBlockBuilder(new BlockBuilderStatus(), positionCount * mapSize);
for (int i = 0; i < positionCount * mapSize; i++) {
int wordLength = ThreadLocalRandom.current().nextInt(5, 10);
valueType.writeSlice(valueBlockBuilder, utf8Slice(randomString(wordLength)));
}
return valueBlockBuilder.build();
}
public static Signature arithmeticExpressionSignature(
ArithmeticBinaryExpression.Type expressionType,
Type returnType,
Type leftType,
Type rightType) {
return internalOperator(
expressionType.name(),
returnType.getTypeSignature(),
leftType.getTypeSignature(),
rightType.getTypeSignature());
}
@Override
protected Type visitArrayConstructor(ArrayConstructor node, AnalysisContext context) {
Type type =
coerceToSingleType(
context, "All ARRAY elements must be the same type: %s", node.getValues());
Type arrayType =
typeManager.getParameterizedType(
ARRAY.getName(), ImmutableList.of(type.getTypeSignature()), ImmutableList.of());
expressionTypes.put(node, arrayType);
return arrayType;
}
@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());
}
@Override
public FunctionInfo specialize(
Map<String, Type> types,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
Type valueType = types.get("T");
Signature signature =
new Signature(NAME, AGGREGATE, valueType.getTypeSignature(), valueType.getTypeSignature());
InternalAggregationFunction aggregation = generateAggregation(valueType);
return new FunctionInfo(signature, getDescription(), aggregation);
}
protected void checkPageSource(
ConnectorPageSource pageSource, List<TestColumn> testColumns, List<Type> types)
throws IOException {
try {
MaterializedResult result = materializeSourceDataStream(SESSION, pageSource, types);
for (MaterializedRow row : result) {
for (int i = 0, testColumnsSize = testColumns.size(); i < testColumnsSize; i++) {
TestColumn testColumn = testColumns.get(i);
Type type = types.get(i);
Object actualValue = row.getField(i);
Object expectedValue = testColumn.getExpectedValue();
if (actualValue == null) {
assertEquals(null, expectedValue, String.format("Expected non-null for column %d", i));
} else if (testColumn.getObjectInspector().getTypeName().equals("float")
|| testColumn.getObjectInspector().getTypeName().equals("double")) {
assertEquals((double) actualValue, (double) expectedValue, EPSILON);
} else if (testColumn.getObjectInspector().getTypeName().equals("date")) {
SqlDate expectedDate = new SqlDate(((Long) expectedValue).intValue());
assertEquals(actualValue, expectedDate);
} else if (testColumn.getObjectInspector().getTypeName().equals("timestamp")) {
SqlTimestamp expectedTimestamp =
new SqlTimestamp((Long) expectedValue, SESSION.getTimeZoneKey());
assertEquals(actualValue, expectedTimestamp);
} else if (testColumn.getObjectInspector().getCategory() == Category.PRIMITIVE) {
if (expectedValue instanceof Slice) {
expectedValue = ((Slice) expectedValue).toStringUtf8();
}
if (actualValue instanceof Slice) {
actualValue = ((Slice) actualValue).toStringUtf8();
}
if (actualValue instanceof SqlVarbinary) {
actualValue = new String(((SqlVarbinary) actualValue).getBytes(), UTF_8);
}
assertEquals(actualValue, expectedValue, String.format("Wrong value for column %d", i));
} else {
BlockBuilder builder = type.createBlockBuilder(new BlockBuilderStatus(), 1);
type.writeObject(builder, expectedValue);
expectedValue = type.getObjectValue(SESSION, builder.build(), 0);
assertEquals(
actualValue,
expectedValue,
String.format("Wrong value for column %s", testColumn.getName()));
}
}
}
} finally {
pageSource.close();
}
}
private Object selectSingleValue(Operator operator) {
Page output = getAtMostOnePage(operator, SOURCE_PAGE);
assertNotNull(output);
assertEquals(output.getPositionCount(), 1);
assertEquals(output.getChannelCount(), 1);
Type type = operator.getTypes().get(0);
Block block = output.getBlock(0);
assertEquals(block.getPositionCount(), 1);
return type.getObjectValue(session.toConnectorSession(), block, 0);
}
public ParquetPageSource(
ParquetReader parquetReader,
ParquetDataSource dataSource,
MessageType fileSchema,
MessageType requestedSchema,
long totalBytes,
Properties splitSchema,
List<HiveColumnHandle> columns,
TupleDomain<HiveColumnHandle> effectivePredicate,
TypeManager typeManager,
boolean useParquetColumnNames) {
checkArgument(totalBytes >= 0, "totalBytes is negative");
requireNonNull(splitSchema, "splitSchema is null");
requireNonNull(columns, "columns is null");
requireNonNull(effectivePredicate, "effectivePredicate is null");
this.parquetReader = parquetReader;
this.dataSource = dataSource;
this.requestedSchema = requestedSchema;
this.totalBytes = totalBytes;
int size = columns.size();
this.constantBlocks = new Block[size];
this.hiveColumnIndexes = new int[size];
ImmutableList.Builder<String> namesBuilder = ImmutableList.builder();
ImmutableList.Builder<Type> typesBuilder = ImmutableList.builder();
for (int columnIndex = 0; columnIndex < size; columnIndex++) {
HiveColumnHandle column = columns.get(columnIndex);
checkState(column.getColumnType() == REGULAR, "column type must be regular");
String name = column.getName();
Type type = typeManager.getType(column.getTypeSignature());
namesBuilder.add(name);
typesBuilder.add(type);
hiveColumnIndexes[columnIndex] = column.getHiveColumnIndex();
if (getParquetType(column, fileSchema, useParquetColumnNames) == null) {
BlockBuilder blockBuilder =
type.createBlockBuilder(new BlockBuilderStatus(), MAX_VECTOR_LENGTH);
for (int i = 0; i < MAX_VECTOR_LENGTH; i++) {
blockBuilder.appendNull();
}
constantBlocks[columnIndex] = blockBuilder.build();
}
}
types = typesBuilder.build();
columnNames = namesBuilder.build();
}
private void coerceType(
AnalysisContext context, Expression expression, Type expectedType, String message) {
Type actualType = process(expression, context);
if (!actualType.equals(expectedType)) {
if (!canCoerce(actualType, expectedType)) {
throw new SemanticException(
TYPE_MISMATCH,
expression,
message + " must evaluate to a %s (actual: %s)",
expectedType,
actualType);
}
expressionCoercions.put(expression, expectedType);
}
}
private boolean processScalarFunction(Method method) throws IllegalAccessException {
ScalarFunction scalarFunction = method.getAnnotation(ScalarFunction.class);
if (scalarFunction == null) {
return false;
}
checkValidMethod(method);
MethodHandle methodHandle = lookup().unreflect(method);
String name = scalarFunction.value();
if (name.isEmpty()) {
name = camelToSnake(method.getName());
}
SqlType returnTypeAnnotation = method.getAnnotation(SqlType.class);
checkArgument(
returnTypeAnnotation != null,
"Method %s return type does not have a @SqlType annotation",
method);
Type returnType = type(typeManager, returnTypeAnnotation);
Signature signature =
new Signature(
name.toLowerCase(ENGLISH),
returnType.getTypeSignature(),
Lists.transform(parameterTypes(typeManager, method), Type::getTypeSignature));
verifyMethodSignature(
method, signature.getReturnType(), signature.getArgumentTypes(), typeManager);
List<Boolean> nullableArguments = getNullableArguments(method);
scalar(
signature,
methodHandle,
scalarFunction.deterministic(),
getDescription(method),
scalarFunction.hidden(),
method.isAnnotationPresent(Nullable.class),
nullableArguments);
for (String alias : scalarFunction.alias()) {
scalar(
signature.withAlias(alias.toLowerCase(ENGLISH)),
methodHandle,
scalarFunction.deterministic(),
getDescription(method),
scalarFunction.hidden(),
method.isAnnotationPresent(Nullable.class),
nullableArguments);
}
return true;
}
@Override
public FunctionInfo specialize(
Map<String, Type> types,
int arity,
TypeManager typeManager,
FunctionRegistry functionRegistry) {
Type type = types.get("T");
TypeSignature typeSignature = type.getTypeSignature();
return operatorInfo(
NOT_EQUAL,
RETURN_TYPE,
ImmutableList.of(typeSignature, typeSignature),
METHOD_HANDLE.bindTo(type),
false,
ImmutableList.of(false, false));
}
@TypeParameter("T")
@SqlType(StandardTypes.BOOLEAN)
@Nullable
public static Boolean contains(
@TypeParameter("T") Type elementType,
@OperatorDependency(
operator = EQUAL,
returnType = StandardTypes.BOOLEAN,
argumentTypes = {"T", "T"})
MethodHandle equals,
@SqlType("array(T)") Block arrayBlock,
@SqlType("T") Slice value) {
boolean foundNull = false;
for (int i = 0; i < arrayBlock.getPositionCount(); i++) {
if (arrayBlock.isNull(i)) {
foundNull = true;
continue;
}
try {
if ((boolean) equals.invokeExact(elementType.getSlice(arrayBlock, i), value)) {
return true;
}
} catch (Throwable t) {
Throwables.propagateIfInstanceOf(t, Error.class);
Throwables.propagateIfInstanceOf(t, PrestoException.class);
throw new PrestoException(INTERNAL_ERROR, t);
}
}
if (foundNull) {
return null;
}
return false;
}
public static void input(Type type, NullableBooleanState state, Block block, int position) {
if (!state.isNull()) {
return;
}
state.setNull(false);
state.setBoolean(type.getBoolean(block, position));
}
@Override
public void deserialize(Block block, int index, NullableDoubleState state) {
state.setNull(block.isNull(index));
if (!state.isNull()) {
state.setDouble(type.getDouble(block, index));
}
}
@TypeParameter("T")
@SqlType(StandardTypes.BIGINT)
public static long arrayPosition(
@TypeParameter("T") Type type,
@OperatorDependency(
operator = EQUAL,
returnType = StandardTypes.BOOLEAN,
argumentTypes = {"T", "T"})
MethodHandle equalMethodHandle,
@SqlType("array(T)") Block array,
@SqlType("T") double element) {
int size = array.getPositionCount();
for (int i = 0; i < size; i++) {
if (!array.isNull(i)) {
double arrayValue = type.getDouble(array, i);
try {
if ((boolean) equalMethodHandle.invokeExact(arrayValue, element)) {
return i + 1; // result is 1-based (instead of 0)
}
} catch (Throwable t) {
Throwables.propagateIfInstanceOf(t, Error.class);
Throwables.propagateIfInstanceOf(t, PrestoException.class);
throw new PrestoException(INTERNAL_ERROR, t);
}
}
}
return 0;
}
public static Signature arrayConstructorSignature(
Type returnType, List<? extends Type> argumentTypes) {
return internalFunction(
ARRAY_CONSTRUCTOR,
returnType.getTypeSignature(),
Lists.transform(argumentTypes, Type::getTypeSignature));
}