diff --git a/CHANGES.md b/CHANGES.md index 7693e65d78a0..aae1d14e820b 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -73,6 +73,8 @@ * The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)). * (Go) Pipeline-level `--resource_hints` are now forwarded to expansion services, so they apply to cross-language transforms as they do in the Java and Python SDKs. The `max_active_bundles_per_worker` hint name is now also accepted by `--resource_hints` ([#23893](https://github.com/apache/beam/issues/23893)). * ReadFromBigQuery now supports Lakehouse runtime catalog (BigLake metastore) tables with `method=DIRECT_READ`, using 4-part `project.catalog.namespace.table` identifiers. Previously `project:catalog.namespace.table` was silently mis-parsed and tables that report no `numBytes` failed to split (Python) ([#39597](https://github.com/apache/beam/issues/39597)). +* (Java) `TextIO`, `FileIO`, and `WriteFiles`: streaming writes can skip the shuffle for large bundles by writing them directly to temp files (`withDirectWriteByteThreshold` or experiment `streaming_file_direct_write_threshold_bytes`). Off by default ([#40533](https://github.com/apache/beam/issues/40533)). +* (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) for configuring AWS service clients such as S3 ([#40528](https://github.com/apache/beam/issues/40528)). * (Python) Added `AwsOptions` with `--aws_*` pipeline options (e.g. `--aws_access_key_id`, `--aws_region_name`, `--aws_endpoint_url`) that set defaults for all AWS service clients. The existing `--s3_*` options in `S3Options` override these defaults for the S3 client only ([#40528](https://github.com/apache/beam/issues/40528)). * (Python) Added `--aws_credentials_secret_manager` so AWS credentials (`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and their `--s3_*` overrides) can be secret specs resolved through a secret manager such as `GoogleCloudSecretManager`, instead of plain-text values ([#40528](https://github.com/apache/beam/issues/40528)). diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java index f1b0c89090d3..d3d6855ec47a 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowRunner.java @@ -180,6 +180,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.hash.Hashing; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.ByteStreams; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.Files; +import org.checkerframework.checker.nullness.qual.Nullable; import org.joda.time.DateTimeUtils; import org.joda.time.DateTimeZone; import org.joda.time.format.DateTimeFormat; @@ -2645,6 +2646,7 @@ static class StreamingShardedWriteFactory // We pick 10 as a default, as it works well with the default number of workers started // by Dataflow. static final int DEFAULT_NUM_SHARDS = 10; + DataflowPipelineWorkerPoolOptions options; StreamingShardedWriteFactory(PipelineOptions options) { @@ -2684,6 +2686,11 @@ static class StreamingShardedWriteFactory replacement = replacement.withWindowedWrites(); } + Long directWriteThreshold = getDirectWriteThreshold(transform); + if (directWriteThreshold != null) { + replacement = replacement.withDirectWriteByteThreshold(directWriteThreshold); + } + if (WriteFilesTranslation.isAutoSharded(transform)) { replacement = replacement.withAutoSharding(); return PTransformReplacement.of( @@ -2693,11 +2700,24 @@ static class StreamingShardedWriteFactory return PTransformReplacement.of( PTransformReplacements.getSingletonMainInput(transform), replacement.withNumShards(numShards)); + } catch (RuntimeException e) { + throw e; } catch (Exception e) { throw new RuntimeException(e); } } + /** + * Returns the direct write threshold set on the original {@link WriteFiles}, else {@code null}. + */ + private @Nullable Long getDirectWriteThreshold(AppliedPTransform transform) { + Object original = transform.getTransform(); + if (original instanceof WriteFiles) { + return ((WriteFiles) original).getDirectWriteByteThreshold(); + } + return null; + } + @Override public Map, ReplacementOutput> mapOutputs( Map, PCollection> outputs, WriteFilesResult newOutput) { diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java index b196aa1714c6..3d05c2486c1e 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/DataflowRunnerTest.java @@ -2149,9 +2149,66 @@ public void testStreamingWriteWithShardingReturnsSameTransform() { assertNull(replacement.getNumShardsProvider()); assertNull(replacement.getComputeNumShards()); + assertNull(replacement.getDirectWriteByteThreshold()); assertTrue(replacement.getWithAutoSharding()); } + @Test + public void testStreamingWriteWithDirectWriteByteThreshold() { + PipelineOptions options = TestPipeline.testingPipelineOptions(); + + TestPipeline p = TestPipeline.fromOptions(options); + StreamingShardedWriteFactory factory = + new StreamingShardedWriteFactory<>(p.getOptions()); + WriteFiles original = + WriteFiles.to(new TestSink(tmpFolder.toString())) + .withRunnerDeterminedSharding() + .withDirectWriteByteThreshold(1000L); + PCollection objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); + AppliedPTransform, WriteFilesResult, WriteFiles> + originalApplication = + AppliedPTransform.of( + "writefiles", + PValues.expandInput(objs), + Collections.emptyMap(), + original, + ResourceHints.create(), + p); + + WriteFiles replacement = + (WriteFiles) + factory.getReplacementTransform(originalApplication).getTransform(); + + assertEquals(Long.valueOf(1000L), replacement.getDirectWriteByteThreshold()); + } + + @Test + public void testStreamingWriteDirectWriteThresholdDefaultNull() { + PipelineOptions options = TestPipeline.testingPipelineOptions(); + + TestPipeline p = TestPipeline.fromOptions(options); + StreamingShardedWriteFactory factory = + new StreamingShardedWriteFactory<>(p.getOptions()); + WriteFiles original = + WriteFiles.to(new TestSink(tmpFolder.toString())).withRunnerDeterminedSharding(); + PCollection objs = (PCollection) p.apply(Create.empty(VoidCoder.of())); + AppliedPTransform, WriteFilesResult, WriteFiles> + originalApplication = + AppliedPTransform.of( + "writefiles", + PValues.expandInput(objs), + Collections.emptyMap(), + original, + ResourceHints.create(), + p); + + WriteFiles replacement = + (WriteFiles) + factory.getReplacementTransform(originalApplication).getTransform(); + + assertNull(replacement.getDirectWriteByteThreshold()); + } + private void verifyMergingStatefulParDoRejected(PipelineOptions options) throws Exception { Pipeline p = Pipeline.create(options); diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java index 3f4289003614..683c167b1a12 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java @@ -1071,6 +1071,8 @@ public static FileNaming relativeFileNaming( abstract @Nullable Duration getBatchMaxBufferingDuration(); + abstract @Nullable Long getDirectWriteByteThreshold(); + abstract @Nullable ErrorHandler getBadRecordErrorHandler(); abstract Builder toBuilder(); @@ -1133,6 +1135,9 @@ abstract Builder setMaxNumWritersPerBundle( abstract Builder setBatchMaxBufferingDuration( @Nullable Duration batchMaxBufferingDuration); + abstract Builder setDirectWriteByteThreshold( + @Nullable Long directWriteByteThreshold); + abstract Builder setBadRecordErrorHandler( @Nullable ErrorHandler badRecordErrorHandler); @@ -1436,6 +1441,13 @@ public Write withBatchMaxBufferingDuration( return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build(); } + /** See {@link WriteFiles#withDirectWriteByteThreshold(long)}. */ + public Write withDirectWriteByteThreshold(long thresholdBytes) { + checkArgument( + thresholdBytes > 0, "thresholdBytes must be positive, but was: %s", thresholdBytes); + return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build(); + } + @VisibleForTesting Contextful> resolveFileNamingFn() { if (getDynamic()) { @@ -1565,6 +1577,9 @@ public WriteFilesResult expand(PCollection input) { if (getBatchMaxBufferingDuration() != null) { writeFiles = writeFiles.withBatchMaxBufferingDuration(getBatchMaxBufferingDuration()); } + if (getDirectWriteByteThreshold() != null) { + writeFiles = writeFiles.withDirectWriteByteThreshold(getDirectWriteByteThreshold()); + } return input.apply(writeFiles); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java index b96465e5cab7..dc7dd1c5d49f 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java @@ -713,6 +713,9 @@ public abstract static class TypedWrite /** Batch max buffering duration for input records. */ abstract @Nullable Duration getBatchMaxBufferingDuration(); + /** Byte threshold above which a streaming bundle is written directly to temporary files. */ + abstract @Nullable Long getDirectWriteByteThreshold(); + /** Whether to write windowed output files. */ abstract boolean getWindowedWrites(); @@ -780,6 +783,9 @@ abstract Builder setNumShards( abstract Builder setBatchMaxBufferingDuration( @Nullable Duration batchMaxBufferingDuration); + abstract Builder setDirectWriteByteThreshold( + @Nullable Long directWriteByteThreshold); + abstract Builder setWindowedWrites(boolean windowedWrites); abstract Builder setAutoSharding(boolean windowedWrites); @@ -926,6 +932,13 @@ public TypedWrite withBatchMaxBufferingDuration( return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build(); } + /** See {@link WriteFiles#withDirectWriteByteThreshold(long)}. */ + public TypedWrite withDirectWriteByteThreshold(long thresholdBytes) { + checkArgument( + thresholdBytes > 0, "thresholdBytes must be positive, but was: %s", thresholdBytes); + return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build(); + } + /** Set the base directory used to generate temporary files. */ public TypedWrite withTempDirectory( ValueProvider tempDirectory) { @@ -1195,6 +1208,9 @@ public WriteFilesResult expand(PCollection input) { if (getBatchMaxBufferingDuration() != null) { write = write.withBatchMaxBufferingDuration(getBatchMaxBufferingDuration()); } + if (getDirectWriteByteThreshold() != null) { + write = write.withDirectWriteByteThreshold(getDirectWriteByteThreshold()); + } return input.apply("WriteFiles", write); } @@ -1388,6 +1404,11 @@ public Write withBatchMaxBufferingDuration(@Nullable Duration batchMaxBufferingD return new Write(inner.withBatchMaxBufferingDuration(batchMaxBufferingDuration)); } + /** See {@link TypedWrite#withDirectWriteByteThreshold(long)}. */ + public Write withDirectWriteByteThreshold(long thresholdBytes) { + return new Write(inner.withDirectWriteByteThreshold(thresholdBytes)); + } + /** * Specify that output filenames are wanted. * diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java index 2fd77a1ad0c6..6f68591024f2 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java @@ -49,6 +49,8 @@ import org.apache.beam.sdk.io.FileBasedSink.WriteOperation; import org.apache.beam.sdk.io.FileBasedSink.Writer; import org.apache.beam.sdk.io.fs.ResourceId; +import org.apache.beam.sdk.metrics.Counter; +import org.apache.beam.sdk.metrics.Metrics; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.ValueProvider; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -79,6 +81,7 @@ import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.MoreFutures; +import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.KV; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollection.IsBounded; @@ -210,6 +213,13 @@ public static WriteFiles> getSideInputs(); public abstract @Nullable ShardingFunction getShardingFunction(); @@ -251,6 +261,9 @@ abstract Builder setBatchSizeBytes( abstract Builder setBatchMaxBufferingDuration( @Nullable Duration batchMaxBufferingDuration); + abstract Builder setDirectWriteByteThreshold( + @Nullable Long directWriteByteThreshold); + abstract Builder setSideInputs( List> sideInputs); @@ -445,6 +458,23 @@ public WriteFiles withBatchMaxBufferingDuration( return toBuilder().setBatchMaxBufferingDuration(batchMaxBufferingDuration).build(); } + /** + * Returns a new {@link WriteFiles} that writes large input bundles directly to temporary files, + * skipping the shuffle before temporary file creation. + * + *

Applies only to unbounded inputs, and is ignored when dynamic destinations use side inputs. + * Records of each bundle are buffered until their encoded size reaches {@code thresholdBytes}; + * bundles that reach it are written straight to temporary files. Smaller bundles keep the + * existing grouped path so output files stay large. Shard numbers are assigned at finalization, + * which may override {@link #withNumShards(int)} (the number of files per window can differ from + * the configured number of shards). + */ + public WriteFiles withDirectWriteByteThreshold( + long thresholdBytes) { + checkArgument(thresholdBytes > 0, "thresholdBytes must be positive, got %s", thresholdBytes); + return toBuilder().setDirectWriteByteThreshold(thresholdBytes).build(); + } + public WriteFiles withSideInputs( List> sideInputs) { return toBuilder().setSideInputs(sideInputs).build(); @@ -583,6 +613,22 @@ public WriteFilesResult expand(PCollection input) { (getComputeNumShards() == null) ? null : input.apply(getComputeNumShards()); boolean fixedSharding = getComputeNumShards() != null || getNumShardsProvider() != null; + Long directWriteThreshold = getDirectWriteByteThreshold(); + if (directWriteThreshold != null + && getSideInputs().isEmpty() + && input.isBounded() == IsBounded.UNBOUNDED) { + // Shard numbers are assigned in finalize across both the direct and grouped files, so any + // fixed shard count is not enforced there. + return input + .apply( + "WriteShardedBundlesWithDirectWrite", + new WriteShardedBundlesWithDirectWrite( + destinationCoder, fileResultCoder, numShardsView, directWriteThreshold)) + .apply( + "FinalizeTempFileBundles", + new FinalizeTempFileBundles( + null, destinationCoder, /* assignShardNumbersInFinalize= */ true)); + } PCollection>> tempFileResults; if (fixedSharding) { tempFileResults = @@ -614,7 +660,9 @@ public WriteFilesResult expand(PCollection input) { } return tempFileResults.apply( - "FinalizeTempFileBundles", new FinalizeTempFileBundles(numShardsView, destinationCoder)); + "FinalizeTempFileBundles", + new FinalizeTempFileBundles( + numShardsView, destinationCoder, /* assignShardNumbersInFinalize= */ false)); } @Override @@ -817,7 +865,6 @@ public void processElement( ProcessContext c, BoundedWindow window, MultiOutputReceiver outputReceiver) throws Exception { getDynamicDestinations().setSideInputAccessorFromProcessContext(c); - PaneInfo paneInfo = c.pane(); // If we are doing windowed writes, we need to ensure that we have separate files for // data in different windows/panes. Similar for dynamic writes, make sure that different // destinations go to different writers. @@ -835,10 +882,10 @@ public void processElement( writer = writers.get(key); if (writer == null) { if (getMaxNumWritersPerBundle() < 0 || writers.size() <= getMaxNumWritersPerBundle()) { - writer = openAndRegisterWriter(key, window, paneInfo, destination); + writer = openAndRegisterWriter(writers, key); } else if (getEvictWritersWhenFull()) { evictOldestWriter(); - writer = openAndRegisterWriter(key, window, paneInfo, destination); + writer = openAndRegisterWriter(writers, key); } else { if (spilledShardNum == UNKNOWN_SHARDNUM) { // Cache the random value so we only call ThreadLocalRandom once per DoFn instance. @@ -865,29 +912,6 @@ public void processElement( writeOrClose(writer, formattedRecord); } - private Writer openAndRegisterWriter( - WriterKey key, - BoundedWindow window, - PaneInfo paneInfo, - DestinationT destination) - throws Exception { - synchronized (writers) { - String uuid = UUID.randomUUID().toString(); - LOG.info( - "Opening writer {} for window {} pane {} destination {}", - uuid, - window, - paneInfo, - destination); - Writer writer = writeOperation.createWriter(); - writer.setDestination(destination); - writer.open(uuid); - writers.put(key, writer); - LOG.debug("Done opening writer"); - return writer; - } - } - private void evictOldestWriter() throws Exception { synchronized (writers) { Iterator, Writer>> iterator = @@ -927,30 +951,56 @@ public void finishBundle(FinishBundleContext c) throws Exception { BoundedWindow window = evictedResult.getWindow(); c.output(evictedResult, window.maxTimestamp(), window); } - for (Map.Entry, Writer> entry : - writers.entrySet()) { - WriterKey key = entry.getKey(); - Writer writer = entry.getValue(); - try { - writer.close(); - } catch (Exception e) { - // If anything goes wrong, make sure to delete the temporary file. - writer.cleanup(); - throw e; - } - BoundedWindow window = key.window; - c.output( - new FileResult<>( - writer.getOutputFile(), UNKNOWN_SHARDNUM, window, key.paneInfo, key.destination), - window.maxTimestamp(), - window); - } + closeAllWritersAndOutputFileResults(writers, c); writers.clear(); evictedFileResults.clear(); } } } + private Writer openAndRegisterWriter( + Map, Writer> writers, + WriterKey key) + throws Exception { + String uuid = UUID.randomUUID().toString(); + LOG.info( + "Opening writer {} for window {} pane {} destination {}", + uuid, + key.window, + key.paneInfo, + key.destination); + Writer writer = writeOperation.createWriter(); + writer.setDestination(key.destination); + writer.open(uuid); + writers.put(key, writer); + LOG.debug("Done opening writer"); + return writer; + } + + private static void closeAllWritersAndOutputFileResults( + Map, Writer> writers, + DoFn>.FinishBundleContext c) + throws Exception { + for (Map.Entry, Writer> entry : + writers.entrySet()) { + WriterKey key = entry.getKey(); + Writer writer = entry.getValue(); + try { + writer.close(); + } catch (Exception e) { + // If anything goes wrong, make sure to delete the temporary file. + writer.cleanup(); + throw e; + } + BoundedWindow window = key.window; + c.output( + new FileResult<>( + writer.getOutputFile(), UNKNOWN_SHARDNUM, window, key.paneInfo, key.destination), + window.maxTimestamp(), + window); + } + } + private static void writeOrClose( Writer writer, OutputT t) throws Exception { try { @@ -1118,6 +1168,367 @@ public PCollection> expand(PCollection input) { } } + /** + * Streaming write that sends large bundles straight to temporary files. + * + *

Small bundles, and records that would exceed the per-bundle writer limit, take the usual + * grouped path ({@link WriteShardedBundlesToTempFiles} when a shard count is set, or the {@link + * GroupIntoBatches} auto-sharded path otherwise). All files then drop their shard number and are + * grouped per destination, so finalization assigns shard numbers across both paths. + */ + private class WriteShardedBundlesWithDirectWrite + extends PTransform, PCollection>>> { + private final Coder destinationCoder; + private final Coder> fileResultCoder; + private final @Nullable PCollectionView numShardsView; + private final long thresholdBytes; + + private WriteShardedBundlesWithDirectWrite( + Coder destinationCoder, + Coder> fileResultCoder, + @Nullable PCollectionView numShardsView, + long thresholdBytes) { + this.destinationCoder = destinationCoder; + this.fileResultCoder = fileResultCoder; + this.numShardsView = numShardsView; + this.thresholdBytes = thresholdBytes; + } + + @Override + public PCollection>> expand(PCollection input) { + TupleTag> directFilesTag = new TupleTag<>("directFiles"); + TupleTag smallBundlesTag = new TupleTag<>("smallBundles"); + + PCollectionTuple routed = + input.apply( + "DirectWriteLargeBundles", + ParDo.of( + new DirectWriteLargeBundlesFn( + smallBundlesTag, thresholdBytes, input.getCoder())) + .withSideInputs(getSideInputs()) + .withOutputTags( + directFilesTag, + TupleTagList.of(ImmutableList.of(smallBundlesTag, BAD_RECORD_TAG)))); + addErrorCollection(routed); + + PCollection smallBundles = routed.get(smallBundlesTag).setCoder(input.getCoder()); + PCollection> groupedFiles; + if (getComputeNumShards() != null || getNumShardsProvider() != null) { + groupedFiles = + smallBundles + .apply( + "WriteShardedBundlesToTempFiles", + new WriteShardedBundlesToTempFiles( + destinationCoder, fileResultCoder, numShardsView)) + .apply( + "DropShardNum", + ParDo.of( + new DoFn, FileResult>() { + @ProcessElement + public void process(ProcessContext c) { + c.output(c.element().withShard(UNKNOWN_SHARDNUM)); + } + })) + .setCoder(fileResultCoder); + } else { + groupedFiles = + writeAutoShardedRecordsToTempFiles( + smallBundles, + destinationCoder, + fileResultCoder, + getBatchSize() != null ? getBatchSize() : FILE_TRIGGERING_RECORD_COUNT, + getBatchSizeBytes() != null ? getBatchSizeBytes() : FILE_TRIGGERING_BYTE_COUNT, + getBatchMaxBufferingDuration() != null + ? getBatchMaxBufferingDuration() + : FILE_TRIGGERING_RECORD_BUFFERING_DURATION); + } + + // When the small path goes through GroupIntoShards (GBK), its trigger transitions to the + // continuation trigger. Set the direct path to the same windowing strategy so + // FlattenTempFiles does not reject mismatched triggers when user pipelines use non-default + // triggers. + PCollection> directFiles = + routed + .get(directFilesTag) + .setCoder(fileResultCoder) + .setWindowingStrategyInternal(groupedFiles.getWindowingStrategy()); + + // Group all temp files per destination so each window/pane firing is finalized together + // and shard numbers are assigned once across both paths. + return PCollectionList.of(directFiles) + .and(groupedFiles) + .apply("FlattenTempFiles", Flatten.pCollections()) + .setCoder(fileResultCoder) + .apply( + "KeyedByDestination", + WithKeys.of( + new SimpleFunction, DestinationT>() { + @Override + public DestinationT apply(FileResult input) { + return input.getDestination(); + } + })) + .setCoder(KvCoder.of(destinationCoder, fileResultCoder)) + .apply("GroupTempFilesByDestination", GroupByKey.create()) + .apply( + "ExtractValuesToList", + ParDo.of( + new DoFn< + KV>>, + List>>() { + @ProcessElement + public void processElement(ProcessContext c) { + // Stamp every file with this firing's pane so files from different + // firings of the same window get distinct final names. + List> result = new ArrayList<>(); + for (FileResult e : c.element().getValue()) { + result.add( + new FileResult<>( + e.getTempFilename(), + UNKNOWN_SHARDNUM, + e.getWindow(), + c.pane(), + e.getDestination())); + } + c.output(result); + } + })) + .setCoder(ListCoder.of(fileResultCoder)); + } + } + + private static class ByteSizeObserver extends ElementByteSizeObserver { + private long byteSize; + + @Override + protected void reportElementSize(long elementByteSize) { + this.byteSize = elementByteSize; + } + } + + /** + * Buffers each bundle's records until their encoded size reaches {@code thresholdBytes}. Bundles + * that reach it are written directly to temporary files. Records of smaller bundles, and records + * that would exceed the per-bundle writer limit, are emitted to {@code smallBundlesTag} at the + * end of the bundle for the grouped path. + */ + private class DirectWriteLargeBundlesFn extends DoFn> { + private final TupleTag smallBundlesTag; + private final long thresholdBytes; + private final Coder inputCoder; + + private final Counter directBundles = Metrics.counter(WriteFiles.class, "directWriteBundles"); + private final Counter groupedBundles = Metrics.counter(WriteFiles.class, "groupedWriteBundles"); + + // Per-bundle state, reset in startBundle. + private transient List> buffer; + private transient long bufferedBytes; + private transient boolean directMode; + private transient Map, Writer> writers; + private transient List> spilled; + + DirectWriteLargeBundlesFn( + TupleTag smallBundlesTag, long thresholdBytes, Coder inputCoder) { + this.smallBundlesTag = smallBundlesTag; + this.thresholdBytes = thresholdBytes; + this.inputCoder = inputCoder; + } + + @StartBundle + public void startBundle() { + buffer = new ArrayList<>(); + bufferedBytes = 0; + directMode = false; + writers = new LinkedHashMap<>(); + spilled = new ArrayList<>(); + } + + @ProcessElement + public void processElement( + ProcessContext c, BoundedWindow window, MultiOutputReceiver outputReceiver) + throws Exception { + getDynamicDestinations().setSideInputAccessorFromProcessContext(c); + BufferedRecord record = + new BufferedRecord<>(c.element(), c.timestamp(), window, c.pane()); + if (directMode) { + write(record, outputReceiver); + return; + } + buffer.add(record); + bufferedBytes += estimateSize(c.element()); + if (bufferedBytes >= thresholdBytes) { + directMode = true; + // Direct writes are disabled when destinations use side inputs (see + // expand()), so resolving buffered records' destinations here is safe. + for (BufferedRecord buffered : buffer) { + write(buffered, outputReceiver); + } + buffer.clear(); + } + } + + private long estimateSize(UserT value) throws Exception { + if (inputCoder.isRegisterByteSizeObserverCheap(value)) { + ByteSizeObserver observer = new ByteSizeObserver(); + inputCoder.registerByteSizeObserver(value, observer); + observer.advance(); + return observer.byteSize; + } + return CoderUtils.encodeToByteArray(inputCoder, value).length; + } + + private void write(BufferedRecord record, MultiOutputReceiver outputReceiver) + throws Exception { + MaybeDestination maybeDestination = + getDestinationWithErrorHandling(record.value, outputReceiver, inputCoder); + if (!maybeDestination.isValid) { + return; + } + DestinationT destination = maybeDestination.destination; + WriterKey key = new WriterKey<>(record.window, record.pane, destination); + Writer writer = writers.get(key); + if (writer == null) { + int maxWriters = + getMaxNumWritersPerBundle() < 0 + ? DEFAULT_MAX_NUM_WRITERS_PER_BUNDLE + : getMaxNumWritersPerBundle(); + if (writers.size() >= maxWriters) { + // Too many open writers in this bundle; send the record to the grouped path. + spilled.add(record); + return; + } + writer = openAndRegisterWriter(writers, key); + } + OutputT formattedRecord = + formatRecordWithErrorHandling(record.value, outputReceiver, inputCoder); + if (formattedRecord == null) { + return; + } + writeOrClose(writer, formattedRecord); + } + + @FinishBundle + public void finishBundle(FinishBundleContext c) throws Exception { + if (directMode) { + directBundles.inc(); + } else if (!buffer.isEmpty()) { + groupedBundles.inc(); + } + closeAllWritersAndOutputFileResults(writers, c); + for (BufferedRecord record : buffer) { + c.output(smallBundlesTag, record.value, record.timestamp, record.window); + } + for (BufferedRecord record : spilled) { + c.output(smallBundlesTag, record.value, record.timestamp, record.window); + } + buffer.clear(); + spilled.clear(); + writers.clear(); + bufferedBytes = 0; + directMode = false; + } + } + + /** A record held by {@link DirectWriteLargeBundlesFn} with the metadata needed to re-emit it. */ + private static class BufferedRecord { + private final T value; + private final Instant timestamp; + private final BoundedWindow window; + private final PaneInfo pane; + + BufferedRecord(T value, Instant timestamp, BoundedWindow window, PaneInfo pane) { + this.value = value; + this.timestamp = timestamp; + this.window = window; + this.pane = pane; + } + } + + private PCollection> writeAutoShardedRecordsToTempFiles( + PCollection input, + Coder destinationCoder, + Coder> fileResultCoder, + int batchSize, + int batchSizeBytes, + Duration maxBufferingDuration) { + // Auto-sharding is achieved via GroupIntoBatches.WithShardedKey which shards, groups and at + // the same time batches the input records. The sharding behavior depends on runners. The + // batching is per window and we also emit the batches if there are a certain number of + // records buffered or they have been buffered for a certain time, controlled by + // FILE_TRIGGERING_RECORD_COUNT and BUFFERING_DURATION respectively. + // + // TODO(https://github.com/apache/beam/issues/20928): The implementation doesn't currently + // work with merging windows. + TupleTag> shardTag = new TupleTag<>("shardTag"); + + PCollectionTuple shardedElements = + input.apply( + "KeyedByDestinationHash", + ParDo.of(new KeyByDestinationHash(input.getCoder(), destinationCoder)) + .withOutputTags(shardTag, TupleTagList.of(BAD_RECORD_TAG))); + addErrorCollection(shardedElements); + + PCollection, Iterable>> shardedInput = + shardedElements + .get(shardTag) + .setCoder(KvCoder.of(VarIntCoder.of(), input.getCoder())) + .apply( + "ShardAndBatch", + GroupIntoBatches.ofSize(batchSize) + .withByteSize(batchSizeBytes) + .withMaxBufferingDuration(maxBufferingDuration) + .withShardedKey()) + .setCoder( + KvCoder.of( + org.apache.beam.sdk.util.ShardedKey.Coder.of(VarIntCoder.of()), + IterableCoder.of(input.getCoder()))); + + TupleTag> writtenRecordsTag = new TupleTag<>("writtenRecords"); + // Write grouped elements to temp files. + PCollectionTuple writtenFiles = + shardedInput + .apply( + "AddDummyShard", + MapElements.via( + new SimpleFunction< + KV, Iterable>, + KV, Iterable>>() { + @Override + public KV, Iterable> apply( + KV, Iterable> input) { + // Add dummy shard since it is required by WriteShardsIntoTempFilesFn. It + // will be dropped after we generate the temp files. + return KV.of( + ShardedKey.of(input.getKey().getKey(), DUMMY_SHARDNUM), + input.getValue()); + } + })) + .setCoder( + KvCoder.of( + ShardedKeyCoder.of(VarIntCoder.of()), IterableCoder.of(input.getCoder()))) + .apply( + "WriteShardsIntoTempFiles", + ParDo.of(new WriteShardsIntoTempFilesFn(input.getCoder())) + .withSideInputs(getSideInputs()) + .withOutputTags(writtenRecordsTag, TupleTagList.of(BAD_RECORD_TAG))); + + addErrorCollection(writtenFiles); + + return writtenFiles + .get(writtenRecordsTag) + .setCoder(fileResultCoder) + .apply( + "DropShardNum", + ParDo.of( + new DoFn, FileResult>() { + @ProcessElement + public void process(ProcessContext c) { + c.output(c.element().withShard(UNKNOWN_SHARDNUM)); + } + })) + .setCoder(fileResultCoder); + } + private class WriteAutoShardedBundlesToTempFiles extends PTransform, PCollection>>> { private final Coder destinationCoder; @@ -1144,84 +1555,14 @@ private WriteAutoShardedBundlesToTempFiles( @Override public PCollection>> expand(PCollection input) { - // Auto-sharding is achieved via GroupIntoBatches.WithShardedKey which shards, groups and at - // the same time batches the input records. The sharding behavior depends on runners. The - // batching is per window and we also emit the batches if there are a certain number of - // records buffered or they have been buffered for a certain time, controlled by - // FILE_TRIGGERING_RECORD_COUNT and BUFFERING_DURATION respectively. - // - // TODO(https://github.com/apache/beam/issues/20928): The implementation doesn't currently - // work with merging windows. - TupleTag> shardTag = new TupleTag<>("shardTag"); - - PCollectionTuple shardedElements = - input.apply( - "KeyedByDestinationHash", - ParDo.of(new KeyByDestinationHash(input.getCoder(), destinationCoder)) - .withOutputTags(shardTag, TupleTagList.of(BAD_RECORD_TAG))); - addErrorCollection(shardedElements); - - PCollection, Iterable>> shardedInput = - shardedElements - .get(shardTag) - .setCoder(KvCoder.of(VarIntCoder.of(), input.getCoder())) - .apply( - "ShardAndBatch", - GroupIntoBatches.ofSize(batchSize) - .withByteSize(batchSizeBytes) - .withMaxBufferingDuration(maxBufferingDuration) - .withShardedKey()) - .setCoder( - KvCoder.of( - org.apache.beam.sdk.util.ShardedKey.Coder.of(VarIntCoder.of()), - IterableCoder.of(input.getCoder()))); - - TupleTag> writtenRecordsTag = new TupleTag<>("writtenRecords"); - // Write grouped elements to temp files. - PCollectionTuple writtenFiles = - shardedInput - .apply( - "AddDummyShard", - MapElements.via( - new SimpleFunction< - KV, Iterable>, - KV, Iterable>>() { - @Override - public KV, Iterable> apply( - KV, Iterable> - input) { - // Add dummy shard since it is required by WriteShardsIntoTempFilesFn. It - // will be dropped after we generate the temp files. - return KV.of( - ShardedKey.of(input.getKey().getKey(), DUMMY_SHARDNUM), - input.getValue()); - } - })) - .setCoder( - KvCoder.of( - ShardedKeyCoder.of(VarIntCoder.of()), IterableCoder.of(input.getCoder()))) - .apply( - "WriteShardsIntoTempFiles", - ParDo.of(new WriteShardsIntoTempFilesFn(input.getCoder())) - .withSideInputs(getSideInputs()) - .withOutputTags(writtenRecordsTag, TupleTagList.of(BAD_RECORD_TAG))); - - addErrorCollection(writtenFiles); - PCollection> tempFiles = - writtenFiles - .get(writtenRecordsTag) - .setCoder(fileResultCoder) - .apply( - "DropShardNum", - ParDo.of( - new DoFn, FileResult>() { - @ProcessElement - public void process(ProcessContext c) { - c.output(c.element().withShard(UNKNOWN_SHARDNUM)); - } - })) - .setCoder(fileResultCoder); + writeAutoShardedRecordsToTempFiles( + input, + destinationCoder, + fileResultCoder, + batchSize, + batchSizeBytes, + maxBufferingDuration); // Group temp file results by destinations again to gather all the results in the same window. // This is needed since we don't have shard idx associated with each temp file so have to rely @@ -1484,11 +1825,16 @@ private class FinalizeTempFileBundles PCollection>>, WriteFilesResult> { private final @Nullable PCollectionView numShardsView; private final Coder destinationCoder; + // When true, ignore the fixed shard count and let finalize number the files 0..K-1. + private final boolean assignShardNumbersInFinalize; private FinalizeTempFileBundles( - @Nullable PCollectionView numShardsView, Coder destinationCoder) { + @Nullable PCollectionView numShardsView, + Coder destinationCoder, + boolean assignShardNumbersInFinalize) { this.numShardsView = numShardsView; this.destinationCoder = destinationCoder; + this.assignShardNumbersInFinalize = assignShardNumbersInFinalize; } @Override @@ -1519,7 +1865,9 @@ private class FinalizeFn public void process(ProcessContext c) throws Exception { getDynamicDestinations().setSideInputAccessorFromProcessContext(c); @Nullable Integer fixedNumShards; - if (numShardsView != null) { + if (assignShardNumbersInFinalize) { + fixedNumShards = null; + } else if (numShardsView != null) { fixedNumShards = c.sideInput(numShardsView); } else if (getNumShardsProvider() != null) { fixedNumShards = getNumShardsProvider().get(); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java index 7a0dfb4822b3..eab7195d0918 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/FileIOTest.java @@ -723,6 +723,13 @@ public void testWriteUnboundedWithCustomBatchSizeBytes() throws IOException { assertTrue(metadata.size() >= 2); } + @Test + public void testWriteWithDirectWriteByteThreshold() { + FileIO.Write write = + FileIO.write().to("/tmp/x").via(TextIO.sink()).withDirectWriteByteThreshold(1000); + assertEquals(Long.valueOf(1000), write.getDirectWriteByteThreshold()); + } + static void checkFileContents(File rootDir, String prefix, List inputs) throws IOException { List outputFiles = Lists.newArrayList(); diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java index f58906e464be..50b1672c6453 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TextIOWriteTest.java @@ -782,4 +782,10 @@ public void testSink() throws Exception { assertEquals(Arrays.asList("header", "a", "b", "c", "footer"), readLinesFromFile(f)); } + + @Test + public void testWriteWithDirectWriteByteThreshold() { + TextIO.Write write = TextIO.write().to("/tmp/x").withDirectWriteByteThreshold(1000); + assertEquals(Long.valueOf(1000), write.inner.getDirectWriteByteThreshold()); + } } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java index 09f3f04d5b09..40d6efb4a671 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/WriteFilesTest.java @@ -54,6 +54,7 @@ import java.util.regex.Matcher; import java.util.regex.Pattern; import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.io.DefaultFilenamePolicy.Params; import org.apache.beam.sdk.io.FileBasedSink.DynamicDestinations; @@ -64,6 +65,10 @@ 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.metrics.MetricNameFilter; +import org.apache.beam.sdk.metrics.MetricQueryResults; +import org.apache.beam.sdk.metrics.MetricResult; +import org.apache.beam.sdk.metrics.MetricsFilter; import org.apache.beam.sdk.options.Description; import org.apache.beam.sdk.options.PipelineOptionsFactoryTest.TestPipelineOptions; import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider; @@ -89,10 +94,12 @@ import org.apache.beam.sdk.transforms.display.DisplayData; import org.apache.beam.sdk.transforms.errorhandling.ErrorHandler.BadRecordErrorHandler; import org.apache.beam.sdk.transforms.errorhandling.ErrorHandlingTestUtils.ErrorSinkTransform; +import org.apache.beam.sdk.transforms.windowing.AfterPane; 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.Repeatedly; import org.apache.beam.sdk.transforms.windowing.Sessions; import org.apache.beam.sdk.transforms.windowing.Window; import org.apache.beam.sdk.util.MimeTypes; @@ -106,6 +113,7 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; import org.apache.commons.compress.utils.Sets; +import org.checkerframework.checker.nullness.qual.Nullable; import org.hamcrest.Matchers; import org.joda.time.Duration; import org.joda.time.Instant; @@ -422,6 +430,250 @@ public void testWithRunnerDeterminedShardingTestStream() throws IOException { getBaseOutputFilename(), elements, Optional.absent(), !write.getWindowedWrites(), null); } + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteAllBundlesLarge() throws IOException { + runDirectWriteTestStream(1L); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteAllBundlesSmall() throws IOException { + runDirectWriteTestStream(Long.MAX_VALUE); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteUnshardedMixedBundles() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + // Two 10-element bundles (~100 bytes each), one 8-element bundle (~88 bytes), and one + // 2-element bundle (~22 bytes, < 50) so both direct and auto-sharded paths run. + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 28), String.class)) + .addElements(elements.get(28), elements.get(29)) + .advanceWatermarkToInfinity(); + + WriteFiles write = + WriteFiles.to(makeSimpleSink()) + .withWindowedWrites() + .withRunnerDeterminedSharding() + .withDirectWriteByteThreshold(50L); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertThat(counterValue(result, "directWriteBundles"), Matchers.greaterThan(0L)); + assertThat(counterValue(result, "groupedWriteBundles"), Matchers.greaterThan(0L)); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteNonDefaultTriggerLargeBundles() throws IOException { + runDirectWriteNonDefaultTrigger(1L); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteNonDefaultTriggerSmallBundles() throws IOException { + runDirectWriteNonDefaultTrigger(Long.MAX_VALUE); + } + + private void runDirectWriteNonDefaultTrigger(long thresholdBytes) throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + FilenamePolicy paneAwarePolicy = + new PaneAwarePerWindowFiles( + getBaseOutputDirectory().resolve("file", StandardResolveOptions.RESOLVE_FILE), + "simple"); + + WriteFiles write = + WriteFiles.to(SimpleSink.makeSimpleSink(getBaseOutputDirectory(), paneAwarePolicy)) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(thresholdBytes); + p.apply(testStream) + .apply( + Window.into(FixedWindows.of(Duration.standardSeconds(10))) + .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(5))) + .discardingFiredPanes() + .withAllowedLateness(Duration.standardMinutes(1))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertDirectWriteMetrics(result, thresholdBytes); + } + + private static class SideInputDestinations extends DynamicDestinations { + private final ResourceId baseOutputDirectory; + private final PCollectionView sideInputView; + + SideInputDestinations(ResourceId baseOutputDirectory, PCollectionView sideInputView) { + this.baseOutputDirectory = baseOutputDirectory; + this.sideInputView = sideInputView; + } + + @Override + public List> getSideInputs() { + return Collections.singletonList(sideInputView); + } + + @Override + public String formatRecord(String record) { + return record; + } + + @Override + public String getDestination(String element) { + return sideInput(sideInputView); + } + + @Override + public String getDefaultDestination() { + return "file"; + } + + @Override + public FilenamePolicy getFilenamePolicy(String destination) { + return new PerWindowFiles( + baseOutputDirectory.resolve(destination, StandardResolveOptions.RESOLVE_FILE), "simple"); + } + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteDisabledWithSideInputs() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + PCollectionView sideInputView = + p.apply("SideInput", Create.of("file")).apply(View.asSingleton()); + SideInputDestinations destinations = + new SideInputDestinations(getBaseOutputDirectory(), sideInputView); + SimpleSink sink = + new SimpleSink<>(getBaseOutputDirectory(), destinations, Compression.UNCOMPRESSED); + + WriteFiles write = + WriteFiles.to(sink).withWindowedWrites().withNumShards(3).withDirectWriteByteThreshold(1L); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertEquals(0L, counterValue(result, "directWriteBundles")); + } + + private void runDirectWriteTestStream(long thresholdBytes) throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + Instant startInstant = new Instant(0L); + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(startInstant) + .addElements(elements.get(0), Iterables.toArray(elements.subList(1, 10), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(5))) + .addElements( + elements.get(10), Iterables.toArray(elements.subList(11, 20), String.class)) + .advanceWatermarkTo(startInstant.plus(Duration.standardSeconds(10))) + .addElements( + elements.get(20), Iterables.toArray(elements.subList(21, 30), String.class)) + .advanceWatermarkToInfinity(); + + WriteFiles write = + WriteFiles.to(makeSimpleSink()) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(thresholdBytes); + p.apply(testStream) + .apply(Window.into(FixedWindows.of(Duration.standardSeconds(10)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + // Shard count is not fixed on the direct-write path, so don't check it. + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertDirectWriteMetrics(result, thresholdBytes); + } + + private static void assertDirectWriteMetrics(PipelineResult result, long thresholdBytes) { + long directBundles = counterValue(result, "directWriteBundles"); + long groupedBundles = counterValue(result, "groupedWriteBundles"); + if (thresholdBytes == 1L) { + assertThat(directBundles, Matchers.greaterThan(0L)); + } else if (thresholdBytes == Long.MAX_VALUE) { + assertEquals(0L, directBundles); + assertThat(groupedBundles, Matchers.greaterThan(0L)); + } + } + + private static long counterValue(PipelineResult result, String name) { + MetricQueryResults metrics = + result + .metrics() + .queryMetrics( + MetricsFilter.builder() + .addNameFilter(MetricNameFilter.named(WriteFiles.class, name)) + .build()); + long total = 0; + for (MetricResult counter : metrics.getCounters()) { + total += counter.getAttempted(); + } + return total; + } + /** Test a WriteFiles transform with an empty PCollection. */ @Test @Category(NeedsRunner.class) @@ -1109,6 +1361,50 @@ public ResourceId unwindowedFilename( } } + private static class PaneAwarePerWindowFiles extends FilenamePolicy { + private static final DateTimeFormatter FORMATTER = ISODateTimeFormat.hourMinuteSecondMillis(); + private final ResourceId baseFilename; + private final String suffix; + + public PaneAwarePerWindowFiles(ResourceId baseFilename, String suffix) { + this.baseFilename = baseFilename; + this.suffix = suffix; + } + + @Override + public ResourceId windowedFilename( + int shardNumber, + int numShards, + BoundedWindow window, + PaneInfo paneInfo, + OutputFileHints outputFileHints) { + DecimalFormat df = new DecimalFormat("0000"); + IntervalWindow intervalWindow = (IntervalWindow) window; + String prefix = + baseFilename.isDirectory() ? "" : firstNonNull(baseFilename.getFilename(), ""); + String filename = + String.format( + "%s%s-%s-pane-%d-%s-of-%s%s%s", + prefix, + FORMATTER.print(intervalWindow.start()), + FORMATTER.print(intervalWindow.end()), + paneInfo.getIndex(), + df.format(shardNumber), + df.format(numShards), + outputFileHints.getSuggestedFilenameSuffix(), + suffix); + return baseFilename + .getCurrentDirectory() + .resolve(filename, StandardResolveOptions.RESOLVE_FILE); + } + + @Override + public @Nullable ResourceId unwindowedFilename( + int shardNumber, int numShards, OutputFileHints outputFileHints) { + throw new UnsupportedOperationException(); + } + } + /** * Same as {@link #runShardedWrite(List, PTransform, String, WriteFiles, BiFunction, boolean)} but * without shard content check. This means content will be checked only globally, that shards