use of com.facebook.presto.thrift.api.udf.ThriftUdfPage in project presto by prestodb.
the class ThriftSqlFunctionExecutor method toSqlFunctionResult.
private SqlFunctionResult toSqlFunctionResult(ThriftUdfResult result, Type returnType) {
ThriftUdfPage page = result.getResult();
switch(page.getPageFormat()) {
case PRESTO_THRIFT:
return new SqlFunctionResult(getOnlyElement(page.getThriftPage().getThriftBlocks()).toBlock(returnType), result.getUdfStats().getTotalCpuTimeMs());
case PRESTO_SERIALIZED:
checkState(blockEncodingSerde != null, "blockEncodingSerde not set");
PagesSerde pagesSerde = new PagesSerde(blockEncodingSerde, Optional.empty(), Optional.empty(), Optional.empty());
return new SqlFunctionResult(pagesSerde.deserialize(page.getPrestoPage().toSerializedPage()).getBlock(0), result.getUdfStats().getTotalCpuTimeMs());
default:
throw new IllegalArgumentException(format("Unknown page format: %s", page.getPageFormat()));
}
}
use of com.facebook.presto.thrift.api.udf.ThriftUdfPage in project presto by prestodb.
the class ThriftSqlFunctionExecutor method executeFunction.
@Override
public CompletableFuture<SqlFunctionResult> executeFunction(String source, RemoteScalarFunctionImplementation functionImplementation, Page input, List<Integer> channels, List<Type> argumentTypes, Type returnType) {
ThriftUdfPage page = buildThriftPage(functionImplementation, input, channels, argumentTypes);
SqlFunctionHandle functionHandle = functionImplementation.getFunctionHandle();
SqlFunctionId functionId = functionHandle.getFunctionId();
ThriftFunctionHandle thriftFunctionHandle = new ThriftFunctionHandle(functionId.getFunctionName().toString(), functionId.getArgumentTypes().stream().map(TypeSignature::toString).collect(toImmutableList()), returnType.toString(), functionHandle.getVersion());
ThriftUdfService thriftUdfService = thriftUdfClient.get(Optional.of(functionImplementation.getLanguage().getLanguage()));
return invokeUdfWithRetry(thriftUdfService, new ThriftUdfRequest(source, thriftFunctionHandle, page)).thenApply(thriftResult -> toSqlFunctionResult(thriftResult, returnType));
}
use of com.facebook.presto.thrift.api.udf.ThriftUdfPage in project presto by prestodb.
the class EchoFirstInputThriftUdfService method invokeUdf.
@Override
public ListenableFuture<ThriftUdfResult> invokeUdf(ThriftUdfRequest request) throws ThriftUdfServiceException {
ThriftUdfPage inputs = request.getInputs();
ThriftUdfPage result;
switch(inputs.getPageFormat()) {
case PRESTO_THRIFT:
PrestoThriftPage thriftPage = inputs.getThriftPage();
if (thriftPage.getThriftBlocks().isEmpty()) {
throw toThriftUdfServiceException(false, EchoFirstInputErrorCode.GENERIC_INTERNAL_ERROR.toErrorCode(), new UnsupportedOperationException("No input to echo"));
}
result = thriftPage(new PrestoThriftPage(ImmutableList.of(thriftPage.getThriftBlocks().get(0)), inputs.getThriftPage().getPositionCount()));
break;
case PRESTO_SERIALIZED:
PagesSerde pagesSerde = new PagesSerde(blockEncodingSerde, Optional.empty(), Optional.empty(), Optional.empty());
Page page = pagesSerde.deserialize(inputs.getPrestoPage().toSerializedPage());
if (page.getChannelCount() == 0) {
throw toThriftUdfServiceException(false, EchoFirstInputErrorCode.GENERIC_INTERNAL_ERROR.toErrorCode(), new UnsupportedOperationException("No input to echo"));
}
result = prestoPage(pagesSerde.serialize(new Page(page.getBlock(0))));
break;
default:
throw new UnsupportedOperationException();
}
return immediateFuture(new ThriftUdfResult(result, new ThriftUdfStats(100)));
}
Aggregations