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