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

Example 6 with UserDefinedFunctionResponse

use of com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionResponse 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 7 with UserDefinedFunctionResponse

use of com.amazonaws.athena.connector.lambda.udf.UserDefinedFunctionResponse 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 8 with UserDefinedFunctionResponse

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

the class UserDefinedFunctionHandler method doHandleRequest.

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);
    }
}

Example 9 with UserDefinedFunctionResponse

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

the class UserDefinedFunctionResponseSerDeTest method deserialize.

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

Example 10 with UserDefinedFunctionResponse

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

the class UserDefinedFunctionResponseSerDeTest method beforeTest.

@Before
public void beforeTest() throws IOException {
    String productField = "product";
    Schema outputSchema = SchemaBuilder.newBuilder().addField(productField, Types.MinorType.INT.getType()).build();
    Block records = allocator.createBlock(outputSchema);
    int num_records = 10;
    for (int i = 0; i < num_records; i++) {
        BlockUtils.setValue(records.getFieldVector(productField), i, i * 2);
    }
    records.setRowCount(num_records);
    expected = new UserDefinedFunctionResponse(records, "test-method");
    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "UserDefinedFunctionResponse.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}