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

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

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

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

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

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

/**
 * 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;
}
 

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

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

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

/**
 * 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());
}
 

/**
 * 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());
}
 

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

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

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

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

/**
 * 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());
}
 

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

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

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

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

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

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

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

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

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

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

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

/**
 * 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.");
    }
}
 

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

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