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

Example 36 with RowSetLoader

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

the class JsonMessageReader method parseAndWrite.

private void parseAndWrite(ConsumerRecord<?, ?> record, byte[] recordArray) {
    stream.setValue(new ByteArrayInputStream(recordArray));
    if (kafkaJsonLoader == null) {
        JsonLoaderOptions jsonLoaderOptions = new JsonLoaderOptions();
        jsonLoaderOptions.allTextMode = readOptions.isAllTextMode();
        jsonLoaderOptions.readNumbersAsDouble = readOptions.isReadNumbersAsDouble();
        jsonLoaderOptions.skipMalformedRecords = readOptions.isSkipInvalidRecords();
        jsonLoaderOptions.allowNanInf = readOptions.isAllowNanInf();
        jsonLoaderOptions.enableEscapeAnyChar = readOptions.isAllowEscapeAnyChar();
        jsonLoaderOptions.skipMalformedDocument = readOptions.isSkipInvalidRecords();
        kafkaJsonLoader = (KafkaJsonLoader) new KafkaJsonLoader.KafkaJsonLoaderBuilder().resultSetLoader(resultSetLoader).standardOptions(negotiator.queryOptions()).options(jsonLoaderOptions).errorContext(negotiator.parentErrorContext()).fromStream(() -> stream).build();
    }
    RowSetLoader rowWriter = resultSetLoader.writer();
    rowWriter.start();
    if (kafkaJsonLoader.parser().next()) {
        writeValue(rowWriter, MetaDataField.KAFKA_TOPIC, record.topic());
        writeValue(rowWriter, MetaDataField.KAFKA_PARTITION_ID, record.partition());
        writeValue(rowWriter, MetaDataField.KAFKA_OFFSET, record.offset());
        writeValue(rowWriter, MetaDataField.KAFKA_TIMESTAMP, record.timestamp());
        writeValue(rowWriter, MetaDataField.KAFKA_MSG_KEY, record.key() != null ? record.key().toString() : null);
        rowWriter.save();
    }
}

Example 37 with RowSetLoader

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

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

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

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