Search in sources :

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();
}
Also used : VariantWriter(org.apache.drill.exec.vector.accessor.VariantWriter) SingleRowSet(org.apache.drill.exec.physical.rowSet.RowSet.SingleRowSet) ObjectWriter(org.apache.drill.exec.vector.accessor.ObjectWriter) ScalarReader(org.apache.drill.exec.vector.accessor.ScalarReader) ArrayReader(org.apache.drill.exec.vector.accessor.ArrayReader) TupleMetadata(org.apache.drill.exec.record.metadata.TupleMetadata) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) ObjectReader(org.apache.drill.exec.vector.accessor.ObjectReader) ArrayWriter(org.apache.drill.exec.vector.accessor.ArrayWriter) ScalarWriter(org.apache.drill.exec.vector.accessor.ScalarWriter) ExtendableRowSet(org.apache.drill.exec.physical.rowSet.RowSet.ExtendableRowSet) VariantReader(org.apache.drill.exec.vector.accessor.VariantReader) SubOperatorTest(org.apache.drill.test.SubOperatorTest) Test(org.junit.Test)

Aggregations

SingleRowSet (org.apache.drill.exec.physical.rowSet.RowSet.SingleRowSet)11 SchemaBuilder (org.apache.drill.exec.record.metadata.SchemaBuilder)11 TupleMetadata (org.apache.drill.exec.record.metadata.TupleMetadata)11 VariantWriter (org.apache.drill.exec.vector.accessor.VariantWriter)11 SubOperatorTest (org.apache.drill.test.SubOperatorTest)11 Test (org.junit.Test)11 ResultSetLoader (org.apache.drill.exec.physical.resultSet.ResultSetLoader)6 RowSetLoader (org.apache.drill.exec.physical.resultSet.RowSetLoader)6 ArrayWriter (org.apache.drill.exec.vector.accessor.ArrayWriter)6 RowSet (org.apache.drill.exec.physical.rowSet.RowSet)5 ExtendableRowSet (org.apache.drill.exec.physical.rowSet.RowSet.ExtendableRowSet)5 ObjectWriter (org.apache.drill.exec.vector.accessor.ObjectWriter)5 VariantReader (org.apache.drill.exec.vector.accessor.VariantReader)5 ObjectReader (org.apache.drill.exec.vector.accessor.ObjectReader)4 ScalarReader (org.apache.drill.exec.vector.accessor.ScalarReader)4 TupleWriter (org.apache.drill.exec.vector.accessor.TupleWriter)4 ScalarWriter (org.apache.drill.exec.vector.accessor.ScalarWriter)3 MinorType (org.apache.drill.common.types.TypeProtos.MinorType)2 ArrayReader (org.apache.drill.exec.vector.accessor.ArrayReader)2 MajorType (org.apache.drill.common.types.TypeProtos.MajorType)1