Search in sources :

Example 1 with FileIO

use of org.apache.beam.sdk.io.FileIO in project beam by apache.

the class ParquetIOIT method writeThenReadAll.

@Test
public void writeThenReadAll() {
    PCollection<String> testFiles = pipeline.apply("Generate sequence", GenerateSequence.from(0).to(numberOfTextLines)).apply("Produce text lines", ParDo.of(new FileBasedIOITHelper.DeterministicallyConstructTestTextLineFn())).apply("Produce Avro records", ParDo.of(new DeterministicallyConstructAvroRecordsFn())).setCoder(AvroCoder.of(SCHEMA)).apply("Gather write start times", ParDo.of(new TimeMonitor<>(PARQUET_NAMESPACE, "writeStart"))).apply("Write Parquet files", FileIO.<GenericRecord>write().via(ParquetIO.sink(SCHEMA)).to(filenamePrefix)).getPerDestinationOutputFilenames().apply("Gather write end times", ParDo.of(new TimeMonitor<>(PARQUET_NAMESPACE, "writeEnd"))).apply("Get file names", Values.create());
    PCollection<String> consolidatedHashcode = testFiles.apply("Find files", FileIO.matchAll()).apply("Read matched files", FileIO.readMatches()).apply("Gather read start time", ParDo.of(new TimeMonitor<>(PARQUET_NAMESPACE, "readStart"))).apply("Read parquet files", ParquetIO.readFiles(SCHEMA)).apply("Gather read end time", ParDo.of(new TimeMonitor<>(PARQUET_NAMESPACE, "readEnd"))).apply("Map records to strings", MapElements.into(strings()).via((SerializableFunction<GenericRecord, String>) record -> String.valueOf(record.get("row")))).apply("Calculate hashcode", Combine.globally(new HashingFn()));
    PAssert.thatSingleton(consolidatedHashcode).isEqualTo(expectedHash);
    testFiles.apply("Delete test files", ParDo.of(new FileBasedIOITHelper.DeleteFileFn()).withSideInputs(consolidatedHashcode.apply(View.asSingleton())));
    PipelineResult result = pipeline.run();
    result.waitUntilFinish();
    collectAndPublishMetrics(result);
}
Also used : TimeMonitor(org.apache.beam.sdk.testutils.metrics.TimeMonitor) GenericRecordBuilder(org.apache.avro.generic.GenericRecordBuilder) FileIO(org.apache.beam.sdk.io.FileIO) BeforeClass(org.junit.BeforeClass) PipelineResult(org.apache.beam.sdk.PipelineResult) MetricsReader(org.apache.beam.sdk.testutils.metrics.MetricsReader) Combine(org.apache.beam.sdk.transforms.Combine) RunWith(org.junit.runner.RunWith) SerializableFunction(org.apache.beam.sdk.transforms.SerializableFunction) View(org.apache.beam.sdk.transforms.View) Timestamp(com.google.cloud.Timestamp) FileBasedIOITHelper(org.apache.beam.sdk.io.common.FileBasedIOITHelper) Function(java.util.function.Function) HashSet(java.util.HashSet) InfluxDBSettings(org.apache.beam.sdk.testutils.publishing.InfluxDBSettings) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) TypeDescriptors.strings(org.apache.beam.sdk.values.TypeDescriptors.strings) NamedTestResult(org.apache.beam.sdk.testutils.NamedTestResult) DoFn(org.apache.beam.sdk.transforms.DoFn) MapElements(org.apache.beam.sdk.transforms.MapElements) GenericRecord(org.apache.avro.generic.GenericRecord) HashingFn(org.apache.beam.sdk.io.common.HashingFn) Schema(org.apache.avro.Schema) FileBasedIOITHelper.readFileBasedIOITPipelineOptions(org.apache.beam.sdk.io.common.FileBasedIOITHelper.readFileBasedIOITPipelineOptions) PAssert(org.apache.beam.sdk.testing.PAssert) FileBasedIOITHelper.appendTimestampSuffix(org.apache.beam.sdk.io.common.FileBasedIOITHelper.appendTimestampSuffix) TimeMonitor(org.apache.beam.sdk.testutils.metrics.TimeMonitor) Set(java.util.Set) GenerateSequence(org.apache.beam.sdk.io.GenerateSequence) Test(org.junit.Test) UUID(java.util.UUID) JUnit4(org.junit.runners.JUnit4) PCollection(org.apache.beam.sdk.values.PCollection) AvroCoder(org.apache.beam.sdk.coders.AvroCoder) IOITMetrics(org.apache.beam.sdk.testutils.metrics.IOITMetrics) Rule(org.junit.Rule) ParDo(org.apache.beam.sdk.transforms.ParDo) FileBasedIOTestPipelineOptions(org.apache.beam.sdk.io.common.FileBasedIOTestPipelineOptions) Values(org.apache.beam.sdk.transforms.Values) FileBasedIOITHelper(org.apache.beam.sdk.io.common.FileBasedIOITHelper) PipelineResult(org.apache.beam.sdk.PipelineResult) GenericRecord(org.apache.avro.generic.GenericRecord) HashingFn(org.apache.beam.sdk.io.common.HashingFn) Test(org.junit.Test)

Aggregations

Timestamp (com.google.cloud.Timestamp)1 HashSet (java.util.HashSet)1 Set (java.util.Set)1 UUID (java.util.UUID)1 Function (java.util.function.Function)1 Schema (org.apache.avro.Schema)1 GenericRecord (org.apache.avro.generic.GenericRecord)1 GenericRecordBuilder (org.apache.avro.generic.GenericRecordBuilder)1 PipelineResult (org.apache.beam.sdk.PipelineResult)1 AvroCoder (org.apache.beam.sdk.coders.AvroCoder)1 FileIO (org.apache.beam.sdk.io.FileIO)1 GenerateSequence (org.apache.beam.sdk.io.GenerateSequence)1 FileBasedIOITHelper (org.apache.beam.sdk.io.common.FileBasedIOITHelper)1 FileBasedIOITHelper.appendTimestampSuffix (org.apache.beam.sdk.io.common.FileBasedIOITHelper.appendTimestampSuffix)1 FileBasedIOITHelper.readFileBasedIOITPipelineOptions (org.apache.beam.sdk.io.common.FileBasedIOITHelper.readFileBasedIOITPipelineOptions)1 FileBasedIOTestPipelineOptions (org.apache.beam.sdk.io.common.FileBasedIOTestPipelineOptions)1 HashingFn (org.apache.beam.sdk.io.common.HashingFn)1 PAssert (org.apache.beam.sdk.testing.PAssert)1 TestPipeline (org.apache.beam.sdk.testing.TestPipeline)1 NamedTestResult (org.apache.beam.sdk.testutils.NamedTestResult)1