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

Example 46 with TupleWriter

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

the class TestResultSetLoaderDicts method testMapValueRequiredFields.

/**
 * Create dict with map value. Then, add columns to the map
 * on the fly. Use required, variable-width columns since
 * those require the most processing and are most likely to
 * fail if anything is out of place.
 */
@Test
public void testMapValueRequiredFields() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addDict("d", MinorType.VARCHAR).mapValue().add("b", MinorType.VARCHAR).resumeDict().resumeSchema().buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    assertEquals(5, rsLoader.schemaVersion());
    final RowSetLoader rootWriter = rsLoader.writer();
    rsLoader.startBatch();
    rootWriter.addRow(10, map("a", mapValue("c1"), "b", mapValue("c2")));
    // Validate first batch
    RowSet actual = fixture.wrap(rsLoader.harvest());
    assertEquals(5, rsLoader.schemaVersion());
    SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(10, map("a", mapValue("c1"), "b", mapValue("c2"))).build();
    RowSetUtilities.verify(expected, actual);
    // Now add columns in the second batch.
    rsLoader.startBatch();
    rootWriter.addRow(20, map("a2", mapValue("c11"), "b2", mapValue("c12"), "c2", mapValue("c13")));
    final DictWriter dictWriter = rootWriter.dict("d");
    final TupleWriter nestedMapWriter = dictWriter.valueWriter().tuple();
    nestedMapWriter.addColumn(SchemaBuilder.columnSchema("c", MinorType.VARCHAR, DataMode.REQUIRED));
    rootWriter.addRow(30, map("a3", mapValue("c21", "d21")));
    // And another set while the write proceeds.
    nestedMapWriter.addColumn(SchemaBuilder.columnSchema("d", MinorType.VARCHAR, DataMode.REQUIRED));
    rootWriter.addRow(40, map("a4", mapValue("c31", "d31", "e31"), "b4", mapValue("c32", "d32", "e32")));
    // Validate second batch
    actual = fixture.wrap(rsLoader.harvest());
    assertEquals(7, rsLoader.schemaVersion());
    final TupleMetadata expectedSchema = new SchemaBuilder().add("a", MinorType.INT).addDict("d", MinorType.VARCHAR).mapValue().add("b", MinorType.VARCHAR).add("c", MinorType.VARCHAR).add("d", MinorType.VARCHAR).resumeDict().resumeSchema().buildSchema();
    expected = fixture.rowSetBuilder(expectedSchema).addRow(20, map("a2", mapValue("c11", "", ""), "b2", mapValue("c12", "", ""), "c2", mapValue("c13", "", ""))).addRow(30, map("a3", mapValue("c21", "d21", ""))).addRow(40, map("a4", mapValue("c31", "d31", "e31"), "b4", mapValue("c32", "d32", "e32"))).build();
    RowSetUtilities.verify(expected, actual);
    rsLoader.close();
}

Example 47 with TupleWriter

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

the class TestResultSetLoaderMaps method testNameSpace.

/**
 * Verify that map name spaces (and implementations) are
 * independent.
 */
@Test
public void testNameSpace() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").add("a", MinorType.INT).addMap("m").add("a", MinorType.INT).resumeMap().resumeSchema().buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    assertFalse(rsLoader.isProjectionEmpty());
    final RowSetLoader rootWriter = rsLoader.writer();
    rsLoader.startBatch();
    // Write a row the way that clients will do.
    final ScalarWriter a1Writer = rootWriter.scalar("a");
    final TupleWriter m1Writer = rootWriter.tuple("m");
    final ScalarWriter a2Writer = m1Writer.scalar("a");
    final TupleWriter m2Writer = m1Writer.tuple("m");
    final ScalarWriter a3Writer = m2Writer.scalar("a");
    rootWriter.start();
    a1Writer.setInt(11);
    a2Writer.setInt(12);
    a3Writer.setInt(13);
    rootWriter.save();
    rootWriter.start();
    a1Writer.setInt(21);
    a2Writer.setInt(22);
    a3Writer.setInt(23);
    rootWriter.save();
    // Try simplified test format
    rootWriter.addRow(31, mapValue(32, mapValue(33)));
    // Verify
    final RowSet actual = fixture.wrap(rsLoader.harvest());
    final SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(11, mapValue(12, mapValue(13))).addRow(21, mapValue(22, mapValue(23))).addRow(31, mapValue(32, mapValue(33))).build();
    RowSetUtilities.verify(expected, actual);
    rsLoader.close();
}

Example 48 with TupleWriter

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

the class TestResultSetLoaderMaps method testEmptyMapAddition.

/**
 * Test adding an empty map to a loader after writing the first row.
 * Then add columns in another batch. Yes, this is a bizarre condition,
 * but we must check it anyway for robustness.
 */
@Test
public void testEmptyMapAddition() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    assertEquals(1, rsLoader.schemaVersion());
    final RowSetLoader rootWriter = rsLoader.writer();
    // Start without the map. Add a map after the first row.
    rsLoader.startBatch();
    rootWriter.addRow(10);
    final int mapIndex = rootWriter.addColumn(SchemaBuilder.columnSchema("m", MinorType.MAP, DataMode.REQUIRED));
    final TupleWriter mapWriter = rootWriter.tuple(mapIndex);
    rootWriter.addRow(20, mapValue()).addRow(30, mapValue());
    RowSet actual = fixture.wrap(rsLoader.harvest());
    assertEquals(2, rsLoader.schemaVersion());
    assertEquals(3, actual.rowCount());
    // Validate first batch
    TupleMetadata expectedSchema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").resumeSchema().buildSchema();
    SingleRowSet expected = fixture.rowSetBuilder(expectedSchema).addRow(10, mapValue()).addRow(20, mapValue()).addRow(30, mapValue()).build();
    RowSetUtilities.verify(expected, actual);
    // Now add another column to the map
    rsLoader.startBatch();
    mapWriter.addColumn(SchemaBuilder.columnSchema("a", MinorType.VARCHAR, DataMode.REQUIRED));
    rootWriter.addRow(40, mapValue("fred")).addRow(50, mapValue("barney"));
    actual = fixture.wrap(rsLoader.harvest());
    assertEquals(3, rsLoader.schemaVersion());
    assertEquals(2, actual.rowCount());
    // Validate first batch
    expectedSchema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").add("a", MinorType.VARCHAR).resumeSchema().buildSchema();
    expected = fixture.rowSetBuilder(expectedSchema).addRow(40, mapValue("fred")).addRow(50, mapValue("barney")).build();
    RowSetUtilities.verify(expected, actual);
    rsLoader.close();
}

Example 49 with TupleWriter

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

the class TestResultSetLoaderUnions method testRepeatedListOfUnion.

/**
 * The repeated list type is way off in the weeds in Drill. It is not fully
 * supported and the semantics are very murky as a result. It is not clear
 * how such a structure fits into SQL or into an xDBC client. Still, it is
 * part of Drill at present and must be supported in the result set loader.
 * <p>
 * This test models using the repeated list as a 2D array of UNION.
 */
@Test
public void testRepeatedListOfUnion() {
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator());
    final RowSetLoader writer = rsLoader.writer();
    // Can write a batch as if this was a repeated Varchar, except
    // that any value can also be null.
    rsLoader.startBatch();
    writer.addColumn(MaterializedField.create("id", Types.required(MinorType.INT)));
    // A union requires a structured column. The only tool to build that a present
    // is the schema builder, so we use that and grab a single column.
    final TupleMetadata dummySchema = new SchemaBuilder().addRepeatedList("list").addArray(MinorType.UNION).resumeSchema().buildSchema();
    writer.addColumn(dummySchema.metadata(0));
    // Sanity check: should be an array of array of variants.
    final ArrayWriter outerArrWriter = writer.array("list");
    assertEquals(ObjectType.ARRAY, outerArrWriter.entryType());
    final ArrayWriter innerArrWriter = outerArrWriter.array();
    assertEquals(ObjectType.VARIANT, innerArrWriter.entryType());
    final VariantWriter variant = innerArrWriter.variant();
    // No types, so all we can do is add a null list, or a list of nulls.
    writer.addRow(1, null).addRow(2, objArray()).addRow(3, objArray(null, null)).addRow(4, objArray(variantArray(), variantArray())).addRow(5, objArray(variantArray(null, null), variantArray(null, null)));
    // Add a String. Now we can create a list of strings and/or nulls.
    variant.addMember(MinorType.VARCHAR);
    assertTrue(variant.hasType(MinorType.VARCHAR));
    writer.addRow(6, objArray(variantArray("fred", "wilma", null), variantArray("barney", "betty", null)));
    // Add a map
    final TupleWriter mapWriter = variant.addMember(MinorType.MAP).tuple();
    mapWriter.addColumn(MetadataUtils.newScalar("first", Types.optional(MinorType.VARCHAR)));
    mapWriter.addColumn(MetadataUtils.newScalar("last", Types.optional(MinorType.VARCHAR)));
    // Add a map-based record
    writer.addRow(7, objArray(variantArray(mapValue("fred", "flintstone"), mapValue("wilma", "flintstone")), variantArray(mapValue("barney", "rubble"), mapValue("betty", "rubble"))));
    // Verify
    final RowSet result = fixture.wrap(rsLoader.harvest());
    final TupleMetadata schema = new SchemaBuilder().add("id", MinorType.INT).addRepeatedList("list").addList().addType(MinorType.VARCHAR).addMap().addNullable("first", MinorType.VARCHAR).addNullable("last", MinorType.VARCHAR).resumeUnion().resumeRepeatedList().resumeSchema().buildSchema();
    final SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(1, null).addRow(2, objArray()).addRow(3, objArray(null, null)).addRow(4, objArray(variantArray(), variantArray())).addRow(5, objArray(variantArray(null, null), variantArray(null, null))).addRow(6, objArray(variantArray("fred", "wilma", null), variantArray("barney", "betty", null))).addRow(7, objArray(variantArray(mapValue("fred", "flintstone"), mapValue("wilma", "flintstone")), variantArray(mapValue("barney", "rubble"), mapValue("betty", "rubble")))).build();
    RowSetUtilities.verify(expected, result);
}

Example 50 with TupleWriter

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

the class TestMapAccessors method testDoubleNestedArray.

/**
 * Test a doubly-nested array of maps.
 */
@Test
public void testDoubleNestedArray() {
    TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addMapArray("m1").add("b", MinorType.INT).addMapArray("m2").add("c", MinorType.INT).addArray("d", MinorType.VARCHAR).resumeMap().resumeSchema().buildSchema();
    RowSetBuilder builder = fixture.rowSetBuilder(schema);
    RowSetWriter rootWriter = builder.writer();
    ScalarWriter aWriter = rootWriter.scalar("a");
    ArrayWriter a1Writer = rootWriter.array("m1");
    TupleWriter m1Writer = a1Writer.tuple();
    ScalarWriter bWriter = m1Writer.scalar("b");
    ArrayWriter a2Writer = m1Writer.array("m2");
    TupleWriter m2Writer = a2Writer.tuple();
    ScalarWriter cWriter = m2Writer.scalar("c");
    ScalarWriter dWriter = m2Writer.array("d").scalar();
    for (int i = 0; i < 5; i++) {
        aWriter.setInt(i);
        for (int j = 0; j < 4; j++) {
            int a1Key = i + 10 + j;
            bWriter.setInt(a1Key);
            for (int k = 0; k < 3; k++) {
                int a2Key = a1Key * 10 + k;
                cWriter.setInt(a2Key);
                for (int l = 0; l < 2; l++) {
                    dWriter.setString("d-" + (a2Key * 10 + l));
                }
                a2Writer.save();
            }
            a1Writer.save();
        }
        rootWriter.save();
    }
    RowSet results = builder.build();
    RowSetReader reader = results.reader();
    ScalarReader aReader = reader.scalar("a");
    ArrayReader a1Reader = reader.array("m1");
    TupleReader m1Reader = a1Reader.tuple();
    ScalarReader bReader = m1Reader.scalar("b");
    ArrayReader a2Reader = m1Reader.array("m2");
    TupleReader m2Reader = a2Reader.tuple();
    ScalarReader cReader = m2Reader.scalar("c");
    ArrayReader dArray = m2Reader.array("d");
    ScalarReader dReader = dArray.scalar();
    for (int i = 0; i < 5; i++) {
        assertTrue(reader.next());
        assertEquals(i, aReader.getInt());
        for (int j = 0; j < 4; j++) {
            assertTrue(a1Reader.next());
            int a1Key = i + 10 + j;
            assertEquals(a1Key, bReader.getInt());
            for (int k = 0; k < 3; k++) {
                assertTrue(a2Reader.next());
                int a2Key = a1Key * 10 + k;
                assertEquals(a2Key, cReader.getInt());
                for (int l = 0; l < 2; l++) {
                    assertTrue(dArray.next());
                    assertEquals("d-" + (a2Key * 10 + l), dReader.getString());
                }
            }
        }
    }
    results.clear();
}