[Python] Infer the right type for slices of fixed-size tuples

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)).

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):

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])

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