Repository navigation
[Python] Infer the right type for slices of fixed-size tuples - #40680
Conversation
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 apache#20899
There was a problem hiding this comment.
📋 Review Summary
This PR fixes a bug in Python bytecode type inference where slicing a fixed-size Tuple retained the full original tuple type instead of inferring the sliced tuple type. The changes cleanly unify slice handling across BUILD_SLICE, BINARY_SLICE, and BINARY_SUBSCR (supporting Python 3.10 through 3.14+) and include comprehensive unit and pipeline tests.
🔍 General Feedback
- The unification of slice handling via
_slice_typeand preserving constant bounds inbuild_sliceasConst(slice(...))is clean, concise, and handles both constant and dynamic slice bounds appropriately across Python versions. - Test coverage in
trivial_inference_test.pyandtyped_pipeline_test.pythoroughly exercises positive/negative bounds, steps, out-of-range slices, dynamic indices, and end-to-endCombinePerKeypipeline type checking.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #40680 +/- ##
=============================================
+ Coverage 56.16% 80.02% +23.85%
=============================================
Files 1124 490 -634
Lines 178336 82825 -95511
Branches 1489 0 -1489
=============================================
- Hits 100170 66278 -33892
+ Misses 75645 16547 -59098
+ Partials 2521 0 -2521
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Assigning reviewers: R: @jrmccluskey for label python. This pull request likely touches a core component ("core" label). Please review with scrutiny. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
jrmccluskey
left a comment
There was a problem hiding this comment.
LGTM, just one note around the CHANGES.md file entry. Thanks!
| * (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)). |
There was a problem hiding this comment.
This is a hair verbose, let's streamline a bit and let the context live in the linked bug.
| * (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)). | |
| * (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:])`. ([#20899](https://github.com/apache/beam/issues/20899)). |
Fixes #20899
Type inference gives a slice of a fixed-size tuple the type of the whole tuple. For an element of type
Tuple[str, str, int],row[1:]is inferred asTuple[str, str, int]instead ofTuple[str, int]. This breaks pipelines, not only the reportedelement_type. Dropping the first field before combining fails at construction:When the sliced PCollection is encoded, for example before a
Reshuffle, the coder for the whole tuple is used and the pipeline fails withTypeError: '<' not supported between instances of 'str' and 'int'forrow[1:], orIndexError: tuple index out of rangeforrow[:2].With this change, a slice of a fixed-size
Tuplewith constant bounds is inferred as the tuple of the selected element types (Tuple[str, int]above,Tuple[float, str, int]for[::-1]onTuple[int, str, float]). A slice with bounds that aren't constants, such asrow[i:], is inferred asTuple[Union[...], ...]of the element types, in line with a variable index (row[i]is inferred asUnion[...]) and tuple unpacking ((*a, *b)). Slices of lists,Tuple[T, ...]andstrkeep the type of the sliced value, as before.Slices compile to different bytecode depending on the Python version:
BUILD_SLICEfollowed by a subscript on 3.10 and 3.11 and for slices with a step,BINARY_SLICEon 3.12 and later, and a constantsliceobject on 3.14.build_slicenow keeps constant bounds as aConstslice, andbinary_subscrandbinary_sliceshare one helper, so all versions infer the same types.Tests:
testGetItemSliceintrivial_inference_test.pyhas new cases for fixed-size tuples.test_inferred_type_of_tuple_sliceintyped_pipeline_test.pyruns the pipeline above.Both fail before the change and pass after it on Python 3.10, 3.11, 3.12, 3.13 and 3.14.
Testing:
pytest apache_beam/typehints/trivial_inference_test.py apache_beam/typehints/typed_pipeline_test.pyon Python 3.10, 3.11, 3.12, 3.13 and 3.14pytest apache_beam/typehints apache_beam/transforms apache_beam/coders apache_beam/pipeline_test.py apache_beam/dataframe/transforms_test.py apache_beam/yaml/yaml_mapping_test.py apache_beam/yaml/yaml_udf_test.pyon Python 3.12yapf --diffandruff checkon the changed filesThank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123orfixes #456).CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.