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