apache / apache/beam

Element type inference doesn't work for multi-output DoFns

Open
#18,957 0 comments 0 reactions 0 assignees View on GitHub
bug core P3 python types
Dominant language
Java
Stars
8.7k
Forks
4.7k
Avg merge
1d 20h
Merged PRs (30d)
196

Description

TLDR: if you have a multi-output DoFn, then the non-main PCollections with incorrectly have their element types set to None. This affects type checking for pipelines involving these PCollections.

Minimal example:
```

import apache_beam as beam

class TripleDoFn(beam.DoFn):
def process(self, elem):
yield_elem

if elem % 2 == 0:
yield beam.pvalue.TaggedOutput('ten_times', elem * 10)
if elem % 3
== 0:
yield beam.pvalue.TaggedOutput('hundred_times', elem * 100)

@beam.typehints.with_input_types(int)
@beam.typehints.with_output_types(int)
class
MultiplyBy(beam.DoFn):
def __init__(self, multiplier):
self._multiplier = multiplier

def
process(self, elem):
return elem * self._multiplier

def main():
with beam.Pipeline() as
p:
x, a, b = (
p
| 'Create' >> beam.Create([1, 2, 3])
| 'TripleDo' >> beam.ParDo(TripleDoFn()).with_outputs(

'ten_times', 'hundred_times', main='main_output'))

_ = a | 'MultiplyBy2' >> beam.ParDo(MultiplyBy(2))

if
__name__ == '__main__':
main()

```

Running this yields the following error:
```

apache_beam.typehints.decorators.TypeCheckError: Type hint violation for 'MultiplyBy2': requires but got None for elem

```

Replacing `a` with `b` yields the same error. Replacing `a` with `x` instead yields the following error:
```

apache_beam.typehints.decorators.TypeCheckError: Type hint violation for 'MultiplyBy2': requires but got Union[TaggedOutput, int] for elem

```

I would expect Beam to correctly infer that `a` and `b` have element types of `int` rather than `None`, and I would also expect Beam to correctly figure out that the element types of `x` are compatible with `int`.

Imported from Jira [BEAM-4132](https://issues.apache.org/jira/browse/BEAM-4132). Original Jira may contain additional context.
Reported by: chuanyu.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.