diff --git a/CHANGES.md b/CHANGES.md index 1087db89cfa2..c1ec83bb5b87 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -178,6 +178,7 @@ * (Python) Added typed decision questions and the `EvaluateDecisions` transform, with an optional Jev model adapter ([#40220](https://github.com/apache/beam/pull/40220)). * `QdrantWriteConfig` can now optionally store chunk content in the point payload via `content_payload_key`, so retrieved points include the chunk text (Python) ([#40476](https://github.com/apache/beam/issues/40476)). * 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`). 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`) 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)). * (Python) Added `AwsSecret` (`AwsSecretsManager` / `type:awssecret`) to retrieve secrets from AWS Secrets Manager, configured from `AwsOptions` (`--aws_*`) with optional per-secret connection overrides ([#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 796aaa960a4c..485d26915808 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 @@ -181,6 +181,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; @@ -2644,6 +2645,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) { @@ -2683,6 +2685,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( @@ -2692,11 +2699,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 2246f8b3ee6d..b30f1520ca79 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 @@ -2189,9 +2189,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 ae0746294077..94f79893a55f 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; @@ -62,6 +64,7 @@ import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Reify; import org.apache.beam.sdk.transforms.Reshuffle; +import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.SimpleFunction; import org.apache.beam.sdk.transforms.Values; import org.apache.beam.sdk.transforms.View; @@ -77,6 +80,7 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindows; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.transforms.windowing.Window; +import org.apache.beam.sdk.util.BundleLifter; import org.apache.beam.sdk.util.CoderUtils; import org.apache.beam.sdk.util.MoreFutures; import org.apache.beam.sdk.values.KV; @@ -210,6 +214,13 @@ public static WriteFiles> getSideInputs(); public abstract @Nullable ShardingFunction getShardingFunction(); @@ -251,6 +262,9 @@ abstract Builder setBatchSizeBytes( abstract Builder setBatchMaxBufferingDuration( @Nullable Duration batchMaxBufferingDuration); + abstract Builder setDirectWriteByteThreshold( + @Nullable Long directWriteByteThreshold); + abstract Builder setSideInputs( List> sideInputs); @@ -441,6 +455,26 @@ 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. + * + *

Records are buffered per window, pane and destination until their encoded size reaches + * {@code thresholdBytes}; those that reach it are written straight to temporary files. Smaller + * groups 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). + * + *

Direct writes are only used for unbounded inputs with non-merging windows whose dynamic + * destinations don't use side inputs. Otherwise the threshold is ignored, a warning is logged, + * and the grouped write path is used. + */ + 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(); @@ -579,6 +613,39 @@ public WriteFilesResult expand(PCollection input) { (getComputeNumShards() == null) ? null : input.apply(getComputeNumShards()); boolean fixedSharding = getComputeNumShards() != null || getNumShardsProvider() != null; + Long directWriteThreshold = getDirectWriteByteThreshold(); + if (directWriteThreshold != null) { + List disabledReasons = new ArrayList<>(); + if (input.isBounded() == IsBounded.BOUNDED) { + disabledReasons.add("the input is bounded"); + } + if (!getSideInputs().isEmpty()) { + disabledReasons.add("dynamic destinations use side inputs"); + } + // Merging windows are only final after the GroupByKey, so direct writes would use + // unmerged windows. + if (input.getWindowingStrategy().needsMerge()) { + disabledReasons.add("the windowing strategy uses merging windows"); + } + if (disabledReasons.isEmpty()) { + // 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)); + } + LOG.warn( + "Direct write is disabled (directWriteByteThreshold={}) because {}. Falling back to the" + + " grouped write path.", + directWriteThreshold, + String.join(" and ", disabledReasons)); + } PCollection>> tempFileResults; if (fixedSharding) { tempFileResults = @@ -603,14 +670,33 @@ public WriteFilesResult expand(PCollection input) { new WriteAutoShardedBundlesToTempFiles( destinationCoder, fileResultCoder, - getBatchSize(), - getBatchSizeBytes(), - getBatchMaxBufferingDuration())); + getBatchSizeOrDefault(), + getBatchSizeBytesOrDefault(), + getBatchMaxBufferingDurationOrDefault())); } } return tempFileResults.apply( - "FinalizeTempFileBundles", new FinalizeTempFileBundles(numShardsView, destinationCoder)); + "FinalizeTempFileBundles", + new FinalizeTempFileBundles( + numShardsView, destinationCoder, /* assignShardNumbersInFinalize= */ false)); + } + + private int getBatchSizeOrDefault() { + Integer batchSize = getBatchSize(); + return batchSize != null ? batchSize : FILE_TRIGGERING_RECORD_COUNT; + } + + private int getBatchSizeBytesOrDefault() { + Integer batchSizeBytes = getBatchSizeBytes(); + return batchSizeBytes != null ? batchSizeBytes : FILE_TRIGGERING_BYTE_COUNT; + } + + private Duration getBatchMaxBufferingDurationOrDefault() { + Duration maxBufferingDuration = getBatchMaxBufferingDuration(); + return maxBufferingDuration != null + ? maxBufferingDuration + : FILE_TRIGGERING_RECORD_BUFFERING_DURATION; } @Override @@ -813,7 +899,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. @@ -831,10 +916,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. @@ -861,29 +946,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 = @@ -923,30 +985,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 { @@ -1114,6 +1202,318 @@ 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 smallBundlesTag = new TupleTag<>("smallBundles"); + TupleTag largeBundlesTag = new TupleTag<>("largeBundles"); + TupleTag> directFilesTag = new TupleTag<>("directFiles"); + TupleTag spilledRecordsTag = new TupleTag<>("spilledRecords"); + + // Evaluate the threshold per window, pane and destination, i.e. per output file. + PCollectionTuple lifted = + input.apply( + "LiftLargeBundles", + BundleLifter.of(smallBundlesTag, largeBundlesTag, thresholdBytes, input.getCoder()) + .withKeyFn(new DestinationKeyFn<>(getDynamicDestinations()))); + + PCollectionTuple routed = + lifted + .get(largeBundlesTag) + .setCoder(input.getCoder()) + .apply( + "DirectWriteLargeBundles", + ParDo.of(new DirectWriteRecordsFn(spilledRecordsTag, input.getCoder())) + .withSideInputs(getSideInputs()) + .withOutputTags( + directFilesTag, + TupleTagList.of(ImmutableList.of(spilledRecordsTag, BAD_RECORD_TAG)))); + addErrorCollection(routed); + + PCollection smallBundles = + PCollectionList.of(lifted.get(smallBundlesTag).setCoder(input.getCoder())) + .and(routed.get(spilledRecordsTag).setCoder(input.getCoder())) + .apply("FlattenSmallBundles", Flatten.pCollections()) + .apply("CountGroupedBundles", ParDo.of(new CountGroupedBundlesFn<>())) + .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, + getBatchSizeOrDefault(), + getBatchSizeBytesOrDefault(), + getBatchMaxBufferingDurationOrDefault()); + } + + // 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. Every file is stamped with this + // firing's pane so files from different firings of the same window get distinct names. + PCollection> allTempFiles = + PCollectionList.of(directFiles) + .and(groupedFiles) + .apply("FlattenTempFiles", Flatten.pCollections()) + .setCoder(fileResultCoder); + return groupTempFilesByDestination( + allTempFiles, + destinationCoder, + fileResultCoder, + "GroupTempFilesByDestination", + /* stampFiringPane= */ true); + } + } + + /** + * Key function for {@link BundleLifter}: the record's destination, so that together with the + * window and pane the threshold is evaluated per output file. Returns {@code null} if the + * destination cannot be computed, which sends the record to the grouped path where the usual + * error handling applies. + */ + private static class DestinationKeyFn + implements SerializableFunction { + private final DynamicDestinations dynamicDestinations; + + DestinationKeyFn(DynamicDestinations dynamicDestinations) { + this.dynamicDestinations = dynamicDestinations; + } + + @Override + public @Nullable Object apply(UserT record) { + try { + // Wrapped so that a null destination (e.g. Void) is still a valid key. + return java.util.Optional.ofNullable(dynamicDestinations.getDestination(record)); + } catch (Exception e) { + return null; + } + } + } + + /** Passes records through, counting the bundles that take the grouped path. */ + private static class CountGroupedBundlesFn extends DoFn { + private final Counter groupedBundles = Metrics.counter(WriteFiles.class, "groupedWriteBundles"); + private transient boolean counted; + + @StartBundle + public void startBundle() { + counted = false; + } + + @ProcessElement + public void processElement(@Element T element, OutputReceiver receiver) { + if (!counted) { + groupedBundles.inc(); + counted = true; + } + receiver.output(element); + } + } + + /** + * Writes the records lifted by {@link BundleLifter} directly to temporary files, one writer per + * window, pane and destination. Records that would exceed the per-bundle writer limit are emitted + * to {@code spilledRecordsTag} for the grouped path. + */ + private class DirectWriteRecordsFn extends DoFn> { + private final TupleTag spilledRecordsTag; + private final Coder inputCoder; + + private final Counter directBundles = Metrics.counter(WriteFiles.class, "directWriteBundles"); + + // Per-bundle state, reset in startBundle. + private transient Map, Writer> writers; + + DirectWriteRecordsFn(TupleTag spilledRecordsTag, Coder inputCoder) { + this.spilledRecordsTag = spilledRecordsTag; + this.inputCoder = inputCoder; + } + + @StartBundle + public void startBundle() { + writers = new LinkedHashMap<>(); + } + + @ProcessElement + public void processElement( + ProcessContext c, BoundedWindow window, MultiOutputReceiver outputReceiver) + throws Exception { + getDynamicDestinations().setSideInputAccessorFromProcessContext(c); + MaybeDestination maybeDestination = + getDestinationWithErrorHandling(c.element(), outputReceiver, inputCoder); + if (!maybeDestination.isValid) { + return; + } + DestinationT destination = maybeDestination.destination; + WriterKey key = new WriterKey<>(window, c.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. + outputReceiver.get(spilledRecordsTag).output(c.element()); + return; + } + writer = openAndRegisterWriter(writers, key); + } + OutputT formattedRecord = + formatRecordWithErrorHandling(c.element(), outputReceiver, inputCoder); + if (formattedRecord == null) { + return; + } + writeOrClose(writer, formattedRecord); + } + + @FinishBundle + public void finishBundle(FinishBundleContext c) throws Exception { + if (!writers.isEmpty()) { + directBundles.inc(); + } + closeAllWritersAndOutputFileResults(writers, c); + writers.clear(); + } + } + + 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; @@ -1125,133 +1525,94 @@ private class WriteAutoShardedBundlesToTempFiles private WriteAutoShardedBundlesToTempFiles( Coder destinationCoder, Coder> fileResultCoder, - Integer batchSize, - Integer batchSizeBytes, + int batchSize, + int batchSizeBytes, Duration maxBufferingDuration) { this.destinationCoder = destinationCoder; this.fileResultCoder = fileResultCoder; - this.batchSize = batchSize != null ? batchSize : FILE_TRIGGERING_RECORD_COUNT; - this.batchSizeBytes = batchSizeBytes != null ? batchSizeBytes : FILE_TRIGGERING_BYTE_COUNT; - this.maxBufferingDuration = - maxBufferingDuration != null - ? maxBufferingDuration - : FILE_TRIGGERING_RECORD_BUFFERING_DURATION; + this.batchSize = batchSize; + this.batchSizeBytes = batchSizeBytes; + this.maxBufferingDuration = maxBufferingDuration; } @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 // on the indexing within a bundle. - return tempFiles - .apply( - "KeyedByDestination", - WithKeys.of( - new SimpleFunction, DestinationT>() { - @Override - public DestinationT apply(FileResult input) { - return input.getDestination(); - } - })) - .setCoder(KvCoder.of(destinationCoder, fileResultCoder)) - .apply(GroupByKey.create()) - .apply( - "ExtractValuesToList", - ParDo.of( - new DoFn< - KV>>, - List>>() { - @ProcessElement - public void processElement( - @Element KV>> element, - ProcessContext c) { - List> result = new ArrayList<>(); - for (FileResult e : element.getValue()) { - result.add(e); - } - c.output(result); - } - })) - .setCoder(ListCoder.of(fileResultCoder)); + // "GroupByKey" matches the default name this step has always had (update compatibility). + return groupTempFilesByDestination( + tempFiles, destinationCoder, fileResultCoder, "GroupByKey", /* stampFiringPane= */ false); + } + } + + /** + * Groups temp file results by destination into one list per destination and window firing. Shared + * by the auto-sharded write and the direct-write path. Produces the steps {@code + * KeyedByDestination}, {@code groupByKeyName} and {@code ExtractValuesToList}. + * + * @param stampFiringPane if true, each file is re-stamped with the firing's pane and an unknown + * shard number so shard numbers can be assigned in finalize across all files of the firing. + */ + private PCollection>> groupTempFilesByDestination( + PCollection> tempFiles, + Coder destinationCoder, + Coder> fileResultCoder, + String groupByKeyName, + boolean stampFiringPane) { + return tempFiles + .apply( + "KeyedByDestination", + WithKeys.of( + new SimpleFunction, DestinationT>() { + @Override + public DestinationT apply(FileResult input) { + return input.getDestination(); + } + })) + .setCoder(KvCoder.of(destinationCoder, fileResultCoder)) + .apply(groupByKeyName, GroupByKey.create()) + .apply("ExtractValuesToList", ParDo.of(new ExtractValuesToListFn<>(stampFiringPane))) + .setCoder(ListCoder.of(fileResultCoder)); + } + + private static class ExtractValuesToListFn + extends DoFn< + KV>>, List>> { + private final boolean stampFiringPane; + + ExtractValuesToListFn(boolean stampFiringPane) { + this.stampFiringPane = stampFiringPane; + } + + @ProcessElement + public void processElement(ProcessContext c) { + List> result = new ArrayList<>(); + for (FileResult e : c.element().getValue()) { + if (stampFiringPane) { + // Stamp every file with this firing's pane so files from different firings of the + // same window get distinct final names. + result.add( + new FileResult<>( + e.getTempFilename(), + UNKNOWN_SHARDNUM, + e.getWindow(), + c.pane(), + e.getDestination())); + } else { + result.add(e); + } + } + c.output(result); } } @@ -1480,11 +1841,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 @@ -1515,7 +1881,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/main/java/org/apache/beam/sdk/util/BundleLifter.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java index 539ef147f173..3de2936673e1 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/BundleLifter.java @@ -17,30 +17,57 @@ */ package org.apache.beam.sdk.util; -import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver; -import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.SerializableFunction; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.util.common.ElementByteSizeObserver; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.TupleTag; import org.apache.beam.sdk.values.TupleTagList; import org.checkerframework.checker.nullness.qual.MonotonicNonNull; import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Duration; +import org.joda.time.Instant; import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** - * A PTransform that buffers elements and outputs them to one of two TupleTags based on the total - * size of the bundle in finish_bundle. + * A PTransform that routes elements to one of two TupleTags depending on how many bytes of them + * arrive together within a bundle. + * + *

Elements are buffered per key. Once the buffered elements of a key reach the threshold, they + * are emitted to the large tag, and all later elements of that key in the same bundle are emitted + * to the large tag directly without buffering. Elements of keys that never reach the threshold are + * emitted to the small tag at the end of the bundle. + * + *

Without a key function, all elements of a bundle share a single key, so the threshold applies + * to the whole bundle. With {@link #withKeyFn}, the key is the element's window, pane and the value + * returned by the key function. Elements for which the key function returns {@code null} are + * emitted to the small tag immediately. + * + *

To bound memory, at most twice the threshold is buffered in total. When that is exceeded, the + * keys with the fewest buffered bytes are emitted to the small tag until at most the threshold + * remains buffered. + * + *

Every element keeps its own timestamp and window. * *

For internal use only; no backwards-compatibility guarantees. * @@ -57,89 +84,275 @@ public class BundleLifter extends PTransform, PCollectionTuple final TupleTag smallBatchTag; final TupleTag largeBatchTag; - final int threshold; + final long threshold; final SerializableFunction elementSizer; + final @Nullable SerializableFunction keyFn; /** - * A DoFn that buffers elements within a bundle and outputs them to different tags in - * finish_bundle based on the total bundle size. + * A DoFn that buffers elements per key within a bundle and outputs them to the large tag once a + * key reaches the threshold, or to the small tag otherwise. * * @param The type of elements being processed. */ static class BundleLiftDoFn extends DoFn { private static final Logger LOG = LoggerFactory.getLogger(BundleLiftDoFn.class); + /** The key shared by all elements when no key function is set. */ + private static final Object UNKEYED = new Object(); + final TupleTag smallBatchTag; final TupleTag largeBatchTag; - final int threshold; + final long threshold; + final long maxBufferedBytes; final SerializableFunction elementSizer; + final @Nullable SerializableFunction keyFn; - private transient @MonotonicNonNull List buffer; - private transient long bundleSizeBytes; - private transient @Nullable MultiOutputReceiver receiver; + // Per-bundle state, reset in startBundle. + private transient @MonotonicNonNull Map> buffers; + private transient @MonotonicNonNull Set liftedKeys; + // Reset in startBundle; not transient so SpotBugs doesn't flag it as unrestored state. + private long bufferedBytes; BundleLiftDoFn( TupleTag smallBatchTag, TupleTag largeBatchTag, - int threshold, + long threshold, SerializableFunction elementSizer) { + this(smallBatchTag, largeBatchTag, threshold, elementSizer, null); + } + + BundleLiftDoFn( + TupleTag smallBatchTag, + TupleTag largeBatchTag, + long threshold, + SerializableFunction elementSizer, + @Nullable SerializableFunction keyFn) { this.smallBatchTag = smallBatchTag; this.largeBatchTag = largeBatchTag; this.threshold = threshold; + this.maxBufferedBytes = threshold > Long.MAX_VALUE / 2 ? Long.MAX_VALUE : 2 * threshold; this.elementSizer = elementSizer; + this.keyFn = keyFn; } @StartBundle public void startBundle() { - buffer = new ArrayList<>(); - receiver = null; - bundleSizeBytes = 0L; + buffers = new LinkedHashMap<>(); + liftedKeys = new HashSet<>(); + bufferedBytes = 0L; } @ProcessElement - public void processElement(@Element T element, MultiOutputReceiver mor) { - if (receiver == null) { - receiver = mor; + public void processElement( + @Element T element, + @Timestamp Instant timestamp, + BoundedWindow window, + PaneInfo pane, + MultiOutputReceiver receiver) { + Map> buffers = + checkStateNotNull(this.buffers, "Buffers should be set by startBundle."); + Set liftedKeys = + checkStateNotNull(this.liftedKeys, "Lifted keys should be set by startBundle."); + + Object key; + if (keyFn == null) { + key = UNKEYED; + } else { + @Nullable Object userKey = keyFn.apply(element); + if (userKey == null) { + receiver.get(smallBatchTag).output(element); + return; + } + key = new WindowedKey(window, pane, userKey); + } + + if (liftedKeys.contains(key)) { + // This key already reached the threshold in this bundle; stream it through. + receiver.get(largeBatchTag).output(element); + return; + } + + KeyBuffer buffer = buffers.get(key); + if (buffer == null) { + buffer = new KeyBuffer<>(); + buffers.put(key, buffer); + } + long elementBytes = elementSizer.apply(element); + buffer.add(new BufferedElement<>(element, timestamp, window, pane), elementBytes); + bufferedBytes += elementBytes; + + if (buffer.sizeBytes >= threshold) { + buffers.remove(key); + bufferedBytes -= buffer.sizeBytes; + liftedKeys.add(key); + LOG.debug( + "Key reached {} estimated bytes; emitting {} buffered elements to tag: '{}'", + buffer.sizeBytes, + buffer.elements.size(), + largeBatchTag.getId()); + emit(buffer, receiver.get(largeBatchTag)); + return; + } + + if (bufferedBytes > maxBufferedBytes) { + evictSmallestKeys(buffers, receiver.get(smallBatchTag)); } - checkArgumentNotNull(buffer, "Buffer should be set by startBundle."); - buffer.add(element); - bundleSizeBytes += elementSizer.apply(element); } @FinishBundle - public void finishBundle() { - checkArgumentNotNull(buffer, "Buffer should be set by startBundle."); - if (buffer.isEmpty()) { - return; + public void finishBundle(FinishBundleContext context) { + Map> buffers = + checkStateNotNull(this.buffers, "Buffers should be set by startBundle."); + for (KeyBuffer buffer : buffers.values()) { + LOG.debug( + "Emitting {} elements of {} estimated bytes to tag: '{}'", + buffer.elements.size(), + buffer.sizeBytes, + smallBatchTag.getId()); + for (BufferedElement e : buffer.elements) { + context.output(smallBatchTag, e.value, e.timestamp, e.window); + } } + buffers.clear(); + checkStateNotNull(this.liftedKeys).clear(); + bufferedBytes = 0L; + } - // Select the target tag based on the bundle size - TupleTag targetTag; - targetTag = (bundleSizeBytes < threshold) ? smallBatchTag : largeBatchTag; - LOG.debug( - "Emitting {} elements of {} estimated bytes to tag: '{}'", - buffer.size(), - bundleSizeBytes, - targetTag.getId()); + /** + * Buffered elements are re-emitted from {@link ProcessElement} with their own timestamps, which + * may be earlier than the current element's. This is safe because they belong to the same + * bundle, so the input watermark is still held at or before them. + */ + @Override + public Duration getAllowedTimestampSkew() { + return Duration.millis(Long.MAX_VALUE); + } - checkArgumentNotNull(receiver, "Receiver should be set by startBundle."); - OutputReceiver taggedOutput = receiver.get(targetTag); + /** + * Emits the smallest keys to the small output until at most {@code threshold} bytes remain + * buffered. Each call frees more than {@code threshold} bytes, so at least that many new bytes + * must be buffered before it runs again, which keeps the sort's cost amortized. + */ + private void evictSmallestKeys(Map> buffers, OutputReceiver small) { + List>> bySize = new ArrayList<>(buffers.entrySet()); + bySize.sort(Comparator.comparingLong(entry -> entry.getValue().sizeBytes)); + for (Map.Entry> entry : bySize) { + if (bufferedBytes <= threshold) { + break; + } + KeyBuffer buffer = entry.getValue(); + buffers.remove(entry.getKey()); + bufferedBytes -= buffer.sizeBytes; + emit(buffer, small); + } + } - for (T element : buffer) { - taggedOutput.output(element); + private static void emit(KeyBuffer buffer, OutputReceiver receiver) { + for (BufferedElement e : buffer.elements) { + receiver.outputWindowedValue( + e.value, e.timestamp, Collections.singletonList(e.window), e.pane); } } } - private BundleLifter(TupleTag smallBatchTag, TupleTag largeBatchTag, int threshold) { - this(smallBatchTag, largeBatchTag, threshold, x -> 1); + /** The buffered elements of one key, in arrival order. */ + private static class KeyBuffer { + private final List> elements = new ArrayList<>(); + private long sizeBytes; + + void add(BufferedElement element, long elementBytes) { + elements.add(element); + sizeBytes += elementBytes; + } + } + + /** A buffered element with the metadata needed to re-emit it. */ + private static class BufferedElement { + private final T value; + private final Instant timestamp; + private final BoundedWindow window; + private final PaneInfo pane; + + BufferedElement(T value, Instant timestamp, BoundedWindow window, PaneInfo pane) { + this.value = value; + this.timestamp = timestamp; + this.window = window; + this.pane = pane; + } + } + + /** A user key scoped to a window and pane. */ + private static class WindowedKey { + private final BoundedWindow window; + private final PaneInfo pane; + private final Object key; + + WindowedKey(BoundedWindow window, PaneInfo pane, Object key) { + this.window = window; + this.pane = pane; + this.key = key; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof WindowedKey)) { + return false; + } + WindowedKey other = (WindowedKey) o; + return window.equals(other.window) && pane.equals(other.pane) && key.equals(other.key); + } + + @Override + public int hashCode() { + return Objects.hash(window, pane, key); + } + } + + /** Sizes elements by their encoded size under the given coder. */ + private static class CoderSizer implements SerializableFunction { + private final Coder coder; + + CoderSizer(Coder coder) { + this.coder = coder; + } + + @Override + public Integer apply(T element) { + try { + long bytes; + if (coder.isRegisterByteSizeObserverCheap(element)) { + SizeObserver observer = new SizeObserver(); + coder.registerByteSizeObserver(element, observer); + observer.advance(); + bytes = observer.bytes; + } else { + bytes = CoderUtils.encodeToByteArray(coder, element).length; + } + return (int) Math.min(bytes, Integer.MAX_VALUE); + } catch (Exception e) { + throw new RuntimeException("Failed to estimate the encoded size of an element", e); + } + } + } + + private static class SizeObserver extends ElementByteSizeObserver { + private long bytes; + + @Override + protected void reportElementSize(long elementByteSize) { + bytes = elementByteSize; + } } private BundleLifter( TupleTag smallBatchTag, TupleTag largeBatchTag, - int threshold, - SerializableFunction elementSizer) { + long threshold, + SerializableFunction elementSizer, + @Nullable SerializableFunction keyFn) { if (smallBatchTag == null || largeBatchTag == null) { throw new IllegalArgumentException("smallBatchTag and largeBatchTag must not be null"); } @@ -154,19 +367,37 @@ private BundleLifter( this.largeBatchTag = largeBatchTag; this.threshold = threshold; this.elementSizer = elementSizer; + this.keyFn = keyFn; } + /** Routes by element count: each element counts as one toward the threshold. */ public static BundleLifter of( - TupleTag smallBatchTag, TupleTag largeBatchTag, int threshold) { - return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold); + TupleTag smallBatchTag, TupleTag largeBatchTag, long threshold) { + return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, x -> 1, null); } + /** Routes by the sizes returned by {@code elementSizer}. */ public static BundleLifter of( TupleTag smallBatchTag, TupleTag largeBatchTag, - int threshold, + long threshold, SerializableFunction elementSizer) { - return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, elementSizer); + return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, elementSizer, null); + } + + /** Routes by each element's encoded size under {@code coder}. */ + public static BundleLifter of( + TupleTag smallBatchTag, TupleTag largeBatchTag, long threshold, Coder coder) { + return new BundleLifter<>( + smallBatchTag, largeBatchTag, threshold, new CoderSizer<>(coder), null); + } + + /** + * Returns a copy that evaluates the threshold separately for each window, pane and key returned + * by {@code keyFn}. Elements for which {@code keyFn} returns {@code null} go to the small tag. + */ + public BundleLifter withKeyFn(SerializableFunction keyFn) { + return new BundleLifter<>(smallBatchTag, largeBatchTag, threshold, elementSizer, keyFn); } @Override @@ -175,7 +406,7 @@ public PCollectionTuple expand(PCollection input) { return input.apply( "BundleLiftDoFn", - ParDo.of(new BundleLiftDoFn<>(smallBatchTag, largeBatchTag, threshold, elementSizer)) + ParDo.of(new BundleLiftDoFn<>(smallBatchTag, largeBatchTag, threshold, elementSizer, keyFn)) .withOutputTags(mainOutputTag, TupleTagList.of(smallBatchTag).and(largeBatchTag))); } } 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 c8d55a8724e7..319fd3fd33f4 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; @@ -82,6 +87,7 @@ 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.SerializableFunction; import org.apache.beam.sdk.transforms.SimpleFunction; import org.apache.beam.sdk.transforms.Top; import org.apache.beam.sdk.transforms.Values; @@ -89,10 +95,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 +114,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 +431,388 @@ 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")); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteDisabledWithMergingWindows() 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(); + + // Merging windows require an explicit shard count, and a 1-byte threshold would otherwise + // write every bundle directly. + WriteFiles write = + WriteFiles.to(makeSimpleSink()) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(1L); + p.apply(testStream) + .apply(Window.into(Sessions.withGapDuration(Duration.standardSeconds(1)))) + .apply(write) + .getPerDestinationOutputFilenames() + .apply(new VerifyFilesExist<>()); + PipelineResult result = p.run(); + + checkFileContents(getBaseOutputFilename(), elements, Optional.absent(), false, null); + assertEquals(0L, counterValue(result, "directWriteBundles")); + } + + /** Routes each element to a destination chosen by {@code destinationOf}, without side inputs. */ + private static class PerElementDestinations extends DynamicDestinations { + private final ResourceId baseOutputDirectory; + private final SerializableFunction destinationOf; + + PerElementDestinations( + ResourceId baseOutputDirectory, SerializableFunction destinationOf) { + this.baseOutputDirectory = baseOutputDirectory; + this.destinationOf = destinationOf; + } + + @Override + public String formatRecord(String record) { + return record; + } + + @Override + public String getDestination(String element) { + return destinationOf.apply(element); + } + + @Override + public String getDefaultDestination() { + return "file"; + } + + @Override + public FilenamePolicy getFilenamePolicy(String destination) { + return new PerWindowFiles( + baseOutputDirectory.resolve(destination, StandardResolveOptions.RESOLVE_FILE), "simple"); + } + } + + private PipelineResult runPerDestinationDirectWrite( + List elements, SerializableFunction destinationOf, long threshold) + throws IOException { + TestStream testStream = + TestStream.create(StringUtf8Coder.of()) + .advanceWatermarkTo(new Instant(0L)) + // A single bundle containing every element. + .addElements( + elements.get(0), + Iterables.toArray(elements.subList(1, elements.size()), String.class)) + .advanceWatermarkToInfinity(); + + SimpleSink sink = + new SimpleSink<>( + getBaseOutputDirectory(), + new PerElementDestinations(getBaseOutputDirectory(), destinationOf), + Compression.UNCOMPRESSED); + WriteFiles write = + WriteFiles.to(sink) + .withWindowedWrites() + .withNumShards(3) + .withDirectWriteByteThreshold(threshold); + 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); + return result; + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteThresholdIsPerDestination() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add("number: " + i); + } + // ~300 bytes in one bundle, but split over 10 destinations of ~30 bytes each, so no output + // file reaches the 50-byte threshold and everything takes the grouped path. + PipelineResult result = + runPerDestinationDirectWrite( + elements, e -> "file" + (Integer.parseInt(e.substring(8)) % 10), 50L); + + assertEquals(0L, counterValue(result, "directWriteBundles")); + assertThat(counterValue(result, "groupedWriteBundles"), Matchers.greaterThan(0L)); + } + + @Test + @Category({NeedsRunner.class, UsesUnboundedPCollections.class, UsesTestStream.class}) + public void testDirectWriteLargeAndSmallDestinationsInOneBundle() throws IOException { + List elements = Lists.newArrayList(); + for (int i = 0; i < 23; ++i) { + elements.add("number: " + i); + } + // 20 elements (~220 bytes) go to one destination and 3 (~33 bytes) to another. With a + // 100-byte threshold, the large destination is written directly and the small one is grouped. + PipelineResult result = + runPerDestinationDirectWrite( + elements, e -> Integer.parseInt(e.substring(8)) < 20 ? "fileLarge" : "fileSmall", 100L); + + assertThat(counterValue(result, "directWriteBundles"), Matchers.greaterThan(0L)); + assertThat(counterValue(result, "groupedWriteBundles"), Matchers.greaterThan(0L)); + } + + 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) @@ -1111,6 +1502,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 diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java index 0467f1383bcf..49772865b7ae 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/util/BundleLifterTest.java @@ -18,12 +18,21 @@ package org.apache.beam.sdk.util; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.hamcrest.Matchers.hasSize; +import java.util.ArrayList; +import java.util.List; +import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.windowing.IntervalWindow; import org.apache.beam.sdk.util.BundleLifter.BundleLiftDoFn; import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.joda.time.Instant; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -96,4 +105,155 @@ public void testLargeBundleWithSizer() throws Exception { assertThat( tester.peekOutputElements(STRING_LARGE), containsInAnyOrder("123", "456", "789", "0")); } + + @Test + public void testStreamsElementsThroughOnceThresholdIsReached() throws Exception { + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 2, x -> 1)); + + tester.startBundle(); + tester.processElement(1); + assertThat(tester.peekOutputElements(INTEGER_LARGE), empty()); + tester.processElement(2); + // The buffered elements are flushed as soon as the threshold is reached... + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(1, 2)); + tester.processElement(3); + // ...and later elements are output directly, before the bundle finishes. + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(1, 2, 3)); + tester.finishBundle(); + + assertThat(tester.peekOutputElements(INTEGER_SMALL), empty()); + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(1, 2, 3)); + } + + @Test + public void testKeyedThresholdIsEvaluatedPerKey() throws Exception { + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1, x -> x % 2)); + + tester.startBundle(); + for (int i = 0; i < 5; i++) { + tester.processElement(i); + } + tester.finishBundle(); + + // Even key has 3 elements and reaches the threshold; odd key has 2 and does not. + assertThat(tester.peekOutputElements(INTEGER_LARGE), containsInAnyOrder(0, 2, 4)); + assertThat(tester.peekOutputElements(INTEGER_SMALL), containsInAnyOrder(1, 3)); + } + + @Test + public void testKeyedThresholdIsEvaluatedPerWindow() throws Exception { + IntervalWindow w1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow w2 = new IntervalWindow(new Instant(10), new Instant(20)); + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1, x -> "key")); + + tester.startBundle(); + tester.processWindowedElement(1, new Instant(1), w1); + tester.processWindowedElement(2, new Instant(2), w1); + tester.processWindowedElement(3, new Instant(11), w2); + tester.processWindowedElement(4, new Instant(12), w2); + tester.finishBundle(); + + // 4 elements share the key, but each window only has 2. + assertThat(tester.peekOutputElements(INTEGER_LARGE), empty()); + assertThat(tester.peekOutputElements(INTEGER_SMALL), containsInAnyOrder(1, 2, 3, 4)); + } + + @Test + public void testUnkeyedThresholdIsEvaluatedPerBundle() throws Exception { + IntervalWindow w1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow w2 = new IntervalWindow(new Instant(10), new Instant(20)); + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1)); + + tester.startBundle(); + tester.processWindowedElement(1, new Instant(1), w1); + tester.processWindowedElement(2, new Instant(2), w1); + tester.processWindowedElement(3, new Instant(11), w2); + tester.processWindowedElement(4, new Instant(12), w2); + tester.finishBundle(); + + // Without a key function the whole bundle counts toward the threshold. + assertThat(tester.peekOutputElements(INTEGER_LARGE), containsInAnyOrder(1, 2, 3, 4)); + assertThat(tester.peekOutputElements(INTEGER_SMALL), empty()); + } + + @Test + public void testNullKeyGoesToSmallOutput() throws Exception { + DoFnTester tester = + DoFnTester.of( + new BundleLiftDoFn<>( + INTEGER_SMALL, INTEGER_LARGE, 1, x -> 1, x -> x < 0 ? null : "key")); + + tester.startBundle(); + tester.processElement(-1); + tester.processElement(5); + tester.finishBundle(); + + assertThat(tester.peekOutputElements(INTEGER_SMALL), contains(-1)); + assertThat(tester.peekOutputElements(INTEGER_LARGE), contains(5)); + } + + @Test + public void testMemoryCapEvictsSmallestKey() throws Exception { + // Threshold 4 caps total buffered bytes at 8. + DoFnTester tester = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 4, x -> 1, x -> x)); + + tester.startBundle(); + for (int i = 0; i < 9; i++) { + tester.processElement(i); + } + // The 9th distinct key pushes the total (9) over the cap (8), so the smallest keys are evicted + // in one pass until at most the threshold (4) remains buffered: 5 keys go out early. + assertThat(tester.peekOutputElements(INTEGER_SMALL), hasSize(5)); + tester.finishBundle(); + + assertThat( + tester.peekOutputElements(INTEGER_SMALL), containsInAnyOrder(0, 1, 2, 3, 4, 5, 6, 7, 8)); + assertThat(tester.peekOutputElements(INTEGER_LARGE), empty()); + } + + @Test + public void testPreservesEachElementsWindowAndTimestamp() throws Exception { + IntervalWindow w1 = new IntervalWindow(new Instant(0), new Instant(10)); + IntervalWindow w2 = new IntervalWindow(new Instant(10), new Instant(20)); + + DoFnTester large = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 2, x -> 1)); + large.startBundle(); + large.processWindowedElement(1, new Instant(5), w1); + large.processWindowedElement(2, new Instant(15), w2); + large.finishBundle(); + assertThat( + describe(large.getMutableOutput(INTEGER_LARGE)), + contains("1@" + new Instant(5) + " in " + w1, "2@" + new Instant(15) + " in " + w2)); + + DoFnTester small = + DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1)); + small.startBundle(); + small.processWindowedElement(1, new Instant(5), w1); + small.processWindowedElement(2, new Instant(15), w2); + small.finishBundle(); + assertThat( + describe(small.getMutableOutput(INTEGER_SMALL)), + contains("1@" + new Instant(5) + " in " + w1, "2@" + new Instant(15) + " in " + w2)); + } + + private static List describe(List> outputs) { + List described = new ArrayList<>(); + for (ValueInSingleWindow output : outputs) { + described.add(output.getValue() + "@" + output.getTimestamp() + " in " + output.getWindow()); + } + return described; + } + + @Test + public void testCoderBasedSizing() { + BundleLifter lifter = + BundleLifter.of(STRING_SMALL, STRING_LARGE, 10L, StringUtf8Coder.of()); + assertThat(lifter.elementSizer.apply("hello"), greaterThanOrEqualTo(5)); + } }