Extract shared code into FileSystems public API

Author: shnapz

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

@@ -26,12 +26,11 @@
 import java.nio.channels.ReadableByteChannel;
 import java.nio.channels.SeekableByteChannel;
 import java.util.ArrayList;
-import java.util.HashSet;
 import java.util.List;
 import java.util.ListIterator;
 import java.util.NoSuchElementException;
 import java.util.concurrent.atomic.AtomicReference;
-import org.apache.beam.sdk.io.FileSystem.LineageLevel;
+import java.util.stream.Collectors;
 import org.apache.beam.sdk.io.fs.EmptyMatchTreatment;
 import org.apache.beam.sdk.io.fs.MatchResult;
 import org.apache.beam.sdk.io.fs.MatchResult.Metadata;
@@ -318,35 +317,10 @@ public final List<? extends FileBasedSource<T>> split(
     }
   }
 
-  /**
-   * Report source Lineage. Due to the size limit of Beam metrics, report full file name or only dir
-   * depend on the number of files.
-   *
-   * <p>- Number of files<=100, report full file paths;
-   *
-   * <p>- Number of directory<=100, report directory names (one level up);
-   *
-   * <p>- Otherwise, report top level only.
-   */
   private static void reportSourceLineage(List<Metadata> expandedFiles) {
-    if (expandedFiles.size() <= 100) {
-      for (Metadata metadata : expandedFiles) {
-        FileSystems.reportSourceLineage(metadata.resourceId());
-      }
-    } else {
-      HashSet<ResourceId> uniqueDirs = new HashSet<>();
-      for (Metadata metadata : expandedFiles) {
-        ResourceId dir = metadata.resourceId().getCurrentDirectory();
-        uniqueDirs.add(dir);
-        if (uniqueDirs.size() > 100) {
-          FileSystems.reportSourceLineage(dir, LineageLevel.TOP_LEVEL);
-          return;
-        }
-      }
-      for (ResourceId uniqueDir : uniqueDirs) {
-        FileSystems.reportSourceLineage(uniqueDir);
-      }
-    }
+    List<ResourceId> resourceIds =
+        expandedFiles.stream().map(Metadata::resourceId).collect(Collectors.toList());
+    FileSystems.reportSourceLineage(resourceIds);
   }
 
   /**

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

@@ -29,6 +29,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -398,6 +399,37 @@ public ResourceId apply(@Nonnull Metadata input) {
         .delete(resourceIdsToDelete);
   }
 
+  /**
+   * Report source {@link Lineage} metrics for multiple resource ids. Due to the size limit of Beam
+   * metrics, report full file name or only dir depend on the number of files.
+   *
+   * <p>- Number of files<=100, report full file paths;
+   *
+   * <p>- Number of directory<=100, report directory names (one level up);
+   *
+   * <p>- Otherwise, report top level only.
+   */
+  public static void reportSourceLineage(List<ResourceId> resourceIds) {
+    if (resourceIds.size() <= 100) {
+      for (ResourceId resourceId : resourceIds) {
+        FileSystems.reportSourceLineage(resourceId);
+      }
+    } else {
+      HashSet<ResourceId> uniqueDirs = new HashSet<>();
+      for (ResourceId resourceId : resourceIds) {
+        ResourceId dir = resourceId.getCurrentDirectory();
+        uniqueDirs.add(dir);
+        if (uniqueDirs.size() > 100) {
+          FileSystems.reportSourceLineage(dir, LineageLevel.TOP_LEVEL);
+          return;
+        }
+      }
+      for (ResourceId uniqueDir : uniqueDirs) {
+        FileSystems.reportSourceLineage(uniqueDir);
+      }
+    }
+  }
+
   /** Report source {@link Lineage} metrics for resource id. */
   public static void reportSourceLineage(ResourceId resourceId) {
     reportSourceLineage(resourceId, LineageLevel.FILE);

sdks/java/io/hadoop-format/src/main/java/org/apache/beam/sdk/io/hadoop/format/HadoopFormatIO.java

@@ -52,7 +52,6 @@
 import org.apache.beam.sdk.coders.CoderRegistry;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.io.BoundedSource;
-import org.apache.beam.sdk.io.FileSystem;
 import org.apache.beam.sdk.io.FileSystems;
 import org.apache.beam.sdk.io.fs.ResourceId;
 import org.apache.beam.sdk.io.hadoop.SerializableConfiguration;
@@ -753,6 +752,7 @@ public List<BoundedSource<KV<K, V>>> split(long desiredBundleSizeBytes, Pipeline
           .collect(Collectors.toList());
     }
 
+    /** Report only file-based sources */
     private void reportSourceLineage(final List<SerializableSplit> inputSplits) {
       List<ResourceId> fileResources = new ArrayList<>();
 
@@ -766,24 +766,7 @@ private void reportSourceLineage(final List<SerializableSplit> inputSplits) {
         }
       }
 
-      if (fileResources.size() <= 100) {
-        for (ResourceId resource : fileResources) {
-          FileSystems.reportSourceLineage(resource);
-        }
-      } else {
-        HashSet<ResourceId> uniqueDirs = new HashSet<>();
-        for (ResourceId resource : fileResources) {
-          ResourceId dir = resource.getCurrentDirectory();
-          uniqueDirs.add(dir);
-          if (uniqueDirs.size() > 100) {
-            FileSystems.reportSourceLineage(dir, FileSystem.LineageLevel.TOP_LEVEL);
-            return;
-          }
-        }
-        for (ResourceId dir : uniqueDirs) {
-          FileSystems.reportSourceLineage(dir);
-        }
-      }
+      FileSystems.reportSourceLineage(fileResources);
     }
 
     @Override