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:]

sdks/python/apache_beam/testing/synthetic_pipeline_test.py

@@ -47,7 +47,8 @@ def input_spec(
     value_size,
     bundle_size_distribution_type='const',
     bundle_size_distribution_param=0,
-    force_initial_num_bundles=0):
+    force_initial_num_bundles=0,
+    algorithm=None):
   return {
       'numRecords': num_records,
       'keySizeBytes': key_size,
@@ -57,6 +58,7 @@ def input_spec(
           'param': bundle_size_distribution_param
       },
       'forceNumInitialBundles': force_initial_num_bundles,
+      'algorithm': algorithm,
   }
 
 
@@ -181,6 +183,21 @@ def assert_size(element, expected_size):
           | beam.Map(lambda elm: elm[1]) | 'value' >> beam.Map(assert_size, 15))
       assert_that(pcoll | beam.combiners.Count.Globally(), equal_to([300]))
 
+  def test_synthetic_source_gen_kv_separately(self):
+    source1 = synthetic_pipeline.SyntheticSource(
+        input_spec(10, 8, 12, algorithm='lcg'))
+    source2 = synthetic_pipeline.SyntheticSource(
+        input_spec(10, 10, 10, algorithm='lcg'))
+    self.assertTrue(source1._gen_kv_separately)
+    self.assertFalse(source2._gen_kv_separately)
+    bytes1 = [
+        k + v for k, v in source1.read(source1.get_range_tracker(None, None))
+    ]
+    bytes2 = [
+        k + v for k, v in source2.read(source2.get_range_tracker(None, None))
+    ]
+    self.assertEqual(bytes1, bytes2)
+
   def test_synthetic_source_split_even(self):
     source = synthetic_pipeline.SyntheticSource(
         input_spec(1000, 1, 1, 'const', 0))