def process(self, element):
if element.primary().operation == WriteMutation._OPERATION_DELETE:
# As delete mutations are not batchable.
yield TaggedOutput(_BatchableFilterFn.OUTPUT_TAG_UNBATCHABLE, element)
else:
mg_info = element.info
if mg_info['byte_size'] > self._max_batch_size_bytes \
or mg_info['cells'] > self._max_number_cells \
or mg_info['rows'] > self._max_number_rows:
yield TaggedOutput(_BatchableFilterFn.OUTPUT_TAG_UNBATCHABLE, element)
else:
yield element
def mux_input(x):
if isinstance(x, int):
yield TaggedOutput('numbers', x)
else:
yield TaggedOutput('letters', x)
def process(self, element, train_size, val_label):
"""Randomly assigns element to training or validation set."""
if random.random() > train_size:
yield TaggedOutput(val_label, element)
else:
yield element
def test_passed_tuple_as_tag(self):
with self.assertRaisesRegex(
TypeError,
r'Attempting to create a TaggedOutput with non-string tag \(1, 2, 3\)'
):
TaggedOutput((1, 2, 3), 'value')
def process(self, element):
e = json.loads(element)
yield TaggedOutput(e.pop('_Dest'), e)