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