Add Error Handlers to File IO and related IOs (TextIO, AvroIO)

sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileIO.java

@@ -61,6 +61,8 @@
 import org.apache.beam.sdk.transforms.Watch.Growth.TerminationCondition;
 import org.apache.beam.sdk.transforms.display.DisplayData;
 import org.apache.beam.sdk.transforms.display.HasDisplayData;
+import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
+import org.apache.beam.sdk.transforms.errorhandling.ErrorHandler;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
@@ -236,6 +238,26 @@
  * destination-dependent: every window/pane for every destination will use the same number of shards
  * specified via {@link Write#withNumShards} or {@link Write#withSharding}.
  *
+ * <h3>Handling Errors</h3>
+ *
+ * <p>When using dynamic destinations, or when using a formatting function to format a record for
+ * writing, it's possible for an individual record to be malformed, causing an exception. By
+ * default, these exceptions are propagated to the runner, and are usually retried, though this
+ * depends on the runner. Alternately, these errors can be routed to another {@link PTransform} by
+ * using {@link Write#withBadRecordErrorHandler(ErrorHandler)}. The ErrorHandler is registered with
+ * the pipeline (see below). See {@link ErrorHandler} for more documentation. Of note, this error
+ * handling only handles errors related to specific records. It does not handle errors related to
+ * connectivity, authorization, etc. as those should be retried by the runner.
+ *
+ * <pre>{@code
+ * PCollection<> records = ...;
+ * PTransform<PCollection<BadRecord>,?> alternateSink = ...;
+ * try (BadRecordErrorHandler<?> handler = pipeline.registerBadRecordErrorHandler(alternateSink) {
+ *    records.apply("Write", FileIO.writeDynamic().otherConfigs()
+ *        .withBadRecordErrorHandler(handler));
+ * }
+ * }</pre>
+ *
  * <h3>Writing custom types to sinks</h3>
  *
  * <p>Normally, when writing a collection of a custom type using a {@link Sink} that takes a
@@ -1016,6 +1038,8 @@ public static FileNaming relativeFileNaming(
 
     abstract boolean getNoSpilling();
 
+    abstract @Nullable ErrorHandler<BadRecord, ?> getBadRecordErrorHandler();
+
     abstract Builder<DestinationT, UserT> toBuilder();
 
     @AutoValue.Builder
@@ -1062,6 +1086,9 @@ abstract Builder<DestinationT, UserT> setSharding(
 
       abstract Builder<DestinationT, UserT> setNoSpilling(boolean noSpilling);
 
+      abstract Builder<DestinationT, UserT> setBadRecordErrorHandler(
+          @Nullable ErrorHandler<BadRecord, ?> badRecordErrorHandler);
+
       abstract Write<DestinationT, UserT> build();
     }
 
@@ -1288,6 +1315,12 @@ public Write<DestinationT, UserT> withNoSpilling() {
       return toBuilder().setNoSpilling(true).build();
     }
 
+    /** See {@link WriteFiles#withBadRecordErrorHandler(ErrorHandler)}. */
+    public Write<DestinationT, UserT> withBadRecordErrorHandler(
+        ErrorHandler<BadRecord, ?> errorHandler) {
+      return toBuilder().setBadRecordErrorHandler(errorHandler).build();
+    }
+
     @VisibleForTesting
     Contextful<Fn<DestinationT, FileNaming>> resolveFileNamingFn() {
       if (getDynamic()) {
@@ -1391,6 +1424,9 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
       if (getNoSpilling()) {
         writeFiles = writeFiles.withNoSpilling();
       }
+      if (getBadRecordErrorHandler() != null) {
+        writeFiles = writeFiles.withBadRecordErrorHandler(getBadRecordErrorHandler());
+      }
       return input.apply(writeFiles);
     }
 

sdks/java/core/src/main/java/org/apache/beam/sdk/io/TextIO.java

@@ -51,6 +51,8 @@
 import org.apache.beam.sdk.transforms.SerializableFunctions;
 import org.apache.beam.sdk.transforms.Watch.Growth.TerminationCondition;
 import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.errorhandling.BadRecord;
+import org.apache.beam.sdk.transforms.errorhandling.ErrorHandler;
 import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PDone;
@@ -176,6 +178,10 @@
  *
  * <p>For backwards compatibility, {@link TextIO} also supports the legacy {@link
  * DynamicDestinations} interface for advanced features via {@link Write#to(DynamicDestinations)}.
+ *
+ * <p>Error handling for records that are malformed can be handled by using {@link
+ * TypedWrite#withBadRecordErrorHandler(ErrorHandler)}. See documentation in
+ * {@link FileIO} for details on usage
  */
 @SuppressWarnings({
   "nullness" // TODO(https://github.com/apache/beam/issues/20497)
@@ -708,6 +714,8 @@ public abstract static class TypedWrite<UserT, DestinationT>
      */
     abstract WritableByteChannelFactory getWritableByteChannelFactory();
 
+    abstract @Nullable ErrorHandler<BadRecord, ?> getBadRecordErrorHandler();
+
     abstract Builder<UserT, DestinationT> toBuilder();
 
     @AutoValue.Builder
@@ -754,6 +762,9 @@ abstract Builder<UserT, DestinationT> setNumShards(
       abstract Builder<UserT, DestinationT> setWritableByteChannelFactory(
           WritableByteChannelFactory writableByteChannelFactory);
 
+      abstract Builder<UserT, DestinationT> setBadRecordErrorHandler(
+          @Nullable ErrorHandler<BadRecord, ?> badRecordErrorHandler);
+
       abstract TypedWrite<UserT, DestinationT> build();
     }
 
@@ -993,6 +1004,12 @@ public TypedWrite<UserT, DestinationT> withNoSpilling() {
       return toBuilder().setNoSpilling(true).build();
     }
 
+    /** See {@link WriteFiles#withBadRecordErrorHandler(ErrorHandler)}. */
+    public TypedWrite<UserT, DestinationT> withBadRecordErrorHandler(
+        ErrorHandler<BadRecord, ?> errorHandler) {
+      return toBuilder().setBadRecordErrorHandler(errorHandler).build();
+    }
+
     /** Don't write any output files if the PCollection is empty. */
     public TypedWrite<UserT, DestinationT> skipIfEmpty() {
       return toBuilder().setSkipIfEmpty(true).build();
@@ -1083,6 +1100,9 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
       if (getNoSpilling()) {
         write = write.withNoSpilling();
       }
+      if (getBadRecordErrorHandler() != null) {
+        write = write.withBadRecordErrorHandler(getBadRecordErrorHandler());
+      }
       if (getSkipIfEmpty()) {
         write = write.withSkipIfEmpty();
       }