Add EmptyMatchTreatment support for SMB reads (#5759)

This change adds comprehensive EmptyMatchTreatment support for Sort Merge Bucket (SMB) operations:

API-level changes:
- Add EmptyMatchTreatment parameter to all SMB IO read operations (JSON, TensorFlow, ParquetAvro, ParquetType)
- Update BucketedInput.of() methods to accept EmptyMatchTreatment parameter
- Maintain backward compatibility with existing code

End-to-end implementation:
- Fix BucketedInput#getOrSampleByteSize to handle empty directories gracefully when EmptyMatchTreatment.ALLOW is used
- Enhance MultiSourceKeyGroupReader to filter out sources without valid metadata, preventing failures with empty directories
- Add comprehensive tests for both API and integration functionality

This addresses the SMB filesystem abstraction bypasses by ensuring EmptyMatchTreatment works throughout the entire SMB pipeline.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

Author: sushmullur

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/AvroSortedBucketIO.java

@@ -38,6 +38,7 @@
 import org.apache.beam.sdk.extensions.smb.SortedBucketSource.Predicate;
 import org.apache.beam.sdk.extensions.smb.SortedBucketTransform.NewBucketMetadataFn;
 import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -218,6 +219,9 @@ public abstract static class Read<T extends IndexedRecord> extends SortedBucketI
     @Nullable
     abstract Predicate<T> getPredicate();
 
+    @Nullable
+    abstract EmptyMatchTreatment getEmptyMatchTreatment();
+
     abstract Builder<T> toBuilder();
 
     @AutoValue.Builder
@@ -236,6 +240,8 @@ abstract static class Builder<T extends IndexedRecord> {
 
       abstract Builder<T> setPredicate(Predicate<T> predicate);
 
+      abstract Builder<T> setEmptyMatchTreatment(EmptyMatchTreatment emptyMatchTreatment);
+
       abstract Read<T> build();
     }
 
@@ -264,6 +270,11 @@ public Read<T> withPredicate(Predicate<T> predicate) {
       return toBuilder().setPredicate(predicate).build();
     }
 
+    /** Specifies the empty match treatment for handling empty directories. */
+    public Read<T> withEmptyMatchTreatment(EmptyMatchTreatment emptyMatchTreatment) {
+      return toBuilder().setEmptyMatchTreatment(emptyMatchTreatment).build();
+    }
+
     @SuppressWarnings("unchecked")
     FileOperations<T> getFileOperations() {
       return AvroFileOperations.of(getDatumFactory(), getSchema());
@@ -278,7 +289,8 @@ public SortedBucketSource.BucketedInput<T> toBucketedInput(
           getInputDirectories(),
           getFilenameSuffix(),
           getFileOperations(),
-          getPredicate());
+          getPredicate(),
+          getEmptyMatchTreatment());
     }
   }
 

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/JsonSortedBucketIO.java

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.extensions.smb.SortedBucketTransform.NewBucketMetadataFn;
 import org.apache.beam.sdk.io.Compression;
 import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -118,6 +119,9 @@ public abstract static class Read extends SortedBucketIO.Read<TableRow> {
     @Nullable
     abstract Predicate<TableRow> getPredicate();
 
+    @Nullable
+    abstract EmptyMatchTreatment getEmptyMatchTreatment();
+
     abstract Builder toBuilder();
 
     @AutoValue.Builder
@@ -134,6 +138,8 @@ abstract static class Builder {
 
       abstract Builder setPredicate(Predicate<TableRow> predicate);
 
+      abstract Builder setEmptyMatchTreatment(EmptyMatchTreatment emptyMatchTreatment);
+
       abstract Read build();
     }
 
@@ -156,6 +162,19 @@ public Read withPredicate(Predicate<TableRow> predicate) {
       return toBuilder().setPredicate(predicate).build();
     }
 
+    /**
+     * Specifies how to handle empty file matches when reading input directories.
+     *
+     * @param emptyMatchTreatment the behavior when file patterns don't match any files. Use
+     *     {@link EmptyMatchTreatment#ALLOW} to proceed with empty inputs, {@link
+     *     EmptyMatchTreatment#DISALLOW} to fail on empty matches, or {@link
+     *     EmptyMatchTreatment#ALLOW_IF_WILDCARD} to allow empty matches only for wildcard patterns.
+     * @return a new {@link Read} instance with the specified empty match treatment
+     */
+    public Read withEmptyMatchTreatment(EmptyMatchTreatment emptyMatchTreatment) {
+      return toBuilder().setEmptyMatchTreatment(emptyMatchTreatment).build();
+    }
+
     FileOperations<TableRow> getFileOperations() {
       return JsonFileOperations.of(getCompression());
     }
@@ -168,7 +187,8 @@ public BucketedInput<TableRow> toBucketedInput(final SortedBucketSource.Keying k
           getInputDirectories(),
           getFilenameSuffix(),
           getFileOperations(),
-          getPredicate());
+          getPredicate(),
+          getEmptyMatchTreatment());
     }
   }
 

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/MultiSourceKeyGroupReader.java

@@ -97,6 +97,7 @@ public MultiSourceKeyGroupReader(
     this.resultSchema = resultSchema;
     this.bucketedInputs =
         sources.stream()
+            .filter(src -> hasValidDirectories(src))
             .map(src -> new BucketIterator<>(src, bucketId, effectiveParallelism, options))
             .collect(Collectors.toList());
     // this only operates on the primary key
@@ -105,6 +106,20 @@ public MultiSourceKeyGroupReader(
             someArbitraryBucketMetadata.rehashBucket(bytes, effectiveParallelism) == bucketId;
   }
 
+  /**
+   * Check if a source has valid directories with metadata files.
+   * This is needed to handle EmptyMatchTreatment.ALLOW cases where some directories might be empty.
+   */
+  private static boolean hasValidDirectories(SortedBucketSource.BucketedInput<?> source) {
+    try {
+      source.getSourceMetadata();
+      return true;
+    } catch (Exception e) {
+      // If we can't get source metadata, the directory is likely empty or invalid
+      return false;
+    }
+  }
+
   public KV<KeyType, CoGbkResult> readNext() {
     advance();
     return head;

scio-smb/src/main/java/org/apache/beam/sdk/extensions/smb/ParquetAvroSortedBucketIO.java

@@ -34,6 +34,7 @@
 import org.apache.beam.sdk.extensions.smb.SortedBucketSource.Predicate;
 import org.apache.beam.sdk.extensions.smb.SortedBucketTransform.NewBucketMetadataFn;
 import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
@@ -222,6 +223,9 @@ public abstract static class Read<T extends IndexedRecord> extends SortedBucketI
     @Nullable
     abstract Schema getProjection();
 
+    @Nullable
+    abstract EmptyMatchTreatment getEmptyMatchTreatment();
+
     abstract Builder<T> toBuilder();
 
     @AutoValue.Builder
@@ -246,6 +250,8 @@ abstract static class Builder<T extends IndexedRecord> {
 
       abstract Builder<T> setProjection(Schema projection);
 
+      abstract Builder<T> setEmptyMatchTreatment(EmptyMatchTreatment emptyMatchTreatment);
+
       abstract Read<T> build();
     }
 
@@ -283,6 +289,19 @@ public Read<T> withProjection(Schema projection) {
       return toBuilder().setProjection(projection).build();
     }
 
+    /**
+     * Specifies how to handle empty file matches when reading input directories.
+     *
+     * @param emptyMatchTreatment the behavior when file patterns don't match any files. Use
+     *     {@link EmptyMatchTreatment#ALLOW} to proceed with empty inputs, {@link
+     *     EmptyMatchTreatment#DISALLOW} to fail on empty matches, or {@link
+     *     EmptyMatchTreatment#ALLOW_IF_WILDCARD} to allow empty matches only for wildcard patterns.
+     * @return a new {@link Read} instance with the specified empty match treatment
+     */
+    public Read<T> withEmptyMatchTreatment(EmptyMatchTreatment emptyMatchTreatment) {
+      return toBuilder().setEmptyMatchTreatment(emptyMatchTreatment).build();
+    }
+
     @Override
     public BucketedInput<T> toBucketedInput(final SortedBucketSource.Keying keying) {
       return BucketedInput.of(
@@ -291,7 +310,8 @@ public BucketedInput<T> toBucketedInput(final SortedBucketSource.Keying keying)
           getInputDirectories(),
           getFilenameSuffix(),
           getFileOperations(),
-          getPredicate());
+          getPredicate(),
+          getEmptyMatchTreatment());
     }
 
     FileOperations<T> getFileOperations() {