Allow dependency declarations for java.

robertwb · robertwb · commit e40304170332 · 2025-03-18T08:55:59.000-07:00
diff --git a/sdks/python/apache_beam/yaml/examples/transforms/elementwise/map_to_fields_with_java_deps.yaml b/sdks/python/apache_beam/yaml/examples/transforms/elementwise/map_to_fields_with_java_deps.yaml
@@ -0,0 +1,57 @@
+# coding=utf-8
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+pipeline:
+  type: chain
+  transforms:
+    - type: Create
+      config:
+        elements:
+          - {sdk: MapReduce, year: 2004}
+          - {sdk: MillWheel, year: 2008}
+          - {sdk: Flume, year: 2010}
+          - {sdk: Dataflow, year: 2014}
+          - {sdk: Apache Beam, year: 2016}
+    - type: MapToFields
+      name: ToRoman
+      config:
+        language: java
+        fields:
+          tool_name: sdk
+          year:
+            callable: |
+              import org.apache.beam.sdk.values.Row;
+              import java.util.function.Function;
+              import com.github.chaosfirebolt.converter.RomanInteger;
+
+              public class MyFunction implements Function<Row, String> {
+                public String apply(Row row) {
+                  return RomanInteger.parse(
+                      String.valueOf(row.getInt64("year"))).toString();
+                }
+              }
+        dependencies:
+          - 'com.github.chaosfirebolt.converter:roman-numeral-converter:2.1.0'
+    - type: LogForTesting
+
+# Expected:
+#  Row(tool_name='MapReduce', year='MMIV')
+#  Row(tool_name='MillWheel', year='MMVIII')
+#  Row(tool_name='Flume', year='MMX')
+#  Row(tool_name='Dataflow', year='MMXIV')
+#  Row(tool_name='Apache Beam', year='MMXVI')
diff --git a/sdks/python/apache_beam/yaml/yaml_provider.py b/sdks/python/apache_beam/yaml/yaml_provider.py
@@ -356,10 +356,13 @@ def cache_artifacts(self):
 
 
 class ExternalJavaProvider(ExternalProvider):
-  def __init__(self, urns, jar_provider):
+  def __init__(self, urns, jar_provider, classpath=None):
     super().__init__(
-        urns, lambda: external.JavaJarExpansionService(jar_provider()))
+        urns,
+        lambda: external.JavaJarExpansionService(
+            jar_provider(), classpath=classpath))
     self._jar_provider = jar_provider
+    self._classpath = classpath
 
   def available(self):
     # pylint: disable=subprocess-run-check
@@ -381,6 +384,15 @@ def try_decode(bs):
   def cache_artifacts(self):
     return [self._jar_provider()]
 
+  def _with_extra_dependencies(self, dependencies: Iterable[str]):
+    jars = sum((
+        external.JavaJarExpansionService._expand_jars(dep)
+        for dep in dependencies), [])
+    return ExternalJavaProvider(
+        self._urns,
+        jar_provider=self._jar_provider,
+        classpath=(list(self._classpath or []) + list(jars)))
+
 
 @ExternalProvider.register_provider_type('python')
 def python(urns, provider_base_path, packages=()):
diff --git a/sdks/python/apache_beam/yaml/yaml_transform.py b/sdks/python/apache_beam/yaml/yaml_transform.py
@@ -368,7 +368,7 @@ def expand(pcolls):
         if pcoll in providers_by_input
     ]
     provider = self.best_provider(spec, input_providers)
-    extra_dependencies = extract_extra_dependencies(spec)
+    extra_dependencies, spec = extract_extra_dependencies(spec)
     if extra_dependencies:
       provider = provider.with_extra_dependencies(frozenset(extra_dependencies))
 
@@ -713,9 +713,13 @@ def extract_name(spec):
 
 def extract_extra_dependencies(spec):
   deps = spec.get('config', {}).get('dependencies', [])
+  if not deps:
+    return [], spec
   if not isinstance(deps, list):
     raise TypeErrorError(f'Dependencies must be a list of strings, got {deps}')
-  return deps
+  return deps, dict(
+      spec,
+      config={k: v for k, v in spec['config'].items() if k != 'dependencies'})
 
 
 def push_windowing_to_roots(spec):
