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

Example 61 with Block

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

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

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

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

Example 65 with Block

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

the class GetSplitsRequestSerDeTest method beforeTest.

@Before
public void beforeTest() throws IOException {
    String yearCol = "year";
    String monthCol = "month";
    String dayCol = "day";
    Schema schema = SchemaBuilder.newBuilder().addField(yearCol, new ArrowType.Int(32, true)).addField(monthCol, new ArrowType.Int(32, true)).addField(dayCol, new ArrowType.Int(32, true)).addField("col2", new ArrowType.Utf8()).addField("col3", Types.MinorType.FLOAT8.getType()).addField("col4", Types.MinorType.FLOAT8.getType()).addField("col5", Types.MinorType.FLOAT8.getType()).build();
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col3", SortedRangeSet.copyOf(Types.MinorType.FLOAT8.getType(), ImmutableList.of(Range.greaterThan(allocator, Types.MinorType.FLOAT8.getType(), -10000D)), false));
    constraintsMap.put("col4", EquatableValueSet.newBuilder(allocator, Types.MinorType.FLOAT8.getType(), false, true).add(1.1D).build());
    constraintsMap.put("col5", new AllOrNoneValueSet(Types.MinorType.FLOAT8.getType(), false, true));
    Constraints constraints = new Constraints(constraintsMap);
    Block partitions = allocator.createBlock(schema);
    int num_partitions = 10;
    for (int i = 0; i < num_partitions; i++) {
        BlockUtils.setValue(partitions.getFieldVector(yearCol), i, 2016 + i);
        BlockUtils.setValue(partitions.getFieldVector(monthCol), i, (i % 12) + 1);
        BlockUtils.setValue(partitions.getFieldVector(dayCol), i, (i % 28) + 1);
    }
    partitions.setRowCount(num_partitions);
    expected = new GetSplitsRequest(federatedIdentity, "test-query-id", "test-catalog", new TableName("test-schema", "test-table"), partitions, ImmutableList.of(yearCol, monthCol, dayCol), constraints, "test-continuation-token");
    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "GetSplitsRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}