[Feature Request]: In Python how to assign types to multi-output tags?

[cozos]

What would you like to happen?

I have a multi-output DoFn:

    class DoFn1:
        def process(self, row) -> Iterable[Union[Dict[str, Any], pvalue.TaggedOutput]]:
            if something:
                yield some_dict(...)
            else:
                yield pvalue.TaggedOutput("bad", ...)

And another DoFn that consumes its outputs

    class DoFn2:
        def process(self, row: Dict[str, Any]) -> Iterable[...]:
            if something:
                yield some_dict(...)
            else:
                yield pvalue.TaggedOutput("bad", ...)

And then when I use it like this:

    pcoll = ...
    pcoll = pcoll | "dofn1" >> beam.ParDo(DoFn1()).with_outputs(
        "bad",
        main="good",
    )
    pcoll["good"] | "dofn2" >> beam.ParDo(DoFn2())

I get an error that looks like this:

apache_beam.typehints.decorators.TypeCheckError: Type hint violation for 'dofn2': requires Dict[str, Any] 
but got Union[Dict[str, Any], TaggedOutput] for row

How can I tell Beam type checking that the good tag gives Dict[str, Any], and bad gives TaggedOutput? Am I forced to propagate the tagged output type hints?

Also asked in SO here if not appropriate for Github: https://stackoverflow.com/questions/76151787/in-apache-beam-dataflow-multi-output-dofns-how-do-you-assign-type-hints-to-spec

I found a related discussion here: #9810

Issue Priority

Priority: 3 (nice-to-have improvement)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[tvalentyn]

Have you tried using .with_input_types ?

beam/sdks/python/apache_beam/transforms/ptransform.py

Line 871 in c31243e

def with_input_types(

[RXminuS]

@tvalentyn I'm not quite understanding what you mean. I think this issue is more referring to the fact that it doesn't really seam possible to set the return type of a expand to be either a Tuple or a TaggedPCollection so that you get typed outputs for each of the keys/indexes.

I'm stuck on this as well breaking both static typing with a warning

WARNING  | apache_beam.typehints.decorators:strip_pcoll_helper:381 - This output type hint will be ignored and not used for type-checking purposes. Typically, output type hints for a PTransform are single (or nested) types wrapped by a PCollection, PDone, or None. Got: Tuple[apache_beam.pvalue.PCollection[XXX], apache_beam.pvalue.PCollection[YYY]] instead.

Nor can I perform runtime type-checking correctly

[tvalentyn]

In the context of example in the description, I am asking whether we tried to replace

    pcoll["good"] | "dofn2" >> beam.ParDo(DoFn2())

to

    pcoll["good"] | "dofn2" >> beam.ParDo(DoFn2()).with_input_types(<correct typehint for elements belonging to good pcoll, which seems to be Dict[str, Any]>)

[tvalentyn]

also cc: @jrmccluskey who is an expert on this topic.

[lazarillo]

Any updates on this? I am having trouble with the same thing as the OP.

[tvalentyn]

@lazarillo have you tried using with_input_types on transforms that consume tagged output?

[lazarillo]

Yes, I have. I think you are correct and it will solve the issue. Unfortunately, I have multiple issues going on, so for the moment, I cannot fully confirm that it worked.

[robmoore]

I tried this approach with beam.io.WriteToCsv without success.

[tvalentyn]

@robmoore could you elaborate please?