Search in sources :

Example 1 with FieldWriter

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

the class JdbcRecordHandler method makeFactory.

/**
 * Create a field extractor for complex List type.
 * @param field Field's metadata information.
 * @return Extractor for the List type.
 */
protected FieldWriterFactory makeFactory(Field field) {
    return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
        Array arrayField = ((ResultSet) context).getArray(field.getName());
        if (!((ResultSet) context).wasNull()) {
            List<Object> fieldValue = new ArrayList<>(Arrays.asList((Object[]) arrayField.getArray()));
            BlockUtils.setComplexValue(vector, rowNum, FieldResolver.DEFAULT, fieldValue);
        }
        return true;
    };
}
Also used : Array(java.sql.Array) ConstraintProjector(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector) ResultSet(java.sql.ResultSet) ArrayList(java.util.ArrayList) FieldVector(org.apache.arrow.vector.FieldVector) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) DecimalExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor) DateDayExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor) VarBinaryExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor) BitExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor) IntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor) Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor) Float8Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) VarCharExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor) Float4Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor) DateMilliExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor) FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriter)

Example 2 with FieldWriter

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

the class GeneratedRowWriter method writeRow.

public boolean writeRow(Block block, int rowNum, Object context) throws Exception {
    checkAndRecompile(block);
    boolean matched = true;
    for (FieldWriter next : fieldWriters) {
        matched &= next.write(context, rowNum);
    }
    return matched;
}
Also used : FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriter) BitFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BitFieldWriter) IntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.IntFieldWriter) Float8FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.Float8FieldWriter) TinyIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.TinyIntFieldWriter) DecimalFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DecimalFieldWriter) DateDayFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DateDayFieldWriter) BigIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BigIntFieldWriter) Float4FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.Float4FieldWriter) SmallIntFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.SmallIntFieldWriter) VarBinaryFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.VarBinaryFieldWriter) DateMilliFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DateMilliFieldWriter) VarCharFieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.VarCharFieldWriter)

Example 3 with FieldWriter

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

the class ExampleRecordHandler method makeFactory.

/**
 * Since GeneratedRowWriter doesn't yet support complex types (STRUCT, LIST) we use this to
 * create our own FieldWriters via customer FieldWriterFactory. In this case we are producing
 * FieldWriters that only work for our exact example schema. This will be enhanced with a more
 * generic solution in a future release.
 */
private FieldWriterFactory makeFactory(Field field, RowContext rowContext) {
    Types.MinorType fieldType = Types.getMinorTypeForArrowType(field.getType());
    switch(fieldType) {
        case LIST:
            Field child = field.getChildren().get(0);
            Types.MinorType childType = Types.getMinorTypeForArrowType(child.getType());
            switch(childType) {
                case LIST:
                    return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
                        UnionListWriter writer = ((ListVector) vector).getWriter();
                        writer.setPosition(rowNum);
                        writer.startList();
                        BaseWriter.ListWriter innerWriter = writer.list();
                        innerWriter.startList();
                        for (int i = 0; i < 3; i++) {
                            byte[] bytes = String.valueOf(1000 + i).getBytes(Charsets.UTF_8);
                            try (ArrowBuf buf = vector.getAllocator().buffer(bytes.length)) {
                                buf.writeBytes(bytes);
                                innerWriter.varChar().writeVarChar(0, (int) (buf.readableBytes()), buf);
                            }
                        }
                        innerWriter.endList();
                        writer.endList();
                        ((ListVector) vector).setNotNull(rowNum);
                        return true;
                    };
                case STRUCT:
                    return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
                        UnionListWriter writer = ((ListVector) vector).getWriter();
                        writer.setPosition(rowNum);
                        writer.startList();
                        BaseWriter.StructWriter structWriter = writer.struct();
                        structWriter.start();
                        byte[] bytes = "chars".getBytes(Charsets.UTF_8);
                        try (ArrowBuf buf = vector.getAllocator().buffer(bytes.length)) {
                            buf.writeBytes(bytes);
                            structWriter.varChar("varchar").writeVarChar(0, (int) (buf.readableBytes()), buf);
                        }
                        structWriter.bigInt("bigint").writeBigInt(100L);
                        structWriter.end();
                        writer.endList();
                        ((ListVector) vector).setNotNull(rowNum);
                        return true;
                    };
                default:
                    throw new IllegalArgumentException("Unsupported type " + childType);
            }
        case MAP:
            return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) -> (FieldWriter) (Object context, int rowNum) -> {
                UnionMapWriter writer = ((MapVector) vector).getWriter();
                writer.setPosition(rowNum);
                writer.startMap();
                writer.startEntry();
                byte[] bytes = "chars".getBytes(Charsets.UTF_8);
                try (ArrowBuf buf = vector.getAllocator().buffer(bytes.length)) {
                    buf.writeBytes(bytes);
                    writer.key().varChar("key").writeVarChar(0, (int) (buf.readableBytes()), buf);
                }
                writer.value().integer("value").writeInt(1001);
                writer.endEntry();
                writer.endMap();
                ((MapVector) vector).setNotNull(rowNum);
                return true;
            };
        default:
            throw new IllegalArgumentException("Unsupported type " + fieldType);
    }
}
Also used : Types(org.apache.arrow.vector.types.Types) BaseWriter(org.apache.arrow.vector.complex.writer.BaseWriter) ArrowBuf(org.apache.arrow.memory.ArrowBuf) ConstraintProjector(com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector) FieldVector(org.apache.arrow.vector.FieldVector) FieldWriter(com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriter) Field(org.apache.arrow.vector.types.pojo.Field) UnionMapWriter(org.apache.arrow.vector.complex.impl.UnionMapWriter) UnionListWriter(org.apache.arrow.vector.complex.impl.UnionListWriter) ListVector(org.apache.arrow.vector.complex.ListVector) BigIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor) DecimalExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor) DateDayExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor) TinyIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor) VarBinaryExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor) BitExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor) IntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor) Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor) Float8Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor) SmallIntExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor) VarCharExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor) Float4Extractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor) DateMilliExtractor(com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor) MapVector(org.apache.arrow.vector.complex.MapVector)

Aggregations

FieldWriter (com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.FieldWriter)3 BigIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BigIntExtractor)2 BitExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.BitExtractor)2 DateDayExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateDayExtractor)2 DateMilliExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DateMilliExtractor)2 DecimalExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.DecimalExtractor)2 Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Extractor)2 Float4Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float4Extractor)2 Float8Extractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.Float8Extractor)2 IntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.IntExtractor)2 SmallIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.SmallIntExtractor)2 TinyIntExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.TinyIntExtractor)2 VarBinaryExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarBinaryExtractor)2 VarCharExtractor (com.amazonaws.athena.connector.lambda.data.writers.extractors.VarCharExtractor)2 ConstraintProjector (com.amazonaws.athena.connector.lambda.domain.predicate.ConstraintProjector)2 FieldVector (org.apache.arrow.vector.FieldVector)2 BigIntFieldWriter (com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BigIntFieldWriter)1 BitFieldWriter (com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.BitFieldWriter)1 DateDayFieldWriter (com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DateDayFieldWriter)1 DateMilliFieldWriter (com.amazonaws.athena.connector.lambda.data.writers.fieldwriters.DateMilliFieldWriter)1