Updates ExpansionService to support dynamically discovering and expanding SchemaTransforms #23413

chamikaramj · 2022-09-29T01:32:30Z

This adds:

Support for dynamically discovering and registering SchemaTransforms in the Java expansion service.
Support for dynamically discovering registered SchemaTransforms from the Python side.
Support for using SchemaTransforms in Python pipelines.

This addresses #23412

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Choose reviewer(s) and mention them in a comment (R: @username).
Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI.

chamikaramj · 2022-09-29T01:36:42Z

cc: @pabloem @johnjcasey @ahmedabu98

codecov · 2022-09-29T02:00:16Z

Codecov Report

Merging #23413 (aa6e6c7) into master (226bc97) will increase coverage by 0.23%.
The diff coverage is 77.58%.

@@            Coverage Diff             @@
##           master   #23413      +/-   ##
==========================================
+ Coverage   73.08%   73.31%   +0.23%     
==========================================
  Files         725      711      -14     
  Lines       96664    96195     -469     
==========================================
- Hits        70647    70530     -117     
+ Misses      24705    24353     -352     
  Partials     1312     1312

Flag	Coverage Δ
python	`83.13% <77.58%> (+0.40%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
sdks/python/apache_beam/transforms/external.py	`78.88% <77.19%> (-0.86%)`	⬇️
sdks/python/apache_beam/portability/common_urns.py	`100.00% <100.00%> (ø)`
...python/apache_beam/examples/complete/distribopt.py	`0.00% <0.00%> (-98.57%)`	⬇️
...hon/apache_beam/examples/cookbook/mergecontacts.py	`23.21% <0.00%> (-73.22%)`	⬇️
...he_beam/examples/cookbook/multiple_output_pardo.py	`37.50% <0.00%> (-58.34%)`	⬇️
...python/apache_beam/examples/wordcount_debugging.py	`38.77% <0.00%> (-57.15%)`	⬇️
...python/apache_beam/examples/dataframe/wordcount.py	`38.46% <0.00%> (-53.85%)`	⬇️
...s/python/apache_beam/examples/wordcount_minimal.py	`40.74% <0.00%> (-51.86%)`	⬇️
sdks/python/apache_beam/examples/wordcount.py	`48.27% <0.00%> (-44.83%)`	⬇️
...apache_beam/examples/cookbook/custom_ptransform.py	`53.65% <0.00%> (-41.58%)`	⬇️
... and 76 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

github-actions · 2022-09-29T03:00:17Z

Assigning reviewers. If you would like to opt out of this review, comment assign to next reviewer:

R: @AnandInguva for label python.
R: @lukecwik for label java.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

chamikaramj · 2022-09-29T19:39:27Z

R: @robertwb

github-actions · 2022-09-29T20:04:25Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/external_transforms.proto

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

robertwb · 2022-10-06T21:08:30Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    return transformProvider;
+  }
+
+  static class RowTransform extends PTransform {


Perhaps at least a 1-line docstring as to why this class exists?

Looking at this a bit more, inputs always come in (and leave) as dicts. Could we just preserve them in this case rather than adding another level of indirection?

Added a docstring explaining why this is needed. May be we could get rid of this in the future if PCollectionRowTuple becomes a portable type that is understood by all SDKs.

So this class basically exists to undo the logic at https://github.com/apache/beam/blob/release-2.42.0/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java#L348 .

I think we should instead override the createInput (and possibly extractOutputs) methods in ExpansionServiceSchemaTransformProvider.

Yeah, that worked! Updated. Thanks.

robertwb · 2022-10-06T21:10:55Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    }
+
+    @Override
+    public POutput expand(PInput input) {


For UI and uniqueness considerations, should we inherit the stage name of its underlying transform?

I don't think we can use the exact same name due to uniqueness requirement of runners but changed the name of this transform to a derivative of the underlying transform.

One can nest transforms of the same name. (In fact, nesting with distinct prefixes is how we get around users having to specify names everywhere.)

RowTransform class was removed. So this is obsolete now.

robertwb · 2022-10-06T21:31:59Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+      }
+      PCollectionRowTuple output = inputRowTuple.apply(this.rowTuplePTransform);
+
+      if (output.getAll().size() > 1) {


Is there any advantage to doing this? (In particular, the name of the single output gets lost.)

This just converts the PCollectionRowTuple to a PCollectionTuple. I think currently we need this conversion.

I meant why can't we just return the PCollectionRowTuple itself? (Possibly we need to update the expansion service code to handle this type, but that should be done anyway.)

RowTransform class was removed. So this is obsolete now.

robertwb · 2022-10-06T21:40:34Z

sdks/python/apache_beam/transforms/external.py

@@ -104,6 +105,43 @@ def payload(self):
    """
    return self.build().SerializeToString()

+  def _get_schema_proto_and_payload(self, ignored_arg_format, *args, **kwargs):


The ignored_arg_format param feels a bit awkward, and it'd be easy to accidentally forget it and grab the first arg as well (and I'm not sure the JavaClassLookupPayloadBuilder.IGNORED_ARG_FORMAT logic should be in this superclass as well).

Instead, could we have the java-class calling one just make a dict using JavaClassLookupPayloadBuilder.IGNORED_ARG_FORMAT for the positional args, and then call shared code to go from a dict to a schema and payload?

robertwb · 2022-10-06T21:42:15Z

sdks/python/apache_beam/transforms/external.py

+        self._expansion_service)
+
+  @staticmethod
+  def discover(expansion_service, regex=None):


For simplicity I think we can skip the regex filtering until we know we need it.

sdks/python/apache_beam/transforms/external.py

pabloem · 2022-10-15T22:40:00Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+  }
+
+  Iterable<org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider> getAllProviders() {
+    return schemaTransformProviders.values();


so if a provider is not (duh) provided, even if the SchemaTransform implementation exists, it would not show up in the Discover call. Is that fine?

I would hope that we can eventually replace writing a Provider with a decorator on the (suitably configured) PTransform class itself, but that's future work.

Ack. Resolving.

chamikaramj · 2022-10-17T19:55:55Z

@pabloem I don't think a SchemaTransform is useful without the provider. Is that not the case ?

BTW I was held up due to some other work but hope to address comments here and push a change soon.

…ding SchemaTransforms

chamikaramj

Thanks. PTAL.

chamikaramj · 2022-10-13T23:52:33Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    return transformProvider;
+  }
+
+  static class RowTransform extends PTransform {


Added a docstring explaining why this is needed. May be we could get rid of this in the future if PCollectionRowTuple becomes a portable type that is understood by all SDKs.

chamikaramj · 2022-10-14T00:12:19Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    }
+
+    @Override
+    public POutput expand(PInput input) {


I don't think we can use the exact same name due to uniqueness requirement of runners but changed the name of this transform to a derivative of the underlying transform.

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

chamikaramj · 2022-10-21T22:21:19Z

sdks/python/apache_beam/transforms/external.py

@@ -104,6 +105,43 @@ def payload(self):
    """
    return self.build().SerializeToString()

+  def _get_schema_proto_and_payload(self, ignored_arg_format, *args, **kwargs):


chamikaramj · 2022-10-21T22:29:04Z

sdks/python/apache_beam/transforms/external.py

+        self._expansion_service)
+
+  @staticmethod
+  def discover(expansion_service, regex=None):


chamikaramj · 2022-10-24T22:47:53Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+      }
+      PCollectionRowTuple output = inputRowTuple.apply(this.rowTuplePTransform);
+
+      if (output.getAll().size() > 1) {


This just converts the PCollectionRowTuple to a PCollectionTuple. I think currently we need this conversion.

chamikaramj · 2022-10-25T18:04:21Z

Retest this please

chamikaramj · 2022-10-28T00:25:02Z

Run Java PreCommit

chamikaramj · 2022-10-28T00:25:20Z

Run Python PreCommit

chamikaramj · 2022-10-31T18:42:19Z

Friendly ping :)

robertwb · 2022-11-04T00:48:13Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+
+    @Override
+    public String getName() {
+      return "RowTransform_of_" + this.rowTuplePTransform.getName();


Does the "RowTransform_of_" prefix add value? Maybe simply drop it, as this will be the name the user sees before drilling down.

RowTransform class was removed. So this is obsolete now.

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/external_transforms.proto

robertwb · 2022-11-04T00:51:41Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    }
+
+    @Override
+    public POutput expand(PInput input) {


One can nest transforms of the same name. (In fact, nesting with distinct prefixes is how we get around users having to specify names everywhere.)

robertwb · 2022-11-04T00:53:42Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+      }
+      PCollectionRowTuple output = inputRowTuple.apply(this.rowTuplePTransform);
+
+      if (output.getAll().size() > 1) {


I meant why can't we just return the PCollectionRowTuple itself? (Possibly we need to update the expansion service code to handle this type, but that should be done anyway.)

robertwb · 2022-11-04T00:59:25Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    return transformProvider;
+  }
+
+  static class RowTransform extends PTransform {


So this class basically exists to undo the logic at https://github.com/apache/beam/blob/release-2.42.0/sdks/java/expansion-service/src/main/java/org/apache/beam/sdk/expansion/service/ExpansionService.java#L348 .

I think we should instead override the createInput (and possibly extractOutputs) methods in ExpansionServiceSchemaTransformProvider.

robertwb · 2022-11-04T01:00:49Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+  }
+
+  Iterable<org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider> getAllProviders() {
+    return schemaTransformProviders.values();


I would hope that we can eventually replace writing a Provider with a decorator on the (suitably configured) PTransform class itself, but that's future work.

robertwb · 2022-11-04T01:02:38Z

sdks/python/apache_beam/transforms/external.py

+    # Expand the transform using the expansion service and the config_row.
+    if self._expansion_service is None:
+      self._expansion_service = BeamJarExpansionService(
+          ':sdks:java:expansion-service:app:shadowJar',


Do we expect many schema transforms to live in this jar, or should we make identifying the expansion service mandatory?

I don't expect many (or any) SchemaTransforms to live in this jar. So I believe we should make specifying an expansion service mandatory.

robertwb · 2022-11-04T01:04:15Z

sdks/python/apache_beam/transforms/external_test.py

+    ComplexType = typing.NamedTuple(
+        "ComplexType", [
+            ("str_sub_field", str),
+            ("int_sub_field", np.int32),


Can't we just use int here? (Does that give us an int64?)

Yeah just "int" works. Updated.

chamikaramj · 2022-11-08T19:17:08Z

Thanks. Hope to respond by the end of the week (have been bit busy due to the release).

chamikaramj

Thanks Robert. PTAL.

model/pipeline/src/main/proto/org/apache/beam/model/pipeline/v1/external_transforms.proto

chamikaramj · 2022-11-13T23:00:00Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    return transformProvider;
+  }
+
+  static class RowTransform extends PTransform {


Yeah, that worked! Updated. Thanks.

chamikaramj · 2022-11-13T23:00:47Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    }
+
+    @Override
+    public POutput expand(PInput input) {


RowTransform class was removed. So this is obsolete now.

chamikaramj · 2022-11-13T23:01:59Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+      }
+      PCollectionRowTuple output = inputRowTuple.apply(this.rowTuplePTransform);
+
+      if (output.getAll().size() > 1) {


RowTransform class was removed. So this is obsolete now.

chamikaramj · 2022-11-13T23:02:08Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+
+    @Override
+    public String getName() {
+      return "RowTransform_of_" + this.rowTuplePTransform.getName();


RowTransform class was removed. So this is obsolete now.

chamikaramj · 2022-11-13T23:02:32Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+  }
+
+  Iterable<org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider> getAllProviders() {
+    return schemaTransformProviders.values();


Ack. Resolving.

chamikaramj · 2022-11-15T02:24:46Z

sdks/python/apache_beam/transforms/external.py

+    # Expand the transform using the expansion service and the config_row.
+    if self._expansion_service is None:
+      self._expansion_service = BeamJarExpansionService(
+          ':sdks:java:expansion-service:app:shadowJar',


I don't expect many (or any) SchemaTransforms to live in this jar. So I believe we should make specifying an expansion service mandatory.

chamikaramj · 2022-11-15T02:27:59Z

sdks/python/apache_beam/transforms/external_test.py

+    ComplexType = typing.NamedTuple(
+        "ComplexType", [
+            ("str_sub_field", str),
+            ("int_sub_field", np.int32),


Yeah just "int" works. Updated.

robertwb · 2022-11-22T00:17:59Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+      return PCollectionRowTuple.of(
+          DEFAULT_INPUT_TAG, (PCollection<Row>) inputs.values().iterator().next());
+    } else {
+      PCollectionRowTuple inputRowTuple = PCollectionRowTuple.empty(p);


Can't we use this for all arities, rather than special-casing the 0 and 1 sizes above?

robertwb · 2022-11-22T00:18:16Z

...main/java/org/apache/beam/sdk/expansion/service/ExpansionServiceSchemaTransformProvider.java

+    } else if (output.getAll().size() == 1) {
+      return ImmutableMap.of("output", output.getAll().values().iterator().next());
+    } else {
+      ImmutableMap.Builder<String, PCollection<?>> pCollectionMap = ImmutableMap.builder();


chamikaramj · 2022-11-23T00:05:21Z

Thanks!

chamikaramj · 2022-11-23T05:16:34Z

Run Java PreCommit

chamikaramj · 2022-11-23T05:16:48Z

Run GoPortable PreCommit

chamikaramj · 2022-11-23T17:03:29Z

Run Java PreCommit

chamikaramj · 2022-11-23T18:41:00Z

Run Java PreCommit

chamikaramj · 2022-11-23T22:21:25Z

Run GoPortable PreCommit

chamikaramj · 2022-11-23T22:21:40Z

Run Java_Examples_Dataflow PreCommit

chamikaramj · 2022-11-23T23:39:38Z

Run GoPortable PreCommit

chamikaramj · 2022-11-24T00:12:07Z

Run Java_Examples_Dataflow_Java11 PreCommit

chamikaramj · 2022-11-24T00:12:17Z

Run Java_Examples_Dataflow_Java17 PreCommit

…ding SchemaTransforms (apache#23413) * Updates ExpansionService to support dynamically discovering and expanding SchemaTransforms * Fixing checker framework errors. * Address reviewer comments * Addressing reviewer comments * Addressing reviewer comments

github-actions bot added java model python labels Sep 29, 2022

chamikaramj requested a review from robertwb September 29, 2022 01:32

github-actions bot added the Next Action: Reviewers label Sep 29, 2022

TheNeuralBit self-requested a review October 4, 2022 21:08

robertwb reviewed Oct 6, 2022

View reviewed changes

pabloem reviewed Oct 15, 2022

View reviewed changes

chamikaramj added 3 commits October 24, 2022 16:42

Updates ExpansionService to support dynamically discovering and expan…

d494711

…ding SchemaTransforms

Fixing checker framework errors.

a0d7f37

Address reviewer comments

c1a90aa

chamikaramj commented Oct 24, 2022

View reviewed changes

robertwb reviewed Nov 4, 2022

View reviewed changes

chamikaramj force-pushed the easy_x-lang_final branch from 1dde786 to a36c489 Compare November 15, 2022 03:28

Addressing reviewer comments

82a0f28

chamikaramj force-pushed the easy_x-lang_final branch from a36c489 to 82a0f28 Compare November 15, 2022 03:31

chamikaramj commented Nov 15, 2022

View reviewed changes

robertwb approved these changes Nov 22, 2022

View reviewed changes

Addressing reviewer comments

aa6e6c7

chamikaramj closed this Nov 23, 2022

chamikaramj deleted the easy_x-lang_final branch November 23, 2022 00:00

chamikaramj restored the easy_x-lang_final branch November 23, 2022 00:01

chamikaramj reopened this Nov 23, 2022

chamikaramj merged commit 23e59af into apache:master Nov 24, 2022

This was referenced Nov 24, 2022

[Feature Request]: Easy Multi-language via a SchemaTransform-aware Expansion Service #23412

Closed

Updates ExpansionService to support dynamically discovering and expanding SchemaTransforms #22802

Closed

Updates ExpansionService to support dynamically discovering and expanding SchemaTransforms #23413

Updates ExpansionService to support dynamically discovering and expanding SchemaTransforms #23413

Conversation

chamikaramj commented Sep 29, 2022

GitHub Actions Tests Status (on master branch)

chamikaramj commented Sep 29, 2022

codecov bot commented Sep 29, 2022 • edited Loading

Codecov Report

github-actions bot commented Sep 29, 2022

chamikaramj commented Sep 29, 2022

github-actions bot commented Sep 29, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chamikaramj commented Oct 17, 2022

chamikaramj left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chamikaramj commented Oct 25, 2022

chamikaramj commented Oct 28, 2022

chamikaramj commented Oct 28, 2022

chamikaramj commented Oct 31, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chamikaramj commented Nov 8, 2022

chamikaramj left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 23, 2022

chamikaramj commented Nov 24, 2022

chamikaramj commented Nov 24, 2022

codecov bot commented Sep 29, 2022 •

edited

Loading