[Feature Request]: Support TypedDict hint in Apache Beam.

[tvalentyn]

What would you like to happen?

CoGroupByKey produces Dict[K, Union[...]], which is narrower than TypedDict. If the user follows the CoGroupByKey with a ParDo where the function takes TypeDict as input, e.g. ParDo(fn(element: TypedDict)), then the type check in Beam fails.

Users having types like:

myTypeDict = TypedDict(
    'myTypeDict', {
        'field_1': List[Field1Type],
        'field_2': List[Field2Type],
    })

are using less specific hints like myLessSpecificDict = Dict(str, Union[List[Field1Type], List[Field2Type]]) in order to satisfy beams typing.

We could have a first-class typehint TypeDict, which may be easier once we deprecate Python 3.7.

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

[AnandInguva]

Update the issue title?

[tvalentyn]

Done, thanks.

[bfontaine]

Why have specific code for TypedDict when you can use the more generic Mapping, which accepts not only TypedDict, but also any kind of dict?

Note: there’s a typo in the issue title: it should be "TypedDict", not "TypeDict".

[aaronsarna]

The value of TypedDict over Mapping is if the different keys are expected to have different values. For example, let's say I did the following:

 to_group = {
  'ints': pcoll1,  # PCollection of (str, int)
  'strs': pcoll2,  # PCollection of (str, str)
}
grouped = to_group | beam.CoGroupByKey()
something = grouped | beam.MapTuple(map_grouped)

If using Mapping, the method would have to start:

def map_grouped(key: str, values: Mapping[str, Iterable[str | int]]):
  ints = values['ints']  # type: Iterable[int]
  strs = values['strs']  # type: Iterable[str]
  ...

since the type checker can't tell which key corresponds with which value type.

If I could instead have a TypedDict, it could be a bit cleaner:

class GroupedValues(TypedDict):
  ints: Iterable[int]
  strs: Iterable[str]

def map_grouped(key: str, values: GroupedValues):
  ints = values['ints']
  strs = values['strs']

and the type checker needs no additional help.

[tvalentyn]

cc: @jrmccluskey