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

Example 11 with UserDefinedFunctionResponse

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

the class ExampleUserDefinedFunctionHandlerTest method testConcatenateMethod.

@Test
public void testConcatenateMethod() throws Exception {
    Schema inputSchema = SchemaBuilder.newBuilder().addListField("list", Types.MinorType.VARCHAR.getType()).build();
    Schema outputSchema = SchemaBuilder.newBuilder().addField("string", Types.MinorType.VARCHAR.getType()).build();
    Block inputRecords = allocator.createBlock(inputSchema);
    inputRecords.setRowCount(1);
    FieldVector fieldVector = inputRecords.getFieldVector("list");
    List<String> value = Lists.newArrayList("a", "b");
    BlockUtils.setComplexValue(fieldVector, 0, FieldResolver.DEFAULT, value);
    UserDefinedFunctionResponse response = runAndAssertSerialization(inputRecords, outputSchema, "concatenate");
    Block outputRecords = response.getRecords();
    assertEquals(1, outputRecords.getRowCount());
    FieldReader fieldReader = outputRecords.getFieldReader("string");
    ArrowValueProjector arrowValueProjector = ProjectorUtils.createArrowValueProjector(fieldReader);
    assertEquals(exampleUserDefinedFunctionHandler.concatenate(Lists.newArrayList("a", "b")), arrowValueProjector.project(0));
}

Example 12 with UserDefinedFunctionResponse

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

the class UserDefinedFunctionHandlerTest method testInvocationWithStructType.

@Test
public void testInvocationWithStructType() throws Exception {
    int rowCount = 20;
    UserDefinedFunctionRequest udfRequest = createUDFRequest(rowCount, Map.class, "test_row_type", true, Map.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);
        Map<String, Object> actual = (Map) UnitTestBlockUtils.getValue(fieldReader, pos);
        Map<String, Object> input = ImmutableMap.of("intVal", pos + 100, "doubleVal", pos + 200.2);
        Map<String, Object> expected = handler.test_row_type(input);
        for (Map.Entry<String, Object> entry : expected.entrySet()) {
            String key = entry.getKey();
            assertTrue(actual.containsKey(key));
            assertEquals(expected.get(key), actual.get(key));
        }
    }
}

Example 13 with UserDefinedFunctionResponse

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

the class UserDefinedFunctionHandlerTest method testInvocationWithListType.

@Test
public void testInvocationWithListType() throws Exception {
    int rowCount = 20;
    UserDefinedFunctionRequest udfRequest = createUDFRequest(rowCount, List.class, "test_list_type", true, List.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);
        List<Integer> result = (List) UnitTestBlockUtils.getValue(fieldReader, pos);
        List<Integer> expected = handler.test_list_type(ImmutableList.of(pos + 100, pos + 200, pos + 300));
        assertArrayEquals(expected.toArray(), result.toArray());
    }
}

Example 14 with UserDefinedFunctionResponse

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

the class UserDefinedFunctionHandler method processScalarFunction.

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