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

Example 6 with UserDefinedFunctionRequest

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

the class CompositeHandler method handleRequest.

/**
 * Handles routing the request to the appropriate Handler, either MetadataHandler or RecordHandler.
 *
 * @param allocator The BlockAllocator to use for Apache Arrow Resources.
 * @param rawReq The request object itself.
 * @param outputStream The OutputStream to which all responses should be written.
 * @param objectMapper The ObjectMapper that can be used for serializing responses.
 * @throws Exception
 * @note that PingRequests are routed to the MetadataHandler even though both MetadataHandler and RecordHandler
 * implemented PingRequest handling.
 */
public final void handleRequest(BlockAllocator allocator, FederationRequest rawReq, OutputStream outputStream, ObjectMapper objectMapper) throws Exception {
    if (rawReq instanceof PingRequest) {
        try (PingResponse response = metadataHandler.doPing((PingRequest) rawReq)) {
            assertNotNull(response);
            objectMapper.writeValue(outputStream, response);
        }
        return;
    }
    if (rawReq instanceof MetadataRequest) {
        metadataHandler.doHandleRequest(allocator, objectMapper, (MetadataRequest) rawReq, outputStream);
    } else if (rawReq instanceof RecordRequest) {
        recordHandler.doHandleRequest(allocator, objectMapper, (RecordRequest) rawReq, outputStream);
    } else if (udfhandler != null && rawReq instanceof UserDefinedFunctionRequest) {
        udfhandler.doHandleRequest(allocator, objectMapper, (UserDefinedFunctionRequest) rawReq, outputStream);
    } else {
        throw new IllegalArgumentException("Unknown request class " + rawReq.getClass());
    }
}

Example 7 with UserDefinedFunctionRequest

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

the class UserDefinedFunctionHandler method handleRequest.

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

Example 8 with UserDefinedFunctionRequest

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

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

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

the class UserDefinedFunctionHandlerTest method createUDFRequest.

private UserDefinedFunctionRequest createUDFRequest(int rowCount, Class returnType, String methodName, boolean nonNullData, Class... argumentTypes) {
    Schema inputSchema = buildSchema(argumentTypes);
    Schema outputSchema = buildSchema(returnType);
    Block block = allocator.createBlock(inputSchema);
    block.setRowCount(rowCount);
    if (nonNullData) {
        writeData(block, rowCount);
    }
    return new UserDefinedFunctionRequest(null, block, outputSchema, methodName, SCALAR);
}