diff --git a/sdks/python/apache_beam/testing/fast_test_utils.pxd b/sdks/python/apache_beam/testing/fast_test_utils.pxd index b35566a579eb..f04e87d760bd 100644 --- a/sdks/python/apache_beam/testing/fast_test_utils.pxd +++ b/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) diff --git a/sdks/python/apache_beam/testing/fast_test_utils.pyx b/sdks/python/apache_beam/testing/fast_test_utils.pyx index d9e4481b5149..471c772ec0c8 100644 --- a/sdks/python/apache_beam/testing/fast_test_utils.pyx +++ b/sdks/python/apache_beam/testing/fast_test_utils.pyx @@ -30,6 +30,10 @@ 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 (self._seed >> (self._bits - 32)) + cpdef void seed(self, libc.stdint.uint64_t seed): self._seed = (seed * self._a + self._c) & self._mask @@ -37,12 +41,10 @@ cdef class LCGenerator(object): self._seed = (seed ^ self._a) & self._mask cpdef libc.stdint.int32_t next_int(self): - self.seed(self._seed) - return (self._seed >> (self._bits - 32)) + return self._next_uint() cpdef libc.stdint.uint32_t next_uint(self): - self.seed(self._seed) - return (self._seed >> (self._bits - 32)) + return self._next_uint() cpdef bytes randbytes(self, int length): cdef libc.stdint.int32_t ints = (length + 3) // 4 @@ -50,11 +52,11 @@ cdef class LCGenerator(object): 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((data + i * 4), (& value), 4) retval = PyBytes_FromStringAndSize(data, length) libc.stdlib.free(data) return retval cpdef double random_sample(self): - return (self.next_uint() >> 8) / (1 << 24) + return (self._next_uint() >> 8) / (1 << 24) diff --git a/sdks/python/apache_beam/testing/synthetic_pipeline.py b/sdks/python/apache_beam/testing/synthetic_pipeline.py index 7070dc04ea02..18a4c37d2313 100644 --- a/sdks/python/apache_beam/testing/synthetic_pipeline.py +++ b/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:] diff --git a/sdks/python/apache_beam/testing/synthetic_pipeline_test.py b/sdks/python/apache_beam/testing/synthetic_pipeline_test.py index 7e973a3ca7d7..176e1a7e8348 100644 --- a/sdks/python/apache_beam/testing/synthetic_pipeline_test.py +++ b/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))