implements RequestStreamHandler

{

private static final Logger logger = LoggerFactory.getLogger(UserDefinedFunctionHandler.class);

private static final int RETURN_COLUMN_COUNT = 1;

//Used to tag log lines generated by this connector for diagnostic purposes when interacting with Athena.

private final String sourceType;

public UserDefinedFunctionHandler(String sourceType)

{

this.sourceType = sourceType;

}

@Override

public final void handleRequest(InputStream inputStream, OutputStream outputStream, Context context)

{

try (BlockAllocator allocator = new BlockAllocatorImpl()) {

ObjectMapper objectMapper = VersionedObjectMapperFactory.create(allocator);

try (FederationRequest rawRequest = objectMapper.readValue(inputStream, FederationRequest.class)) {

if (rawRequest instanceof PingRequest) {

try (PingResponse response = doPing((PingRequest) rawRequest)) {

assertNotNull(response);

objectMapper.writeValue(outputStream, response);

}

return;

}

if (!(rawRequest instanceof UserDefinedFunctionRequest)) {

throw new RuntimeException("Expected a UserDefinedFunctionRequest but found "

+ rawRequest.getClass());

}

doHandleRequest(allocator, objectMapper, (UserDefinedFunctionRequest) rawRequest, outputStream);

}

catch (Exception ex) {

throw (ex instanceof RuntimeException) ? (RuntimeException) ex : new RuntimeException(ex);

}

}

}

protected final void doHandleRequest(BlockAllocator allocator,

ObjectMapper objectMapper,

UserDefinedFunctionRequest req,

OutputStream outputStream)

throws Exception

{

logger.info("doHandleRequest: request[{}]", req);

try (UserDefinedFunctionResponse response = processFunction(allocator, req)) {

logger.info("doHandleRequest: response[{}]", response);

assertNotNull(response);

objectMapper.writeValue(outputStream, response);

}

}

@VisibleForTesting

UserDefinedFunctionResponse processFunction(BlockAllocator allocator, UserDefinedFunctionRequest req)

throws Exception

{

UserDefinedFunctionType functionType = req.getFunctionType();

switch (functionType) {

case SCALAR:

return processScalarFunction(allocator, req);

default:

throw new UnsupportedOperationException("Unsupported function type " + functionType);

}

}

private UserDefinedFunctionResponse processScalarFunction(BlockAllocator allocator, UserDefinedFunctionRequest req)

throws Exception

{

Method udfMethod = extractScalarFunctionMethod(req);

Block inputRecords = req.getInputRecords();

Schema outputSchema = req.getOutputSchema();

Block outputRecords = processRows(allocator, udfMethod, inputRecords, outputSchema);

return new UserDefinedFunctionResponse(outputRecords, udfMethod.getName());

}

/**

* Processes a group by rows. This method takes in a block of data (containing multiple rows), process them and

* returns multiple rows of the output column in a block.

* <p>

* UDF methods are invoked row-by-row in a for loop. Arrow values are converted to Java Objects and then passed into

* the UDF java method. This is not very efficient because we might potentially be doing a lot of data copying.

* Advanced users could choose to override this method and directly deal with Arrow data to achieve better

* performance.

*

* @param allocator arrow memory allocator

* @param udfMethod the extracted java method matching the User-Defined-Function defined in Athena.

* @param inputRecords input data in Arrow format

* @param outputSchema output data schema in Arrow format

* @return output data in Arrow format

*/

protected Block processRows(BlockAllocator allocator, Method udfMethod, Block inputRecords, Schema outputSchema)

throws Exception

{

int rowCount = inputRecords.getRowCount();

List<ArrowValueProjector> valueProjectors = Lists.newArrayList();

for (Field field : inputRecords.getFields()) {

FieldReader fieldReader = inputRecords.getFieldReader(field.getName());

ArrowValueProjector arrowValueProjector = ProjectorUtils.createArrowValueProjector(fieldReader);

valueProjectors.add(arrowValueProjector);

}

Field outputField = outputSchema.getFields().get(0);

GeneratedRowWriter outputRowWriter = createOutputRowWriter(outputField, valueProjectors, udfMethod);

Block outputRecords = allocator.createBlock(outputSchema);

outputRecords.setRowCount(rowCount);

try {

for (int rowNum = 0; rowNum < rowCount; ++rowNum) {

outputRowWriter.writeRow(outputRecords, rowNum, rowNum);

}

}

catch (Throwable t) {

try {

outputRecords.close();

}

catch (Exception e) {

logger.error("Error closing output block", e);

}

throw t;

}

return outputRecords;

}

/**

* Use reflection to find tha java method that maches the UDF function defined in Athena SQL.

*

* @param req UDF request

* @return java method matching the UDF defined in Athena query.

*/

private Method extractScalarFunctionMethod(UserDefinedFunctionRequest req)

{

String methodName = req.getMethodName();

Class[] argumentTypes = extractJavaTypes(req.getInputRecords().getSchema());

Class[] returnTypes = extractJavaTypes(req.getOutputSchema());

checkState(returnTypes.length == RETURN_COLUMN_COUNT,

String.format("Expecting %d return columns, found %d in method signature.",

RETURN_COLUMN_COUNT, returnTypes.length));

Class returnType = returnTypes[0];

Method udfMethod;

try {

udfMethod = this.getClass().getMethod(methodName, argumentTypes);

logger.info(String.format("Found UDF method %s with input types [%s] and output types [%s]",

methodName, Arrays.toString(argumentTypes), returnType.getName()));

}

catch (NoSuchMethodException e) {

String msg = "Failed to find UDF method. " + e.getMessage()

+ " Please make sure the method name contains only lowercase and the method signature (name and"

+ " argument types) in Lambda matches the function signature defined in SQL.";

throw new RuntimeException(msg, e);

}

if (!returnType.equals(udfMethod.getReturnType())) {

throw new IllegalArgumentException("signature return type " + returnType +

" does not match udf implementation return type " + udfMethod.getReturnType());

}

return udfMethod;

}

private Class[] extractJavaTypes(Schema schema)

{

Class[] types = new Class[schema.getFields().size()];

List<Field> fields = schema.getFields();

for (int i = 0; i < fields.size(); ++i) {

Types.MinorType minorType = Types.getMinorTypeForArrowType(fields.get(i).getType());

types[i] = BlockUtils.getJavaType(minorType);

}

return types;

}

private final PingResponse doPing(PingRequest request)

{

PingResponse response = new PingResponse(request.getCatalogName(), request.getQueryId(), sourceType, CAPABILITIES, SERDE_VERSION);

try {

onPing(request);

}

catch (Exception ex) {

logger.warn("doPing: encountered an exception while delegating onPing.", ex);

}

return response;

}

protected void onPing(PingRequest request)

{

//NoOp

}

private void assertNotNull(FederationResponse response)

{

if (response == null) {

throw new RuntimeException("Response was null");

}

}

private GeneratedRowWriter createOutputRowWriter(Field outputField, List<ArrowValueProjector> valueProjectors, Method udfMethod)

{

GeneratedRowWriter.RowWriterBuilder builder = GeneratedRowWriter.newBuilder();

Extractor extractor = makeExtractor(outputField, valueProjectors, udfMethod);

if (extractor != null) {

builder.withExtractor(outputField.getName(), extractor);

}

else {

builder.withFieldWriterFactory(outputField.getName(), makeFactory(outputField, valueProjectors, udfMethod));

}

return builder.build();

}

/**

* Creates an Extractor for the given outputField.

* @param outputField outputField

* @param valueProjectors projectors that we use to read input data.

* @param udfMethod

* @return

*/

private Extractor makeExtractor(Field outputField, List<ArrowValueProjector> valueProjectors, Method udfMethod)

{

Types.MinorType fieldType = Types.getMinorTypeForArrowType(outputField.getType());

Object[] arguments = new Object[valueProjectors.size()];

switch (fieldType) {

case INT:

return (IntExtractor) (Object inputRowNum, NullableIntHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (int) result;

}

};

case DATEMILLI:

return (DateMilliExtractor) (Object inputRowNum, NullableDateMilliHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = ((LocalDateTime) result).atZone(BlockUtils.UTC_ZONE_ID).toInstant().toEpochMilli();

}

};

case DATEDAY:

return (DateDayExtractor) (Object inputRowNum, NullableDateDayHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (int) ((LocalDate) result).toEpochDay();

}

};

case TINYINT:

return (TinyIntExtractor) (Object inputRowNum, NullableTinyIntHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (byte) result;

}

};

case SMALLINT:

return (SmallIntExtractor) (Object inputRowNum, NullableSmallIntHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (short) result;

}

};

case FLOAT4:

return (Float4Extractor) (Object inputRowNum, NullableFloat4Holder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (float) result;

}

};

case FLOAT8:

return (Float8Extractor) (Object inputRowNum, NullableFloat8Holder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (double) result;

}

};

case DECIMAL:

return (DecimalExtractor) (Object inputRowNum, NullableDecimalHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = ((BigDecimal) result);

}

};

case BIT:

return (BitExtractor) (Object inputRowNum, NullableBitHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = ((boolean) result) ? 1 : 0;

}

};

case BIGINT:

return (BigIntExtractor) (Object inputRowNum, NullableBigIntHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (long) result;

}

};

case VARCHAR:

return (VarCharExtractor) (Object inputRowNum, NullableVarCharHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = ((String) result);

}

};

case VARBINARY:

return (VarBinaryExtractor) (Object inputRowNum, NullableVarBinaryHolder dst) ->

{

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

if (result == null) {

dst.isSet = 0;

}

else {

dst.isSet = 1;

dst.value = (byte[]) result;

}

};

default:

return null;

}

}

private FieldWriterFactory makeFactory(Field field, List<ArrowValueProjector> valueProjectors, Method udfMethod)

{

Object[] arguments = new Object[valueProjectors.size()];

Types.MinorType fieldType = Types.getMinorTypeForArrowType(field.getType());

switch (fieldType) {

case LIST:

case STRUCT:

return (FieldVector vector, Extractor extractor, ConstraintProjector ignored) ->

(Object inputRowNum, int outputRowNum) -> {

Object result = invokeMethod(udfMethod, arguments, (int) inputRowNum, valueProjectors);

BlockUtils.setComplexValue(vector, outputRowNum, FieldResolver.DEFAULT, result);

return true; // push-down does not apply in UDFs

};

default:

throw new IllegalArgumentException("Unsupported type " + fieldType);

}

}

private Object invokeMethod(Method udfMethod,

Object[] arguments,

int inputRowNum,

List<ArrowValueProjector> valueProjectors)

{

for (int col = 0; col < valueProjectors.size(); ++col) {

arguments[col] = valueProjectors.get(col).project(inputRowNum);

}

try {

return udfMethod.invoke(this, arguments);

}

catch (IllegalAccessException e) {

throw new RuntimeException(e);

}

catch (InvocationTargetException e) {

if (Objects.isNull(e)) {

throw new RuntimeException(e);

}

throw new RuntimeException(e.getCause());

}

catch (IllegalArgumentException e) {

String msg = String.format("%s. Expected function types %s, got types %s",

e.getMessage(),

Arrays.stream(udfMethod.getParameterTypes()).map(clazz -> clazz.getName()).collect(Collectors.toList()),

Arrays.stream(arguments).map(arg -> arg.getClass().getName()).collect(Collectors.toList()));

throw new RuntimeException(msg, e);

}

}