From 6e6be214c414542313aff2c098fc451ce5889e80 Mon Sep 17 00:00:00 2001
From: Valentyn Tymofieiev <valentyn@google.com>
Date: Wed, 22 Jul 2026 07:11:55 -0700
Subject: [PATCH] Support Python's --files_to_stage option in Java xlang.

---
 .../python/PythonExternalTransform.java       | 34 ++++++++++++++++++
 .../python/PythonExternalTransformTest.java   | 36 +++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java b/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java
index fb44f8272154..316252eda104 100644
--- a/sdks/java/extensions/python/src/main/java/org/apache/beam/sdk/extensions/python/PythonExternalTransform.java
+++ b/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) {
@@ -510,6 +534,11 @@ public OutputT expand(InputT input) {
         // * It was explicitly requested.
         // * Python executable is not available in the system but Docker is available.
         if (useTransformService || (!pythonAvailable && dockerAvailable)) {
+          if (!filesToStage.isEmpty()) {
+            throw new UnsupportedOperationException(
+                "Staging local files is not supported when using the Docker-based transform service. "
+                    + "Please use the default local Python-based expansion service instead.");
+          }
           // A unique project name ensures that this expansion gets a dedicated instance of the
           // transform service.
           String projectName = UUID.randomUUID().toString();
@@ -550,6 +579,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())
diff --git a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/PythonExternalTransformTest.java b/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/PythonExternalTransformTest.java
index 30fe0b90f397..665168c6f4e5 100644
--- a/sdks/java/extensions/python/src/test/java/org/apache/beam/sdk/extensions/python/PythonExternalTransformTest.java
+++ b/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 =
@@ -393,4 +404,29 @@ public void testLoopbackEnvironmentWithPythonExternalTransform() {
             .apply(Keys.create());
     PAssert.that(output).containsInAnyOrder("A", "B");
   }
+
+  @Test
+  public void testFilesToStageWithTransformServiceThrows() {
+    PythonExternalTransformOptions options =
+        PipelineOptionsFactory.create().as(PythonExternalTransformOptions.class);
+    options.setUseTransformService(true);
+
+    Pipeline p = Pipeline.create(options);
+
+    PythonExternalTransform<PCollection<String>, PCollection<String>> transform =
+        PythonExternalTransform.<PCollection<String>, PCollection<String>>from("DummyTransform")
+            .withFilesToStage(ImmutableList.of("file1.py"));
+
+    PCollection<String> input = p.apply(Create.of("a"));
+
+    try {
+      input.apply(transform);
+      org.junit.Assert.fail("Expected UnsupportedOperationException");
+    } catch (UnsupportedOperationException e) {
+      assertTrue(
+          e.getMessage()
+              .contains(
+                  "Staging local files is not supported when using the Docker-based transform service"));
+    }
+  }
 }