Examples with Block com.amazonaws.athena.connector.lambda.data.Block used on opensource projects

Example 1 with Block

use of com.amazonaws.athena.connector.lambda.data.Block 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 Block

use of com.amazonaws.athena.connector.lambda.data.Block 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 Block

use of com.amazonaws.athena.connector.lambda.data.Block 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 Block

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

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