You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "liferoad (via GitHub)" <gi...@apache.org> on 2025/08/02 14:19:08 UTC

[PR] Fix ml transform [beam]

liferoad opened a new pull request, #35774:
URL: https://github.com/apache/beam/pull/35774

   **Please** add a meaningful description for your change here
   
   ------------------------
   
   Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
   
    - [ ] 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](https://www.apache.org/licenses/icla.pdf).
   
   See the [Contributor Guide](https://beam.apache.org/contribute) for more tips on [how to make review process smoother](https://github.com/apache/beam/blob/master/CONTRIBUTING.md#make-the-reviewers-job-easier).
   
   To check the build health, please visit [https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md](https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md)
   
   GitHub Actions Tests Status (on master branch)
   ------------------------------------------------------------------------------------------------
   [![Build python source distribution and wheels](https://github.com/apache/beam/actions/workflows/build_wheels.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Build+python+source+distribution+and+wheels%22+branch%3Amaster+event%3Aschedule)
   [![Python tests](https://github.com/apache/beam/actions/workflows/python_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Python+Tests%22+branch%3Amaster+event%3Aschedule)
   [![Java tests](https://github.com/apache/beam/actions/workflows/java_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Java+Tests%22+branch%3Amaster+event%3Aschedule)
   [![Go tests](https://github.com/apache/beam/actions/workflows/go_tests.yml/badge.svg?event=schedule&&?branch=master)](https://github.com/apache/beam/actions?query=workflow%3A%22Go+tests%22+branch%3Amaster+event%3Aschedule)
   
   See [CI.md](https://github.com/apache/beam/blob/master/CI.md) for more information about GitHub Actions CI or the [workflows README](https://github.com/apache/beam/blob/master/.github/workflows/README.md) to see a list of phrases to trigger workflows.
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "derrickaw (via GitHub)" <gi...@apache.org>.
derrickaw commented on PR #35774:
URL: https://github.com/apache/beam/pull/35774#issuecomment-3162138303

   For #35769


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "charlespnh (via GitHub)" <gi...@apache.org>.
charlespnh commented on PR #35774:
URL: https://github.com/apache/beam/pull/35774#issuecomment-3182006307

   Does this PR also address https://github.com/apache/beam/issues/35768, i.e. making tft install optional now?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "liferoad (via GitHub)" <gi...@apache.org>.
liferoad commented on code in PR #35774:
URL: https://github.com/apache/beam/pull/35774#discussion_r2250001330


##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
 
+  if transforms and any(t.get('type') == 'SentenceTransformerEmbeddings'
+                        for t in transforms):

Review Comment:
   we can also add `InferenceAPIEmbeddings`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #35774:
URL: https://github.com/apache/beam/pull/35774#discussion_r2267988918


##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -512,13 +503,36 @@ def ml_transform(
     raise ValueError(
         'tensorflow-transform must be installed to use this MLTransform')
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
-  # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
-  # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result_ml_transform = MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
 
-
-if tft is not None:
+  if transforms and any(t.get('type', '').endswith('Embeddings')
+                        for t in transforms):
+    from apache_beam.typehints import List
+    try:
+      if pcoll.element_type:
+        new_fields = named_fields_from_element_type(pcoll.element_type)
+        columns_to_change = set()
+        for t_spec in transforms:
+          if t_spec.get('type', '').endswith('Embeddings'):
+            columns_to_change.update(
+                t_spec.get('config', {}).get('columns', []))
+
+        final_fields = []
+        for name, typ in new_fields:
+          if name in columns_to_change:
+            final_fields.append((name, List[float]))
+          else:
+            final_fields.append((name, typ))
+        output_schema = RowTypeConstraint.from_fields(final_fields)
+        return pcoll | result_ml_transform.with_output_types(output_schema)
+    except TypeError:
+      # If we can't get a schema, just return the result.
+      pass
+  return pcoll | result_ml_transform

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   This block of code can be refactored for better readability and efficiency. The current implementation iterates over `transforms` twice: once in `any()` and again in the `for` loop. This can be optimized by iterating only once to find all embedding transforms and then processing them. Using comprehensions can also make the code more concise and Pythonic.
   
   ```python
     if transforms:
       embedding_transforms = [
           t for t in transforms if t.get('type', '').endswith('Embeddings')
       ]
       if embedding_transforms:
         from apache_beam.typehints import List
         try:
           if pcoll.element_type:
             columns_to_change = {
                 col
                 for t_spec in embedding_transforms
                 for col in t_spec.get('config', {}).get('columns', [])
             }
             new_fields = named_fields_from_element_type(pcoll.element_type)
             final_fields = [
                 (name, List[float] if name in columns_to_change else typ)
                 for name, typ in new_fields
             ]
             output_schema = RowTypeConstraint.from_fields(final_fields)
             return pcoll | result_ml_transform.with_output_types(output_schema)
         except TypeError:
           # If we can't get a schema, just return the result.
           pass
     return pcoll | result_ml_transform
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "liferoad (via GitHub)" <gi...@apache.org>.
liferoad commented on code in PR #35774:
URL: https://github.com/apache/beam/pull/35774#discussion_r2267985235


##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
 
+  if transforms and any(t.get('type') == 'SentenceTransformerEmbeddings'
+                        for t in transforms):
+    from apache_beam.typehints import List
+    try:
+      if pcoll.element_type:
+        new_fields = named_fields_from_element_type(pcoll.element_type)
+        columns_to_change = set()
+        for t_spec in transforms:
+          if t_spec.get('type') == 'SentenceTransformerEmbeddings':
+            columns_to_change.update(
+                t_spec.get('config', {}).get('columns', []))
+
+        final_fields = []
+        for name, typ in new_fields:
+          if name in columns_to_change:
+            final_fields.append((name, List[float]))
+          else:
+            final_fields.append((name, typ))
+        output_schema = RowTypeConstraint.from_fields(final_fields)
+        return result | beam.Map(lambda x: x).with_output_types(output_schema)

Review Comment:
   Good idea. Done.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on code in PR #35774:
URL: https://github.com/apache/beam/pull/35774#discussion_r2260544316


##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
 
+  if transforms and any(t.get('type') == 'SentenceTransformerEmbeddings'
+                        for t in transforms):
+    from apache_beam.typehints import List
+    try:
+      if pcoll.element_type:
+        new_fields = named_fields_from_element_type(pcoll.element_type)
+        columns_to_change = set()
+        for t_spec in transforms:
+          if t_spec.get('type') == 'SentenceTransformerEmbeddings':
+            columns_to_change.update(
+                t_spec.get('config', {}).get('columns', []))
+
+        final_fields = []
+        for name, typ in new_fields:
+          if name in columns_to_change:
+            final_fields.append((name, List[float]))
+          else:
+            final_fields.append((name, typ))
+        output_schema = RowTypeConstraint.from_fields(final_fields)
+        return result | beam.Map(lambda x: x).with_output_types(output_schema)

Review Comment:
   Rather than applying an extra identity map, could we store the transform above with something like:
   
   ```
   to_be_applied = MLTransform(
         write_artifact_location=write_artifact_location,
         read_artifact_location=read_artifact_location,
         transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
   ```
   
   Then here we could:
   
   ```
   return result | to_be_applied.with_output_types(output_schema)
   ```



##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
 
+  if transforms and any(t.get('type') == 'SentenceTransformerEmbeddings'

Review Comment:
   We can probably get rid of the tft check above (or at least modify it) since you don't need tft installed for embeddings



##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
 
+  if transforms and any(t.get('type') == 'SentenceTransformerEmbeddings'
+                        for t in transforms):

Review Comment:
   Can we do a generic match of anything ending in `Embeddings`?



##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])

Review Comment:
   Will _config_to_obj work? I don't think it will right now.



##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])

Review Comment:
   Oh I see, it will work but only if tft is installed. Again, so we should probably also address that so it is not required.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "liferoad (via GitHub)" <gi...@apache.org>.
liferoad commented on code in PR #35774:
URL: https://github.com/apache/beam/pull/35774#discussion_r2249263183


##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -514,11 +509,36 @@ def ml_transform(
   options.YamlOptions.check_enabled(pcoll.pipeline, 'ML')
   # TODO(robertwb): Perhaps _config_to_obj could be pushed into MLTransform
   # itself for better cross-language support?
-  return pcoll | MLTransform(
+  result = pcoll | MLTransform(
       write_artifact_location=write_artifact_location,
       read_artifact_location=read_artifact_location,
       transforms=[_config_to_obj(t) for t in transforms] if transforms else [])
 
+  if transforms and any(t.get('type') == 'SentenceTransformerEmbeddings'
+                        for t in transforms):

Review Comment:
   this is ugly.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "charlespnh (via GitHub)" <gi...@apache.org>.
charlespnh commented on PR #35774:
URL: https://github.com/apache/beam/pull/35774#issuecomment-3188570639

   > > Does this PR also address #35768, i.e. making tft install optional now?
   > 
   > It should be. Let me know if you still have the issue.
   
   Can confirm it works now. Thanks a lot!
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "liferoad (via GitHub)" <gi...@apache.org>.
liferoad merged PR #35774:
URL: https://github.com/apache/beam/pull/35774


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "damccorm (via GitHub)" <gi...@apache.org>.
damccorm commented on code in PR #35774:
URL: https://github.com/apache/beam/pull/35774#discussion_r2269873106


##########
sdks/python/apache_beam/yaml/yaml_ml.py:
##########
@@ -512,13 +503,36 @@ def ml_transform(
     raise ValueError(

Review Comment:
   We need to remove this check as well (will also fix linting issues)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #35774:
URL: https://github.com/apache/beam/pull/35774#issuecomment-3146951821

   Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #35774:
URL: https://github.com/apache/beam/pull/35774#issuecomment-3148501639

   Assigning reviewers:
   
   R: @shunping for label python.
   
   
   Note: If you would like to opt out of this review, comment `assign to next reviewer`.
   
   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).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "liferoad (via GitHub)" <gi...@apache.org>.
liferoad commented on PR #35774:
URL: https://github.com/apache/beam/pull/35774#issuecomment-3176845857

   /gemini review


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] Fix ml transform [beam]

Posted by "liferoad (via GitHub)" <gi...@apache.org>.
liferoad commented on PR #35774:
URL: https://github.com/apache/beam/pull/35774#issuecomment-3183917016

   > Does this PR also address #35768, i.e. making tft install optional now?
   
   It should be. Let me know if you still have the issue.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@beam.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org