Examples with TupleWriter org.apache.drill.exec.vector.accessor.TupleWriter used on opensource projects

Example 51 with TupleWriter

use of org.apache.drill.exec.vector.accessor.TupleWriter in project drill by apache.

the class TestMapAccessors method testBasics.

@Test
public void testBasics() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").add("c", MinorType.INT).add("d", MinorType.VARCHAR).resumeSchema().add("e", MinorType.VARCHAR).buildSchema();
    RowSetBuilder builder = fixture.rowSetBuilder(schema);
    RowSetWriter rootWriter = builder.writer();
    // Verify structure and schema
    final TupleMetadata actualSchema = rootWriter.tupleSchema();
    assertEquals(3, actualSchema.size());
    assertTrue(actualSchema.metadata(1).isMap());
    assertEquals(2, actualSchema.metadata("m").tupleSchema().size());
    assertEquals(2, actualSchema.column("m").getChildren().size());
    // Write a row the way that clients will do.
    final ScalarWriter aWriter = rootWriter.scalar("a");
    final TupleWriter mWriter = rootWriter.tuple("m");
    final ScalarWriter cWriter = mWriter.scalar("c");
    final ScalarWriter dWriter = mWriter.scalar("d");
    final ScalarWriter eWriter = rootWriter.scalar("e");
    aWriter.setInt(10);
    cWriter.setInt(110);
    dWriter.setString("fred");
    eWriter.setString("pebbles");
    rootWriter.save();
    // Write another using the test-time conveniences
    rootWriter.addRow(20, mapValue(210, "barney"), "bam-bam");
    RowSet result = builder.build();
    assertEquals(2, result.rowCount());
    // Validate internal structure.
    VectorContainer container = result.container();
    assertEquals(3, container.getNumberOfColumns());
    ValueVector v = container.getValueVector(1).getValueVector();
    assertTrue(v instanceof MapVector);
    MapVector mv = (MapVector) v;
    assertEquals(2, mv.getAccessor().getValueCount());
    // Validate data. Do so using the readers to avoid verifying
    // using the very mechanisms we want to test.
    RowSetReader rootReader = result.reader();
    final ScalarReader aReader = rootReader.scalar("a");
    final TupleReader mReader = rootReader.tuple("m");
    final ScalarReader cReader = mReader.scalar("c");
    final ScalarReader dReader = mReader.scalar("d");
    final ScalarReader eReader = rootReader.scalar("e");
    rootReader.next();
    assertEquals(10, aReader.getInt());
    assertEquals(110, cReader.getInt());
    assertEquals("fred", dReader.getString());
    assertEquals("pebbles", eReader.getString());
    rootReader.next();
    assertEquals(20, aReader.getInt());
    assertEquals(210, cReader.getInt());
    assertEquals("barney", dReader.getString());
    assertEquals("bam-bam", eReader.getString());
    // Verify using the convenience methods.
    final SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(10, mapValue(110, "fred"), "pebbles").addRow(20, mapValue(210, "barney"), "bam-bam").build();
    new RowSetComparison(expected).verify(result);
    // Test that the row set rebuilds its internal structure from
    // a vector container.
    RowSet wrapped = fixture.wrap(result.container());
    RowSetUtilities.verify(expected, wrapped);
}

Example 52 with TupleWriter

use of org.apache.drill.exec.vector.accessor.TupleWriter in project drill by apache.

the class TestRowSet method testMapStructure.

/**
 * Test a simple map structure at the top level of a row.
 *
 * @throws VectorOverflowException should never occur
 */
@Test
public void testMapStructure() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").addArray("b", MinorType.INT).resumeSchema().buildSchema();
    final ExtendableRowSet rowSet = fixture.rowSet(schema);
    final RowSetWriter writer = rowSet.writer();
    // Map and Int
    // Test Invariants
    assertEquals(ObjectType.SCALAR, writer.column("a").type());
    assertEquals(ObjectType.SCALAR, writer.column(0).type());
    assertEquals(ObjectType.TUPLE, writer.column("m").type());
    assertEquals(ObjectType.TUPLE, writer.column(1).type());
    assertSame(writer.column(1).tuple(), writer.tuple(1));
    final TupleWriter mapWriter = writer.column(1).tuple();
    assertEquals(ObjectType.SCALAR, mapWriter.column("b").array().entry().type());
    assertEquals(ObjectType.SCALAR, mapWriter.column("b").array().entryType());
    final ScalarWriter aWriter = writer.column("a").scalar();
    final ScalarWriter bWriter = writer.column("m").tuple().column("b").array().entry().scalar();
    assertSame(bWriter, writer.tuple(1).array(0).scalar());
    assertEquals(ValueType.INTEGER, bWriter.valueType());
    try {
        writer.column(1).scalar();
        fail();
    } catch (final UnsupportedOperationException e) {
    // Expected
    }
    try {
        writer.column(1).array();
        fail();
    } catch (final UnsupportedOperationException e) {
    // Expected
    }
    // Write data
    aWriter.setInt(10);
    bWriter.setInt(11);
    bWriter.setInt(12);
    writer.save();
    aWriter.setInt(20);
    bWriter.setInt(21);
    bWriter.setInt(22);
    writer.save();
    aWriter.setInt(30);
    bWriter.setInt(31);
    bWriter.setInt(32);
    writer.save();
    // Finish the row set and get a reader.
    final SingleRowSet actual = writer.done();
    final RowSetReader reader = actual.reader();
    assertEquals(ObjectType.SCALAR, reader.column("a").type());
    assertEquals(ObjectType.SCALAR, reader.column(0).type());
    assertEquals(ObjectType.TUPLE, reader.column("m").type());
    assertEquals(ObjectType.TUPLE, reader.column(1).type());
    assertSame(reader.column(1).tuple(), reader.tuple(1));
    final ScalarReader aReader = reader.column(0).scalar();
    final TupleReader mReader = reader.column(1).tuple();
    final ArrayReader bArray = mReader.column("b").array();
    assertEquals(ObjectType.SCALAR, bArray.entryType());
    final ScalarReader bReader = bArray.scalar();
    assertEquals(ValueType.INTEGER, bReader.valueType());
    // Row 1: (10, {[11, 12]})
    assertTrue(reader.next());
    assertEquals(10, aReader.getInt());
    assertFalse(mReader.isNull());
    assertTrue(bArray.next());
    assertFalse(bReader.isNull());
    assertEquals(11, bReader.getInt());
    assertTrue(bArray.next());
    assertFalse(bReader.isNull());
    assertEquals(12, bReader.getInt());
    assertFalse(bArray.next());
    // Row 2: (20, {[21, 22]})
    assertTrue(reader.next());
    assertEquals(20, aReader.getInt());
    assertFalse(mReader.isNull());
    assertTrue(bArray.next());
    assertEquals(21, bReader.getInt());
    assertTrue(bArray.next());
    assertEquals(22, bReader.getInt());
    // Row 3: (30, {[31, 32]})
    assertTrue(reader.next());
    assertEquals(30, aReader.getInt());
    assertFalse(mReader.isNull());
    assertTrue(bArray.next());
    assertEquals(31, bReader.getInt());
    assertTrue(bArray.next());
    assertEquals(32, bReader.getInt());
    assertFalse(reader.next());
    // Verify that the map accessor's value count was set.
    final MapVector mapVector = (MapVector) actual.container().getValueVector(1).getValueVector();
    assertEquals(actual.rowCount(), mapVector.getAccessor().getValueCount());
    final SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(10, objArray(intArray(11, 12))).addRow(20, objArray(intArray(21, 22))).addRow(30, objArray(intArray(31, 32))).build();
    RowSetUtilities.verify(expected, actual);
}

Example 53 with TupleWriter

use of org.apache.drill.exec.vector.accessor.TupleWriter in project drill by apache.

the class TestRowSet method testRepeatedMapStructure.

@Test
public void testRepeatedMapStructure() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addMapArray("m").add("b", MinorType.INT).add("c", MinorType.INT).resumeSchema().buildSchema();
    final ExtendableRowSet rowSet = fixture.rowSet(schema);
    final RowSetWriter writer = rowSet.writer();
    // Map and Int
    // Pick out components and lightly test. (Assumes structure
    // tested earlier is still valid, so no need to exhaustively
    // test again.)
    assertEquals(ObjectType.SCALAR, writer.column("a").type());
    assertEquals(ObjectType.ARRAY, writer.column("m").type());
    final ArrayWriter maWriter = writer.column(1).array();
    assertEquals(ObjectType.TUPLE, maWriter.entryType());
    final TupleWriter mapWriter = maWriter.tuple();
    assertEquals(ObjectType.SCALAR, mapWriter.column("b").type());
    assertEquals(ObjectType.SCALAR, mapWriter.column("c").type());
    final ScalarWriter aWriter = writer.column("a").scalar();
    final ScalarWriter bWriter = mapWriter.scalar("b");
    final ScalarWriter cWriter = mapWriter.scalar("c");
    assertEquals(ValueType.INTEGER, aWriter.valueType());
    assertEquals(ValueType.INTEGER, bWriter.valueType());
    assertEquals(ValueType.INTEGER, cWriter.valueType());
    // Write data
    aWriter.setInt(10);
    bWriter.setInt(101);
    cWriter.setInt(102);
    // Advance to next array position
    maWriter.save();
    bWriter.setInt(111);
    cWriter.setInt(112);
    maWriter.save();
    writer.save();
    aWriter.setInt(20);
    bWriter.setInt(201);
    cWriter.setInt(202);
    maWriter.save();
    bWriter.setInt(211);
    cWriter.setInt(212);
    maWriter.save();
    writer.save();
    aWriter.setInt(30);
    bWriter.setInt(301);
    cWriter.setInt(302);
    maWriter.save();
    bWriter.setInt(311);
    cWriter.setInt(312);
    maWriter.save();
    writer.save();
    // Finish the row set and get a reader.
    final SingleRowSet actual = writer.done();
    final RowSetReader reader = actual.reader();
    // Verify reader structure
    assertEquals(ObjectType.SCALAR, reader.column("a").type());
    assertEquals(ObjectType.ARRAY, reader.column("m").type());
    final ArrayReader maReader = reader.column(1).array();
    assertEquals(ObjectType.TUPLE, maReader.entryType());
    final TupleReader mapReader = maReader.tuple();
    assertEquals(ObjectType.SCALAR, mapReader.column("b").type());
    assertEquals(ObjectType.SCALAR, mapReader.column("c").type());
    final ScalarReader aReader = reader.column("a").scalar();
    final ScalarReader bReader = mapReader.scalar("b");
    final ScalarReader cReader = mapReader.scalar("c");
    assertEquals(ValueType.INTEGER, aReader.valueType());
    assertEquals(ValueType.INTEGER, bReader.valueType());
    assertEquals(ValueType.INTEGER, cReader.valueType());
    // Row 1: Use iterator-like accessors
    assertTrue(reader.next());
    assertEquals(10, aReader.getInt());
    // Array itself is not null
    assertFalse(maReader.isNull());
    assertTrue(maReader.next());
    // Tuple 0 is not null
    assertFalse(mapReader.isNull());
    assertEquals(101, mapReader.scalar(0).getInt());
    assertEquals(102, mapReader.scalar(1).getInt());
    assertTrue(maReader.next());
    assertEquals(111, mapReader.scalar(0).getInt());
    assertEquals(112, mapReader.scalar(1).getInt());
    // Row 2: use explicit positioning,
    // but access scalars through the map reader.
    assertTrue(reader.next());
    assertEquals(20, aReader.getInt());
    maReader.setPosn(0);
    assertEquals(201, mapReader.scalar(0).getInt());
    assertEquals(202, mapReader.scalar(1).getInt());
    maReader.setPosn(1);
    assertEquals(211, mapReader.scalar(0).getInt());
    assertEquals(212, mapReader.scalar(1).getInt());
    // Row 3: use scalar accessor
    assertTrue(reader.next());
    assertEquals(30, aReader.getInt());
    assertTrue(maReader.next());
    assertEquals(301, bReader.getInt());
    assertEquals(302, cReader.getInt());
    assertTrue(maReader.next());
    assertEquals(311, bReader.getInt());
    assertEquals(312, cReader.getInt());
    assertFalse(reader.next());
    // Verify that the map accessor's value count was set.
    final RepeatedMapVector mapVector = (RepeatedMapVector) actual.container().getValueVector(1).getValueVector();
    assertEquals(3, mapVector.getAccessor().getValueCount());
    // Verify the readers and writers again using the testing tools.
    final SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(10, objArray(objArray(101, 102), objArray(111, 112))).addRow(20, objArray(objArray(201, 202), objArray(211, 212))).addRow(30, objArray(objArray(301, 302), objArray(311, 312))).build();
    RowSetUtilities.verify(expected, actual);
}

Example 54 with TupleWriter

use of org.apache.drill.exec.vector.accessor.TupleWriter in project drill by apache.

the class TestVariantAccessors method testListOfMaps.

/**
 * List of maps. Like a repeated map, but each list entry can be
 * null.
 */
@Test
public void testListOfMaps() {
    final TupleMetadata schema = new SchemaBuilder().addList("list").addMap().addNullable("a", MinorType.INT).addNullable("b", MinorType.VARCHAR).resumeUnion().resumeSchema().buildSchema();
    final ExtendableRowSet rowSet = fixture.rowSet(schema);
    final RowSetWriter writer = rowSet.writer();
    {
        final ObjectWriter listObj = writer.column("list");
        assertEquals(ObjectType.ARRAY, listObj.type());
        final ArrayWriter listArray = listObj.array();
        final ObjectWriter itemObj = listArray.entry();
        assertEquals(ObjectType.TUPLE, itemObj.type());
        final TupleWriter mapWriter = itemObj.tuple();
        final ScalarWriter aWriter = mapWriter.scalar("a");
        final ScalarWriter bWriter = mapWriter.scalar("b");
        // First row:
        // {1, "fred"}, null, {3, null}
        aWriter.setInt(1);
        bWriter.setString("fred");
        listArray.save();
        // Can't mark the map as null. Instead, we simply skip
        // the map and the contained nullable members will automatically
        // back-fill each entry with a null value.
        listArray.save();
        aWriter.setInt(3);
        bWriter.setNull();
        listArray.save();
        writer.save();
        // Second row: null
        writer.save();
        // Third row: {null, "dino"}
        aWriter.setNull();
        bWriter.setString("dino");
        listArray.save();
        writer.save();
        // Fourth row: empty array. Note that there is no trigger
        // to say that the column is not null, so we have to do it
        // explicitly.
        listArray.setNull(false);
        writer.save();
        // Last row: {4, "pebbles"}
        aWriter.setInt(4);
        bWriter.setString("pebbles");
        listArray.save();
        writer.save();
    }
    final SingleRowSet result = writer.done();
    assertEquals(5, result.rowCount());
    {
        final RowSetReader reader = result.reader();
        final ObjectReader listObj = reader.column("list");
        assertEquals(ObjectType.ARRAY, listObj.type());
        final ArrayReader listArray = listObj.array();
        assertEquals(ObjectType.TUPLE, listArray.entry().type());
        final TupleReader mapReader = listArray.tuple();
        final ScalarReader aReader = mapReader.scalar("a");
        final ScalarReader bReader = mapReader.scalar("b");
        // First row:
        // {1, "fred"}, null, {3, null}
        assertTrue(reader.next());
        assertFalse(listArray.isNull());
        assertFalse(mapReader.isNull());
        assertEquals(3, listArray.size());
        assertTrue(listArray.next());
        assertFalse(aReader.isNull());
        assertEquals(1, aReader.getInt());
        assertFalse(bReader.isNull());
        assertEquals("fred", bReader.getString());
        assertTrue(listArray.next());
        // Awkward: the map has no null state, but its
        // members do.
        assertTrue(aReader.isNull());
        assertTrue(bReader.isNull());
        assertTrue(listArray.next());
        assertFalse(aReader.isNull());
        assertEquals(3, aReader.getInt());
        assertTrue(bReader.isNull());
        assertFalse(listArray.next());
        // Second row: null
        assertTrue(reader.next());
        assertTrue(listArray.isNull());
        assertEquals(0, listArray.size());
        // Third row: {null, "dino"}
        assertTrue(reader.next());
        assertFalse(listArray.isNull());
        assertEquals(1, listArray.size());
        assertTrue(listArray.next());
        assertTrue(aReader.isNull());
        assertFalse(bReader.isNull());
        assertEquals("dino", bReader.getString());
        assertFalse(listArray.next());
        // Fourth row: empty array.
        assertTrue(reader.next());
        assertFalse(listArray.isNull());
        assertEquals(0, listArray.size());
        assertFalse(listArray.next());
        // Last row: {4, "pebbles"}
        assertTrue(reader.next());
        assertFalse(listArray.isNull());
        assertEquals(1, listArray.size());
        assertTrue(listArray.next());
        assertEquals(4, aReader.getInt());
        assertEquals("pebbles", bReader.getString());
        assertFalse(listArray.next());
        assertFalse(reader.next());
    }
    result.clear();
}

Example 55 with TupleWriter

use of org.apache.drill.exec.vector.accessor.TupleWriter in project drill by apache.

the class HDF5BatchReader method writeAttributes.

/**
 * Gets the attributes for an HDF5 datapath. These attributes are projected as
 * a map in select * queries when the defaultPath is null.
 *
 * @param rowWriter
 *          the row to which the data will be written
 * @param record
 *          the record for the attributes
 */
private void writeAttributes(TupleWriter rowWriter, HDF5DrillMetadata record) {
    Map<String, HDF5Attribute> attribs = getAttributes(record.getPath());
    Iterator<Map.Entry<String, HDF5Attribute>> entries = attribs.entrySet().iterator();
    int index = rowWriter.tupleSchema().index("attributes");
    if (index == -1) {
        index = rowWriter.addColumn(SchemaBuilder.columnSchema("attributes", TypeProtos.MinorType.MAP, TypeProtos.DataMode.REQUIRED));
    }
    TupleWriter mapWriter = rowWriter.tuple(index);
    while (entries.hasNext()) {
        Map.Entry<String, HDF5Attribute> entry = entries.next();
        String key = entry.getKey();
        HDF5Attribute attrib = entry.getValue();
        switch(attrib.getDataType()) {
            case BIT:
                writeBooleanColumn(mapWriter, key, (Boolean) attrib.getValue());
                break;
            case BIGINT:
                writeLongColumn(mapWriter, key, (Long) attrib.getValue());
                break;
            case INT:
                writeIntColumn(mapWriter, key, (Integer) attrib.getValue());
                break;
            case SMALLINT:
                writeSmallIntColumn(mapWriter, key, (Short) attrib.getValue());
                break;
            case TINYINT:
                writeByteColumn(mapWriter, key, (Byte) attrib.getValue());
                break;
            case FLOAT8:
                writeFloat8Column(mapWriter, key, (Double) attrib.getValue());
                break;
            case FLOAT4:
                writeFloat4Column(mapWriter, key, (Float) attrib.getValue());
                break;
            case VARCHAR:
                writeStringColumn(mapWriter, key, (String) attrib.getValue());
                break;
            case TIMESTAMP:
                writeTimestampColumn(mapWriter, key, (Long) attrib.getValue());
                break;
            case GENERIC_OBJECT:
                // This is the case for HDF5 enums
                String enumText = attrib.getValue().toString();
                writeStringColumn(mapWriter, key, enumText);
                break;
            default:
                throw new IllegalStateException(attrib.getDataType().name());
        }
    }
}