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

Example 96 with ScalarWriter

use of org.apache.drill.exec.vector.accessor.ScalarWriter 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 97 with ScalarWriter

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

the class TestVariantAccessors method testListOfUnions.

/**
 * Test a union list.
 */
@Test
public void testListOfUnions() {
    final TupleMetadata schema = new SchemaBuilder().addList("list").addType(MinorType.INT).addType(MinorType.VARCHAR).resumeSchema().buildSchema();
    final ExtendableRowSet rowSet = fixture.rowSet(schema);
    final RowSetWriter writer = rowSet.writer();
    {
        final ObjectWriter listObj = writer.column(0);
        assertEquals(ObjectType.ARRAY, listObj.type());
        final ArrayWriter listArray = listObj.array();
        final ObjectWriter itemObj = listArray.entry();
        assertEquals(ObjectType.VARIANT, itemObj.type());
        final VariantWriter variant = itemObj.variant();
        final ScalarWriter intWriter = variant.scalar(MinorType.INT);
        final ScalarWriter strWriter = variant.scalar(MinorType.VARCHAR);
        // First row: (1, "two", 3)
        variant.setType(MinorType.INT);
        intWriter.setInt(1);
        listArray.save();
        variant.setType(MinorType.VARCHAR);
        strWriter.setString("two");
        listArray.save();
        variant.setType(MinorType.INT);
        intWriter.setInt(3);
        listArray.save();
        writer.save();
        // Second row: null
        writer.save();
        // Third row: 4, null, "six", null int, null string
        variant.setType(MinorType.INT);
        intWriter.setInt(4);
        listArray.save();
        variant.setNull();
        listArray.save();
        variant.setType(MinorType.VARCHAR);
        strWriter.setString("six");
        listArray.save();
        variant.setType(MinorType.INT);
        intWriter.setNull();
        listArray.save();
        variant.setType(MinorType.VARCHAR);
        intWriter.setNull();
        listArray.save();
        writer.save();
        // Fourth row: empty array.
        listArray.setNull(false);
        writer.save();
        // Fifth row: 9
        variant.setType(MinorType.INT);
        intWriter.setInt(9);
        listArray.save();
        writer.save();
    }
    final SingleRowSet result = writer.done();
    assertEquals(5, result.rowCount());
    {
        final RowSetReader reader = result.reader();
        final ObjectReader listObj = reader.column(0);
        assertEquals(ObjectType.ARRAY, listObj.type());
        final ArrayReader listArray = listObj.array();
        assertEquals(ObjectType.VARIANT, listArray.entry().type());
        final VariantReader variant = listArray.variant();
        final ScalarReader intReader = variant.scalar(MinorType.INT);
        final ScalarReader strReader = variant.scalar(MinorType.VARCHAR);
        // First row: (1, "two", 3)
        assertTrue(reader.next());
        assertFalse(listArray.isNull());
        assertEquals(3, listArray.size());
        assertTrue(listArray.next());
        assertEquals(MinorType.INT, variant.dataType());
        assertFalse(intReader.isNull());
        assertTrue(strReader.isNull());
        assertEquals(1, intReader.getInt());
        assertEquals(1, variant.scalar().getInt());
        assertTrue(listArray.next());
        assertEquals(MinorType.VARCHAR, variant.dataType());
        assertTrue(intReader.isNull());
        assertFalse(strReader.isNull());
        assertEquals("two", strReader.getString());
        assertEquals("two", variant.scalar().getString());
        assertTrue(listArray.next());
        assertEquals(MinorType.INT, variant.dataType());
        assertEquals(3, intReader.getInt());
        assertFalse(listArray.next());
        // Second row: null
        assertTrue(reader.next());
        assertTrue(listArray.isNull());
        assertEquals(0, listArray.size());
        // Third row: 4, null, "six", null int, null string
        assertTrue(reader.next());
        assertEquals(5, listArray.size());
        assertTrue(listArray.next());
        assertEquals(4, intReader.getInt());
        assertTrue(listArray.next());
        assertTrue(variant.isNull());
        assertTrue(listArray.next());
        assertEquals("six", strReader.getString());
        assertTrue(listArray.next());
        assertEquals(MinorType.INT, variant.dataType());
        assertTrue(intReader.isNull());
        assertTrue(listArray.next());
        assertEquals(MinorType.VARCHAR, variant.dataType());
        assertTrue(strReader.isNull());
        assertFalse(listArray.next());
        // Fourth row: empty array.
        assertTrue(reader.next());
        assertFalse(listArray.isNull());
        assertEquals(0, listArray.size());
        assertFalse(listArray.next());
        // Fifth row: 9
        assertTrue(reader.next());
        assertEquals(1, listArray.size());
        assertTrue(listArray.next());
        assertEquals(9, intReader.getInt());
        assertFalse(listArray.next());
        assertFalse(reader.next());
    }
    result.clear();
}

Example 98 with ScalarWriter

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

the class HDF5BatchReader method writeSmallIntColumn.

/**
 * Helper function to write a 2D short list
 * @param rowWriter the row to which the data will be written
 * @param name the name of the outer list
 * @param list the list of data
 */
private void writeSmallIntColumn(TupleWriter rowWriter, String name, short[] list) {
    int index = rowWriter.tupleSchema().index(name);
    if (index == -1) {
        ColumnMetadata colSchema = MetadataUtils.newScalar(name, MinorType.SMALLINT, TypeProtos.DataMode.REPEATED);
        index = rowWriter.addColumn(colSchema);
    }
    ScalarWriter arrayWriter = rowWriter.column(index).array().scalar();
    int maxElements = Math.min(list.length, PREVIEW_ROW_LIMIT);
    for (int i = 0; i < maxElements; i++) {
        arrayWriter.setInt(list[i]);
    }
}

Example 99 with ScalarWriter

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

the class HDF5BatchReader method writeLongListColumn.

private void writeLongListColumn(TupleWriter rowWriter, String name, long[] list) {
    int index = rowWriter.tupleSchema().index(name);
    if (index == -1) {
        ColumnMetadata colSchema = MetadataUtils.newScalar(name, TypeProtos.MinorType.BIGINT, TypeProtos.DataMode.REPEATED);
        index = rowWriter.addColumn(colSchema);
    }
    ScalarWriter arrayWriter = rowWriter.column(index).array().scalar();
    int maxElements = Math.min(list.length, PREVIEW_ROW_LIMIT);
    for (int i = 0; i < maxElements; i++) {
        arrayWriter.setLong(list[i]);
    }
}

Example 100 with ScalarWriter

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

the class HDF5BatchReader method writeStringListColumn.

private void writeStringListColumn(TupleWriter rowWriter, String name, String[] list) {
    int index = rowWriter.tupleSchema().index(name);
    if (index == -1) {
        ColumnMetadata colSchema = MetadataUtils.newScalar(name, TypeProtos.MinorType.VARCHAR, TypeProtos.DataMode.REPEATED);
        index = rowWriter.addColumn(colSchema);
    }
    ScalarWriter arrayWriter = rowWriter.column(index).array().scalar();
    int maxElements = Math.min(list.length, PREVIEW_ROW_LIMIT);
    for (int i = 0; i < maxElements; i++) {
        arrayWriter.setString(list[i]);
    }
}