Create a varint32 coder and used it for RowCoder (#34354)

* Create a varint32 coder and used it for RowCoder

* Fix checks

* Handle a race when multiple process downloading jar at the same time

Author: Abacn

.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json

@@ -1,4 +1,4 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run",
-  "modification": 5
+  "modification": 4
 }

sdks/python/apache_beam/coders/coder_impl.pxd

@@ -130,6 +130,11 @@ cdef class VarIntCoderImpl(StreamCoderImpl):
   cpdef bytes encode(self, value)
 
 
+cdef class VarInt32CoderImpl(StreamCoderImpl):
+  @cython.locals(ivalue=libc.stdint.int32_t)
+  cpdef bytes encode(self, value)
+
+
 cdef class SingletonCoderImpl(CoderImpl):
   cdef object _value
 

sdks/python/apache_beam/coders/coder_impl.py

@@ -974,6 +974,37 @@ def estimate_size(self, value, nested=False):
     return get_varint_size(value)
 
 
+class VarInt32CoderImpl(StreamCoderImpl):
+  """For internal use only; no backwards-compatibility guarantees.
+
+  A coder for int32 objects."""
+  def encode_to_stream(self, value, out, nested):
+    # type: (int, create_OutputStream, bool) -> None
+    out.write_var_int32(value)
+
+  def decode_from_stream(self, in_stream, nested):
+    # type: (create_InputStream, bool) -> int
+    return in_stream.read_var_int32()
+
+  def encode(self, value):
+    ivalue = value  # type cast
+    if 0 <= ivalue < len(small_ints):
+      return small_ints[ivalue]
+    return StreamCoderImpl.encode(self, value)
+
+  def decode(self, encoded):
+    if len(encoded) == 1:
+      i = ord(encoded)
+      if 0 <= i < 128:
+        return i
+    return StreamCoderImpl.decode(self, encoded)
+
+  def estimate_size(self, value, nested=False):
+    # type: (Any, bool) -> int
+    # Note that VarInts are encoded the same way regardless of nesting.
+    return get_varint_size(value & 0xFFFFFFFF)
+
+
 class SingletonCoderImpl(CoderImpl):
   """For internal use only; no backwards-compatibility guarantees.
 

sdks/python/apache_beam/coders/coders.py

@@ -629,7 +629,7 @@ def __repr__(self):
 
 
 class VarIntCoder(FastCoder):
-  """Variable-length integer coder."""
+  """Variable-length integer coder  matches Java SDK's VarLongCoder."""
   def _create_impl(self):
     return coder_impl.VarIntCoderImpl()
 
@@ -650,6 +650,25 @@ def __hash__(self):
 Coder.register_structured_urn(common_urns.coders.VARINT.urn, VarIntCoder)
 
 
+class VarInt32Coder(FastCoder):
+  """Variable-length integer coder matches Java SDK's VarIntCoder."""
+  def _create_impl(self):
+    return coder_impl.VarInt32CoderImpl()
+
+  def is_deterministic(self):
+    # type: () -> bool
+    return True
+
+  def to_type_hint(self):
+    return int
+
+  def __eq__(self, other):
+    return type(self) == type(other)
+
+  def __hash__(self):
+    return hash(type(self))
+
+
 class BigEndianShortCoder(FastCoder):
   """A coder used for big-endian int16 values."""
   def _create_impl(self):

sdks/python/apache_beam/coders/coders_test_common.py

@@ -318,6 +318,20 @@ def test_varint_coder(self):
             for k in range(0, int(math.log(MAX_64_BIT_INT)))
         ])
 
+  def test_varint32_coder(self):
+    # Small ints.
+    self.check_coder(coders.VarInt32Coder(), *range(-10, 10))
+    # Multi-byte encoding starts at 128
+    self.check_coder(coders.VarInt32Coder(), *range(120, 140))
+    # Large values
+    MAX_32_BIT_INT = 0x7fffffff
+    self.check_coder(
+        coders.VarIntCoder(),
+        *[
+            int(math.pow(-1, k) * math.exp(k))
+            for k in range(0, int(math.log(MAX_32_BIT_INT)))
+        ])
+
   def test_float_coder(self):
     self.check_coder(
         coders.FloatCoder(), *[float(0.1 * x) for x in range(-100, 100)])

sdks/python/apache_beam/coders/row_coder.py

@@ -33,6 +33,7 @@
 from apache_beam.coders.coders import SinglePrecisionFloatCoder
 from apache_beam.coders.coders import StrUtf8Coder
 from apache_beam.coders.coders import TimestampCoder
+from apache_beam.coders.coders import VarInt32Coder
 from apache_beam.coders.coders import VarIntCoder
 from apache_beam.portability import common_urns
 from apache_beam.portability.api import schema_pb2
@@ -142,8 +143,10 @@ def _coder_from_type(field_type):
 def _nonnull_coder_from_type(field_type):
   type_info = field_type.WhichOneof("type_info")
   if type_info == "atomic_type":
-    if field_type.atomic_type in (schema_pb2.INT32, schema_pb2.INT64):
+    if field_type.atomic_type == schema_pb2.INT64:
       return VarIntCoder()
+    elif field_type.atomic_type == schema_pb2.INT32:
+      return VarInt32Coder()
     if field_type.atomic_type == schema_pb2.INT16:
       return BigEndianShortCoder()
     elif field_type.atomic_type == schema_pb2.FLOAT:

sdks/python/apache_beam/coders/row_coder_test.py

@@ -203,6 +203,29 @@ def test_create_row_coder_from_schema(self):
     for test_case in self.PEOPLE:
       self.assertEqual(test_case, coder.decode(coder.encode(test_case)))
 
+  def test_row_coder_negative_varint(self):
+    schema = schema_pb2.Schema(
+        id="negative",
+        fields=[
+            schema_pb2.Field(
+                name="i64",
+                type=schema_pb2.FieldType(atomic_type=schema_pb2.INT64)),
+            schema_pb2.Field(
+                name="i32",
+                type=schema_pb2.FieldType(atomic_type=schema_pb2.INT32))
+        ])
+    coder = RowCoder(schema)
+    Negative = typing.NamedTuple(
+        "Negative", [
+            ("i64", np.int64),
+            ("i32", np.int32),
+        ])
+    test_cases = [
+        Negative(-1, -1023), Negative(-1023, -1), Negative(-2**63, -2**31)
+    ]
+    for test_case in test_cases:
+      self.assertEqual(test_case, coder.decode(coder.encode(test_case)))
+
   @unittest.skip(
       "https://github.com/apache/beam/issues/19696 - Overflow behavior in "
       "VarIntCoder is currently inconsistent")

sdks/python/apache_beam/coders/slow_stream.py

@@ -58,6 +58,9 @@ def write_var_int64(self, v: int) -> None:
       if not v:
         break
 
+  def write_var_int32(self, v: int) -> None:
+    self.write_var_int64(int(v) & 0xFFFFFFFF)
+
   def write_bigendian_int64(self, v):
     self.write(struct.pack('>q', v))
 
@@ -156,6 +159,10 @@ def read_var_int64(self):
       result -= 1 << 64
     return result
 
+  def read_var_int32(self):
+    v = self.read_var_int64()
+    return struct.unpack('<i', struct.pack('<I', v))[0]
+
   def read_bigendian_int64(self):
     return struct.unpack('>q', self.read(8))[0]
 

sdks/python/apache_beam/coders/stream.pxd

@@ -26,6 +26,7 @@ cdef class OutputStream(object):
   cpdef write(self, bytes b, bint nested=*)
   cpdef write_byte(self, unsigned char val)
   cpdef write_var_int64(self, libc.stdint.int64_t v)
+  cpdef write_var_int32(self, libc.stdint.int32_t v)
   cpdef write_bigendian_int64(self, libc.stdint.int64_t signed_v)
   cpdef write_bigendian_uint64(self, libc.stdint.uint64_t signed_v)
   cpdef write_bigendian_int32(self, libc.stdint.int32_t signed_v)
@@ -43,6 +44,8 @@ cdef class ByteCountingOutputStream(OutputStream):
   cdef size_t count
 
   cpdef write(self, bytes b, bint nested=*)
+  cpdef write_var_int64(self, libc.stdint.int64_t val)
+  cpdef write_var_int32(self, libc.stdint.int32_t val)
   cpdef write_byte(self, unsigned char val)
   cpdef write_bigendian_int64(self, libc.stdint.int64_t val)
   cpdef write_bigendian_uint64(self, libc.stdint.uint64_t val)
@@ -61,6 +64,7 @@ cdef class InputStream(object):
   cpdef bytes read(self, size_t len)
   cpdef long read_byte(self) except? -1
   cpdef libc.stdint.int64_t read_var_int64(self) except? -1
+  cpdef libc.stdint.int32_t read_var_int32(self) except? -1
   cpdef libc.stdint.int64_t read_bigendian_int64(self) except? -1
   cpdef libc.stdint.uint64_t read_bigendian_uint64(self) except? -1
   cpdef libc.stdint.int32_t read_bigendian_int32(self) except? -1

sdks/python/apache_beam/coders/stream.pyx

@@ -73,6 +73,11 @@ cdef class OutputStream(object):
       if not v:
         break
 
+  cpdef write_var_int32(self, libc.stdint.int32_t signed_v):
+    """Encode an int using variable-length encoding to a stream."""
+    cdef libc.stdint.int64_t v = signed_v & <libc.stdint.int64_t>(0xFFFFFFFF)
+    self.write_var_int64(v)
+
   cpdef write_bigendian_int64(self, libc.stdint.int64_t signed_v):
     self.write_bigendian_uint64(signed_v)
 
@@ -91,7 +96,7 @@ cdef class OutputStream(object):
 
   cpdef write_bigendian_int32(self, libc.stdint.int32_t signed_v):
     cdef libc.stdint.uint32_t v = signed_v
-    if  self.buffer_size < self.pos + 4:
+    if self.buffer_size < self.pos + 4:
       self.extend(4)
     self.data[self.pos    ] = <unsigned char>(v >> 24)
     self.data[self.pos + 1] = <unsigned char>(v >> 16)
@@ -151,6 +156,12 @@ cdef class ByteCountingOutputStream(OutputStream):
   cpdef write_var_int64(self, libc.stdint.int64_t signed_v):
     self.count += get_varint_size(signed_v)
 
+  cpdef write_var_int32(self, libc.stdint.int32_t signed_v):
+    if signed_v < 0:
+      self.count += 5
+    else:
+      self.count += get_varint_size(signed_v)
+
   cpdef write_byte(self, unsigned char _):
     self.count += 1
 
@@ -225,6 +236,11 @@ cdef class InputStream(object):
 
     return result
 
+  cpdef libc.stdint.int32_t read_var_int32(self) except? -1:
+    """Decode a variable-length encoded int32 from a stream."""
+    cdef libc.stdint.int64_t v = self.read_var_int64()
+    return <libc.stdint.int32_t>(v);
+
   cpdef libc.stdint.int64_t read_bigendian_int64(self) except? -1:
     return self.read_bigendian_uint64()