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

Example 56 with ResultSetLoader

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

the class ExcelBatchReader method open.

@Override
public boolean open(FileSchemaNegotiator negotiator) {
    split = negotiator.split();
    errorContext = negotiator.parentErrorContext();
    ResultSetLoader loader = negotiator.build();
    rowWriter = loader.writer();
    openFile(negotiator);
    if (negotiator.hasProvidedSchema()) {
        TupleMetadata providedSchema = negotiator.providedSchema();
        logger.debug("Found inline schema");
        // Add Implicit columns to schema
        SchemaBuilder builderForProvidedSchema = new SchemaBuilder();
        builderForProvidedSchema.addAll(providedSchema);
        TupleMetadata finalSchema = builderForProvidedSchema.build();
        buildColumnWritersFromProvidedSchema(finalSchema);
        // Add schema to file negotiator
        logger.debug("Metadata added to provided schema.");
        addMetadataToSchema(builderForProvidedSchema);
        // Build column writer array
        negotiator.tableSchema(finalSchema, true);
    } else {
        defineSchema(negotiator);
    }
    return true;
}

Example 57 with ResultSetLoader

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

the class TestResultSetLoaderTorture method doTortureTest.

private void doTortureTest() {
    TupleMetadata schema = new SchemaBuilder().add("n0", MinorType.INT).addMap("m1").addNullable("n1", MinorType.INT).addMapArray("m2").addNullable("n2", MinorType.INT).addNullable("s2", MinorType.VARCHAR).addMap("m3").addNullable("n3", MinorType.INT).addArray("s3", MinorType.VARCHAR).resumeMap().resumeMap().resumeSchema().buildSchema();
    ResultSetLoaderImpl.ResultSetOptions options = new ResultSetOptionBuilder().rowCountLimit(ValueVector.MAX_ROW_COUNT).readerSchema(schema).build();
    ResultSetLoader rsLoader = new ResultSetLoaderImpl(fixture.allocator(), options);
    RowSetLoader rootWriter = rsLoader.writer();
    TestSetup setup = new TestSetup();
    BatchWriter batchWriter = new BatchWriter(setup, rootWriter);
    int totalRowCount = 0;
    ReadState readState = new ReadState();
    for (int batchCount = 0; batchCount < 10; batchCount++) {
        rsLoader.startBatch();
        batchWriter.writeBatch();
        // Now the hard part. Verify the above batch.
        RowSet result = fixture.wrap(rsLoader.harvest());
        // result.print();
        // Should have overflowed
        int savedCount = batchWriter.rowCount();
        assertEquals(savedCount, result.rowCount());
        totalRowCount += savedCount;
        assertEquals(totalRowCount, rsLoader.totalRowCount());
        assertEquals(batchCount + 1, rsLoader.batchCount());
        BatchReader reader = new BatchReader(setup, result.reader(), readState);
        reader.verify();
        result.clear();
    }
    // Last row overflow row
    {
        rsLoader.startBatch();
        // Use this to visualize a string buffer. There is also a method
        // to visualize offset vectors. These two are the most pesky vectors
        // to get right.
        // VectorPrinter.printStrings((VarCharVector) ((NullableVarCharVector) ((AbstractScalarWriter) batchWriter.s2Writer).vector()).getValuesVector(), 0, 8);
        RowSet result = fixture.wrap(rsLoader.harvest());
        // Use this here, or earlier, when things go amiss and you need
        // to see what the actual results might be.
        // result.print();
        totalRowCount++;
        assertEquals(totalRowCount, rsLoader.totalRowCount());
        BatchReader reader = new BatchReader(setup, result.reader(), readState);
        reader.verify();
        result.clear();
    }
    rsLoader.close();
}

Example 58 with ResultSetLoader

use of org.apache.drill.exec.physical.resultSet.ResultSetLoader 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 59 with ResultSetLoader

use of org.apache.drill.exec.physical.resultSet.ResultSetLoader 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 60 with ResultSetLoader

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