[Cherrypick] [AvroIO, FileIO] Exposing evictWritersWhenFull in AvroIO (#40336)

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

@@ -396,6 +396,7 @@ public static <InputT> Write<Void, InputT> write() {
         .setIgnoreWindowing(false)
         .setAutoSharding(false)
         .setNoSpilling(false)
+        .setEvictWritersWhenFull(false)
         .build();
   }
 
@@ -410,6 +411,7 @@ public static <DestT, InputT> Write<DestT, InputT> writeDynamic() {
         .setIgnoreWindowing(false)
         .setAutoSharding(false)
         .setNoSpilling(false)
+        .setEvictWritersWhenFull(false)
         .build();
   }
 
@@ -1059,6 +1061,8 @@ public static FileNaming relativeFileNaming(
 
     abstract boolean getNoSpilling();
 
+    abstract boolean getEvictWritersWhenFull();
+
     abstract @Nullable Integer getMaxNumWritersPerBundle();
 
     abstract @Nullable Integer getBatchSize();
@@ -1117,6 +1121,8 @@ abstract Builder<DestinationT, UserT> setSharding(
 
       abstract Builder<DestinationT, UserT> setNoSpilling(boolean noSpilling);
 
+      abstract Builder<DestinationT, UserT> setEvictWritersWhenFull(boolean evictWritersWhenFull);
+
       abstract Builder<DestinationT, UserT> setMaxNumWritersPerBundle(
           @Nullable Integer maxNumWritersPerBundle);
 
@@ -1366,6 +1372,11 @@ public Write<DestinationT, UserT> withNoSpilling() {
       return toBuilder().setNoSpilling(true).build();
     }
 
+    /** See {@link WriteFiles#withEvictWritersWhenFull()}. */
+    public Write<DestinationT, UserT> withEvictWritersWhenFull() {
+      return toBuilder().setEvictWritersWhenFull(true).build();
+    }
+
     /**
      * Set the maximum number of writers created in a bundle before spilling to shuffle. See {@link
      * WriteFiles#withMaxNumWritersPerBundle()}.
@@ -1511,6 +1522,7 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
       resolvedSpec.setIgnoreWindowing(getIgnoreWindowing());
       resolvedSpec.setAutoSharding(getAutoSharding());
       resolvedSpec.setNoSpilling(getNoSpilling());
+      resolvedSpec.setEvictWritersWhenFull(getEvictWritersWhenFull());
       if (getMaxNumWritersPerBundle() != null) {
         resolvedSpec.setMaxNumWritersPerBundle(getMaxNumWritersPerBundle());
       }
@@ -1535,6 +1547,9 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
       if (getNoSpilling()) {
         writeFiles = writeFiles.withNoSpilling();
       }
+      if (getEvictWritersWhenFull()) {
+        writeFiles = writeFiles.withEvictWritersWhenFull();
+      }
       if (getMaxNumWritersPerBundle() != null) {
         writeFiles = writeFiles.withMaxNumWritersPerBundle(getMaxNumWritersPerBundle());
       }

sdks/java/extensions/avro/src/main/java/org/apache/beam/sdk/extensions/avro/io/AvroIO.java

@@ -615,6 +615,7 @@ private static <UserT, OutputT> TypedWrite.Builder<UserT, Void, OutputT> default
         .setMetadata(ImmutableMap.of())
         .setWindowedWrites(false)
         .setNoSpilling(false)
+        .setEvictWritersWhenFull(false)
         .setSyncInterval(DataFileConstants.DEFAULT_SYNC_INTERVAL);
   }
 
@@ -1426,6 +1427,8 @@ public abstract static class TypedWrite<UserT, DestinationT, OutputT>
 
     abstract boolean getNoSpilling();
 
+    abstract boolean getEvictWritersWhenFull();
+
     abstract @Nullable Integer getMaxNumWritersPerBundle();
 
     abstract @Nullable FilenamePolicy getFilenamePolicy();
@@ -1488,6 +1491,9 @@ public Builder<UserT, DestinationT, OutputT> setGenericRecords(boolean genericRe
 
       abstract Builder<UserT, DestinationT, OutputT> setNoSpilling(boolean noSpilling);
 
+      abstract Builder<UserT, DestinationT, OutputT> setEvictWritersWhenFull(
+          boolean evictWritersWhenFull);
+
       abstract Builder<UserT, DestinationT, OutputT> setMaxNumWritersPerBundle(
           @Nullable Integer maxNumWritersPerBundle);
 
@@ -1698,6 +1704,11 @@ public TypedWrite<UserT, DestinationT, OutputT> withNoSpilling() {
       return toBuilder().setNoSpilling(true).build();
     }
 
+    /** See {@link WriteFiles#withEvictWritersWhenFull()}. */
+    public TypedWrite<UserT, DestinationT, OutputT> withEvictWritersWhenFull() {
+      return toBuilder().setEvictWritersWhenFull(true).build();
+    }
+
     /** See {@link WriteFiles#withMaxNumWritersPerBundle()}. */
     public TypedWrite<UserT, DestinationT, OutputT> withMaxNumWritersPerBundle(
         @Nullable Integer maxNumWritersPerBundle) {
@@ -1819,6 +1830,9 @@ public WriteFilesResult<DestinationT> expand(PCollection<UserT> input) {
       if (getBadRecordErrorHandler() != null) {
         write = write.withBadRecordErrorHandler(getBadRecordErrorHandler());
       }
+      if (getEvictWritersWhenFull()) {
+        write = write.withEvictWritersWhenFull();
+      }
       return input.apply("Write", write);
     }