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

Example 1 with UserDefinedFunctionResponse

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

the class ExampleUserDefinedFunctionHandlerTest method testGetDefaultValueIfNullMethod.

@Test
public void testGetDefaultValueIfNullMethod() throws Exception {
    Schema inputSchema = SchemaBuilder.newBuilder().addField("input", Types.MinorType.BIGINT.getType()).build();
    Schema outputSchema = SchemaBuilder.newBuilder().addField("output", Types.MinorType.BIGINT.getType()).build();
    Block inputRecords = allocator.createBlock(inputSchema);
    inputRecords.setRowCount(2);
    BigIntVector fieldVector = (BigIntVector) inputRecords.getFieldVector("input");
    fieldVector.setSafe(0, 123l);
    fieldVector.setNull(1);
    UserDefinedFunctionResponse response = runAndAssertSerialization(inputRecords, outputSchema, "get_default_value_if_null");
    Block outputRecords = response.getRecords();
    assertEquals(2, outputRecords.getRowCount());
    FieldReader fieldReader = outputRecords.getFieldReader("output");
    ArrowValueProjector arrowValueProjector = ProjectorUtils.createArrowValueProjector(fieldReader);
    assertEquals(exampleUserDefinedFunctionHandler.get_default_value_if_null(123l), arrowValueProjector.project(0));
    assertEquals(exampleUserDefinedFunctionHandler.get_default_value_if_null(null), arrowValueProjector.project(1));
}

Example 2 with UserDefinedFunctionResponse

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

the class ExampleUserDefinedFunctionHandlerTest method testMultiplyMethod.

@Test
public void testMultiplyMethod() throws Exception {
    Schema inputSchema = SchemaBuilder.newBuilder().addField("factor1", Types.MinorType.INT.getType()).addField("factor2", Types.MinorType.INT.getType()).build();
    Schema outputSchema = SchemaBuilder.newBuilder().addField("product", Types.MinorType.INT.getType()).build();
    Block inputRecords = allocator.createBlock(inputSchema);
    inputRecords.setRowCount(1);
    IntVector inputVector1 = (IntVector) inputRecords.getFieldVector("factor1");
    IntVector inputVector2 = (IntVector) inputRecords.getFieldVector("factor2");
    inputVector1.setSafe(0, 2);
    inputVector2.setSafe(0, 3);
    UserDefinedFunctionResponse response = runAndAssertSerialization(inputRecords, outputSchema, "multiply");
    Block outputRecords = response.getRecords();
    assertEquals(1, outputRecords.getRowCount());
    FieldReader fieldReader = outputRecords.getFieldReader("product");
    ArrowValueProjector arrowValueProjector = ProjectorUtils.createArrowValueProjector(fieldReader);
    assertEquals(exampleUserDefinedFunctionHandler.multiply(2, 3), arrowValueProjector.project(0));
}

Example 3 with UserDefinedFunctionResponse

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

the class ExampleUserDefinedFunctionHandlerTest method testToJsonMethod.

@Test
public void testToJsonMethod() throws Exception {
    Schema inputSchema = SchemaBuilder.newBuilder().addStructField("struct").addChildField("struct", "int", Types.MinorType.INT.getType()).addChildField("struct", "double", Types.MinorType.FLOAT8.getType()).addChildField("struct", "string", Types.MinorType.VARCHAR.getType()).build();
    Schema outputSchema = SchemaBuilder.newBuilder().addField("json", Types.MinorType.VARCHAR.getType()).build();
    Block inputRecords = allocator.createBlock(inputSchema);
    inputRecords.setRowCount(1);
    FieldVector fieldVector = inputRecords.getFieldVector("struct");
    Map<String, Object> struct = new HashMap<>();
    struct.put("int", 10);
    struct.put("double", 2.3);
    struct.put("string", "test_string");
    BlockUtils.setComplexValue(fieldVector, 0, FieldResolver.DEFAULT, struct);
    UserDefinedFunctionResponse response = runAndAssertSerialization(inputRecords, outputSchema, "to_json");
    Block outputRecords = response.getRecords();
    assertEquals(1, outputRecords.getRowCount());
    FieldReader fieldReader = outputRecords.getFieldReader("json");
    ArrowValueProjector arrowValueProjector = ProjectorUtils.createArrowValueProjector(fieldReader);
    assertEquals(exampleUserDefinedFunctionHandler.to_json(struct), arrowValueProjector.project(0));
}

Example 4 with UserDefinedFunctionResponse

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

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