Examples with RowSetReader org.apache.drill.test.rowSet.RowSetReader used on opensource projects

Example 26 with RowSetReader

use of org.apache.drill.test.rowSet.RowSetReader in project drill by axbaretto.

the class TestResultSetLoaderOverflow method testArrayOverflowWithOtherArrays.

/**
 * Test the complete set of array overflow cases:
 * <ul>
 * <li>Array a is written before the column that has overflow,
 * and must be copied, in its entirety, to the overflow row.</li>
 * <li>Column b causes the overflow.</li>
 * <li>Column c is written after the overflow, and should go
 * to the look-ahead row.</li>
 * <li>Column d is written for a while, then has empties before
 * the overflow row, but is written in the overflow row.<li>
 * <li>Column e is like d, but is not written in the overflow
 * row.</li>
 */
@Test
public void testArrayOverflowWithOtherArrays() {
    TupleMetadata schema = new SchemaBuilder().addArray("a", MinorType.INT).addArray("b", MinorType.VARCHAR).addArray("c", MinorType.INT).addArray("d", MinorType.INT).buildSchema();
    ResultSetOptions options = new OptionBuilder().setRowCountLimit(ValueVector.MAX_ROW_COUNT).setSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    // Fill batch with rows of with a single array, three values each. Tack on
    // a suffix to each so we can be sure the proper data is written and moved
    // to the overflow batch.
    byte[] value = new byte[512];
    Arrays.fill(value, (byte) 'X');
    String strValue = new String(value, Charsets.UTF_8);
    int aCount = 3;
    int bCount = 11;
    int cCount = 5;
    int dCount = 7;
    int cCutoff = ValueVector.MAX_BUFFER_SIZE / value.length / bCount / 2;
    ScalarWriter aWriter = rootWriter.array("a").scalar();
    ScalarWriter bWriter = rootWriter.array("b").scalar();
    ScalarWriter cWriter = rootWriter.array("c").scalar();
    ScalarWriter dWriter = rootWriter.array("d").scalar();
    int count = 0;
    rsLoader.startBatch();
    while (rootWriter.start()) {
        if (rootWriter.rowCount() == 2952) {
            count = count + 0;
        }
        for (int i = 0; i < aCount; i++) {
            aWriter.setInt(count * aCount + i);
        }
        for (int i = 0; i < bCount; i++) {
            String cellValue = strValue + (count * bCount + i);
            bWriter.setString(cellValue);
        }
        if (count < cCutoff) {
            for (int i = 0; i < cCount; i++) {
                cWriter.setInt(count * cCount + i);
            }
        }
        if (count < cCutoff || rootWriter.isFull()) {
            for (int i = 0; i < dCount; i++) {
                dWriter.setInt(count * dCount + i);
            }
        }
        rootWriter.save();
        count++;
    }
    // Verify
    RowSet result = fixture.wrap(rsLoader.harvest());
    assertEquals(count - 1, result.rowCount());
    RowSetReader reader = result.reader();
    ScalarElementReader aReader = reader.array("a").elements();
    ScalarElementReader bReader = reader.array("b").elements();
    ScalarElementReader cReader = reader.array("c").elements();
    ScalarElementReader dReader = reader.array("d").elements();
    while (reader.next()) {
        int rowId = reader.rowIndex();
        assertEquals(aCount, aReader.size());
        for (int i = 0; i < aCount; i++) {
            assertEquals(rowId * aCount + i, aReader.getInt(i));
        }
        assertEquals(bCount, bReader.size());
        for (int i = 0; i < bCount; i++) {
            String cellValue = strValue + (rowId * bCount + i);
            assertEquals(cellValue, bReader.getString(i));
        }
        if (rowId < cCutoff) {
            assertEquals(cCount, cReader.size());
            for (int i = 0; i < cCount; i++) {
                assertEquals(rowId * cCount + i, cReader.getInt(i));
            }
            assertEquals(dCount, dReader.size());
            for (int i = 0; i < dCount; i++) {
                assertEquals(rowId * dCount + i, dReader.getInt(i));
            }
        } else {
            assertEquals(0, cReader.size());
            assertEquals(0, dReader.size());
        }
    }
    result.clear();
    int firstCount = count - 1;
    // One row is in the batch. Write more, skipping over the
    // initial few values for columns c and d. Column d has a
    // roll-over value, c has an empty roll-over.
    rsLoader.startBatch();
    for (int j = 0; j < 5; j++) {
        rootWriter.start();
        for (int i = 0; i < aCount; i++) {
            aWriter.setInt(count * aCount + i);
        }
        for (int i = 0; i < bCount; i++) {
            String cellValue = strValue + (count * bCount + i);
            bWriter.setString(cellValue);
        }
        if (j > 3) {
            for (int i = 0; i < cCount; i++) {
                cWriter.setInt(count * cCount + i);
            }
            for (int i = 0; i < dCount; i++) {
                dWriter.setInt(count * dCount + i);
            }
        }
        rootWriter.save();
        count++;
    }
    result = fixture.wrap(rsLoader.harvest());
    assertEquals(6, result.rowCount());
    reader = result.reader();
    aReader = reader.array("a").elements();
    bReader = reader.array("b").elements();
    cReader = reader.array("c").elements();
    dReader = reader.array("d").elements();
    int j = 0;
    while (reader.next()) {
        int rowId = firstCount + reader.rowIndex();
        assertEquals(aCount, aReader.size());
        for (int i = 0; i < aCount; i++) {
            assertEquals("Index " + i, rowId * aCount + i, aReader.getInt(i));
        }
        assertEquals(bCount, bReader.size());
        for (int i = 0; i < bCount; i++) {
            String cellValue = strValue + (rowId * bCount + i);
            assertEquals(cellValue, bReader.getString(i));
        }
        if (j > 4) {
            assertEquals(cCount, cReader.size());
            for (int i = 0; i < cCount; i++) {
                assertEquals(rowId * cCount + i, cReader.getInt(i));
            }
        } else {
            assertEquals(0, cReader.size());
        }
        if (j == 0 || j > 4) {
            assertEquals(dCount, dReader.size());
            for (int i = 0; i < dCount; i++) {
                assertEquals(rowId * dCount + i, dReader.getInt(i));
            }
        } else {
            assertEquals(0, dReader.size());
        }
        j++;
    }
    result.clear();
    rsLoader.close();
}

Example 27 with RowSetReader

use of org.apache.drill.test.rowSet.RowSetReader in project drill by axbaretto.

the class TestResultSetLoaderOverflow method testOverflowWithNullables.

@Test
public void testOverflowWithNullables() {
    TupleMetadata schema = new SchemaBuilder().add("n", MinorType.INT).addNullable("a", MinorType.VARCHAR).addNullable("b", MinorType.VARCHAR).addNullable("c", MinorType.VARCHAR).buildSchema();
    ResultSetOptions options = new OptionBuilder().setRowCountLimit(ValueVector.MAX_ROW_COUNT).setSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    rsLoader.startBatch();
    byte[] value = new byte[512];
    Arrays.fill(value, (byte) 'X');
    int count = 0;
    while (!rootWriter.isFull()) {
        rootWriter.start();
        rootWriter.scalar(0).setInt(count);
        rootWriter.scalar(1).setNull();
        rootWriter.scalar(2).setBytes(value, value.length);
        rootWriter.scalar(3).setNull();
        rootWriter.save();
        count++;
    }
    // Result should exclude the overflow row
    RowSet result = fixture.wrap(rsLoader.harvest());
    assertEquals(count - 1, result.rowCount());
    RowSetReader reader = result.reader();
    while (reader.next()) {
        assertEquals(reader.rowIndex(), reader.scalar(0).getInt());
        assertTrue(reader.scalar(1).isNull());
        assertTrue(Arrays.equals(value, reader.scalar(2).getBytes()));
        assertTrue(reader.scalar(3).isNull());
    }
    result.clear();
    // Next batch should start with the overflow row
    rsLoader.startBatch();
    result = fixture.wrap(rsLoader.harvest());
    reader = result.reader();
    assertEquals(1, result.rowCount());
    assertTrue(reader.next());
    assertEquals(count - 1, reader.scalar(0).getInt());
    assertTrue(reader.scalar(1).isNull());
    assertTrue(Arrays.equals(value, reader.scalar(2).getBytes()));
    assertTrue(reader.scalar(3).isNull());
    result.clear();
    rsLoader.close();
}

Example 28 with RowSetReader

use of org.apache.drill.test.rowSet.RowSetReader in project drill by axbaretto.

the class TestResultSetSchemaChange method testSchemaChangeWithOverflow.

/**
 * Test a schema change on the row that overflows. If the
 * new column is added after overflow, it will appear as
 * a schema-change in the following batch. This is fine as
 * we are essentially time-shifting: pretending that the
 * overflow row was written in the next batch (which, in
 * fact, it is: that's what overflow means.)
 */
@Test
public void testSchemaChangeWithOverflow() {
    ResultSetOptions options = new OptionBuilder().setRowCountLimit(ValueVector.MAX_ROW_COUNT).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    rootWriter.addColumn(SchemaBuilder.columnSchema("a", MinorType.VARCHAR, DataMode.REQUIRED));
    rsLoader.startBatch();
    byte[] value = new byte[512];
    Arrays.fill(value, (byte) 'X');
    int count = 0;
    while (!rootWriter.isFull()) {
        rootWriter.start();
        rootWriter.scalar(0).setBytes(value, value.length);
        if (rootWriter.isFull()) {
            rootWriter.addColumn(SchemaBuilder.columnSchema("b", MinorType.INT, DataMode.OPTIONAL));
            rootWriter.scalar(1).setInt(count);
            // Add a Varchar to ensure its offset fiddling is done properly
            rootWriter.addColumn(SchemaBuilder.columnSchema("c", MinorType.VARCHAR, DataMode.OPTIONAL));
            rootWriter.scalar(2).setString("c-" + count);
            // Allow adding a required column at this point.
            // (Not intuitively obvious that this should work; we back-fill
            // with zeros.)
            rootWriter.addColumn(SchemaBuilder.columnSchema("d", MinorType.INT, DataMode.REQUIRED));
        }
        rootWriter.save();
        count++;
    }
    // Result should include only the first column.
    BatchSchema expectedSchema = new SchemaBuilder().add("a", MinorType.VARCHAR).build();
    RowSet result = fixture.wrap(rsLoader.harvest());
    assertTrue(result.batchSchema().isEquivalent(expectedSchema));
    assertEquals(count - 1, result.rowCount());
    result.clear();
    assertEquals(1, rsLoader.schemaVersion());
    // Double check: still can add a required column after
    // starting the next batch. (No longer in overflow state.)
    rsLoader.startBatch();
    rootWriter.addColumn(SchemaBuilder.columnSchema("e", MinorType.INT, DataMode.REQUIRED));
    // Next batch should start with the overflow row, including
    // the column added at the end of the previous batch, after
    // overflow.
    result = fixture.wrap(rsLoader.harvest());
    assertEquals(5, rsLoader.schemaVersion());
    assertEquals(1, result.rowCount());
    expectedSchema = new SchemaBuilder(expectedSchema).addNullable("b", MinorType.INT).addNullable("c", MinorType.VARCHAR).add("d", MinorType.INT).add("e", MinorType.INT).build();
    assertTrue(result.batchSchema().isEquivalent(expectedSchema));
    RowSetReader reader = result.reader();
    reader.next();
    assertEquals(count - 1, reader.scalar(1).getInt());
    assertEquals("c-" + (count - 1), reader.scalar(2).getString());
    assertEquals(0, reader.scalar("d").getInt());
    assertEquals(0, reader.scalar("e").getInt());
    result.clear();
    rsLoader.close();
}

Example 29 with RowSetReader

use of org.apache.drill.test.rowSet.RowSetReader in project drill by axbaretto.

the class TestResultSetLoaderMapArray method testOmittedValues.

/**
 * Check that the "fill-empties" logic descends down into
 * a repeated map.
 */
@Test
public void testOmittedValues() {
    TupleMetadata schema = new SchemaBuilder().add("id", MinorType.INT).addMapArray("m").addNullable("a", MinorType.INT).addNullable("b", MinorType.VARCHAR).resumeSchema().buildSchema();
    ResultSetLoaderImpl.ResultSetOptions options = new OptionBuilder().setSchema(schema).setRowCountLimit(ValueVector.MAX_ROW_COUNT).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    int mapSkip = 5;
    int entrySkip = 3;
    int rowCount = 1000;
    int entryCount = 10;
    rsLoader.startBatch();
    ArrayWriter maWriter = rootWriter.array("m");
    TupleWriter mWriter = maWriter.tuple();
    for (int i = 0; i < rowCount; i++) {
        rootWriter.start();
        rootWriter.scalar(0).setInt(i);
        if (i % mapSkip != 0) {
            for (int j = 0; j < entryCount; j++) {
                if (j % entrySkip != 0) {
                    mWriter.scalar(0).setInt(i * entryCount + j);
                    mWriter.scalar(1).setString("b-" + i + "." + j);
                }
                maWriter.save();
            }
        }
        rootWriter.save();
    }
    RowSet result = fixture.wrap(rsLoader.harvest());
    assertEquals(rowCount, result.rowCount());
    RowSetReader reader = result.reader();
    ArrayReader maReader = reader.array("m");
    TupleReader mReader = maReader.tuple();
    for (int i = 0; i < rowCount; i++) {
        assertTrue(reader.next());
        assertEquals(i, reader.scalar(0).getInt());
        if (i % mapSkip == 0) {
            assertEquals(0, maReader.size());
            continue;
        }
        assertEquals(entryCount, maReader.size());
        for (int j = 0; j < entryCount; j++) {
            maReader.setPosn(j);
            if (j % entrySkip == 0) {
                assertTrue(mReader.scalar(0).isNull());
                assertTrue(mReader.scalar(1).isNull());
            } else {
                assertFalse(mReader.scalar(0).isNull());
                assertFalse(mReader.scalar(1).isNull());
                assertEquals(i * entryCount + j, mReader.scalar(0).getInt());
                assertEquals("b-" + i + "." + j, mReader.scalar(1).getString());
            }
        }
    }
    result.clear();
    rsLoader.close();
}

Example 30 with RowSetReader

use of org.apache.drill.test.rowSet.RowSetReader in project drill by axbaretto.

the class TestResultSetLoaderMapArray method testDoubleNestedArray.

/**
 * Test a doubly-nested arrays 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();
    ResultSetLoaderImpl.ResultSetOptions options = new OptionBuilder().setSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    rsLoader.startBatch();
    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++) {
        rootWriter.start();
        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 = fixture.wrap(rsLoader.harvest());
    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");
    ScalarElementReader dReader = m2Reader.elements("d");
    for (int i = 0; i < 5; i++) {
        reader.next();
        assertEquals(i, aReader.getInt());
        for (int j = 0; j < 4; j++) {
            a1Reader.setPosn(j);
            int a1Key = i + 10 + j;
            assertEquals(a1Key, bReader.getInt());
            for (int k = 0; k < 3; k++) {
                a2Reader.setPosn(k);
                int a2Key = a1Key * 10 + k;
                assertEquals(a2Key, cReader.getInt());
                for (int l = 0; l < 2; l++) {
                    assertEquals("d-" + (a2Key * 10 + l), dReader.getString(l));
                }
            }
        }
    }
    rsLoader.close();
}