Examples with UnionMapWriter org.apache.arrow.vector.complex.impl.UnionMapWriter used on opensource projects

Example 1 with UnionMapWriter

use of org.apache.arrow.vector.complex.impl.UnionMapWriter 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);
    }
}

Example 2 with UnionMapWriter

use of org.apache.arrow.vector.complex.impl.UnionMapWriter in project aws-athena-query-federation by awslabs.

the class BlockUtils method setComplexValue.

/**
 * Used to set complex values (Struct, List, etc...) on the provided FieldVector.
 *
 * @param vector The FieldVector into which we should write the provided value.
 * @param pos The row number that the value should be written to.
 * @param resolver The FieldResolver that can be used to map your value to the complex type (mostly for Structs, Maps).
 * @param value The value to write.
 * @note This method incurs more Object overhead (heap churn) than using Arrow's native interface. Users of this Utility
 * should weigh their performance needs vs. the readability / ease of use.
 */
public static void setComplexValue(FieldVector vector, int pos, FieldResolver resolver, Object value) {
    if (vector instanceof MapVector) {
        UnionMapWriter writer = ((MapVector) vector).getWriter();
        writer.setPosition(pos);
        writeMap(vector.getAllocator(), writer, vector.getField(), pos, value, resolver);
        writer.endMap();
    } else if (vector instanceof ListVector) {
        if (value != null) {
            UnionListWriter writer = ((ListVector) vector).getWriter();
            writer.setPosition(pos);
            writeList(vector.getAllocator(), writer, vector.getField(), pos, ((List) value), resolver);
            ((ListVector) vector).setNotNull(pos);
        }
    } else if (vector instanceof StructVector) {
        StructWriter writer = ((StructVector) vector).getWriter();
        writer.setPosition(pos);
        writeStruct(vector.getAllocator(), writer, vector.getField(), pos, value, resolver);
    } else {
        throw new RuntimeException("Unsupported 'Complex' vector " + vector.getClass().getSimpleName() + " for field " + vector.getField().getName());
    }
}