org.apache.arrow.vector.types.Types Java Examples

The following examples show how to use org.apache.arrow.vector.types.Types. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: ElasticsearchQueryUtilsTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void getExclusiveEquitableValuesPredicate()
{
    logger.info("getExclusiveEquitableValuesPredicate - enter");

    constraintsMap.put("age", EquatableValueSet.newBuilder(allocator, Types.MinorType.INT.getType(),
            false, true).addAll(ImmutableList.of(20, 25, 30, 35)).build());
    String expectedPredicate = "NOT age:(20 OR 25 OR 30 OR 35)";

    // Get the actual predicate and compare to the expected one.
    QueryBuilder builder = ElasticsearchQueryUtils.getQuery(constraintsMap);
    String actualPredicate = builder.queryName();

    logger.info("Predicates - Expected: {}, Actual: {}", expectedPredicate, actualPredicate);
    assertEquals("Predicates do not match", expectedPredicate, actualPredicate);

    logger.info("getExclusiveEquitableValuesPredicate - exit");
}
 
Example #2
Source File: BlockTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void constrainedBlockTest()
        throws Exception
{
    Schema schema = SchemaBuilder.newBuilder()
            .addIntField("col1")
            .addIntField("col2")
            .build();

    Block block = allocator.createBlock(schema);

    ValueSet col1Constraint = EquatableValueSet.newBuilder(allocator, Types.MinorType.INT.getType(), true, false)
            .add(10).build();
    Constraints constraints = new Constraints(Collections.singletonMap("col1", col1Constraint));
    try (ConstraintEvaluator constraintEvaluator = new ConstraintEvaluator(allocator, schema, constraints)) {
        block.constrain(constraintEvaluator);
        assertTrue(block.setValue("col1", 0, 10));
        assertTrue(block.offerValue("col1", 0, 10));
        assertFalse(block.setValue("col1", 0, 11));
        assertFalse(block.offerValue("col1", 0, 11));
        assertTrue(block.offerValue("unkown_col", 0, 10));
    }
}
 
Example #3
Source File: ExampleUserDefinedFunctionHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@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 #4
Source File: ExampleUserDefinedFunctionHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@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 #5
Source File: HbaseSchemaUtilsTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void coerceTypeNativeTest()
{
    boolean isNative = true;
    assertEquals("asf", coerceType(isNative, Types.MinorType.VARCHAR.getType(), "asf".getBytes()));
    assertEquals("2.0", coerceType(isNative, Types.MinorType.VARCHAR.getType(), "2.0".getBytes()));
    assertEquals(1, coerceType(isNative, Types.MinorType.INT.getType(), toBytes(isNative, 1)));
    assertEquals(-1, coerceType(isNative, Types.MinorType.INT.getType(), toBytes(isNative, -1)));
    assertEquals(1L, coerceType(isNative, Types.MinorType.BIGINT.getType(), toBytes(isNative, 1L)));
    assertEquals(-1L, coerceType(isNative, Types.MinorType.BIGINT.getType(), toBytes(isNative, -1L)));
    assertEquals(1.1F, coerceType(isNative, Types.MinorType.FLOAT4.getType(), toBytes(isNative, 1.1F)));
    assertEquals(-1.1F, coerceType(isNative, Types.MinorType.FLOAT4.getType(), toBytes(isNative, -1.1F)));
    assertEquals(1.1D, coerceType(isNative, Types.MinorType.FLOAT8.getType(), toBytes(isNative, 1.1D)));
    assertEquals(-1.1D, coerceType(isNative, Types.MinorType.FLOAT8.getType(), toBytes(isNative, -1.1D)));
    assertArrayEquals("-1.1".getBytes(), (byte[]) coerceType(isNative, Types.MinorType.VARBINARY.getType(), "-1.1".getBytes()));
}
 
Example #6
Source File: DDBTypeUtils.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
/**
 * Coerces the raw value from DynamoDB to normalized type
 * @param value raw value from DynamoDB
 * @param field Arrow field from table schema
 * @param fieldType Corresponding MinorType for field
 * @param recordMetadata DDBRecordMetadata object containing any metadata that is passed from schema metadata
 * @return coerced value to normalized type
 */
public static Object coerceValueToExpectedType(Object value, Field field,
                                               Types.MinorType fieldType, DDBRecordMetadata recordMetadata)
{
    if (!recordMetadata.isContainsCoercibleType()) {
        return value;
    }
    if (DDBRecordMetadata.isDateTimeFieldType(fieldType) && (value instanceof String || value instanceof BigDecimal)) {
        String dateTimeFormat = recordMetadata.getDateTimeFormat(field.getName());
        if (value instanceof String && StringUtils.isEmpty(dateTimeFormat)) {
            logger.info("Date format not in cache for column {}. Trying to infer format...", field.getName());
            dateTimeFormat = DateTimeFormatterUtil.inferDateTimeFormat((String) value);
            if (StringUtils.isNotEmpty(dateTimeFormat)) {
                logger.info("Adding datetime format {} for column {} to cache", dateTimeFormat, field.getName());
                recordMetadata.setDateTimeFormat(field.getName(), dateTimeFormat);
            }
        }
        value = coerceDateTimeToExpectedType(value, fieldType, dateTimeFormat, recordMetadata.getDefaultTimeZone());
    }
    else if (!fieldType.equals(Types.MinorType.DECIMAL) && value instanceof BigDecimal) {
        value = DDBTypeUtils.coerceDecimalToExpectedType((BigDecimal) value, fieldType);
    }
    return value;
}
 
Example #7
Source File: AllOrNoneValueSetTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void testAll()
        throws Exception
{
    AllOrNoneValueSet valueSet = AllOrNoneValueSet.all(Types.MinorType.INT.getType());
    assertEquals(valueSet.getType(), Types.MinorType.INT.getType());
    assertFalse(valueSet.isNone());
    assertTrue(valueSet.isAll());
    assertFalse(valueSet.isSingleValue());
    assertTrue(valueSet.containsValue(Marker.exactly(allocator, Types.MinorType.INT.getType(), 0)));

    try {
        valueSet.getSingleValue();
        fail();
    }
    catch (Exception ignored) {
    }
}
 
Example #8
Source File: ElasticsearchQueryUtilsTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void getNoneValuePredicate()
{
    logger.info("getNoneValuePredicate - enter");

    constraintsMap.put("number", new AllOrNoneValueSet(Types.MinorType.INT.getType(), false, false));
    String expectedPredicate = "(NOT _exists_:number)";

    // Get the actual predicate and compare to the expected one.
    QueryBuilder builder = ElasticsearchQueryUtils.getQuery(constraintsMap);
    String actualPredicate = builder.queryName();

    logger.info("Predicates - Expected: {}, Actual: {}", expectedPredicate, actualPredicate);
    assertEquals("Predicates do not match", expectedPredicate, actualPredicate);

    logger.info("getNoneValuePredicate - exit");
}
 
Example #9
Source File: JdbcArrowTypeConverterTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void toArrowType()
{
    Assert.assertEquals(Types.MinorType.BIT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BIT, 0, 0));
    Assert.assertEquals(Types.MinorType.BIT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BOOLEAN, 0, 0));
    Assert.assertEquals(Types.MinorType.TINYINT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.TINYINT, 0, 0));
    Assert.assertEquals(Types.MinorType.SMALLINT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.SMALLINT, 0, 0));
    Assert.assertEquals(Types.MinorType.INT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.INTEGER, 0, 0));
    Assert.assertEquals(Types.MinorType.BIGINT.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BIGINT, 0, 0));
    Assert.assertEquals(Types.MinorType.FLOAT4.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.REAL, 0, 0));
    Assert.assertEquals(Types.MinorType.FLOAT4.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.FLOAT, 0, 0));
    Assert.assertEquals(Types.MinorType.FLOAT8.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.DOUBLE, 0, 0));
    Assert.assertEquals(new ArrowType.Decimal(5, 3), JdbcArrowTypeConverter.toArrowType(java.sql.Types.DECIMAL, 5, 3));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.CHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.NCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.VARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.NVARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.LONGVARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARCHAR.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.LONGNVARCHAR, 0, 0));
    Assert.assertEquals(Types.MinorType.VARBINARY.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.BINARY, 0, 0));
    Assert.assertEquals(Types.MinorType.VARBINARY.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.VARBINARY, 0, 0));
    Assert.assertEquals(Types.MinorType.VARBINARY.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.LONGVARBINARY, 0, 0));
    Assert.assertEquals(Types.MinorType.DATEMILLI.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.DATE, 0, 0));
    Assert.assertEquals(Types.MinorType.TIMEMILLI.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.TIME, 0, 0));
    Assert.assertEquals(Types.MinorType.TIMESTAMPMILLI.getType(), JdbcArrowTypeConverter.toArrowType(java.sql.Types.TIMESTAMP, 0, 0));
}
 
Example #10
Source File: RedisMetadataHandler.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
/**
 * Retrieves the schema for the request Table from Glue then enriches that result with Redis specific
 * metadata and columns.
 */
@Override
public GetTableResponse doGetTable(BlockAllocator blockAllocator, GetTableRequest request)
        throws Exception
{
    GetTableResponse response = super.doGetTable(blockAllocator, request);

    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder();
    response.getSchema().getFields().forEach((Field field) ->
            schemaBuilder.addField(field.getName(), field.getType(), field.getChildren())
    );

    response.getSchema().getCustomMetadata().entrySet().forEach((Map.Entry<String, String> meta) ->
            schemaBuilder.addMetadata(meta.getKey(), meta.getValue()));

    schemaBuilder.addField(KEY_COLUMN_NAME, Types.MinorType.VARCHAR.getType());

    return new GetTableResponse(response.getCatalogName(), response.getTableName(), schemaBuilder.build());
}
 
Example #11
Source File: HbaseMetadataHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
/**
 * TODO: Add more types.
 */
@Test
public void doGetTable()
        throws Exception
{
    List<Result> results = TestUtils.makeResults();

    ResultScanner mockScanner = mock(ResultScanner.class);
    when(mockScanner.iterator()).thenReturn(results.iterator());

    when(mockClient.scanTable(anyObject(), any(Scan.class), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        ResultProcessor processor = (ResultProcessor) invocationOnMock.getArguments()[2];
        return processor.scan(mockScanner);
    });

    GetTableRequest req = new GetTableRequest(IDENTITY, QUERY_ID, DEFAULT_CATALOG, TABLE_NAME);
    GetTableResponse res = handler.doGetTable(allocator, req);
    logger.info("doGetTable - {}", res);

    Schema expectedSchema = TestUtils.makeSchema()
            .addField(HbaseSchemaUtils.ROW_COLUMN_NAME, Types.MinorType.VARCHAR.getType())
            .build();

    assertEquals(expectedSchema.getFields().size(), res.getSchema().getFields().size());
}
 
Example #12
Source File: UserDefinedFunctionResponseSerDeTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Before
public void beforeTest()
        throws IOException
{
    String productField = "product";

    Schema outputSchema = SchemaBuilder.newBuilder()
            .addField(productField, Types.MinorType.INT.getType())
            .build();

    Block records = allocator.createBlock(outputSchema);
    int num_records = 10;
    for (int i = 0; i < num_records; i++) {
        BlockUtils.setValue(records.getFieldVector(productField), i, i * 2);
    }
    records.setRowCount(num_records);

    expected = new UserDefinedFunctionResponse(records, "test-method");

    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "UserDefinedFunctionResponse.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}
 
Example #13
Source File: ElasticsearchQueryUtilsTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void getWhitelistedEquitableValuesPredicate()
{
    logger.info("getWhitelistedEquitableValuesPredicate - enter");

    constraintsMap.put("age", EquatableValueSet.newBuilder(allocator, Types.MinorType.INT.getType(),
            true, true).addAll(ImmutableList.of(20, 25, 30, 35)).build());
    String expectedPredicate = "age:(20 OR 25 OR 30 OR 35)";

    // Get the actual predicate and compare to the expected one.
    QueryBuilder builder = ElasticsearchQueryUtils.getQuery(constraintsMap);
    String actualPredicate = builder.queryName();

    logger.info("Predicates - Expected: {}, Actual: {}", expectedPredicate, actualPredicate);
    assertEquals("Predicates do not match", expectedPredicate, actualPredicate);

    logger.info("getWhitelistedEquitableValuesPredicate - exit");
}
 
Example #14
Source File: MarkerTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void testUnbounded()
        throws Exception
{
    assertTrue(Marker.lowerUnbounded(allocator, Types.MinorType.INT.getType()).isLowerUnbounded());
    assertFalse(Marker.lowerUnbounded(allocator, Types.MinorType.INT.getType()).isUpperUnbounded());
    assertTrue(Marker.upperUnbounded(allocator, Types.MinorType.INT.getType()).isUpperUnbounded());
    assertFalse(Marker.upperUnbounded(allocator, Types.MinorType.INT.getType()).isLowerUnbounded());

    assertFalse(Marker.below(allocator, Types.MinorType.INT.getType(), 1).isLowerUnbounded());
    assertFalse(Marker.below(allocator, Types.MinorType.INT.getType(), 1).isUpperUnbounded());
    assertFalse(Marker.exactly(allocator, Types.MinorType.INT.getType(), 1).isLowerUnbounded());
    assertFalse(Marker.exactly(allocator, Types.MinorType.INT.getType(), 1).isUpperUnbounded());
    assertFalse(Marker.above(allocator, Types.MinorType.INT.getType(), 1).isLowerUnbounded());
    assertFalse(Marker.above(allocator, Types.MinorType.INT.getType(), 1).isUpperUnbounded());
}
 
Example #15
Source File: ElasticsearchQueryUtilsTest.java    From aws-athena-query-federation with Apache License 2.0 6 votes vote down vote up
@Test
public void getRangePredicateTest()
{
    logger.info("getRangePredicateTest - enter");

    constraintsMap.put("year", SortedRangeSet.copyOf(Types.MinorType.INT.getType(),
            ImmutableList.of(
                    Range.lessThan(allocator, Types.MinorType.INT.getType(), 1950),
                    Range.equal(allocator, Types.MinorType.INT.getType(), 1952),
                    Range.range(allocator, Types.MinorType.INT.getType(),
                            1955, false, 1972, true),
                    Range.equal(allocator, Types.MinorType.INT.getType(), 1996),
                    Range.greaterThanOrEqual(allocator, Types.MinorType.INT.getType(), 2010)),
            false));
    String expectedPredicate = "(_exists_:year) AND year:((<1950) OR (>1955 AND <=1972) OR (>=2010) OR 1952 OR 1996)";

    // Get the actual predicate and compare to the expected one.
    QueryBuilder builder = ElasticsearchQueryUtils.getQuery(constraintsMap);
    String actualPredicate = builder.queryName();

    logger.info("Predicates - Expected: {}, Actual: {}", expectedPredicate, actualPredicate);
    assertEquals("Predicates do not match", expectedPredicate, actualPredicate);

    logger.info("getRangePredicateTest - exit");
}
 
Example #16
Source File: ElasticsearchFieldResolver.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
/**
 * Allows for coercion of a list of values where the returned types do not match the schema.
 * Multiple fields in Elasticsearch can be returned as a string, numeric (Integer, Long, Double), or null.
 * @param field is the field that we are coercing the value into.
 * @param fieldValue is the list of value to coerce
 * @return the coerced list of value.
 * @throws RuntimeException if the fieldType is not a LIST or the fieldValue is instanceof Map (STRUCT).
 */
protected Object coerceListField(Field field, Object fieldValue)
        throws RuntimeException
{
    Types.MinorType fieldType = Types.getMinorTypeForArrowType(field.getType());

    switch (fieldType) {
        case LIST:
            Field childField = field.getChildren().get(0);
            if (fieldValue instanceof List) {
                // Both fieldType and fieldValue are lists => Return as a new list of values, applying coercion
                // where necessary in order to match the type of the field being mapped into.
                List<Object> coercedValues = new ArrayList<>();
                ((List) fieldValue).forEach(value ->
                        coercedValues.add(coerceField(childField, value)));
                return coercedValues;
            }
            else if (!(fieldValue instanceof Map)) {
                // This is an abnormal case where the fieldType was defined as a list in the schema,
                // however, the fieldValue returns as a single value => Return as a list of a single value
                // applying coercion where necessary in order to match the type of the field being mapped into.
                return Collections.singletonList(coerceField(childField, fieldValue));
            }
            break;
        default:
            break;
    }

    throw new RuntimeException("Invalid field value encountered in Document for field: " + field.toString() +
            ",value: " + fieldValue.toString());
}
 
Example #17
Source File: UserDefinedFunctionRequestSerDeTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Before
public void beforeTest()
        throws IOException
{
    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);

    expected = new UserDefinedFunctionRequest(federatedIdentity,
            inputRecords,
            outputSchema,
            "test-method",
            UserDefinedFunctionType.SCALAR);


    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "UserDefinedFunctionRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}
 
Example #18
Source File: GetTableLayoutRequestSerDeTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Before
public void beforeTest()
        throws IOException
{
    Schema schema = SchemaBuilder.newBuilder()
            .addField("year", new ArrowType.Int(32, true))
            .addField("month", new ArrowType.Int(32, true))
            .addField("day", 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);

    expected = new GetTableLayoutRequest(federatedIdentity,
            "test-query-id",
            "test-catalog",
            new TableName("test-schema", "test-table"),
            constraints,
            schema,
            ImmutableSet.of("year", "month", "day"));

    String expectedSerDeFile = utils.getResourceOrFail("serde/v2", "GetTableLayoutRequest.json");
    expectedSerDeText = utils.readAllAsString(expectedSerDeFile).trim();
}
 
Example #19
Source File: AllOrNoneValueSetTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void testOverlaps()
        throws Exception
{
    AllOrNoneValueSet all = AllOrNoneValueSet.all(Types.MinorType.INT.getType());
    AllOrNoneValueSet none = AllOrNoneValueSet.none(Types.MinorType.INT.getType());

    assertTrue(all.overlaps(allocator, all));
    assertFalse(all.overlaps(allocator, none));
    assertFalse(none.overlaps(allocator, all));
    assertFalse(none.overlaps(allocator, none));
}
 
Example #20
Source File: MarkerSerializationTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void serializationTest()
        throws IOException
{
    logger.info("serializationTest - enter");

    ObjectMapper serializer = ObjectMapperFactory.create(new BlockAllocatorImpl());

    int expectedValue = 1024;
    Marker expectedMarker = Marker.exactly(allocator, Types.MinorType.INT.getType(), expectedValue);

    ByteArrayOutputStream out = new ByteArrayOutputStream();
    serializer.writeValue(out, expectedMarker);

    ObjectMapper deserializer = ObjectMapperFactory.create(allocator);

    Marker actualMarker = deserializer.readValue(new ByteArrayInputStream(out.toByteArray()), Marker.class);

    assertEquals(expectedMarker.getSchema().getCustomMetadata(), actualMarker.getSchema().getCustomMetadata());
    assertEquals(expectedMarker.getSchema().getFields(), actualMarker.getSchema().getFields());
    assertEquals(expectedMarker.getBound(), actualMarker.getBound());
    assertEquals(expectedMarker.getValue(), actualMarker.getValue());
    assertEquals(expectedValue, actualMarker.getValue());
    assertEquals(false, actualMarker.isNullValue());

    logger.info("serializationTest - exit");
}
 
Example #21
Source File: AllOrNoneValueSetTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void testContains()
        throws Exception
{
    AllOrNoneValueSet all = AllOrNoneValueSet.all(Types.MinorType.INT.getType());
    AllOrNoneValueSet none = AllOrNoneValueSet.none(Types.MinorType.INT.getType());

    assertTrue(all.contains(allocator, all));
    assertTrue(all.contains(allocator, none));
    assertFalse(none.contains(allocator, all));
    assertTrue(none.contains(allocator, none));
}
 
Example #22
Source File: AllOrNoneValueSetTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void testSubtract()
        throws Exception
{
    AllOrNoneValueSet all = AllOrNoneValueSet.all(Types.MinorType.INT.getType());
    AllOrNoneValueSet none = AllOrNoneValueSet.none(Types.MinorType.INT.getType());

    assertEquals(all.subtract(allocator, all), none);
    assertEquals(all.subtract(allocator, none), all);
    assertEquals(none.subtract(allocator, all), none);
    assertEquals(none.subtract(allocator, none), none);
}
 
Example #23
Source File: DocDBFieldResolver.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Override
public Object getFieldValue(Field field, Object value)
{
    Types.MinorType minorType = Types.getMinorTypeForArrowType(field.getType());
    if (minorType == Types.MinorType.LIST) {
        return TypeUtils.coerce(field, ((Document) value).get(field.getName()));
    }
    else if (value instanceof Document) {
        Object rawVal = ((Document) value).get(field.getName());
        return TypeUtils.coerce(field, rawVal);
    }
    throw new RuntimeException("Expected LIST or Document type but found " + minorType);
}
 
Example #24
Source File: GlueFieldLexerTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void lexTest()
{
    logger.info("lexTest: enter");

    Field field = GlueFieldLexer.lex("testField", INPUT1);

    logger.info("lexTest: {}", field);
    assertEquals("testField", field.getName());
    assertEquals(Types.MinorType.STRUCT, Types.getMinorTypeForArrowType(field.getType()));
    assertEquals(3, field.getChildren().size());

    List<Field> level1 = field.getChildren();
    assertEquals("street_address", level1.get(0).getName());
    assertEquals(Types.MinorType.STRUCT, Types.getMinorTypeForArrowType(level1.get(0).getType()));
    assertEquals(3, level1.get(0).getChildren().size());

    List<Field> level2 = level1.get(0).getChildren();
    assertEquals("street_number", level2.get(0).getName());
    assertEquals(Types.MinorType.INT, Types.getMinorTypeForArrowType(level2.get(0).getType()));
    assertEquals(0, level2.get(0).getChildren().size());
    assertEquals("street_name", level2.get(1).getName());
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(level2.get(1).getType()));
    assertEquals(0, level2.get(1).getChildren().size());
    assertEquals("street_type", level2.get(2).getName());
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(level2.get(2).getType()));
    assertEquals(0, level2.get(2).getChildren().size());

    assertEquals("country", level1.get(1).getName());
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(level1.get(1).getType()));
    assertEquals(0, level1.get(1).getChildren().size());

    assertEquals("postal_code", level1.get(2).getName());
    assertEquals(Types.MinorType.LIST, Types.getMinorTypeForArrowType(level1.get(2).getType()));
    assertEquals(1, level1.get(2).getChildren().size());
    assertEquals(Types.MinorType.VARCHAR, Types.getMinorTypeForArrowType(level1.get(2).getChildren().get(0).getType()));

    logger.info("lexTest: exit");
}
 
Example #25
Source File: PostGreSqlMetadataHandler.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Override
public Schema getPartitionSchema(final String catalogName)
{
    SchemaBuilder schemaBuilder = SchemaBuilder.newBuilder()
            .addField(BLOCK_PARTITION_SCHEMA_COLUMN_NAME, Types.MinorType.VARCHAR.getType())
            .addField(BLOCK_PARTITION_COLUMN_NAME, Types.MinorType.VARCHAR.getType());
    return schemaBuilder.build();
}
 
Example #26
Source File: ExampleUserDefinedFunctionHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@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 #27
Source File: ListArrowValueProjector.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
public ListArrowValueProjector(FieldReader listReader)
{
    this.listReader = requireNonNull(listReader, "listReader is null");

    List<Field> children = listReader.getField().getChildren();
    if (children.size() != 1) {
        throw new RuntimeException("Unexpected number of children for ListProjector field "
                + listReader.getField().getName());
    }
    Types.MinorType minorType = Types.getMinorTypeForArrowType(children.get(0).getType());
    projection = createValueProjection(minorType);
}
 
Example #28
Source File: HbaseSchemaUtils.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
/**
 * Helper that can coerce the given HBase value to the requested Apache Arrow type.
 *
 * @param isNative If True, the HBase value is stored using native bytes. If False, the value is serialized as a String.
 * @param type The Apache Arrow Type that the value should be coerced to before returning.
 * @param value The HBase value to coerce.
 * @return The coerced value which is now allowed with the provided Apache Arrow type.
 */
public static Object coerceType(boolean isNative, ArrowType type, byte[] value)
{
    if (value == null) {
        return null;
    }

    Types.MinorType minorType = Types.getMinorTypeForArrowType(type);
    switch (minorType) {
        case VARCHAR:
            return Bytes.toString(value);
        case INT:
            return isNative ? ByteBuffer.wrap(value).getInt() : Integer.parseInt(Bytes.toString(value));
        case BIGINT:
            return isNative ? ByteBuffer.wrap(value).getLong() : Long.parseLong(Bytes.toString(value));
        case FLOAT4:
            return isNative ? ByteBuffer.wrap(value).getFloat() : Float.parseFloat(Bytes.toString(value));
        case FLOAT8:
            return isNative ? ByteBuffer.wrap(value).getDouble() : Double.parseDouble(Bytes.toString(value));
        case BIT:
            if (isNative) {
                return (value[0] != 0);
            }
            else {
                return Boolean.parseBoolean(Bytes.toString(value));
            }
        case VARBINARY:
            return value;
        default:
            throw new IllegalArgumentException(type + " with minorType[" + minorType + "] is not supported.");
    }
}
 
Example #29
Source File: GlueFieldLexer.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
private static Field lexComplex(String name, GlueTypeParser.Token startToken, GlueTypeParser parser, BaseTypeMapper mapper)
{
    FieldBuilder fieldBuilder;

    logger.debug("lexComplex: enter - {}", name);
    if (startToken.getMarker() != GlueTypeParser.FIELD_START) {
        throw new RuntimeException("Parse error, expected " + GlueTypeParser.FIELD_START
                + " but found " + startToken.getMarker());
    }

    if (startToken.getValue().toLowerCase().equals(STRUCT)) {
        fieldBuilder = FieldBuilder.newBuilder(name, Types.MinorType.STRUCT.getType());
    }
    else if (startToken.getValue().toLowerCase().equals(LIST)) {
        GlueTypeParser.Token arrayType = parser.next();
        return FieldBuilder.newBuilder(name, Types.MinorType.LIST.getType())
                .addField(mapper.getField(name, arrayType.getValue())).build();
    }
    else {
        throw new RuntimeException("Unexpected start type " + startToken.getValue());
    }

    while (parser.hasNext() && parser.currentToken().getMarker() != GlueTypeParser.FIELD_END) {
        Field child = lex(parser.next(), parser, mapper);
        fieldBuilder.addField(child);
    }
    parser.next();

    logger.debug("lexComplex: exit - {}", name);
    return fieldBuilder.build();
}
 
Example #30
Source File: ElasticsearchMetadataHandlerTest.java    From aws-athena-query-federation with Apache License 2.0 5 votes vote down vote up
@Test
public void convertFieldTest()
{
    logger.info("convertFieldTest: enter");

    handler = new ElasticsearchMetadataHandler(awsGlue, new LocalKeyFactory(), awsSecretsManager, amazonAthena,
            "spill-bucket", "spill-prefix", domainMapProvider, clientFactory, 10);

    Field field = handler.convertField("myscaled", "SCALED_FLOAT(10.51)");

    assertEquals("myscaled", field.getName());
    assertEquals("10.51", field.getMetadata().get("scaling_factor"));

    field = handler.convertField("myscaledlist", "ARRAY<SCALED_FLOAT(100)>");

    assertEquals("myscaledlist", field.getName());
    assertEquals(Types.MinorType.LIST.getType(), field.getType());
    assertEquals("100", field.getChildren().get(0).getMetadata().get("scaling_factor"));

    field = handler.convertField("myscaledstruct", "STRUCT<myscaledstruct:SCALED_FLOAT(10.0)>");

    assertEquals(Types.MinorType.STRUCT.getType(), field.getType());
    assertEquals("myscaledstruct", field.getChildren().get(0).getName());
    assertEquals("10.0", field.getChildren().get(0).getMetadata().get("scaling_factor"));

    logger.info("convertFieldTest: exit");
}