Support Python's --files_to_stage option in Java xlang.

Author: tvalentyn

sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java

@@ -81,6 +81,7 @@ public class PythonExternalTransform<InputT extends PInput, OutputT extends POut
 
   private String expansionService;
   private List<String> extraPackages;
+  private List<String> filesToStage;
 
   // We preseve the order here since Schema's care about order of fields but the order will not
   // matter when applying kwargs at the Python side.
@@ -98,6 +99,7 @@ private PythonExternalTransform(String fullyQualifiedName, String expansionServi
     this.fullyQualifiedName = fullyQualifiedName;
     this.expansionService = expansionService;
     this.extraPackages = new ArrayList<>();
+    this.filesToStage = new ArrayList<>();
     this.kwargsMap = new TreeMap<>();
     this.typeHints = new HashMap<>();
     // TODO(https://github.com/apache/beam/issues/21567): remove a default type hint for
@@ -304,6 +306,28 @@ public PythonExternalTransform<InputT, OutputT> withExtraPackages(List<String> e
     return this;
   }
 
+  /**
+   * Specifies that the given Python files should be staged for the Python worker.
+   *
+   * @param filesToStage a list of local paths to Python files or data files.
+   * @return updated wrapper for the cross-language transform.
+   */
+  public PythonExternalTransform<InputT, OutputT> withFilesToStage(List<String> filesToStage) {
+    if (filesToStage.isEmpty()) {
+      return this;
+    }
+    Preconditions.checkState(
+        Strings.isNullOrEmpty(expansionService),
+        "Files to stage only apply to auto-started expansion service.");
+    this.filesToStage = filesToStage;
+    return this;
+  }
+
+  @VisibleForTesting
+  List<String> getFilesToStage() {
+    return filesToStage;
+  }
+
   @VisibleForTesting
   Row buildOrGetKwargsRow() {
     if (providedKwargsRow != null) {
@@ -550,6 +574,11 @@ public OutputT expand(InputT input) {
           if (requirementsFile != null) {
             args.add("--requirements_file=" + requirementsFile.getAbsolutePath());
           }
+          if (!filesToStage.isEmpty()) {
+            for (String file : filesToStage) {
+              args.add("--files_to_stage=" + file);
+            }
+          }
           PythonService service =
               new PythonService(
                       "apache_beam.runners.portability.expansion_service_main", args.build())

sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/PythonExternalTransformTest.java

@@ -84,6 +84,17 @@ public void pythonTransformWithDependencies() {
     }
   }
 
+  @Test
+  public void testFilesToStage() {
+    PythonExternalTransform<?, ?> transform =
+        PythonExternalTransform
+            .<PCollection<KV<String, String>>, PCollection<KV<String, Iterable<String>>>>from(
+                "DummyTransform")
+            .withFilesToStage(ImmutableList.of("file1.py", "file2.py"));
+
+    assertEquals(ImmutableList.of("file1.py", "file2.py"), transform.getFilesToStage());
+  }
+
   @Test
   public void generateArgsEmpty() {
     PythonExternalTransform<?, ?> transform =