Optimize Cython LCGenerator and SyntheticSource KV generation

sdks/python/apache_beam/testing/fast_test_utils.pxd

@@ -25,6 +25,7 @@ cdef class LCGenerator(object):
   cdef libc.stdint.uint64_t _mask
   cdef libc.stdint.uint64_t _seed
 
+  cdef inline libc.stdint.uint32_t _next_uint(self) noexcept nogil
   cpdef void seed(self, libc.stdint.uint64_t seed)
   cpdef void seed_jdk(self, libc.stdint.uint64_t seed)
   cpdef libc.stdint.int32_t next_int(self)

sdks/python/apache_beam/testing/fast_test_utils.pyx

@@ -30,31 +30,33 @@ cdef class LCGenerator(object):
     self._mask = (1ull << self._bits) - 1
     self.seed(0)
 
+  cdef inline libc.stdint.uint32_t _next_uint(self) noexcept nogil:
+    self._seed = (self._seed * self._a + self._c) & self._mask
+    return <libc.stdint.uint32_t>(self._seed >> (self._bits - 32))
+
   cpdef void seed(self, libc.stdint.uint64_t seed):
     self._seed = (seed * self._a + self._c) & self._mask
 
   cpdef void seed_jdk(self, libc.stdint.uint64_t seed):
     self._seed = (seed ^ self._a) & self._mask
 
   cpdef libc.stdint.int32_t next_int(self):
-    self.seed(self._seed)
-    return <libc.stdint.int32_t>(self._seed >> (self._bits - 32))
+    return <libc.stdint.int32_t>self._next_uint()
 
   cpdef libc.stdint.uint32_t next_uint(self):
-    self.seed(self._seed)
-    return <libc.stdint.uint32_t>(self._seed >> (self._bits - 32))
+    return self._next_uint()
 
   cpdef bytes randbytes(self, int length):
     cdef libc.stdint.int32_t ints = (length + 3) // 4
     cdef char* data = <char*>libc.stdlib.malloc(ints * 4)
     cdef libc.stdint.uint32_t value
     cdef libc.stdint.int32_t i
     for i in range(0, ints, 1):
-      value = self.next_uint()
+      value = self._next_uint()
       libc.string.memcpy(<void*>(data + i * 4), <void*>(& value), 4)
     retval = PyBytes_FromStringAndSize(data, length)
     libc.stdlib.free(data)
     return retval
 
   cpdef double random_sample(self):
-    return <double>(self.next_uint() >> 8) / <double>(1 << 24)
+    return <double>(self._next_uint() >> 8) / <double>(1 << 24)

sdks/python/apache_beam/testing/synthetic_pipeline.py

@@ -403,6 +403,9 @@ def maybe_parse_byte_size(s):
           'Unknown algorithm for input_spec: %s. Supported '
           'algorithms are "builtin" and "lcg".',
           self.gen_algo)
+    # optimization to reduce overhead of copying bytes
+    self._gen_kv_separately = (
+        self.gen_algo == 'lcg' and self._key_size % 4 == 0)
 
   @property
   def element_size(self):
@@ -469,6 +472,10 @@ def _gen_kv_pair(self, generator, index):
           seed=index % self._num_hot_keys, algorithm=self.gen_algo)
       bytes_ = generator_hot.randbytes(self._key_size), generator.randbytes(
         self._value_size)
+    elif self._gen_kv_separately:
+      bytes_ = (
+          generator.randbytes(self._key_size),
+          generator.randbytes(self._value_size))
     else:
       bytes_ = generator.randbytes(self.element_size)
       bytes_ = bytes_[:self._key_size], bytes_[self._key_size:]