Fix Python VarIntCoder OverflowError on uint64 values (#39047)

* [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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: AviKndr

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)
 
 

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,22 +1058,26 @@ 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):
     # type: (create_InputStream, bool) -> int
     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
 
 

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