| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.sdk.io; |
| |
| import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem; |
| import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.includesDisplayDataFor; |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.MoreObjects.firstNonNull; |
| import static org.hamcrest.Matchers.containsInAnyOrder; |
| import static org.hamcrest.Matchers.containsString; |
| import static org.hamcrest.Matchers.equalTo; |
| import static org.hamcrest.Matchers.instanceOf; |
| import static org.hamcrest.Matchers.is; |
| import static org.hamcrest.Matchers.not; |
| import static org.hamcrest.Matchers.nullValue; |
| import static org.junit.Assert.assertEquals; |
| import static org.junit.Assert.assertFalse; |
| import static org.junit.Assert.assertThat; |
| import static org.junit.Assert.assertTrue; |
| |
| import java.io.BufferedReader; |
| import java.io.File; |
| import java.io.IOException; |
| import java.nio.file.Files; |
| import java.text.DecimalFormat; |
| import java.util.ArrayList; |
| import java.util.Arrays; |
| import java.util.Collections; |
| import java.util.List; |
| import java.util.Set; |
| import java.util.concurrent.ThreadLocalRandom; |
| import java.util.function.BiFunction; |
| import java.util.regex.Matcher; |
| import java.util.regex.Pattern; |
| import org.apache.beam.sdk.Pipeline; |
| import org.apache.beam.sdk.coders.StringUtf8Coder; |
| import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params; |
| import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations; |
| import org.apache.beam.sdk.io.FileBasedSink.FilenamePolicy; |
| import org.apache.beam.sdk.io.FileBasedSink.OutputFileHints; |
| import org.apache.beam.sdk.io.SimpleSink.SimpleWriter; |
| import org.apache.beam.sdk.io.fs.EmptyMatchTreatment; |
| import org.apache.beam.sdk.io.fs.MatchResult.Metadata; |
| import org.apache.beam.sdk.io.fs.ResolveOptions.StandardResolveOptions; |
| import org.apache.beam.sdk.io.fs.ResourceId; |
| import org.apache.beam.sdk.options.Description; |
| import org.apache.beam.sdk.options.PipelineOptionsFactoryTest.TestPipelineOptions; |
| import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; |
| import org.apache.beam.sdk.testing.NeedsRunner; |
| import org.apache.beam.sdk.testing.TestPipeline; |
| import org.apache.beam.sdk.testing.UsesUnboundedPCollections; |
| import org.apache.beam.sdk.transforms.Create; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.Flatten; |
| import org.apache.beam.sdk.transforms.GroupByKey; |
| import org.apache.beam.sdk.transforms.MapElements; |
| import org.apache.beam.sdk.transforms.PTransform; |
| import org.apache.beam.sdk.transforms.ParDo; |
| import org.apache.beam.sdk.transforms.SimpleFunction; |
| import org.apache.beam.sdk.transforms.Top; |
| import org.apache.beam.sdk.transforms.Values; |
| import org.apache.beam.sdk.transforms.View; |
| import org.apache.beam.sdk.transforms.display.DisplayData; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.transforms.windowing.FixedWindows; |
| import org.apache.beam.sdk.transforms.windowing.IntervalWindow; |
| import org.apache.beam.sdk.transforms.windowing.PaneInfo; |
| import org.apache.beam.sdk.transforms.windowing.Sessions; |
| import org.apache.beam.sdk.transforms.windowing.Window; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.PCollection; |
| import org.apache.beam.sdk.values.PCollection.IsBounded; |
| import org.apache.beam.sdk.values.PCollectionView; |
| import org.apache.beam.sdk.values.PDone; |
| import org.apache.beam.sdk.values.ShardedKey; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Optional; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists; |
| import org.apache.commons.compress.utils.Sets; |
| import org.hamcrest.Matchers; |
| import org.joda.time.Duration; |
| import org.joda.time.format.DateTimeFormatter; |
| import org.joda.time.format.ISODateTimeFormat; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.experimental.categories.Category; |
| import org.junit.rules.ExpectedException; |
| import org.junit.rules.TemporaryFolder; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| |
| /** Tests for the WriteFiles PTransform. */ |
| @RunWith(JUnit4.class) |
| public class WriteFilesTest { |
| @Rule public TemporaryFolder tmpFolder = new TemporaryFolder(); |
| @Rule public final TestPipeline p = TestPipeline.create(); |
| @Rule public ExpectedException thrown = ExpectedException.none(); |
| |
| @SuppressWarnings("unchecked") // covariant cast |
| private static final PTransform<PCollection<String>, PCollection<String>> IDENTITY_MAP = |
| (PTransform) |
| MapElements.via( |
| new SimpleFunction<String, String>() { |
| @Override |
| public String apply(String input) { |
| return input; |
| } |
| }); |
| |
| private static final PTransform<PCollection<String>, PCollectionView<Integer>> |
| SHARDING_TRANSFORM = |
| new PTransform<PCollection<String>, PCollectionView<Integer>>() { |
| @Override |
| public PCollectionView<Integer> expand(PCollection<String> input) { |
| return null; |
| } |
| }; |
| |
| private static class WindowAndReshuffle<T> extends PTransform<PCollection<T>, PCollection<T>> { |
| private final Window<T> window; |
| |
| public WindowAndReshuffle(Window<T> window) { |
| this.window = window; |
| } |
| |
| private static class AddArbitraryKey<T> extends DoFn<T, KV<Integer, T>> { |
| |
| @ProcessElement |
| public void processElement(ProcessContext c) { |
| c.output(KV.of(ThreadLocalRandom.current().nextInt(), c.element())); |
| } |
| } |
| |
| private static class RemoveArbitraryKey<T> extends DoFn<KV<Integer, Iterable<T>>, T> { |
| |
| @ProcessElement |
| public void processElement(ProcessContext c) { |
| for (T s : c.element().getValue()) { |
| c.output(s); |
| } |
| } |
| } |
| |
| @Override |
| public PCollection<T> expand(PCollection<T> input) { |
| return input |
| .apply(window) |
| .apply(ParDo.of(new AddArbitraryKey<>())) |
| .apply(GroupByKey.create()) |
| .apply(ParDo.of(new RemoveArbitraryKey<>())); |
| } |
| } |
| |
| private static class VerifyFilesExist<DestinationT> |
| extends PTransform<PCollection<KV<DestinationT, String>>, PDone> { |
| @Override |
| public PDone expand(PCollection<KV<DestinationT, String>> input) { |
| input |
| .apply(Values.create()) |
| .apply(FileIO.matchAll().withEmptyMatchTreatment(EmptyMatchTreatment.DISALLOW)); |
| return PDone.in(input.getPipeline()); |
| } |
| } |
| |
| private String getBaseOutputFilename() { |
| return getBaseOutputDirectory().resolve("file", StandardResolveOptions.RESOLVE_FILE).toString(); |
| } |
| |
| /** Test a WriteFiles transform with a PCollection of elements. */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testWrite() throws IOException { |
| List<String> inputs = |
| Arrays.asList( |
| "Critical canary", |
| "Apprehensive eagle", |
| "Intimidating pigeon", |
| "Pedantic gull", |
| "Frisky finch"); |
| runWrite(inputs, IDENTITY_MAP, getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); |
| } |
| |
| /** Test that WriteFiles with an empty input still produces one shard. */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testEmptyWrite() throws IOException { |
| runWrite( |
| Collections.emptyList(), |
| IDENTITY_MAP, |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink())); |
| checkFileContents( |
| getBaseOutputFilename(), |
| Collections.emptyList(), |
| Optional.of(1), |
| true /* expectRemovedTempDirectory */); |
| } |
| |
| /** |
| * Test that WriteFiles with a configured number of shards produces the desired number of shards |
| * even when there are many elements. |
| */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testShardedWrite() throws IOException { |
| runShardedWrite( |
| Arrays.asList("one", "two", "three", "four", "five", "six"), |
| IDENTITY_MAP, |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink())); |
| } |
| |
| private ResourceId getBaseOutputDirectory() { |
| return LocalResources.fromFile(tmpFolder.getRoot(), true) |
| .resolve("output", StandardResolveOptions.RESOLVE_DIRECTORY); |
| } |
| |
| private SimpleSink<Void> makeSimpleSink() { |
| FilenamePolicy filenamePolicy = |
| new PerWindowFiles( |
| getBaseOutputDirectory().resolve("file", StandardResolveOptions.RESOLVE_FILE), |
| "simple"); |
| return SimpleSink.makeSimpleSink(getBaseOutputDirectory(), filenamePolicy); |
| } |
| |
| @Test |
| @Category(NeedsRunner.class) |
| public void testCustomShardedWrite() throws IOException { |
| // Flag to validate that the pipeline options are passed to the Sink |
| WriteOptions options = TestPipeline.testingPipelineOptions().as(WriteOptions.class); |
| options.setTestFlag("test_value"); |
| Pipeline p = TestPipeline.create(options); |
| |
| List<String> inputs = new ArrayList<>(); |
| // Prepare timestamps for the elements. |
| List<Long> timestamps = new ArrayList<>(); |
| for (long i = 0; i < 1000; i++) { |
| inputs.add(Integer.toString(3)); |
| timestamps.add(i + 1); |
| } |
| |
| SimpleSink<Void> sink = makeSimpleSink(); |
| WriteFiles<String, ?, String> write = WriteFiles.to(sink).withSharding(new LargestInt()); |
| p.apply(Create.timestamped(inputs, timestamps).withCoder(StringUtf8Coder.of())) |
| .apply(IDENTITY_MAP) |
| .apply(write) |
| .getPerDestinationOutputFilenames() |
| .apply(new VerifyFilesExist<>()); |
| |
| p.run(); |
| |
| checkFileContents( |
| getBaseOutputFilename(), inputs, Optional.of(3), true /* expectRemovedTempDirectory */); |
| } |
| |
| /** |
| * Test that WriteFiles with a configured number of shards produces the desired number of shard |
| * even when there are too few elements. |
| */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testExpandShardedWrite() throws IOException { |
| runShardedWrite( |
| Arrays.asList("one", "two", "three", "four", "five", "six"), |
| IDENTITY_MAP, |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink()).withNumShards(20)); |
| } |
| |
| /** |
| * Test that WriteFiles with a configured sharding function distributes elements according to |
| * assigned key. Test use simple sharding which co-locate same elements into same shard. |
| */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testCustomShardingFunctionWrite() throws IOException { |
| runShardedWrite( |
| Arrays.asList("1", "2", "3", "1", "2", "3"), |
| IDENTITY_MAP, |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink()) |
| .withNumShards(4) |
| .withShardingFunction(new TestShardingFunction()), |
| new BiFunction<Integer, List<String>, Void>() { |
| @Override |
| public Void apply(Integer shardNumber, List<String> shardContent) { |
| // Function creates only shards 1,2,3 but WriteFiles will ensure shard 0 with empty |
| // content will be created as well |
| if (shardNumber == 0) { |
| assertTrue(shardContent.isEmpty()); |
| } else { |
| assertEquals( |
| Lists.newArrayList(shardNumber.toString(), shardNumber.toString()), shardContent); |
| } |
| return null; |
| } |
| }); |
| } |
| |
| /** Test a WriteFiles transform with an empty PCollection. */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testWriteWithEmptyPCollection() throws IOException { |
| List<String> inputs = new ArrayList<>(); |
| runWrite(inputs, IDENTITY_MAP, getBaseOutputFilename(), WriteFiles.to(makeSimpleSink())); |
| } |
| |
| /** Test a WriteFiles with a windowed PCollection. */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testWriteWindowed() throws IOException { |
| List<String> inputs = |
| Arrays.asList( |
| "Critical canary", |
| "Apprehensive eagle", |
| "Intimidating pigeon", |
| "Pedantic gull", |
| "Frisky finch"); |
| runWrite( |
| inputs, |
| new WindowAndReshuffle<>(Window.into(FixedWindows.of(Duration.millis(2)))), |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink())); |
| } |
| |
| /** Test a WriteFiles with sessions. */ |
| @Test |
| @Category(NeedsRunner.class) |
| public void testWriteWithSessions() throws IOException { |
| List<String> inputs = |
| Arrays.asList( |
| "Critical canary", |
| "Apprehensive eagle", |
| "Intimidating pigeon", |
| "Pedantic gull", |
| "Frisky finch"); |
| |
| runWrite( |
| inputs, |
| new WindowAndReshuffle<>(Window.into(Sessions.withGapDuration(Duration.millis(1)))), |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink())); |
| } |
| |
| @Test |
| @Category(NeedsRunner.class) |
| public void testWriteSpilling() throws IOException { |
| List<String> inputs = Lists.newArrayList(); |
| for (int i = 0; i < 100; ++i) { |
| inputs.add("mambo_number_" + i); |
| } |
| runWrite( |
| inputs, |
| Window.into(FixedWindows.of(Duration.millis(1))), |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink()).withMaxNumWritersPerBundle(2).withWindowedWrites()); |
| } |
| |
| @Test |
| @Category(NeedsRunner.class) |
| public void testWriteNoSpilling() throws IOException { |
| List<String> inputs = Lists.newArrayList(); |
| for (int i = 0; i < 100; ++i) { |
| inputs.add("mambo_number_" + i); |
| } |
| runWrite( |
| inputs, |
| Window.into(FixedWindows.of(Duration.millis(1))), |
| getBaseOutputFilename(), |
| WriteFiles.to(makeSimpleSink()) |
| .withMaxNumWritersPerBundle(2) |
| .withWindowedWrites() |
| .withNoSpilling()); |
| } |
| |
| @Test |
| public void testBuildWrite() { |
| SimpleSink<Void> sink = makeSimpleSink(); |
| WriteFiles<String, ?, String> write = WriteFiles.to(sink).withNumShards(3); |
| assertThat((SimpleSink<Void>) write.getSink(), is(sink)); |
| PTransform<PCollection<String>, PCollectionView<Integer>> originalSharding = |
| write.getComputeNumShards(); |
| |
| assertThat(write.getComputeNumShards(), is(nullValue())); |
| assertThat(write.getNumShardsProvider(), instanceOf(StaticValueProvider.class)); |
| assertThat(write.getNumShardsProvider().get(), equalTo(3)); |
| assertThat(write.getComputeNumShards(), equalTo(originalSharding)); |
| |
| WriteFiles<String, ?, ?> write2 = write.withSharding(SHARDING_TRANSFORM); |
| assertThat((SimpleSink<Void>) write2.getSink(), is(sink)); |
| assertThat(write2.getComputeNumShards(), equalTo(SHARDING_TRANSFORM)); |
| // original unchanged |
| |
| WriteFiles<String, ?, ?> writeUnsharded = write2.withRunnerDeterminedSharding(); |
| assertThat(writeUnsharded.getComputeNumShards(), nullValue()); |
| assertThat(write.getComputeNumShards(), equalTo(originalSharding)); |
| } |
| |
| @Test |
| public void testDisplayData() { |
| DynamicDestinations<String, Void, String> dynamicDestinations = |
| DynamicFileDestinations.constant( |
| DefaultFilenamePolicy.fromParams( |
| new Params() |
| .withBaseFilename( |
| getBaseOutputDirectory() |
| .resolve("file", StandardResolveOptions.RESOLVE_FILE)) |
| .withShardTemplate("-SS-of-NN"))); |
| SimpleSink<Void> sink = |
| new SimpleSink<Void>( |
| getBaseOutputDirectory(), dynamicDestinations, Compression.UNCOMPRESSED) { |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| builder.add(DisplayData.item("foo", "bar")); |
| } |
| }; |
| WriteFiles<String, ?, String> write = WriteFiles.to(sink); |
| |
| DisplayData displayData = DisplayData.from(write); |
| |
| assertThat(displayData, hasDisplayItem("sink", sink.getClass())); |
| assertThat(displayData, includesDisplayDataFor("sink", sink)); |
| } |
| |
| @Test |
| @Category(NeedsRunner.class) |
| public void testUnboundedNeedsWindowed() { |
| thrown.expect(IllegalArgumentException.class); |
| thrown.expectMessage( |
| "Must use windowed writes when applying WriteFiles to an unbounded PCollection"); |
| |
| SimpleSink<Void> sink = makeSimpleSink(); |
| p.apply(Create.of("foo")).setIsBoundedInternal(IsBounded.UNBOUNDED).apply(WriteFiles.to(sink)); |
| p.run(); |
| } |
| |
| @Test |
| @Category(NeedsRunner.class) |
| public void testUnboundedWritesNeedSharding() { |
| thrown.expect(IllegalArgumentException.class); |
| thrown.expectMessage( |
| "When applying WriteFiles to an unbounded PCollection, " |
| + "must specify number of output shards explicitly"); |
| |
| SimpleSink<Void> sink = makeSimpleSink(); |
| p.apply(Create.of("foo")) |
| .setIsBoundedInternal(IsBounded.UNBOUNDED) |
| .apply(WriteFiles.to(sink).withWindowedWrites()); |
| p.run(); |
| } |
| |
| // Test DynamicDestinations class. Expects user values to be string-encoded integers. |
| // Stores the integer mod 5 as the destination, and uses that in the file prefix. |
| static class TestDestinations extends DynamicDestinations<String, Integer, String> { |
| private ResourceId baseOutputDirectory; |
| |
| TestDestinations(ResourceId baseOutputDirectory) { |
| this.baseOutputDirectory = baseOutputDirectory; |
| } |
| |
| @Override |
| public String formatRecord(String record) { |
| return "record_" + record; |
| } |
| |
| @Override |
| public Integer getDestination(String element) { |
| return Integer.valueOf(element) % 5; |
| } |
| |
| @Override |
| public Integer getDefaultDestination() { |
| return 0; |
| } |
| |
| @Override |
| public FilenamePolicy getFilenamePolicy(Integer destination) { |
| return new PerWindowFiles( |
| baseOutputDirectory.resolve("file_" + destination, StandardResolveOptions.RESOLVE_FILE), |
| "simple"); |
| } |
| } |
| |
| @Test |
| @Category(NeedsRunner.class) |
| public void testDynamicDestinationsBounded() throws Exception { |
| testDynamicDestinationsHelper(true, false); |
| } |
| |
| @Test |
| @Category({NeedsRunner.class, UsesUnboundedPCollections.class}) |
| public void testDynamicDestinationsUnbounded() throws Exception { |
| testDynamicDestinationsHelper(false, false); |
| } |
| |
| @Test |
| @Category(NeedsRunner.class) |
| public void testDynamicDestinationsFillEmptyShards() throws Exception { |
| testDynamicDestinationsHelper(true, true); |
| } |
| |
| private void testDynamicDestinationsHelper(boolean bounded, boolean emptyShards) |
| throws IOException { |
| TestDestinations dynamicDestinations = new TestDestinations(getBaseOutputDirectory()); |
| SimpleSink<Integer> sink = |
| new SimpleSink<>(getBaseOutputDirectory(), dynamicDestinations, Compression.UNCOMPRESSED); |
| |
| // Flag to validate that the pipeline options are passed to the Sink. |
| WriteOptions options = TestPipeline.testingPipelineOptions().as(WriteOptions.class); |
| options.setTestFlag("test_value"); |
| Pipeline p = TestPipeline.create(options); |
| |
| final int numInputs = 100; |
| List<String> inputs = Lists.newArrayList(); |
| for (int i = 0; i < numInputs; ++i) { |
| inputs.add(Integer.toString(i)); |
| } |
| // Prepare timestamps for the elements. |
| List<Long> timestamps = new ArrayList<>(); |
| for (long i = 0; i < inputs.size(); i++) { |
| timestamps.add(i + 1); |
| } |
| // If emptyShards==true make numShards larger than the number of elements per destination. |
| // This will force every destination to generate some empty shards. |
| int numShards = emptyShards ? 2 * numInputs / 5 : 2; |
| WriteFiles<String, Integer, String> writeFiles = WriteFiles.to(sink).withNumShards(numShards); |
| |
| PCollection<String> input = p.apply(Create.timestamped(inputs, timestamps)); |
| WriteFilesResult<Integer> res; |
| if (!bounded) { |
| input.setIsBoundedInternal(IsBounded.UNBOUNDED); |
| input = input.apply(Window.into(FixedWindows.of(Duration.standardDays(1)))); |
| res = input.apply(writeFiles.withWindowedWrites()); |
| } else { |
| res = input.apply(writeFiles); |
| } |
| res.getPerDestinationOutputFilenames().apply(new VerifyFilesExist<>()); |
| p.run(); |
| |
| for (int i = 0; i < 5; ++i) { |
| ResourceId base = |
| getBaseOutputDirectory().resolve("file_" + i, StandardResolveOptions.RESOLVE_FILE); |
| List<String> expected = Lists.newArrayList(); |
| for (int j = i; j < numInputs; j += 5) { |
| expected.add("record_" + j); |
| } |
| checkFileContents( |
| base.toString(), |
| expected, |
| Optional.of(numShards), |
| bounded /* expectRemovedTempDirectory */); |
| } |
| } |
| |
| @Test |
| public void testShardedDisplayData() { |
| DynamicDestinations<String, Void, String> dynamicDestinations = |
| DynamicFileDestinations.constant( |
| DefaultFilenamePolicy.fromParams( |
| new Params() |
| .withBaseFilename( |
| getBaseOutputDirectory() |
| .resolve("file", StandardResolveOptions.RESOLVE_FILE)) |
| .withShardTemplate("-SS-of-NN"))); |
| SimpleSink<Void> sink = |
| new SimpleSink<Void>( |
| getBaseOutputDirectory(), dynamicDestinations, Compression.UNCOMPRESSED) { |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| builder.add(DisplayData.item("foo", "bar")); |
| } |
| }; |
| WriteFiles<String, ?, String> write = WriteFiles.to(sink).withNumShards(1); |
| DisplayData displayData = DisplayData.from(write); |
| assertThat(displayData, hasDisplayItem("sink", sink.getClass())); |
| assertThat(displayData, includesDisplayDataFor("sink", sink)); |
| assertThat(displayData, hasDisplayItem("numShards", 1)); |
| } |
| |
| @Test |
| public void testCustomShardStrategyDisplayData() { |
| DynamicDestinations<String, Void, String> dynamicDestinations = |
| DynamicFileDestinations.constant( |
| DefaultFilenamePolicy.fromParams( |
| new Params() |
| .withBaseFilename( |
| getBaseOutputDirectory() |
| .resolve("file", StandardResolveOptions.RESOLVE_FILE)) |
| .withShardTemplate("-SS-of-NN"))); |
| SimpleSink<Void> sink = |
| new SimpleSink<Void>( |
| getBaseOutputDirectory(), dynamicDestinations, Compression.UNCOMPRESSED) { |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| builder.add(DisplayData.item("foo", "bar")); |
| } |
| }; |
| WriteFiles<String, ?, String> write = |
| WriteFiles.to(sink) |
| .withSharding( |
| new PTransform<PCollection<String>, PCollectionView<Integer>>() { |
| @Override |
| public PCollectionView<Integer> expand(PCollection<String> input) { |
| return null; |
| } |
| |
| @Override |
| public void populateDisplayData(DisplayData.Builder builder) { |
| builder.add(DisplayData.item("spam", "ham")); |
| } |
| }); |
| DisplayData displayData = DisplayData.from(write); |
| assertThat(displayData, hasDisplayItem("sink", sink.getClass())); |
| assertThat(displayData, includesDisplayDataFor("sink", sink)); |
| assertThat(displayData, hasDisplayItem("spam", "ham")); |
| } |
| |
| /** |
| * Performs a WriteFiles transform and verifies the WriteFiles transform calls the appropriate |
| * methods on a test sink in the correct order, as well as verifies that the elements of a |
| * PCollection are written to the sink. |
| */ |
| private void runWrite( |
| List<String> inputs, |
| PTransform<PCollection<String>, PCollection<String>> transform, |
| String baseName, |
| WriteFiles<String, ?, String> write) |
| throws IOException { |
| runShardedWrite(inputs, transform, baseName, write); |
| } |
| |
| private static class PerWindowFiles extends FilenamePolicy { |
| private static final DateTimeFormatter FORMATTER = ISODateTimeFormat.hourMinuteSecondMillis(); |
| private final ResourceId baseFilename; |
| private final String suffix; |
| |
| public PerWindowFiles(ResourceId baseFilename, String suffix) { |
| this.baseFilename = baseFilename; |
| this.suffix = suffix; |
| } |
| |
| public String filenamePrefixForWindow(IntervalWindow window) { |
| String prefix = |
| baseFilename.isDirectory() ? "" : firstNonNull(baseFilename.getFilename(), ""); |
| return String.format( |
| "%s%s-%s", prefix, FORMATTER.print(window.start()), FORMATTER.print(window.end())); |
| } |
| |
| @Override |
| public ResourceId windowedFilename( |
| int shardNumber, |
| int numShards, |
| BoundedWindow window, |
| PaneInfo paneInfo, |
| OutputFileHints outputFileHints) { |
| DecimalFormat df = new DecimalFormat("0000"); |
| IntervalWindow intervalWindow = (IntervalWindow) window; |
| String filename = |
| String.format( |
| "%s-%s-of-%s%s%s", |
| filenamePrefixForWindow(intervalWindow), |
| df.format(shardNumber), |
| df.format(numShards), |
| outputFileHints.getSuggestedFilenameSuffix(), |
| suffix); |
| return baseFilename |
| .getCurrentDirectory() |
| .resolve(filename, StandardResolveOptions.RESOLVE_FILE); |
| } |
| |
| @Override |
| public ResourceId unwindowedFilename( |
| int shardNumber, int numShards, OutputFileHints outputFileHints) { |
| DecimalFormat df = new DecimalFormat("0000"); |
| String prefix = |
| baseFilename.isDirectory() ? "" : firstNonNull(baseFilename.getFilename(), ""); |
| String filename = |
| String.format( |
| "%s-%s-of-%s%s%s", |
| prefix, |
| df.format(shardNumber), |
| df.format(numShards), |
| outputFileHints.getSuggestedFilenameSuffix(), |
| suffix); |
| return baseFilename |
| .getCurrentDirectory() |
| .resolve(filename, StandardResolveOptions.RESOLVE_FILE); |
| } |
| } |
| |
| /** |
| * Same as {@link #runShardedWrite(List, PTransform, String, WriteFiles, BiFunction)} but without |
| * shard content check. This means content will be checked only globally, that shards together |
| * contains written input and not content per shard |
| */ |
| private void runShardedWrite( |
| List<String> inputs, |
| PTransform<PCollection<String>, PCollection<String>> transform, |
| String baseName, |
| WriteFiles<String, ?, String> write) |
| throws IOException { |
| runShardedWrite(inputs, transform, baseName, write, null); |
| } |
| |
| /** |
| * Performs a WriteFiles transform with the desired number of shards. Verifies the WriteFiles |
| * transform calls the appropriate methods on a test sink in the correct order, as well as |
| * verifies that the elements of a PCollection are written to the sink. If numConfiguredShards is |
| * not null, also verifies that the output number of shards is correct. |
| */ |
| private void runShardedWrite( |
| List<String> inputs, |
| PTransform<PCollection<String>, PCollection<String>> transform, |
| String baseName, |
| WriteFiles<String, ?, String> write, |
| BiFunction<Integer, List<String>, Void> shardContentChecker) |
| throws IOException { |
| // Flag to validate that the pipeline options are passed to the Sink |
| WriteOptions options = TestPipeline.testingPipelineOptions().as(WriteOptions.class); |
| options.setTestFlag("test_value"); |
| Pipeline p = TestPipeline.create(options); |
| |
| // Prepare timestamps for the elements. |
| List<Long> timestamps = new ArrayList<>(); |
| for (long i = 0; i < inputs.size(); i++) { |
| timestamps.add(i + 1); |
| } |
| p.apply(Create.timestamped(inputs, timestamps).withCoder(StringUtf8Coder.of())) |
| .apply(transform) |
| .apply(write) |
| .getPerDestinationOutputFilenames() |
| .apply(new VerifyFilesExist<>()); |
| p.run(); |
| |
| Optional<Integer> numShards = |
| (write.getNumShardsProvider() != null && !write.getWindowedWrites()) |
| ? Optional.of(write.getNumShardsProvider().get()) |
| : Optional.absent(); |
| checkFileContents(baseName, inputs, numShards, !write.getWindowedWrites(), shardContentChecker); |
| } |
| |
| static void checkFileContents( |
| String baseName, |
| List<String> inputs, |
| Optional<Integer> numExpectedShards, |
| boolean expectRemovedTempDirectory) |
| throws IOException { |
| checkFileContents(baseName, inputs, numExpectedShards, expectRemovedTempDirectory, null); |
| } |
| |
| static void checkFileContents( |
| String baseName, |
| List<String> inputs, |
| Optional<Integer> numExpectedShards, |
| boolean expectRemovedTempDirectory, |
| BiFunction<Integer, List<String>, Void> shardContentChecker) |
| throws IOException { |
| List<File> outputFiles = Lists.newArrayList(); |
| final String pattern = baseName + "*"; |
| List<Metadata> metadata = |
| FileSystems.match(Collections.singletonList(pattern)).get(0).metadata(); |
| for (Metadata meta : metadata) { |
| outputFiles.add(new File(meta.resourceId().toString())); |
| } |
| assertFalse("Should have produced at least 1 output file", outputFiles.isEmpty()); |
| Pattern shardPattern = Pattern.compile("(\\d{4})-of-\\d{4}"); |
| if (numExpectedShards.isPresent()) { |
| assertEquals(numExpectedShards.get().intValue(), outputFiles.size()); |
| |
| Set<String> expectedShards = Sets.newHashSet(); |
| DecimalFormat df = new DecimalFormat("0000"); |
| for (int i = 0; i < numExpectedShards.get(); i++) { |
| expectedShards.add( |
| String.format("%s-of-%s", df.format(i), df.format(numExpectedShards.get()))); |
| } |
| |
| Set<String> outputShards = Sets.newHashSet(); |
| for (File file : outputFiles) { |
| Matcher matcher = shardPattern.matcher(file.getName()); |
| assertTrue(matcher.find()); |
| assertTrue(outputShards.add(matcher.group())); |
| } |
| assertEquals(expectedShards, outputShards); |
| } |
| |
| List<String> actual = Lists.newArrayList(); |
| for (File outputFile : outputFiles) { |
| List<String> actualShard = Lists.newArrayList(); |
| try (BufferedReader reader = Files.newBufferedReader(outputFile.toPath(), Charsets.UTF_8)) { |
| for (; ; ) { |
| String line = reader.readLine(); |
| if (line == null) { |
| break; |
| } |
| if (!line.equals(SimpleWriter.HEADER) && !line.equals(SimpleWriter.FOOTER)) { |
| actualShard.add(line); |
| } |
| } |
| } |
| if (shardContentChecker != null) { |
| Matcher matcher = shardPattern.matcher(outputFile.getName()); |
| matcher.find(); |
| int shardNumber = Integer.parseInt(matcher.group(1)); |
| shardContentChecker.apply(shardNumber, actualShard); |
| } |
| actual.addAll(actualShard); |
| } |
| assertThat(actual, containsInAnyOrder(inputs.toArray())); |
| if (expectRemovedTempDirectory) { |
| assertThat( |
| Lists.newArrayList(new File(baseName).getParentFile().list()), |
| Matchers.everyItem(not(containsString(FileBasedSink.TEMP_DIRECTORY_PREFIX)))); |
| } |
| } |
| |
| /** Options for test, exposed for PipelineOptionsFactory. */ |
| public interface WriteOptions extends TestPipelineOptions { |
| @Description("Test flag and value") |
| String getTestFlag(); |
| |
| void setTestFlag(String value); |
| } |
| |
| /** |
| * Outputs the largest integer in a {@link PCollection} into a {@link PCollectionView}. The input |
| * {@link PCollection} must be convertible to integers via {@link Integer#valueOf(String)} |
| */ |
| private static class LargestInt |
| extends PTransform<PCollection<String>, PCollectionView<Integer>> { |
| @Override |
| public PCollectionView<Integer> expand(PCollection<String> input) { |
| return input |
| .apply( |
| ParDo.of( |
| new DoFn<String, Integer>() { |
| @ProcessElement |
| public void toInteger(ProcessContext ctxt) { |
| ctxt.output(Integer.valueOf(ctxt.element())); |
| } |
| })) |
| .apply(Top.largest(1)) |
| .apply(Flatten.iterables()) |
| .apply(View.asSingleton()); |
| } |
| } |
| |
| private static class TestShardingFunction implements ShardingFunction<String, Void> { |
| @Override |
| public ShardedKey<Integer> assignShardKey(Void destination, String element, int shardCount) |
| throws Exception { |
| int shard = Integer.parseInt(element); |
| return ShardedKey.of(0, shard); |
| } |
| } |
| } |
