From e6bac9449bea2476e9969489d0307ca13fad98a0 Mon Sep 17 00:00:00 2001 From: Avi Kondareddy Date: Sat, 20 Jun 2026 15:59:36 -0700 Subject: [PATCH 1/3] [BEAM] Fix Python VarIntCoder OverflowError on uint64 values The Cython write_var_int64/get_varint_size stream methods take a signed int64_t parameter. A Python int in the unsigned 64-bit range [2**63, 2**64) -- a uint64 -- is converted to that signed parameter at the call boundary and rejected with an OverflowError before the method body runs, even though the body already operates on the unsigned bit pattern and the VarInt wire encoding is well-defined. Fold such values to the signed int64 with the identical bit pattern (and thus identical VarInt encoding) before handing them to the stream. This matches Java's signed VarIntCoder on the wire; decoding remains signed. Values past 64 bits are left unchanged and still overflow downstream, preserving the coder's documented 64-bit limit. Adds test_varint_coder_uint64 covering no-overflow encoding, wire equivalence to the signed twin, size estimation, signed decode, and the still-raising out-of-range case (guarded on the compiled implementation). Co-Authored-By: Claude Opus 4.8 --- sdks/python/apache_beam/coders/coder_impl.py | 24 +++++++++++++++---- .../apache_beam/coders/coders_test_common.py | 20 ++++++++++++++++ 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py index 0bded25e05d2..08ae4cb55d1b 100644 --- a/sdks/python/apache_beam/coders/coder_impl.py +++ b/sdks/python/apache_beam/coders/coder_impl.py @@ -90,6 +90,18 @@ is_compiled = False fits_in_64_bits = lambda x: -(1 << 63) <= x <= (1 << 63) - 1 + +def _as_signed_int64(value): + # type: (int) -> int + + """Folds a uint64 to the signed int64 with the same bits (and VarInt + encoding), which the Cython int64_t stream params accept. Larger values pass + through and still overflow downstream.""" + if (1 << 63) <= value < (1 << 64): + return int(value) - (1 << 64) + return value + + if TYPE_CHECKING or SLOW_STREAM: from .slow_stream import ByteCountingOutputStream from .slow_stream import InputStream as create_InputStream @@ -1044,12 +1056,14 @@ class VarIntCoderImpl(StreamCoderImpl): def encode_to_stream(self, value, out, nested): # type: (int, create_OutputStream, bool) -> None try: - out.write_var_int64(value) + # Fold uint64 values into signed int64 so Cython doesn't overflow. + out.write_var_int64(_as_signed_int64(value)) except OverflowError as e: raise OverflowError( f"Integer value '{value}' is out of the encodable range for " - f"VarIntCoder. This coder is limited to values that fit " - f"within a 64-bit signed integer (-(2**63) to 2**63 - 1). " + f"VarIntCoder. This coder is limited to 64-bit integers: the " + f"signed range -(2**63) to 2**63 - 1, plus unsigned values up to " + f"2**64 - 1 which share the same wire encoding. " f"Original error: {e}") from e def decode_from_stream(self, in_stream, nested): @@ -1073,11 +1087,11 @@ def estimate_size(self, value, nested=False): # type: (Any, bool) -> int # Note that VarInts are encoded the same way regardless of nesting. try: - return get_varint_size(value) + return get_varint_size(_as_signed_int64(value)) except OverflowError as e: raise OverflowError( f"Cannot estimate size for integer value '{value}'. " - f"Value is out of the range for VarIntCoder (64-bit signed integer). " + f"Value is out of the range for VarIntCoder (64-bit integer). " f"Original error: {e}") from e diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py index 422d494b61c7..0fa6bac6bb01 100644 --- a/sdks/python/apache_beam/coders/coders_test_common.py +++ b/sdks/python/apache_beam/coders/coders_test_common.py @@ -37,6 +37,7 @@ from parameterized import param from parameterized import parameterized +from apache_beam.coders import coder_impl from apache_beam.coders import coders from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message from apache_beam.coders import typecoders @@ -437,6 +438,25 @@ def test_varint_coder(self): for k in range(0, int(math.log(MAX_64_BIT_INT))) ]) + def test_varint_coder_uint64(self): + # uint64 values [2**63, 2**64) must encode like the signed int64 with the + # same bits instead of overflowing Cython's int64_t. Decoding is signed, + # matching Java's VarIntCoder. + coder = coders.VarIntCoder() + impl = coder.get_impl() + for v in [1 << 63, (1 << 63) + 12345, (1 << 64) - 1]: + signed_twin = v - (1 << 64) + encoded = coder.encode(v) + self.assertEqual(encoded, coder.encode(signed_twin)) + self.assertEqual(impl.estimate_size(v), len(encoded)) + self.assertEqual(coder.decode(encoded), signed_twin) + + # Values past 64 bits stay out of range (only the Cython stream enforces it). + if coder_impl.is_compiled: + for v in [1 << 64, (1 << 70)]: + with self.assertRaises(OverflowError): + coder.encode(v) + def test_varint32_coder(self): # Small ints. self.check_coder(coders.VarInt32Coder(), *range(-10, 10)) From 093e6a216f01491a4d0addb6c23a40050025ac30 Mon Sep 17 00:00:00 2001 From: Avi Kondareddy Date: Sun, 28 Jun 2026 12:01:29 -0700 Subject: [PATCH 2/3] Enforce 64-bit range in _as_signed_int64 for path parity Address review feedback: raise OverflowError in _as_signed_int64 for values outside the 64-bit range (reusing fits_in_64_bits) so the pure-Python path matches the Cython int64_t overflow instead of silently encoding out-of-range values. Both encode_to_stream and estimate_size already wrap OverflowError, so the user-facing message is unchanged. Drop the is_compiled guard in test_varint_coder_uint64 so the out-of-range case runs on both paths, add the negative lower-bound case (-(2**63) - 1), and remove the now-unused coder_impl import. Co-Authored-By: Claude Opus 4.8 --- sdks/python/apache_beam/coders/coder_impl.py | 6 ++++-- sdks/python/apache_beam/coders/coders_test_common.py | 10 ++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py index 08ae4cb55d1b..d001b6cb2481 100644 --- a/sdks/python/apache_beam/coders/coder_impl.py +++ b/sdks/python/apache_beam/coders/coder_impl.py @@ -95,10 +95,12 @@ def _as_signed_int64(value): # type: (int) -> int """Folds a uint64 to the signed int64 with the same bits (and VarInt - encoding), which the Cython int64_t stream params accept. Larger values pass - through and still overflow downstream.""" + encoding), which the Cython int64_t stream params accept. Values outside the + 64-bit range raise here so the pure-Python path matches Cython's overflow.""" if (1 << 63) <= value < (1 << 64): return int(value) - (1 << 64) + if not fits_in_64_bits(value): + raise OverflowError("%d is out of range for a 64-bit integer." % value) return value diff --git a/sdks/python/apache_beam/coders/coders_test_common.py b/sdks/python/apache_beam/coders/coders_test_common.py index 0fa6bac6bb01..939e44aff468 100644 --- a/sdks/python/apache_beam/coders/coders_test_common.py +++ b/sdks/python/apache_beam/coders/coders_test_common.py @@ -37,7 +37,6 @@ from parameterized import param from parameterized import parameterized -from apache_beam.coders import coder_impl from apache_beam.coders import coders from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message from apache_beam.coders import typecoders @@ -451,11 +450,10 @@ def test_varint_coder_uint64(self): self.assertEqual(impl.estimate_size(v), len(encoded)) self.assertEqual(coder.decode(encoded), signed_twin) - # Values past 64 bits stay out of range (only the Cython stream enforces it). - if coder_impl.is_compiled: - for v in [1 << 64, (1 << 70)]: - with self.assertRaises(OverflowError): - coder.encode(v) + # Values outside the 64-bit range stay out of range on both paths. + for v in [1 << 64, (1 << 70), -(1 << 63) - 1]: + with self.assertRaises(OverflowError): + coder.encode(v) def test_varint32_coder(self): # Small ints. From a841c6483311bb4bddeb595cee73841d3cb2ad4e Mon Sep 17 00:00:00 2001 From: Avi Kondareddy Date: Sun, 28 Jun 2026 13:07:08 -0700 Subject: [PATCH 3/3] Fix uint64 overflow in VarIntCoderImpl.encode fast path The compiled encode() cast value to a C int64_t (ivalue, typed via coder_impl.pxd) for the small-ints fast path, which overflowed on uint64 values before reaching encode_to_stream's fold -- the cause of the test_varint_coder_uint64 CI failure on the Python PreCommit suites. Do the small-ints check with a plain Python-object comparison so it can't overflow; non-small values (including uint64 and out-of-range ints) fall through to StreamCoderImpl.encode -> encode_to_stream, where _as_signed_int64 folds them and raises the wrapped OverflowError for genuinely out-of-range values. Drop the now-unused ivalue int64_t local hint from coder_impl.pxd. Co-Authored-By: Claude Opus 4.8 --- sdks/python/apache_beam/coders/coder_impl.pxd | 1 - sdks/python/apache_beam/coders/coder_impl.py | 8 +++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/sdks/python/apache_beam/coders/coder_impl.pxd b/sdks/python/apache_beam/coders/coder_impl.pxd index e64177e6fd34..89730c17fe1d 100644 --- a/sdks/python/apache_beam/coders/coder_impl.pxd +++ b/sdks/python/apache_beam/coders/coder_impl.pxd @@ -135,7 +135,6 @@ cdef class TimestampCoderImpl(StreamCoderImpl): cdef list small_ints cdef class VarIntCoderImpl(StreamCoderImpl): - @cython.locals(ivalue=libc.stdint.int64_t) cpdef bytes encode(self, value) diff --git a/sdks/python/apache_beam/coders/coder_impl.py b/sdks/python/apache_beam/coders/coder_impl.py index d001b6cb2481..e6b1fb6c37cd 100644 --- a/sdks/python/apache_beam/coders/coder_impl.py +++ b/sdks/python/apache_beam/coders/coder_impl.py @@ -1073,9 +1073,11 @@ def decode_from_stream(self, in_stream, nested): return in_stream.read_var_int64() def encode(self, value): - ivalue = value # type cast - if 0 <= ivalue < len(small_ints): - return small_ints[ivalue] + # Compare as a Python object: a uint64 value overflows the int64_t cast + # the compiled fast path used to do here. Non-small values (including + # uint64) fall through to encode_to_stream, which folds them. + if 0 <= value < len(small_ints): + return small_ints[value] return StreamCoderImpl.encode(self, value) def decode(self, encoded):