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

Example 11 with RowSetReader

use of org.apache.drill.exec.physical.rowSet.RowSetReader in project drill by apache.

the class TestResultSetReader method testBasics.

@Test
public void testBasics() {
    PullResultSetReader rsReader = new PullResultSetReaderImpl(new BatchGenerator(10, 2, 1));
    // Start state
    try {
        rsReader.reader();
        fail();
    } catch (IllegalStateException e) {
    // Expected
    }
    // Ask for schema. Does an implicit next.
    assertEquals(SCHEMA1, rsReader.schema());
    assertEquals(1, rsReader.schemaVersion());
    // Move to the first batch.
    // (Don't need to do a full reader test, that is already done
    // elsewhere.)
    assertTrue(rsReader.next());
    assertEquals(1, rsReader.schemaVersion());
    RowSetReader reader1;
    {
        RowSetReader reader = rsReader.reader();
        reader1 = reader;
        assertTrue(reader.next());
        assertEquals(1, reader.scalar("id").getInt());
        assertEquals("Row1", reader.scalar("name").getString());
    }
    // Second batch, same schema.
    assertTrue(rsReader.next());
    assertEquals(1, rsReader.schemaVersion());
    {
        RowSetReader reader = rsReader.reader();
        assertSame(reader1, reader);
        reader1 = reader;
        assertTrue(reader.next());
        assertEquals(11, reader.scalar("id").getInt());
        assertEquals("Row11", reader.scalar("name").getString());
    }
    // Batch with new schema
    assertTrue(rsReader.next());
    assertEquals(2, rsReader.schemaVersion());
    {
        RowSetReader reader = rsReader.reader();
        assertNotSame(reader1, reader);
        reader1 = reader;
        assertTrue(reader.next());
        assertEquals(21, reader.scalar("id").getInt());
        assertEquals("Row21", reader.scalar("name").getString());
        assertEquals(210, reader.scalar("amount").getInt());
    }
    assertFalse(rsReader.next());
    rsReader.close();
}

Example 12 with RowSetReader

use of org.apache.drill.exec.physical.rowSet.RowSetReader in project drill by apache.

the class TestResultSetLoaderMaps method testMapOverflowWithNewColumn.

/**
 * Test the case in which a new column is added during the overflow row. Unlike
 * the top-level schema case, internally we must create a copy of the map, and
 * move vectors across only when the result is to include the schema version
 * of the target column. For overflow, the new column is added after the
 * first batch; it is added in the second batch that contains the overflow
 * row in which the column was added.
 */
@Test
public void testMapOverflowWithNewColumn() {
    final TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).addMap("m").add("b", MinorType.INT).add("c", MinorType.VARCHAR).resumeSchema().buildSchema();
    final ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().readerSchema(schema).rowCountLimit(ValueVector.MAX_ROW_COUNT).build();
    final ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    assertEquals(4, rsLoader.schemaVersion());
    final RowSetLoader rootWriter = rsLoader.writer();
    // Can't use the shortcut to populate rows when doing a schema
    // change.
    final ScalarWriter aWriter = rootWriter.scalar("a");
    final TupleWriter mWriter = rootWriter.tuple("m");
    final ScalarWriter bWriter = mWriter.scalar("b");
    final ScalarWriter cWriter = mWriter.scalar("c");
    final byte[] value = new byte[512];
    Arrays.fill(value, (byte) 'X');
    int count = 0;
    rsLoader.startBatch();
    while (!rootWriter.isFull()) {
        rootWriter.start();
        aWriter.setInt(count);
        bWriter.setInt(count * 10);
        cWriter.setBytes(value, value.length);
        if (rootWriter.isFull()) {
            // Overflow just occurred. Add another column.
            mWriter.addColumn(SchemaBuilder.columnSchema("d", MinorType.INT, DataMode.OPTIONAL));
            mWriter.scalar("d").setInt(count * 100);
        }
        rootWriter.save();
        count++;
    }
    // Result set should include the original columns, but not d.
    RowSet result = fixture.wrap(rsLoader.harvest());
    assertEquals(4, rsLoader.schemaVersion());
    assertTrue(schema.isEquivalent(result.schema()));
    final BatchSchema expectedSchema = new BatchSchema(SelectionVectorMode.NONE, schema.toFieldList());
    assertTrue(expectedSchema.isEquivalent(result.batchSchema()));
    // Use a reader to validate row-by-row. Too large to create an expected
    // result set.
    RowSetReader reader = result.reader();
    TupleReader mapReader = reader.tuple("m");
    int rowId = 0;
    while (reader.next()) {
        assertEquals(rowId, reader.scalar("a").getInt());
        assertEquals(rowId * 10, mapReader.scalar("b").getInt());
        assertTrue(Arrays.equals(value, mapReader.scalar("c").getBytes()));
        rowId++;
    }
    result.clear();
    // Next batch should start with the overflow row
    rsLoader.startBatch();
    assertEquals(1, rootWriter.rowCount());
    result = fixture.wrap(rsLoader.harvest());
    assertEquals(1, result.rowCount());
    reader = result.reader();
    mapReader = reader.tuple("m");
    while (reader.next()) {
        assertEquals(rowId, reader.scalar("a").getInt());
        assertEquals(rowId * 10, mapReader.scalar("b").getInt());
        assertTrue(Arrays.equals(value, mapReader.scalar("c").getBytes()));
        assertEquals(rowId * 100, mapReader.scalar("d").getInt());
    }
    result.clear();
    rsLoader.close();
}

Example 13 with RowSetReader

use of org.apache.drill.exec.physical.rowSet.RowSetReader in project drill by apache.

the class TestResultSetLoaderOverflow method testMissingArrayValues.

/**
 * Test the case that an array has "missing values" before the overflow.
 */
@Test
public void testMissingArrayValues() {
    TupleMetadata schema = new SchemaBuilder().add("a", MinorType.INT).add("b", MinorType.VARCHAR).addArray("c", MinorType.INT).buildSchema();
    ResultSetOptions options = new ResultSetOptionBuilder().rowCountLimit(ValueVector.MAX_ROW_COUNT).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    byte[] value = new byte[512];
    Arrays.fill(value, (byte) 'X');
    int blankAfter = ValueVector.MAX_BUFFER_SIZE / 512 * 2 / 3;
    ScalarWriter cWriter = rootWriter.array("c").scalar();
    rsLoader.startBatch();
    int rowId = 0;
    while (rootWriter.start()) {
        rootWriter.scalar("a").setInt(rowId);
        rootWriter.scalar("b").setBytes(value, value.length);
        if (rowId < blankAfter) {
            for (int i = 0; i < 3; i++) {
                cWriter.setInt(rowId * 3 + i);
            }
        }
        rootWriter.save();
        rowId++;
    }
    VectorContainer container = rsLoader.harvest();
    BatchValidator.validate(container);
    RowSet result = fixture.wrap(container);
    assertEquals(rowId - 1, result.rowCount());
    RowSetReader reader = result.reader();
    ArrayReader cArray = reader.array("c");
    ScalarReader cReader = cArray.scalar();
    while (reader.next()) {
        assertEquals(reader.offset(), reader.scalar("a").getInt());
        assertArrayEquals(value, reader.scalar("b").getBytes());
        if (reader.offset() < blankAfter) {
            assertEquals(3, cArray.size());
            for (int i = 0; i < 3; i++) {
                assertTrue(cArray.next());
                assertEquals(reader.offset() * 3 + i, cReader.getInt());
            }
        } else {
            assertEquals(0, cArray.size());
        }
    }
    result.clear();
    rsLoader.close();
}

Example 14 with RowSetReader

use of org.apache.drill.exec.physical.rowSet.RowSetReader in project drill by apache.

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 ResultSetOptionBuilder().rowCountLimit(ValueVector.MAX_ROW_COUNT).readerSchema(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
    {
        VectorContainer container = rsLoader.harvest();
        BatchValidator.validate(container);
        RowSet result = fixture.wrap(container);
        assertEquals(count - 1, result.rowCount());
        RowSetReader reader = result.reader();
        while (reader.next()) {
            assertEquals(reader.offset(), reader.scalar(0).getInt());
            assertTrue(reader.scalar(1).isNull());
            assertArrayEquals(value, reader.scalar(2).getBytes());
            assertTrue(reader.scalar(3).isNull());
        }
        result.clear();
    }
    // Next batch should start with the overflow row
    rsLoader.startBatch();
    {
        VectorContainer container = rsLoader.harvest();
        BatchValidator.validate(container);
        RowSet result = fixture.wrap(container);
        RowSetReader reader = result.reader();
        assertEquals(1, result.rowCount());
        assertTrue(reader.next());
        assertEquals(count - 1, reader.scalar(0).getInt());
        assertTrue(reader.scalar(1).isNull());
        assertArrayEquals(value, reader.scalar(2).getBytes());
        assertTrue(reader.scalar(3).isNull());
        result.clear();
    }
    rsLoader.close();
}

Example 15 with RowSetReader

use of org.apache.drill.exec.physical.rowSet.RowSetReader in project drill by apache.

the class TestResultSetLoaderOverflow method testSizeLimitOnArray.

/**
 * Test a row with a single array column which overflows. Verifies
 * that all the fiddly bits about offset vectors and so on works
 * correctly. Run this test (the simplest case) if you change anything
 * about the array handling code.
 */
@Test
public void testSizeLimitOnArray() {
    TupleMetadata schema = new SchemaBuilder().addArray("s", MinorType.VARCHAR).buildSchema();
    ResultSetOptions options = new ResultSetOptionBuilder().rowCountLimit(ValueVector.MAX_ROW_COUNT).readerSchema(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.
    rsLoader.startBatch();
    byte[] value = new byte[473];
    Arrays.fill(value, (byte) 'X');
    String strValue = new String(value, Charsets.UTF_8);
    int valuesPerArray = 13;
    int count = 0;
    {
        int rowSize = 0;
        int totalSize = 0;
        while (rootWriter.start()) {
            totalSize += rowSize;
            rowSize = 0;
            ScalarWriter array = rootWriter.array(0).scalar();
            for (int i = 0; i < valuesPerArray; i++) {
                String cellValue = strValue + (count + 1) + "." + i;
                array.setString(cellValue);
                rowSize += cellValue.length();
            }
            rootWriter.save();
            count++;
        }
        // Row count should include the overflow row.
        int expectedCount = count - 1;
        // Size without overflow row should fit in the vector, size
        // with overflow should not.
        assertTrue(totalSize <= ValueVector.MAX_BUFFER_SIZE);
        assertTrue(totalSize + rowSize > ValueVector.MAX_BUFFER_SIZE);
        // Result should exclude the overflow row. Last row
        // should hold the last full array.
        VectorContainer container = rsLoader.harvest();
        BatchValidator.validate(container);
        RowSet result = fixture.wrap(container);
        assertEquals(expectedCount, result.rowCount());
        RowSetReader reader = result.reader();
        reader.setPosition(expectedCount - 1);
        ArrayReader arrayReader = reader.array(0);
        ScalarReader strReader = arrayReader.scalar();
        assertEquals(valuesPerArray, arrayReader.size());
        for (int i = 0; i < valuesPerArray; i++) {
            assertTrue(arrayReader.next());
            String cellValue = strValue + (count - 1) + "." + i;
            assertEquals(cellValue, strReader.getString());
        }
        result.clear();
    }
    // Next batch should start with the overflow row.
    // The only row in this next batch should be the whole
    // array being written at the time of overflow.
    {
        rsLoader.startBatch();
        assertEquals(1, rootWriter.rowCount());
        assertEquals(count, rsLoader.totalRowCount());
        VectorContainer container = rsLoader.harvest();
        BatchValidator.validate(container);
        RowSet result = fixture.wrap(container);
        assertEquals(1, result.rowCount());
        RowSetReader reader = result.reader();
        reader.next();
        ArrayReader arrayReader = reader.array(0);
        ScalarReader strReader = arrayReader.scalar();
        assertEquals(valuesPerArray, arrayReader.size());
        for (int i = 0; i < valuesPerArray; i++) {
            assertTrue(arrayReader.next());
            String cellValue = strValue + count + "." + i;
            assertEquals(cellValue, strReader.getString());
        }
        result.clear();
    }
    rsLoader.close();
}