def test_state_backed_iterable_coder(self): # pylint: disable=global-variable-undefined # required for pickling by reference global state state = {} def iterable_state_write(values, element_coder_impl): token = b'state_token_%d' % len(state) state[token] = [element_coder_impl.encode(e) for e in values] return token def iterable_state_read(token, element_coder_impl): return [element_coder_impl.decode(s) for s in state[token]] coder = coders.StateBackedIterableCoder( coders.VarIntCoder(), read_state=iterable_state_read, write_state=iterable_state_write, write_state_threshold=1) context = pipeline_context.PipelineContext( iterable_state_read=iterable_state_read, iterable_state_write=iterable_state_write) self.check_coder(coder, [1, 2, 3], context=context, test_size_estimation=False) # Ensure that state was actually used. self.assertNotEqual(state, {}) self.check_coder(coders.TupleCoder((coder, coder)), ([1], [2, 3]), context=context, test_size_estimation=False)
def test_state_backed_iterable_coder(self): # pylint: disable=global-variable-undefined # required for pickling by reference global state state = {} def iterable_state_write(values, element_coder_impl): token = b'state_token_%d' % len(state) state[token] = [element_coder_impl.encode(e) for e in values] return token def iterable_state_read(token, element_coder_impl): return [element_coder_impl.decode(s) for s in state[token]] coder = coders.StateBackedIterableCoder( coders.VarIntCoder(), read_state=iterable_state_read, write_state=iterable_state_write, write_state_threshold=1) # Note: do not use check_coder # see https://github.com/cloudpipe/cloudpickle/issues/452 self._observe(coder) self.assertEqual([1, 2, 3], coder.decode(coder.encode([1, 2, 3]))) # Ensure that state was actually used. self.assertNotEqual(state, {}) tupleCoder = coders.TupleCoder((coder, coder)) self._observe(tupleCoder) self.assertEqual(([1], [2, 3]), tupleCoder.decode(tupleCoder.encode(([1], [2, 3]))))