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 extraPackages; + private List 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 withExtraPackages(List 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 withFilesToStage(List 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 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>>>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> transform = + PythonExternalTransform., PCollection>from("DummyTransform") + .withFilesToStage(ImmutableList.of("file1.py")); + + PCollection 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")); + } + } }