Search in sources :

Example 6 with ListWriter

use of org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter in project drill by axbaretto.

the class TestRepeated method listOfList.

// 
// @Test
// public void repeatedMap() {
// 
// /**
// * We're going to try to create an object that looks like:
// *
// *  {
// *    a: [
// *      {x: 1, y: 2}
// *      {x: 2, y: 1}
// *    ]
// *  }
// *
// */
// MapVector v = new MapVector("", allocator);
// ComplexWriter writer = new ComplexWriterImpl("col", v);
// 
// MapWriter map = writer.rootAsMap();
// 
// map.start();
// ListWriter list = map.list("a");
// MapWriter inner = list.map();
// 
// IntHolder holder = new IntHolder();
// IntWriter xCol = inner.integer("x");
// IntWriter yCol = inner.integer("y");
// 
// inner.start();
// 
// holder.value = 1;
// xCol.write(holder);
// holder.value = 2;
// yCol.write(holder);
// 
// inner.end();
// 
// inner.start();
// 
// holder.value = 2;
// xCol.write(holder);
// holder.value = 1;
// yCol.write(holder);
// 
// inner.end();
// 
// IntWriter numCol = map.integer("nums");
// holder.value = 14;
// numCol.write(holder);
// 
// map.end();
// 
// 
// assertTrue(writer.ok());
// 
// System.out.println(v.getAccessor().getObject(0));
// 
// }
@Test
public void listOfList() throws Exception {
    /**
     * We're going to try to create an object that looks like:
     *
     *  {
     *    a: [
     *      [1,2,3],
     *      [2,3,4]
     *    ],
     *    nums: 14,
     *    b: [
     *      { c: 1 },
     *      { c: 2 , x: 15}
     *    ]
     *  }
     */
    final MapVector mapVector = new MapVector("", allocator, null);
    final ComplexWriterImpl writer = new ComplexWriterImpl("col", mapVector);
    writer.allocate();
    {
        final MapWriter map = writer.rootAsMap();
        final ListWriter list = map.list("a");
        list.startList();
        final ListWriter innerList = list.list();
        final IntWriter innerInt = innerList.integer();
        innerList.startList();
        final IntHolder holder = new IntHolder();
        holder.value = 1;
        innerInt.write(holder);
        holder.value = 2;
        innerInt.write(holder);
        holder.value = 3;
        innerInt.write(holder);
        innerList.endList();
        innerList.startList();
        holder.value = 4;
        innerInt.write(holder);
        holder.value = 5;
        innerInt.write(holder);
        innerList.endList();
        list.endList();
        final IntWriter numCol = map.integer("nums");
        holder.value = 14;
        numCol.write(holder);
        final MapWriter repeatedMap = map.list("b").map();
        repeatedMap.start();
        holder.value = 1;
        repeatedMap.integer("c").write(holder);
        repeatedMap.end();
        repeatedMap.start();
        holder.value = 2;
        repeatedMap.integer("c").write(holder);
        final BigIntHolder h = new BigIntHolder();
        h.value = 15;
        repeatedMap.bigInt("x").write(h);
        repeatedMap.end();
        map.end();
    }
    {
        writer.setPosition(1);
        final MapWriter map = writer.rootAsMap();
        final ListWriter list = map.list("a");
        list.startList();
        final ListWriter innerList = list.list();
        final IntWriter innerInt = innerList.integer();
        innerList.startList();
        final IntHolder holder = new IntHolder();
        holder.value = -1;
        innerInt.write(holder);
        holder.value = -2;
        innerInt.write(holder);
        holder.value = -3;
        innerInt.write(holder);
        innerList.endList();
        innerList.startList();
        holder.value = -4;
        innerInt.write(holder);
        holder.value = -5;
        innerInt.write(holder);
        innerList.endList();
        list.endList();
        final IntWriter numCol = map.integer("nums");
        holder.value = -28;
        numCol.write(holder);
        final MapWriter repeatedMap = map.list("b").map();
        repeatedMap.start();
        holder.value = -1;
        repeatedMap.integer("c").write(holder);
        repeatedMap.end();
        repeatedMap.start();
        holder.value = -2;
        repeatedMap.integer("c").write(holder);
        final BigIntHolder h = new BigIntHolder();
        h.value = -30;
        repeatedMap.bigInt("x").write(h);
        repeatedMap.end();
        map.end();
    }
    final ObjectWriter ow = new ObjectMapper().writer().withDefaultPrettyPrinter();
    System.out.println("Map of Object[0]: " + ow.writeValueAsString(mapVector.getAccessor().getObject(0)));
    System.out.println("Map of Object[1]: " + ow.writeValueAsString(mapVector.getAccessor().getObject(1)));
    final ByteArrayOutputStream stream = new ByteArrayOutputStream();
    final JsonWriter jsonWriter = new JsonWriter(stream, true, true);
    final FieldReader reader = mapVector.getChild("col", MapVector.class).getReader();
    reader.setPosition(0);
    jsonWriter.write(reader);
    reader.setPosition(1);
    jsonWriter.write(reader);
    System.out.print("Json Read: ");
    System.out.println(new String(stream.toByteArray(), Charsets.UTF_8));
    writer.close();
}
Also used : MapWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter.MapWriter) ObjectWriter(com.fasterxml.jackson.databind.ObjectWriter) ByteArrayOutputStream(java.io.ByteArrayOutputStream) JsonWriter(org.apache.drill.exec.vector.complex.fn.JsonWriter) BigIntHolder(org.apache.drill.exec.expr.holders.BigIntHolder) ListWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter) IntHolder(org.apache.drill.exec.expr.holders.IntHolder) BigIntHolder(org.apache.drill.exec.expr.holders.BigIntHolder) ComplexWriterImpl(org.apache.drill.exec.vector.complex.impl.ComplexWriterImpl) FieldReader(org.apache.drill.exec.vector.complex.reader.FieldReader) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) MapVector(org.apache.drill.exec.vector.complex.MapVector) Test(org.junit.Test)

Example 7 with ListWriter

use of org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter in project drill by apache.

the class VectorTest method testRepeatedMapWriter.

@Test
public void testRepeatedMapWriter() throws Exception {
    // The writers will create a nullable int inner vector.
    MaterializedField field = MaterializedField.create("repeated_map", Types.repeated(TypeProtos.MinorType.MAP));
    try (RepeatedMapVector v = new RepeatedMapVector(field, allocator, null)) {
        SpecialMapVector mapVector = new SpecialMapVector(v);
        @SuppressWarnings("resource") SingleMapWriter mapRoot = new SingleMapWriter(mapVector, null, false);
        ListWriter lw = mapRoot.list("repeated_map");
        MapWriter mw = lw.map();
        // Record 1: [10, 20]
        lw.setPosition(0);
        lw.startList();
        mw.start();
        IntWriter iw = mw.integer("inner");
        iw.writeInt(10);
        mw.end();
        mw.start();
        iw.writeInt(20);
        mw.end();
        lw.endList();
        // Record 2: [30, 40, 50]
        lw.setPosition(1);
        lw.startList();
        mw.start();
        iw.writeInt(30);
        mw.end();
        mw.start();
        iw.writeInt(40);
        mw.end();
        mw.start();
        iw.writeInt(50);
        mw.end();
        lw.endList();
        // Record 3: []
        lw.setPosition(2);
        lw.startList();
        lw.endList();
        v.getMutator().setValueCount(3);
        assertEquals(3, v.getAccessor().getValueCount());
        UInt4Vector.Accessor oa = v.getOffsetVector().getAccessor();
        assertEquals(4, oa.getValueCount());
        assertEquals(0, oa.get(0));
        assertEquals(2, oa.get(1));
        assertEquals(5, oa.get(2));
        assertEquals(5, oa.get(3));
        // Past end
        assertEquals(0, oa.get(4));
        NullableIntVector inner = v.addOrGet("inner", Types.optional(TypeProtos.MinorType.INT), NullableIntVector.class);
        UInt1Vector.Accessor ba = inner.getBitsVector().getAccessor();
        assertEquals(1, ba.get(0));
        assertEquals(1, ba.get(1));
        assertEquals(1, ba.get(2));
        assertEquals(1, ba.get(3));
        assertEquals(1, ba.get(4));
        // Past end
        assertEquals(0, ba.get(5));
        NullableIntVector.Accessor ia = inner.getAccessor();
        assertEquals(10, ia.get(0));
        assertEquals(20, ia.get(1));
        assertEquals(30, ia.get(2));
        assertEquals(40, ia.get(3));
        assertEquals(50, ia.get(4));
    }
}
Also used : SingleMapWriter(org.apache.drill.exec.vector.complex.impl.SingleMapWriter) MapWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter.MapWriter) SingleMapWriter(org.apache.drill.exec.vector.complex.impl.SingleMapWriter) RepeatedMapVector(org.apache.drill.exec.vector.complex.RepeatedMapVector) IntWriter(org.apache.drill.exec.vector.complex.writer.IntWriter) MaterializedField(org.apache.drill.exec.record.MaterializedField) ListWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter) Test(org.junit.Test) BaseTest(org.apache.drill.test.BaseTest)

Example 8 with ListWriter

use of org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter in project drill by apache.

the class MapUtility method writeToMapFromReader.

public static void writeToMapFromReader(FieldReader fieldReader, MapWriter mapWriter, String fieldName, String caller) {
    try {
        MajorType valueMajorType = fieldReader.getType();
        MinorType valueMinorType = valueMajorType.getMinorType();
        WriterExtractor extractor = new WriterExtractor(fieldName, valueMajorType, mapWriter, fieldReader instanceof UnionReader);
        switch(valueMinorType) {
            case TINYINT:
                fieldReader.copyAsValue(extractor.get(ListWriter::tinyInt, MapWriter::tinyInt, UnionVectorWriter::tinyInt));
                break;
            case SMALLINT:
                fieldReader.copyAsValue(extractor.get(ListWriter::smallInt, MapWriter::smallInt, UnionVectorWriter::smallInt));
                break;
            case BIGINT:
                fieldReader.copyAsValue(extractor.get(ListWriter::bigInt, MapWriter::bigInt, UnionVectorWriter::bigInt));
                break;
            case INT:
                fieldReader.copyAsValue(extractor.get(ListWriter::integer, MapWriter::integer, UnionVectorWriter::integer));
                break;
            case UINT1:
                fieldReader.copyAsValue(extractor.get(ListWriter::uInt1, MapWriter::uInt1, UnionVectorWriter::uInt1));
                break;
            case UINT2:
                fieldReader.copyAsValue(extractor.get(ListWriter::uInt2, MapWriter::uInt2, UnionVectorWriter::uInt2));
                break;
            case UINT4:
                fieldReader.copyAsValue(extractor.get(ListWriter::uInt4, MapWriter::uInt4, UnionVectorWriter::uInt4));
                break;
            case UINT8:
                fieldReader.copyAsValue(extractor.get(ListWriter::uInt8, MapWriter::uInt8, UnionVectorWriter::uInt8));
                break;
            case DECIMAL9:
                fieldReader.copyAsValue((Decimal9Writer) extractor.get(ListWriter::decimal9, MapWriter::decimal9, UnionVectorWriter::decimal9));
                break;
            case DECIMAL18:
                fieldReader.copyAsValue((Decimal18Writer) extractor.get(ListWriter::decimal18, MapWriter::decimal18, UnionVectorWriter::decimal18));
                break;
            case DECIMAL28SPARSE:
                fieldReader.copyAsValue((Decimal28SparseWriter) extractor.get(ListWriter::decimal28Sparse, MapWriter::decimal28Sparse, UnionVectorWriter::decimal28Sparse));
                break;
            case DECIMAL38SPARSE:
                fieldReader.copyAsValue((Decimal38SparseWriter) extractor.get(ListWriter::decimal38Sparse, MapWriter::decimal38Sparse, UnionVectorWriter::decimal38Sparse));
                break;
            case VARDECIMAL:
                fieldReader.copyAsValue((VarDecimalWriter) extractor.get(lw -> lw.varDecimal(valueMajorType.getPrecision(), valueMajorType.getScale()), (mw, fn) -> mw.varDecimal(fn, valueMajorType.getPrecision(), valueMajorType.getScale()), uw -> uw.varDecimal(valueMajorType.getPrecision(), valueMajorType.getScale())));
                break;
            case DATE:
                fieldReader.copyAsValue(extractor.get(ListWriter::date, MapWriter::date, UnionVectorWriter::date));
                break;
            case TIME:
                fieldReader.copyAsValue(extractor.get(ListWriter::time, MapWriter::time, UnionVectorWriter::time));
                break;
            case TIMESTAMP:
                fieldReader.copyAsValue(extractor.get(ListWriter::timeStamp, MapWriter::timeStamp, UnionVectorWriter::timeStamp));
                break;
            case INTERVAL:
                fieldReader.copyAsValue(extractor.get(ListWriter::interval, MapWriter::interval, UnionVectorWriter::interval));
                break;
            case INTERVALDAY:
                fieldReader.copyAsValue(extractor.get(ListWriter::intervalDay, MapWriter::intervalDay, UnionVectorWriter::intervalDay));
                break;
            case INTERVALYEAR:
                fieldReader.copyAsValue(extractor.get(ListWriter::intervalYear, MapWriter::intervalYear, UnionVectorWriter::intervalYear));
                break;
            case FLOAT4:
                fieldReader.copyAsValue(extractor.get(ListWriter::float4, MapWriter::float4, UnionVectorWriter::float4));
                break;
            case FLOAT8:
                fieldReader.copyAsValue(extractor.get(ListWriter::float8, MapWriter::float8, UnionVectorWriter::float8));
                break;
            case BIT:
                fieldReader.copyAsValue(extractor.get(ListWriter::bit, MapWriter::bit, UnionVectorWriter::bit));
                break;
            case VARCHAR:
                fieldReader.copyAsValue(extractor.get(ListWriter::varChar, MapWriter::varChar, UnionVectorWriter::varChar));
                break;
            case VARBINARY:
                fieldReader.copyAsValue(extractor.get(ListWriter::varBinary, MapWriter::varBinary, UnionVectorWriter::varBinary));
                break;
            case MAP:
                fieldReader.copyAsValue(extractor.get(ListWriter::map, MapWriter::map, UnionVectorWriter::map));
                break;
            case LIST:
                fieldReader.copyAsValue(mapWriter.list(fieldName).list());
                break;
            case DICT:
                fieldReader.copyAsValue(extractor.get(ListWriter::dict, MapWriter::dict, UnionVectorWriter::dict));
                break;
            default:
                throw new DrillRuntimeException(String.format("%s does not support input of type: %s", caller, valueMinorType));
        }
    } catch (ClassCastException e) {
        final MaterializedField field = fieldReader.getField();
        throw new DrillRuntimeException(String.format(TYPE_MISMATCH_ERROR, caller, field.getName(), field.getType()));
    }
}
Also used : MapWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter.MapWriter) UnionVectorWriter(org.apache.drill.exec.vector.complex.impl.UnionVectorWriter) MajorType(org.apache.drill.common.types.TypeProtos.MajorType) ListWriter(org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter) MinorType(org.apache.drill.common.types.TypeProtos.MinorType) MaterializedField(org.apache.drill.exec.record.MaterializedField) UnionReader(org.apache.drill.exec.vector.complex.impl.UnionReader) DrillRuntimeException(org.apache.drill.common.exceptions.DrillRuntimeException)

Aggregations

ListWriter (org.apache.drill.exec.vector.complex.writer.BaseWriter.ListWriter)8 MapWriter (org.apache.drill.exec.vector.complex.writer.BaseWriter.MapWriter)7 BigIntHolder (org.apache.drill.exec.expr.holders.BigIntHolder)4 IntHolder (org.apache.drill.exec.expr.holders.IntHolder)4 BaseWriter (org.apache.drill.exec.vector.complex.writer.BaseWriter)4 DrillBuf (io.netty.buffer.DrillBuf)3 DrillRuntimeException (org.apache.drill.common.exceptions.DrillRuntimeException)3 PathSegment (org.apache.drill.common.expression.PathSegment)3 SchemaPath (org.apache.drill.common.expression.SchemaPath)3 OutputMutator (org.apache.drill.exec.physical.impl.OutputMutator)3 SingleMapWriter (org.apache.drill.exec.vector.complex.impl.SingleMapWriter)3 ArrayList (java.util.ArrayList)2 Collection (java.util.Collection)2 Collections (java.util.Collections)2 Iterator (java.util.Iterator)2 List (java.util.List)2 Optional (java.util.Optional)2 BiFunction (java.util.function.BiFunction)2 Function (java.util.function.Function)2 Supplier (java.util.function.Supplier)2