From d5454852e20558c33fc9ee8a6a87c448ea8018e8 Mon Sep 17 00:00:00 2001 From: JanakiRamChittineni <341781521+raams95@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:50:03 -0500 Subject: [PATCH] [Python] Infer the right type for slices of fixed-size tuples Type inference gave a slice of a fixed-size Tuple the type of the whole tuple, so for example row[1:] on Tuple[str, str, int] was inferred as Tuple[str, str, int]. A following CombinePerKey then failed the type check, and encoding the slice used the coder for the whole tuple. Slices with constant bounds now select the element types, and other slices of a fixed-size tuple give Tuple[Union[...], ...]. BUILD_SLICE keeps constant bounds as a Const slice, so Python 3.10 to 3.14 infer the same types. Fixes #20899 --- CHANGES.md | 1 + sdks/python/apache_beam/typehints/opcodes.py | 33 ++++++++++++++++--- .../typehints/trivial_inference_test.py | 17 ++++++++++ .../typehints/typed_pipeline_test.py | 11 +++++++ 4 files changed, 57 insertions(+), 5 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index d06940d12cf4..03522eb7a536 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -104,6 +104,7 @@ * (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()` instead of silently dropping them ([#40519](https://github.com/apache/beam/issues/40519)). * (Python) BigQuery streaming inserts now reset the retry budget and backoff for each batch instead of sharing them across a bundle ([#39995](https://github.com/apache/beam/pull/39995)). * (Python) State write failures reported by portable runners now fail the bundle instead of silently losing state updates ([#39992](https://github.com/apache/beam/pull/39992)). +* (Python) Type inference now gives the right type for a slice of a fixed-size tuple, such as `row[1:]` in `beam.Map(lambda row: row[1:])`. The slice used to keep the type of the whole tuple, which could make a following `CombinePerKey` fail with a `TypeCheckError` or encode elements with the wrong coder ([#20899](https://github.com/apache/beam/issues/20899)). * (Go) Fixed a data race on the Prism runner's artifact cache map in JobServices ([#32656](https://github.com/apache/beam/issues/32656)). * (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)). * (Go) Wrap last dataWriter/timerWriter send errors on Close ([#40405](https://github.com/apache/beam/issues/40405)). diff --git a/sdks/python/apache_beam/typehints/opcodes.py b/sdks/python/apache_beam/typehints/opcodes.py index c5dda817d0d2..a4dbb869b113 100644 --- a/sdks/python/apache_beam/typehints/opcodes.py +++ b/sdks/python/apache_beam/typehints/opcodes.py @@ -222,6 +222,22 @@ def binary_true_divide(state, arg): binary_subtract = inplace_subtract = symmetric_binary_op +def _slice_type(base, index): + """Returns the type of base[index] for an indexable base and a slice index. + + The index is a Const if the bounds of the slice are known. + """ + if isinstance(base, typehints.TupleHint.TupleConstraint): + # Slicing changes the length of a fixed-size tuple. + if isinstance(index, Const): + try: + return Tuple[base.tuple_types[index.value]] + except (TypeError, ValueError): + pass + return Tuple[element_type(base), ...] + return base + + def binary_subscr(state, unused_arg): index = state.stack.pop() base = Const.unwrap(state.stack.pop()) @@ -237,7 +253,7 @@ def binary_subscr(state, unused_arg): base, typehints.IndexableTypeConstraint): # The slice is treated as a const in 3.14, using this instead of # BINARY_SLICE - out = base + out = _slice_type(base, index) else: out = element_type(base) state.stack.append(out) @@ -259,13 +275,16 @@ def store_subscr(unused_state, unused_args): def binary_slice(state, args): - _ = state.stack.pop() - _ = state.stack.pop() + stop = state.stack.pop() + start = state.stack.pop() base = Const.unwrap(state.stack.pop()) if base is str: out = base elif isinstance(base, typehints.IndexableTypeConstraint): - out = base + if isinstance(start, Const) and isinstance(stop, Const): + out = _slice_type(base, Const(slice(start.value, stop.value))) + else: + out = _slice_type(base, slice) else: out = element_type(base) state.stack.append(out) @@ -667,7 +686,11 @@ def make_closure(state, arg): def build_slice(state, arg): - state.stack[-arg:] = [slice] # a slice object + bounds = state.stack[-arg:] + if all(isinstance(b, Const) for b in bounds): + state.stack[-arg:] = [Const(slice(*(b.value for b in bounds)))] + else: + state.stack[-arg:] = [slice] # a slice object def to_bool(state, arg): diff --git a/sdks/python/apache_beam/typehints/trivial_inference_test.py b/sdks/python/apache_beam/typehints/trivial_inference_test.py index dcb0bac97e80..f704fb2db26f 100644 --- a/sdks/python/apache_beam/typehints/trivial_inference_test.py +++ b/sdks/python/apache_beam/typehints/trivial_inference_test.py @@ -112,6 +112,23 @@ def testGetItemSlice(self): typehints.List[int], lambda v: v[::-1], [typehints.List[int]]) self.assertReturnType( typehints.Tuple[int], lambda v: v[::-1], [typehints.Tuple[int]]) + self.assertReturnType( + typehints.Tuple[str, int], lambda v: v[1:], + [typehints.Tuple[int, str, int]]) + self.assertReturnType( + typehints.Tuple[int, str], lambda v: v[:-1], + [typehints.Tuple[int, str, float]]) + self.assertReturnType( + typehints.Tuple[float, str, int], lambda v: v[::-1], + [typehints.Tuple[int, str, float]]) + self.assertReturnType( + typehints.Tuple[int, float], lambda v: v[0:3:2], + [typehints.Tuple[int, str, float]]) + self.assertReturnType( + typehints.Tuple[()], lambda v: v[5:], [typehints.Tuple[int, str]]) + self.assertReturnType( + typehints.Tuple[typehints.Union[int, str], ...], lambda v, i: v[i:], + [typehints.Tuple[int, str, int], int]) self.assertReturnType(str, lambda v: v[::-1], [str]) self.assertReturnType(typehints.Any, lambda v: v[::-1], [typehints.Any]) self.assertReturnType(typehints.Any, lambda v: v[::-1], [object]) diff --git a/sdks/python/apache_beam/typehints/typed_pipeline_test.py b/sdks/python/apache_beam/typehints/typed_pipeline_test.py index c97a1977e9c3..29cb97dd7872 100644 --- a/sdks/python/apache_beam/typehints/typed_pipeline_test.py +++ b/sdks/python/apache_beam/typehints/typed_pipeline_test.py @@ -461,6 +461,17 @@ def fn(element: int) -> typehints.Optional[int]: result = [1, 2, 3] | beam.Map(fn) self.assertCountEqual([None, 2, 3], result) + def test_inferred_type_of_tuple_slice(self): + # https://github.com/apache/beam/issues/20899 + rows = [('day1', 'a', 1), ('day2', 'a', 2), ('day1', 'b', 3)] + with TestPipeline() as p: + result = ( + p + | beam.Create(rows).with_output_types(typing.Tuple[str, str, int]) + | beam.Map(lambda row: row[1:]) + | beam.CombinePerKey(sum)) + assert_that(result, equal_to([('a', 3), ('b', 3)])) + def test_typed_flatmap(self): def fn(element: int) -> typehints.Iterable[int]: yield element * 2