Search in sources :

Example 6 with SchemaNegotiator

use of org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator in project drill by apache.

the class TestMockRowReader method testColumnRepeat.

/**
 * Test a repeated column.
 */
@Test
public void testColumnRepeat() {
    int rowCount = 10;
    MockTableDef.MockColumn[] cols = new MockTableDef.MockColumn[] { new MockTableDef.MockColumn("a", MinorType.INT, DataMode.REQUIRED, null, null, null, null, 3, null), new MockTableDef.MockColumn("b", MinorType.VARCHAR, DataMode.REQUIRED, 10, null, null, null, null, null) };
    MockTableDef.MockScanEntry entry = new MockTableDef.MockScanEntry(rowCount, true, null, null, cols);
    MockSubScanPOP config = new MockSubScanPOP("dummy", true, Collections.singletonList(entry));
    ManagedReader<SchemaNegotiator> reader = new ExtendedMockBatchReader(entry);
    List<ManagedReader<SchemaNegotiator>> readers = Collections.singletonList(reader);
    // Create options and the scan operator
    ScanFixture mockBatch = buildScan(config, readers);
    ScanOperatorExec scan = mockBatch.scanOp;
    // First batch: build schema. The reader helps: it returns an
    // empty first batch.
    assertTrue(scan.buildSchema());
    TupleMetadata expectedSchema = new SchemaBuilder().add("a1", MinorType.INT).add("a2", MinorType.INT).add("a3", MinorType.INT).add("b", MinorType.VARCHAR, 10).build();
    BatchSchema expectedBatchSchema = new BatchSchema(SelectionVectorMode.NONE, expectedSchema.toFieldList());
    assertTrue(expectedBatchSchema.isEquivalent(scan.batchAccessor().schema()));
    assertEquals(0, scan.batchAccessor().rowCount());
    // Next call, return with data.
    assertTrue(scan.next());
    assertTrue(expectedBatchSchema.isEquivalent(scan.batchAccessor().schema()));
    assertEquals(rowCount, scan.batchAccessor().rowCount());
    scan.batchAccessor().release();
    // EOF
    assertFalse(scan.next());
    mockBatch.close();
}
Also used : ScanFixture(org.apache.drill.exec.physical.impl.scan.ScanTestUtils.ScanFixture) ManagedReader(org.apache.drill.exec.physical.impl.scan.framework.ManagedReader) BatchSchema(org.apache.drill.exec.record.BatchSchema) ScanOperatorExec(org.apache.drill.exec.physical.impl.scan.ScanOperatorExec) TupleMetadata(org.apache.drill.exec.record.metadata.TupleMetadata) SchemaBuilder(org.apache.drill.exec.record.metadata.SchemaBuilder) SchemaNegotiator(org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator) SubOperatorTest(org.apache.drill.test.SubOperatorTest) Test(org.junit.Test) UnlikelyTest(org.apache.drill.categories.UnlikelyTest)

Example 7 with SchemaNegotiator

use of org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator in project drill by apache.

the class TestMockRowReader method testBatchSize.

/**
 * Verify limit on individual batch size (limiting row count per batch).
 */
@Test
public void testBatchSize() {
    int rowCount = 20;
    int batchSize = 10;
    MockTableDef.MockColumn[] cols = new MockTableDef.MockColumn[] { new MockTableDef.MockColumn("a", MinorType.INT, DataMode.REQUIRED, null, null, null, null, null, null), new MockTableDef.MockColumn("b", MinorType.VARCHAR, DataMode.REQUIRED, 10, null, null, null, null, null) };
    MockTableDef.MockScanEntry entry = new MockTableDef.MockScanEntry(rowCount, true, batchSize, null, cols);
    MockSubScanPOP config = new MockSubScanPOP("dummy", true, Collections.singletonList(entry));
    ManagedReader<SchemaNegotiator> reader = new ExtendedMockBatchReader(entry);
    List<ManagedReader<SchemaNegotiator>> readers = Collections.singletonList(reader);
    // Create options and the scan operator
    ScanFixture mockBatch = buildScan(config, readers);
    ScanOperatorExec scan = mockBatch.scanOp;
    // First batch: build schema. The reader helps: it returns an
    // empty first batch.
    assertTrue(scan.buildSchema());
    assertEquals(0, scan.batchAccessor().rowCount());
    // Next call, return with data, limited by batch size.
    assertTrue(scan.next());
    assertEquals(batchSize, scan.batchAccessor().rowCount());
    scan.batchAccessor().release();
    assertTrue(scan.next());
    assertEquals(batchSize, scan.batchAccessor().rowCount());
    scan.batchAccessor().release();
    // EOF
    assertFalse(scan.next());
    mockBatch.close();
}
Also used : ScanFixture(org.apache.drill.exec.physical.impl.scan.ScanTestUtils.ScanFixture) ManagedReader(org.apache.drill.exec.physical.impl.scan.framework.ManagedReader) ScanOperatorExec(org.apache.drill.exec.physical.impl.scan.ScanOperatorExec) SchemaNegotiator(org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator) SubOperatorTest(org.apache.drill.test.SubOperatorTest) Test(org.junit.Test) UnlikelyTest(org.apache.drill.categories.UnlikelyTest)

Example 8 with SchemaNegotiator

use of org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator in project drill by apache.

the class AvroMessageReader method init.

@Override
public void init(SchemaNegotiator negotiator, ReadOptions readOptions, KafkaStoragePlugin plugin) {
    Properties kafkaConsumerProps = plugin.getConfig().getKafkaConsumerProps();
    Map<String, Object> propertiesMap = kafkaConsumerProps.entrySet().stream().collect(Collectors.toMap(e -> e.getKey().toString(), Map.Entry::getValue));
    deserializer = new KafkaAvroDeserializer(null, propertiesMap);
    TupleMetadata providedSchema = negotiator.providedSchema();
    loader = negotiator.build();
    AvroColumnConverterFactory factory = new AvroColumnConverterFactory(providedSchema);
    converter = factory.getRootConverter(providedSchema, new TupleSchema(), loader.writer());
    String keyDeserializer = kafkaConsumerProps.getProperty(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG);
    deserializeKey = keyDeserializer != null && keyDeserializer.equals(KafkaAvroDeserializer.class.getName());
}
Also used : ColumnConverter(org.apache.drill.exec.record.ColumnConverter) SchemaNegotiator(org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator) MetadataUtils(org.apache.drill.exec.record.metadata.MetadataUtils) UserException(org.apache.drill.common.exceptions.UserException) LoggerFactory(org.slf4j.LoggerFactory) RowSetLoader(org.apache.drill.exec.physical.resultSet.RowSetLoader) ReadOptions(org.apache.drill.exec.store.kafka.ReadOptions) Map(java.util.Map) ColumnMetadata(org.apache.drill.exec.record.metadata.ColumnMetadata) TupleSchema(org.apache.drill.exec.record.metadata.TupleSchema) ByteArrayDeserializer(org.apache.kafka.common.serialization.ByteArrayDeserializer) GenericRecord(org.apache.avro.generic.GenericRecord) AvroColumnConverterFactory(org.apache.drill.exec.store.avro.AvroColumnConverterFactory) Schema(org.apache.avro.Schema) Logger(org.slf4j.Logger) Properties(java.util.Properties) ResultSetLoader(org.apache.drill.exec.physical.resultSet.ResultSetLoader) ConsumerConfig(org.apache.kafka.clients.consumer.ConsumerConfig) IOException(java.io.IOException) TupleMetadata(org.apache.drill.exec.record.metadata.TupleMetadata) Collectors(java.util.stream.Collectors) KafkaAvroDeserializer(io.confluent.kafka.serializers.KafkaAvroDeserializer) KafkaStoragePlugin(org.apache.drill.exec.store.kafka.KafkaStoragePlugin) TypeProtos(org.apache.drill.common.types.TypeProtos) MetaDataField(org.apache.drill.exec.store.kafka.MetaDataField) ConsumerRecord(org.apache.kafka.clients.consumer.ConsumerRecord) KafkaConsumer(org.apache.kafka.clients.consumer.KafkaConsumer) AvroColumnConverterFactory(org.apache.drill.exec.store.avro.AvroColumnConverterFactory) TupleMetadata(org.apache.drill.exec.record.metadata.TupleMetadata) KafkaAvroDeserializer(io.confluent.kafka.serializers.KafkaAvroDeserializer) Properties(java.util.Properties) Map(java.util.Map) TupleSchema(org.apache.drill.exec.record.metadata.TupleSchema)

Aggregations

SchemaNegotiator (org.apache.drill.exec.physical.impl.scan.framework.SchemaNegotiator)8 ScanFixture (org.apache.drill.exec.physical.impl.scan.ScanTestUtils.ScanFixture)7 Test (org.junit.Test)7 UnlikelyTest (org.apache.drill.categories.UnlikelyTest)5 ScanOperatorExec (org.apache.drill.exec.physical.impl.scan.ScanOperatorExec)5 ManagedReader (org.apache.drill.exec.physical.impl.scan.framework.ManagedReader)5 SubOperatorTest (org.apache.drill.test.SubOperatorTest)5 TupleMetadata (org.apache.drill.exec.record.metadata.TupleMetadata)4 UserException (org.apache.drill.common.exceptions.UserException)3 BatchSchema (org.apache.drill.exec.record.BatchSchema)3 SchemaBuilder (org.apache.drill.exec.record.metadata.SchemaBuilder)3 KafkaAvroDeserializer (io.confluent.kafka.serializers.KafkaAvroDeserializer)1 IOException (java.io.IOException)1 HashMap (java.util.HashMap)1 Map (java.util.Map)1 Properties (java.util.Properties)1 Collectors (java.util.stream.Collectors)1 Schema (org.apache.avro.Schema)1 GenericRecord (org.apache.avro.generic.GenericRecord)1 TypeProtos (org.apache.drill.common.types.TypeProtos)1