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 0bded25e05d2..e6b1fb6c37cd 100644 --- a/sdks/python/apache_beam/coders/coder_impl.py +++ b/sdks/python/apache_beam/coders/coder_impl.py @@ -90,6 +90,20 @@ 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. 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 + + if TYPE_CHECKING or SLOW_STREAM: from .slow_stream import ByteCountingOutputStream from .slow_stream import InputStream as create_InputStream @@ -1044,12 +1058,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): @@ -1057,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): @@ -1073,11 +1091,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..939e44aff468 100644 --- a/sdks/python/apache_beam/coders/coders_test_common.py +++ b/sdks/python/apache_beam/coders/coders_test_common.py @@ -437,6 +437,24 @@ 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 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. self.check_coder(coders.VarInt32Coder(), *range(-10, 10))