jrmccluskey commented on code in PR #34897:
URL: https://github.com/apache/beam/pull/34897#discussion_r2082135603


##########
sdks/python/apache_beam/ml/anomaly/transforms.py:
##########
@@ -37,13 +38,16 @@
 from apache_beam.ml.anomaly.specifiable import Specifiable
 from apache_beam.ml.inference.base import RunInference
 from apache_beam.transforms.userstate import ReadModifyWriteStateSpec
+from apache_beam.typehints.typehints import TupleConstraint
 
 KeyT = TypeVar('KeyT')
-TempKeyT = TypeVar('TempKeyT', bound=int)
-InputT = tuple[KeyT, beam.Row]
-KeyedInputT = tuple[KeyT, tuple[TempKeyT, beam.Row]]
-KeyedOutputT = tuple[KeyT, tuple[TempKeyT, AnomalyResult]]
-OutputT = tuple[KeyT, AnomalyResult]
+TempKeyT = TypeVar('TempKeyT', bound=str)
+InputT = beam.Row
+OutputT = AnomalyResult

Review Comment:
   Nit: these type aliases aren't strictly necessary given that they're 
specific single types, but since you use them to compose the tuple types 
they're fine to keep



##########
sdks/python/apache_beam/ml/anomaly/transforms.py:
##########
@@ -600,20 +610,43 @@ def expand(
     #
     # We select uuid.uuid1() for its inclusion of node information, making it
     # more suitable for parallel execution environments.
-    add_temp_key_fn: Callable[[InputT], KeyedInputT] \
-        = lambda e: (e[0], (str(uuid.uuid1()), e[1]))
-    keyed_input = (input | "Add temp key" >> beam.Map(add_temp_key_fn))
 
+    if isinstance(input.element_type, TupleConstraint):
+      keyed_input = input
+    else:
+      # Add a None key if the input is unkeyed.
+      keyed_input = input | beam.WithKeys(None)

Review Comment:
   Nit: I'd shy away from a `None` key and use a real-valued constant. Once 
again not wholly necessary, just a thought



-- 
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-unsubscr...@beam.apache.org

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

Reply via email to