Examples with UserDefinedFunctionRequest com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest used on opensource projects

Example 1 with UserDefinedFunctionRequest

use of com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest in project aws-athena-query-federation by awslabs.

the class ExampleUserDefinedFunctionHandlerTest method runAndAssertSerialization.

private UserDefinedFunctionResponse runAndAssertSerialization(Block inputRecords, Schema outputSchema, String methodName) throws IOException {
    UserDefinedFunctionRequest request = new UserDefinedFunctionRequest(IdentityUtil.fakeIdentity(), inputRecords, outputSchema, methodName, UserDefinedFunctionType.SCALAR);
    ObjectMapperUtil.assertSerialization(request);
    ByteArrayOutputStream out = new ByteArrayOutputStream();
    mapper.writeValue(out, request);
    ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(out.toByteArray());
    ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
    exampleUserDefinedFunctionHandler.handleRequest(byteArrayInputStream, outputStream, null);
    UserDefinedFunctionResponse udfResponse = (UserDefinedFunctionResponse) mapper.readValue(outputStream.toByteArray(), FederationResponse.class);
    ObjectMapperUtil.assertSerialization(udfResponse);
    return udfResponse;
}

Example 2 with UserDefinedFunctionRequest

use of com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest in project aws-athena-query-federation by awslabs.

the class UserDefinedFunctionHandlerTest method testMethodNotFound.

@Test
public void testMethodNotFound() {
    int rowCount = 20;
    UserDefinedFunctionRequest udfRequest = createUDFRequest(rowCount, Integer.class, "method_that_does_not_exsit", true, Integer.class, Integer.class);
    try {
        UserDefinedFunctionResponse udfResponse = handler.processFunction(allocator, udfRequest);
        fail("Expected function to fail due to method not found, but succeeded.");
    } catch (Exception e) {
        assertTrue(e.getCause() instanceof NoSuchMethodException);
    }
}

Example 3 with UserDefinedFunctionRequest

use of com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest in project aws-athena-query-federation by awslabs.

the class UserDefinedFunctionHandlerTest method testInvocationWithNullVAlue.

@Test
public void testInvocationWithNullVAlue() throws Exception {
    int rowCount = 20;
    UserDefinedFunctionRequest udfRequest = createUDFRequest(rowCount, Boolean.class, "test_scalar_function_with_null_value", false, Integer.class);
    UserDefinedFunctionResponse udfResponse = handler.processFunction(allocator, udfRequest);
    Block responseBlock = udfResponse.getRecords();
    assertEquals(1, responseBlock.getFieldReaders().size());
    assertEquals(rowCount, responseBlock.getRowCount());
    FieldReader fieldReader = responseBlock.getFieldReaders().get(0);
    for (int pos = 0; pos < rowCount; ++pos) {
        fieldReader.setPosition(pos);
        assertTrue(fieldReader.isSet());
        Boolean expected = handler.test_scalar_function_with_null_value(null);
        Boolean actual = fieldReader.readBoolean();
        assertEquals(expected, actual);
    }
}

Example 4 with UserDefinedFunctionRequest

use of com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest in project aws-athena-query-federation by awslabs.

the class UserDefinedFunctionHandlerTest method testInvocationWithBasicType.

@Test
public void testInvocationWithBasicType() throws Exception {
    int rowCount = 20;
    UserDefinedFunctionRequest udfRequest = createUDFRequest(rowCount, Integer.class, "test_scalar_udf", true, Integer.class, Integer.class);
    UserDefinedFunctionResponse udfResponse = handler.processFunction(allocator, udfRequest);
    Block responseBlock = udfResponse.getRecords();
    assertEquals(1, responseBlock.getFieldReaders().size());
    assertEquals(rowCount, responseBlock.getRowCount());
    FieldReader fieldReader = responseBlock.getFieldReaders().get(0);
    for (int pos = 0; pos < rowCount; ++pos) {
        fieldReader.setPosition(pos);
        int val = (int) UnitTestBlockUtils.getValue(fieldReader, pos);
        int expected = handler.test_scalar_udf(pos + 100, pos + 100);
        assertEquals(expected, val);
    }
}

Example 5 with UserDefinedFunctionRequest

use of com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionRequest in project aws-athena-query-federation by awslabs.

the class UserDefinedFunctionRequestSerDeTest method deserialize.

@Test
public void deserialize() throws IOException {
    logger.info("deserialize: enter");
    InputStream input = new ByteArrayInputStream(expectedSerDeText.getBytes());
    UserDefinedFunctionRequest actual = (UserDefinedFunctionRequest) mapper.readValue(input, FederationRequest.class);
    logger.info("deserialize: deserialized[{}]", actual);
    assertEquals(expected, actual);
    logger.info("deserialize: exit");
}