Examples with RowSetLoader org.apache.drill.exec.physical.resultSet.RowSetLoader used on opensource projects

Example 86 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader in project drill by apache.

the class TestResultSetLoaderRepeatedList method test2DLateSchema.

@Test
public void test2DLateSchema() {
    final TupleMetadata schema = new SchemaBuilder().add("id", MinorType.INT).addRepeatedList("list2").addArray(MinorType.VARCHAR).resumeSchema().buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    final RowSetLoader writer = rsLoader.writer();
    // Add columns dynamically
    writer.addColumn(schema.metadata(0));
    writer.addColumn(schema.metadata(1).cloneEmpty());
    // Yes, this is ugly. The whole repeated array idea is awkward.
    // The only place it is used at present is in JSON where the
    // awkwardness is mixed in with JSON complexity.
    // Consider improving this API in the future.
    ((RepeatedListWriter) writer.array(1)).defineElement(schema.metadata(1).childSchema());
    do2DTest(schema, rsLoader);
    rsLoader.close();
}

Example 87 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader in project drill by apache.

the class TestResultSetLoaderRepeatedList method do2DTest.

private void do2DTest(TupleMetadata schema, ResultSetLoader rsLoader) {
    final RowSetLoader writer = rsLoader.writer();
    // Sanity check of writer structure
    assertEquals(2, writer.size());
    final ObjectWriter listObj = writer.column("list2");
    assertEquals(ObjectType.ARRAY, listObj.type());
    final ArrayWriter listWriter = listObj.array();
    assertEquals(ObjectType.ARRAY, listWriter.entryType());
    final ArrayWriter innerWriter = listWriter.array();
    assertEquals(ObjectType.SCALAR, innerWriter.entryType());
    final ScalarWriter strWriter = innerWriter.scalar();
    assertEquals(ValueType.STRING, strWriter.valueType());
    // Sanity test of schema
    final TupleMetadata rowSchema = writer.tupleSchema();
    assertEquals(2, rowSchema.size());
    final ColumnMetadata listSchema = rowSchema.metadata(1);
    assertEquals(MinorType.LIST, listSchema.type());
    assertEquals(DataMode.REPEATED, listSchema.mode());
    assertTrue(listSchema instanceof RepeatedListColumnMetadata);
    assertEquals(StructureType.MULTI_ARRAY, listSchema.structureType());
    assertNotNull(listSchema.childSchema());
    final ColumnMetadata elementSchema = listSchema.childSchema();
    assertEquals(listSchema.name(), elementSchema.name());
    assertEquals(MinorType.VARCHAR, elementSchema.type());
    assertEquals(DataMode.REPEATED, elementSchema.mode());
    // Write values
    rsLoader.startBatch();
    writer.addRow(1, objArray(strArray("a", "b"), strArray("c", "d"))).addRow(2, objArray(strArray("e"), strArray(), strArray("f", "g", "h"))).addRow(3, objArray()).addRow(4, objArray(strArray(), strArray("i"), strArray()));
    // Verify the values.
    // (Relies on the row set level repeated list tests having passed.)
    final RowSet expected = fixture.rowSetBuilder(schema).addRow(1, objArray(strArray("a", "b"), strArray("c", "d"))).addRow(2, objArray(strArray("e"), strArray(), strArray("f", "g", "h"))).addRow(3, objArray()).addRow(4, objArray(strArray(), strArray("i"), strArray())).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
}

Example 88 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader in project drill by apache.

the class TestResultSetLoaderUnions method testSimpleList.

/**
 * Test for the case of a list defined to contain exactly one type.
 * Relies on the row set tests to verify that the single type model
 * works for lists. Here we test that the ResultSetLoader put the
 * pieces together correctly.
 */
@Test
public void testSimpleList() {
    // Schema with a list declared with one type, not expandable
    final TupleMetadata schema = new SchemaBuilder().add("id", MinorType.INT).addList("list").addType(MinorType.VARCHAR).resumeSchema().buildSchema();
    schema.metadata("list").variantSchema().becomeSimple();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    final RowSetLoader writer = rsLoader.writer();
    // Sanity check: should be an array of Varchar because we said the
    // types within the list is not expandable.
    final ArrayWriter arrWriter = writer.array("list");
    assertEquals(ObjectType.SCALAR, arrWriter.entryType());
    final ScalarWriter strWriter = arrWriter.scalar();
    assertEquals(ValueType.STRING, strWriter.valueType());
    // Can write a batch as if this was a repeated Varchar, except
    // that any value can also be null.
    rsLoader.startBatch();
    writer.addRow(1, strArray("fred", "barney")).addRow(2, null).addRow(3, strArray("wilma", "betty", "pebbles"));
    // Verify
    final SingleRowSet expected = fixture.rowSetBuilder(schema).addRow(1, strArray("fred", "barney")).addRow(2, null).addRow(3, strArray("wilma", "betty", "pebbles")).build();
    RowSetUtilities.verify(expected, fixture.wrap(rsLoader.harvest()));
}

Example 89 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader in project drill by apache.

the class TestResultSetLoaderUnions method testUnionOverflow.

@Test
public void testUnionOverflow() {
    final TupleMetadata schema = new SchemaBuilder().add("id", MinorType.INT).addUnion("u").addType(MinorType.INT).addType(MinorType.VARCHAR).resumeSchema().buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().rowCountLimit(ValueVector.MAX_ROW_COUNT).readerSchema(schema).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    final RowSetLoader writer = rsLoader.writer();
    // Fill the batch with enough data to cause overflow.
    // Fill even rows with a Varchar, odd rows with an int.
    // Data must be large enough to cause overflow before 32K rows
    // (the half that get strings.
    // 16 MB / 32 K = 512 bytes
    // Make a bit bigger to overflow early.
    final int strLength = 600;
    final byte[] value = new byte[strLength - 6];
    Arrays.fill(value, (byte) 'X');
    final String strValue = new String(value, Charsets.UTF_8);
    int count = 0;
    rsLoader.startBatch();
    while (!writer.isFull()) {
        if (count % 2 == 0) {
            writer.addRow(count, String.format("%s%06d", strValue, count));
        } else {
            writer.addRow(count, count * 10);
        }
        count++;
    }
    // Number of rows should be driven by vector size.
    // Our row count should include the overflow row
    final int expectedCount = ValueVector.MAX_BUFFER_SIZE / strLength * 2;
    assertEquals(expectedCount + 1, count);
    // Loader's row count should include only "visible" rows
    assertEquals(expectedCount, writer.rowCount());
    // Total count should include invisible and look-ahead rows.
    assertEquals(expectedCount + 1, rsLoader.totalRowCount());
    // Result should exclude the overflow row
    RowSet result = fixture.wrap(rsLoader.harvest());
    assertEquals(expectedCount, result.rowCount());
    // Verify the data.
    RowSetReader reader = result.reader();
    int readCount = 0;
    while (reader.next()) {
        assertEquals(readCount, reader.scalar(0).getInt());
        if (readCount % 2 == 0) {
            assertEquals(String.format("%s%06d", strValue, readCount), reader.variant(1).scalar().getString());
        } else {
            assertEquals(readCount * 10, reader.variant(1).scalar().getInt());
        }
        readCount++;
    }
    assertEquals(readCount, result.rowCount());
    result.clear();
    // Write a few more rows to verify the overflow row.
    rsLoader.startBatch();
    for (int i = 0; i < 1000; i++) {
        if (count % 2 == 0) {
            writer.addRow(count, String.format("%s%06d", strValue, count));
        } else {
            writer.addRow(count, count * 10);
        }
        count++;
    }
    result = fixture.wrap(rsLoader.harvest());
    assertEquals(1001, result.rowCount());
    final int startCount = readCount;
    reader = result.reader();
    while (reader.next()) {
        assertEquals(readCount, reader.scalar(0).getInt());
        if (readCount % 2 == 0) {
            assertEquals(String.format("%s%06d", strValue, readCount), reader.variant(1).scalar().getString());
        } else {
            assertEquals(readCount * 10, reader.variant(1).scalar().getInt());
        }
        readCount++;
    }
    assertEquals(readCount - startCount, result.rowCount());
    result.clear();
    rsLoader.close();
}

Example 90 with RowSetLoader

use of org.apache.drill.exec.physical.resultSet.RowSetLoader 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);
}