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

Example 11 with VariantWriter

use of org.apache.drill.exec.vector.accessor.VariantWriter 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();
}