This is an automated email from the ASF dual-hosted git repository.
derrickaw pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new c2474960db7 [yaml] - expand jinja method (#38547)
c2474960db7 is described below
commit c2474960db705c7ac8a39a766be94893a1953667
Author: Derrick Williams <[email protected]>
AuthorDate: Mon Jun 22 17:46:00 2026 -0400
[yaml] - expand jinja method (#38547)
* expand jinja
* minor tweaks
* fix lint
* fix gemini comments
* add more documentation
---
sdks/python/apache_beam/yaml/examples/README.md | 4 +
sdks/python/apache_beam/yaml/main.py | 7 +-
sdks/python/apache_beam/yaml/yaml_provider.py | 11 +-
sdks/python/apache_beam/yaml/yaml_transform.py | 56 ++++++++--
.../python/apache_beam/yaml/yaml_transform_test.py | 119 +++++++++++++++++++++
.../www/site/content/en/documentation/sdks/yaml.md | 34 ++++++
6 files changed, 222 insertions(+), 9 deletions(-)
diff --git a/sdks/python/apache_beam/yaml/examples/README.md
b/sdks/python/apache_beam/yaml/examples/README.md
index 55fd19bd8c4..5282951c696 100644
--- a/sdks/python/apache_beam/yaml/examples/README.md
+++ b/sdks/python/apache_beam/yaml/examples/README.md
@@ -256,6 +256,10 @@ Jinja `% include` directive:
-
[wordCountInclude.yaml](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/yaml/examples/transforms/jinja/include/wordCountInclude.yaml)
-
[Instructions](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/yaml/examples/transforms/jinja/include/README.md)
on how to run the pipeline.
+Jinja template inheritance (`{% extends %}`):
+-
[wordCountInheritance.yaml](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/yaml/examples/transforms/jinja/inheritance/wordCountInheritance.yaml)
+-
[Instructions](https://github.com/apache/beam/blob/master/sdks/python/apache_beam/yaml/examples/transforms/jinja/inheritance/README.md)
on how to run the pipeline.
+
### ML
diff --git a/sdks/python/apache_beam/yaml/main.py
b/sdks/python/apache_beam/yaml/main.py
index dc928dec794..804798b82e0 100644
--- a/sdks/python/apache_beam/yaml/main.py
+++ b/sdks/python/apache_beam/yaml/main.py
@@ -235,8 +235,13 @@ def _build_pipeline_yaml_from_argv(argv):
argv = _preparse_jinja_flags(argv)
known_args, pipeline_args = _parse_arguments(argv)
pipeline_template = _pipeline_spec_from_args(known_args)
+
+ search_paths = []
+ if known_args.yaml_pipeline_file:
+ search_paths.append(FileSystems.split(known_args.yaml_pipeline_file)[0])
+
pipeline_yaml = yaml_transform.expand_jinja(
- pipeline_template, known_args.jinja_variables or {})
+ pipeline_template, known_args.jinja_variables or {}, search_paths)
return known_args, pipeline_args, pipeline_template, pipeline_yaml
diff --git a/sdks/python/apache_beam/yaml/yaml_provider.py
b/sdks/python/apache_beam/yaml/yaml_provider.py
index ce5a6e32058..324ae0c2e73 100755
--- a/sdks/python/apache_beam/yaml/yaml_provider.py
+++ b/sdks/python/apache_beam/yaml/yaml_provider.py
@@ -469,10 +469,14 @@ class ExternalPythonProvider(ExternalProvider):
@ExternalProvider.register_provider_type('yaml')
class YamlProvider(Provider):
- def __init__(self, transforms: Mapping[str, Mapping[str, Any]]):
+ def __init__(
+ self,
+ transforms: Mapping[str, Mapping[str, Any]],
+ provider_base_path: Optional[str] = None):
if not isinstance(transforms, dict):
raise ValueError('Transform mapping must be a dict.')
self._transforms = transforms
+ self._provider_base_path = provider_base_path
def available(self):
return True
@@ -524,7 +528,10 @@ class YamlProvider(Provider):
else:
body_str = yaml.safe_dump(SafeLineLoader.strip_metadata(body))
# Now re-parse resolved templatization.
- body = yaml.load(expand_jinja(body_str, args), Loader=SafeLineLoader)
+ search_paths = [FileSystems.split(self._provider_base_path)[0]
+ ] if self._provider_base_path else []
+ body = yaml.load(
+ expand_jinja(body_str, args, search_paths), Loader=SafeLineLoader)
if (body.get('type') == 'chain' and 'input' not in body and
spec.get('requires_inputs', True)):
body['input'] = 'input'
diff --git a/sdks/python/apache_beam/yaml/yaml_transform.py
b/sdks/python/apache_beam/yaml/yaml_transform.py
index 2b745babad0..a4bb9144f8e 100644
--- a/sdks/python/apache_beam/yaml/yaml_transform.py
+++ b/sdks/python/apache_beam/yaml/yaml_transform.py
@@ -1391,19 +1391,63 @@ def preprocess(spec, verbose=False,
known_transforms=None):
return spec
+def strip_leading_comments(source: str) -> str:
+ lines = source.splitlines(keepends=True)
+ stripped_lines = []
+ in_leading_comments = True
+ for line in lines:
+ stripped_line = line.lstrip()
+ if in_leading_comments:
+ if stripped_line.startswith('#') or not stripped_line:
+ continue
+ else:
+ in_leading_comments = False
+ stripped_lines.append(line)
+ return "".join(stripped_lines)
+
+
class _BeamFileIOLoader(jinja2.BaseLoader):
+ def __init__(self, search_paths=()):
+ self.search_paths = list(search_paths)
+
def get_source(self, environment, path):
- with FileSystems.open(path) as fin:
- source = fin.read().decode()
- return source, path, lambda: True
+ candidates = [path]
+ if FileSystems.get_scheme(path) is None and not os.path.isabs(path):
+ for search_path in self.search_paths:
+ candidates.append(FileSystems.join(search_path, path))
+
+ for candidate in candidates:
+ try:
+ exists = FileSystems.exists(candidate)
+ except Exception:
+ exists = False
+
+ if exists:
+ with FileSystems.open(candidate) as fin:
+ source = fin.read().decode()
+ return strip_leading_comments(source), candidate, lambda: True
+
+ raise jinja2.TemplateNotFound(path)
def expand_jinja(
- jinja_template: str, jinja_variables: Mapping[str, Any]) -> str:
+ jinja_template: str,
+ jinja_variables: Mapping[str, Any],
+ search_paths: Iterable[str] = ()) -> str:
+ beam_root_dir = os.path.dirname(
+ os.path.dirname(os.path.abspath(beam.__file__)))
+
+ all_search_paths = list(search_paths)
+ if beam_root_dir not in all_search_paths:
+ all_search_paths.append(beam_root_dir)
+ if '.' not in all_search_paths:
+ all_search_paths.append('.')
+
return ( # keep formatting
jinja2.Environment(
- undefined=jinja2.StrictUndefined, loader=_BeamFileIOLoader())
- .from_string(jinja_template)
+ undefined=jinja2.StrictUndefined,
+ loader=_BeamFileIOLoader(all_search_paths))
+ .from_string(strip_leading_comments(jinja_template))
.render(datetime=datetime, **jinja_variables))
diff --git a/sdks/python/apache_beam/yaml/yaml_transform_test.py
b/sdks/python/apache_beam/yaml/yaml_transform_test.py
index 192af63a987..4cfe625fc9c 100644
--- a/sdks/python/apache_beam/yaml/yaml_transform_test.py
+++ b/sdks/python/apache_beam/yaml/yaml_transform_test.py
@@ -23,12 +23,16 @@ import shutil
import tempfile
import unittest
+import yaml
+
import apache_beam as beam
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to
from apache_beam.utils import python_callable
from apache_beam.yaml import yaml_provider
+from apache_beam.yaml.yaml_transform import SafeLineLoader
from apache_beam.yaml.yaml_transform import YamlTransform
+from apache_beam.yaml.yaml_transform import expand_jinja
try:
import jsonschema
@@ -1467,6 +1471,65 @@ class TestExternalYamlProvider(unittest.TestCase):
''',
providers=merged_providers)
+ def test_provider_with_jinja_imports(self):
+ # Create a macro file in the same temp directory as the provider
+ macro_path = os.path.join(self.temp_dir, 'my_macros.yaml')
+ with open(macro_path, 'w') as f:
+ f.write(
+ """
+{%- macro power_expr(var, n) -%}
+{{ var }} ** {{ n }}
+{%- endmacro -%}
+""")
+
+ # Create a provider that imports and uses the macro
+ templated_provider_path = os.path.join(
+ self.temp_dir, 'templated_provider.yaml')
+ with open(templated_provider_path, 'w') as f:
+ f.write(
+ """
+- type: yaml
+ transforms:
+ CustomPower:
+ config_schema:
+ properties:
+ n: {type: integer}
+ body: |
+ type: MapToFields
+ config:
+ language: python
+ append: true
+ fields:
+ power: "{% import 'my_macros.yaml' as m %}{{
m.power_expr('element', n) }}"
+""")
+
+ loaded_providers = yaml_provider.load_providers(templated_provider_path)
+ test_providers = yaml_provider.InlineProvider(TEST_PROVIDERS)
+ merged_providers = yaml_provider.merge_providers(
+ loaded_providers, [test_providers])
+
+ with beam.Pipeline(options=beam.options.pipeline_options.PipelineOptions(
+ pickle_library='cloudpickle')) as p:
+ results = p | YamlTransform(
+ '''
+ type: composite
+ transforms:
+ - type: Create
+ config:
+ elements: [2, 3]
+ - type: CustomPower
+ input: Create
+ config:
+ n: 3
+ output: CustomPower
+ ''',
+ providers=merged_providers)
+
+ assert_that(
+ results,
+ equal_to(
+ [beam.Row(element=2, power=8), beam.Row(element=3, power=27)]))
+
@beam.transforms.ptransform.annotate_yaml
class LinearTransform(beam.PTransform):
@@ -1481,6 +1544,62 @@ class LinearTransform(beam.PTransform):
return pcoll | beam.Map(lambda x: a * x.element + b)
+class TestYamlExpandJinja(unittest.TestCase):
+ def setUp(self):
+ self.temp_dir = tempfile.mkdtemp()
+ # Create a macro file with leading comments (license header)
+ self.macro_path = os.path.join(self.temp_dir, 'my_macros.yaml')
+ with open(self.macro_path, 'w') as f:
+ f.write(
+ """# coding=utf-8
+# Licensed to the Apache Software Foundation...
+# Some leading comment line
+
+{%- macro add_n(val, n) -%}
+{{ val }} + {{ n }}
+{%- endmacro -%}
+""")
+
+ # Create a pipeline template that includes/imports the macro
+ self.pipeline_path = os.path.join(self.temp_dir, 'my_pipeline.yaml')
+ with open(self.pipeline_path, 'w') as f:
+ f.write(
+ """# coding=utf-8
+# Licensed to the Apache Software Foundation...
+
+{% import 'my_macros.yaml' as macros %}
+type: composite
+transforms:
+ - type: Create
+ config:
+ elements: [1, 2, 3]
+ - type: MapToFields
+ config:
+ language: python
+ fields:
+ result: {{ macros.add_n('element', 10) }}
+""")
+
+ def tearDown(self):
+ shutil.rmtree(self.temp_dir)
+
+ def test_expand_jinja_with_leading_comments_and_imports(self):
+ # Read the pipeline template
+ with open(self.pipeline_path, 'r') as f:
+ template_content = f.read()
+
+ # Expand the jinja using our temp_dir as a search path
+ expanded = expand_jinja(template_content, {}, [self.temp_dir])
+
+ # Parse the expanded YAML
+ parsed = yaml.load(expanded, Loader=SafeLineLoader)
+
+ # Verify the comment-stripping and import resolution was successful
+ self.assertEqual(parsed['type'], 'composite')
+ self.assertEqual(
+ parsed['transforms'][1]['config']['fields']['result'], 'element + 10')
+
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
unittest.main()
diff --git a/website/www/site/content/en/documentation/sdks/yaml.md
b/website/www/site/content/en/documentation/sdks/yaml.md
index 589679d3e50..a1fda30ddb6 100644
--- a/website/www/site/content/en/documentation/sdks/yaml.md
+++ b/website/www/site/content/en/documentation/sdks/yaml.md
@@ -812,6 +812,40 @@ pipeline:
This pipeline can be run with the same command as in the `% include` example
above.
+You can also use template inheritance (`{% extends %}` and `{% block %}`) to
define a base pipeline layout and allow child pipelines to override specific
blocks of transforms:
+
+<PATH_TO_YOUR_REPO>/base_pipeline.yaml
+```yaml
+pipeline:
+ type: chain
+ transforms:
+ - type: ReadFromText
+ config:
+ path: {{ input_path }}
+
+# Injection point for child templates
+{% block custom_steps %}
+{% endblock %}
+
+ - type: WriteToText
+ config:
+ path: {{ output_path }}
+```
+
+<PATH_TO_YOUR_REPO>/child_pipeline.yaml
+```yaml
+{% extends '<PATH_TO_YOUR_REPO>/base_pipeline.yaml' %}
+
+{% block custom_steps %}
+ - type: Filter
+ config:
+ language: python
+ keep: 'row.count > 10'
+{% endblock %}
+```
+
+This pipeline can be run using the exact same command structure as above.
+
There are many more ways to import and even use template inheritance using
Jinja as seen
[here](https://jinja.palletsprojects.com/en/stable/templates/#import)
and [here](https://jinja.palletsprojects.com/en/stable/templates/#inheritance).
