Fix Python VarIntCoder OverflowError on uint64 values

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