This is an automated email from the ASF dual-hosted git repository.
robertwb 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 210ede6 [BEAM-12769] Java-emulating external transform. (#15546)
210ede6 is described below
commit 210ede679845a68039e7792b6bd19f8b488efb24
Author: Robert Bradshaw <rober...@google.com>
AuthorDate: Fri Sep 24 17:38:30 2021 -0700
[BEAM-12769] Java-emulating external transform. (#15546)
---
sdks/python/apache_beam/transforms/external.py | 46 ++++++++++++++++++++--
.../python/apache_beam/transforms/external_test.py | 37 +++++++++++++++++
2 files changed, 79 insertions(+), 4 deletions(-)
diff --git a/sdks/python/apache_beam/transforms/external.py
b/sdks/python/apache_beam/transforms/external.py
index 16840c4..f8e6ddc 100644
--- a/sdks/python/apache_beam/transforms/external.py
+++ b/sdks/python/apache_beam/transforms/external.py
@@ -236,8 +236,7 @@ class JavaClassLookupPayloadBuilder(PayloadBuilder):
:param args: parameter values of the constructor.
:param kwargs: parameter names and values of the constructor.
"""
- if (self._constructor_method or self._constructor_param_args or
- self._constructor_param_kwargs):
+ if self._has_constructor():
raise ValueError(
'Constructor or constructor method can only be specified once')
@@ -254,8 +253,7 @@ class JavaClassLookupPayloadBuilder(PayloadBuilder):
:param args: parameter values of the constructor method.
:param kwargs: parameter names and values of the constructor method.
"""
- if (self._constructor_method or self._constructor_param_args or
- self._constructor_param_kwargs):
+ if self._has_constructor():
raise ValueError(
'Constructor or constructor method can only be specified once')
@@ -276,6 +274,46 @@ class JavaClassLookupPayloadBuilder(PayloadBuilder):
"""
self._builder_methods_and_params[method_name] = (args, kwargs)
+ def _has_constructor(self):
+ return (
+ self._constructor_method or self._constructor_param_args or
+ self._constructor_param_kwargs)
+
+
+class JavaExternalTransform(ptransform.PTransform):
+ """A proxy for Java-implemented external transforms.
+
+ One builds these transforms just as one would in Java.
+ """
+ def __init__(self, class_name, expansion_service=None):
+ self._payload_builder = JavaClassLookupPayloadBuilder(class_name)
+ self._expansion_service = None
+
+ def __call__(self, *args, **kwargs):
+ self._payload_builder.with_constructor(*args, **kwargs)
+ return self
+
+ def __getattr__(self, name):
+ # Don't try to emulate special methods.
+ if name.startswith('__') and name.endswith('__'):
+ return super().__getattr__(name)
+
+ def construct(*args, **kwargs):
+ if self._payload_builder._has_constructor():
+ builder_method = self._payload_builder.add_builder_method
+ else:
+ builder_method = self._payload_builder.with_constructor_method
+ builder_method(name, *args, **kwargs)
+ return self
+
+ return construct
+
+ def expand(self, pcolls):
+ return pcolls | ExternalTransform(
+ common_urns.java_class_lookup,
+ self._payload_builder.build(),
+ self._expansion_service)
+
class AnnotationBasedPayloadBuilder(SchemaBasedPayloadBuilder):
"""
diff --git a/sdks/python/apache_beam/transforms/external_test.py
b/sdks/python/apache_beam/transforms/external_test.py
index 8673ed5..a528384 100644
--- a/sdks/python/apache_beam/transforms/external_test.py
+++ b/sdks/python/apache_beam/transforms/external_test.py
@@ -39,6 +39,7 @@ from apache_beam.testing.util import equal_to
from apache_beam.transforms.external import AnnotationBasedPayloadBuilder
from apache_beam.transforms.external import ImplicitSchemaPayloadBuilder
from apache_beam.transforms.external import JavaClassLookupPayloadBuilder
+from apache_beam.transforms.external import JavaExternalTransform
from apache_beam.transforms.external import NamedTupleBasedPayloadBuilder
from apache_beam.typehints import typehints
from apache_beam.typehints.native_type_compatibility import
convert_to_beam_type
@@ -509,6 +510,42 @@ class JavaClassLookupPayloadBuilderTest(unittest.TestCase):
with self.assertRaises(ValueError):
payload_builder.with_constructor('def')
+ def test_implicit_builder_with_constructor(self):
+ constructor_transform = (
+ JavaExternalTransform('org.pkg.MyTransform')('abc').withIntProperty(5))
+
+ payload_bytes = constructor_transform._payload_builder.payload()
+ payload_from_bytes = proto_utils.parse_Bytes(
+ payload_bytes, JavaClassLookupPayload)
+ self.assertEqual('org.pkg.MyTransform', payload_from_bytes.class_name)
+ self._verify_row(
+ payload_from_bytes.constructor_schema,
+ payload_from_bytes.constructor_payload, {'ignore0': 'abc'})
+ builder_method = payload_from_bytes.builder_methods[0]
+ self.assertEqual('withIntProperty', builder_method.name)
+ self._verify_row(
+ builder_method.schema, builder_method.payload, {'ignore0': 5})
+
+ def test_implicit_builder_with_constructor_method(self):
+ constructor_transform = JavaExternalTransform('org.pkg.MyTransform').of(
+ str_field='abc').withProperty(int_field=1234).build()
+
+ payload_bytes = constructor_transform._payload_builder.payload()
+ payload_from_bytes = proto_utils.parse_Bytes(
+ payload_bytes, JavaClassLookupPayload)
+ self.assertEqual('of', payload_from_bytes.constructor_method)
+ self._verify_row(
+ payload_from_bytes.constructor_schema,
+ payload_from_bytes.constructor_payload, {'str_field': 'abc'})
+ with_property_method = payload_from_bytes.builder_methods[0]
+ self.assertEqual('withProperty', with_property_method.name)
+ self._verify_row(
+ with_property_method.schema,
+ with_property_method.payload, {'int_field': 1234})
+ build_method = payload_from_bytes.builder_methods[1]
+ self.assertEqual('build', build_method.name)
+ self._verify_row(build_method.schema, build_method.payload, {})
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
