def test_splits_record_batch(self):
feature1 = pa.array([1.0])
feature2 = pa.array([2.0])
feature3 = pa.array([3.0])
record_batch = pa.RecordBatch.from_arrays(
[feature1, feature2, feature3], ['a', 'b', 'c'])
sliced_record_batch = ('slice_key', record_batch)
partitioner = mock.create_autospec(
feature_partition_util.ColumnHasher(0))
partitioner.assign.side_effect = [99, 43, 99]
# Verify we saw the right features.
partitions = list(
feature_partition_util.generate_feature_partitions(
sliced_record_batch, partitioner, frozenset([])))
self.assertCountEqual(
[mock.call('a'), mock.call('b'),
mock.call('c')], partitioner.assign.call_args_list)
# Verify we got the right output slices.
expected = {
('slice_key', 99):
pa.RecordBatch.from_arrays([feature1, feature3], ['a', 'c']),
('slice_key', 43):
pa.RecordBatch.from_arrays([feature2], ['b']),
}
self.assertCountEqual(expected.keys(), [x[0] for x in partitions])
for key, partitioned_record_batch in partitions:
expected_batch = expected[key]
test_util.make_arrow_record_batches_equal_fn(
self, [expected_batch])([partitioned_record_batch])
def test_csv_decoder_with_schema(self):
input_lines = ['1,1,2.0,hello',
'5,5,12.34,world']
column_names = ['int_feature_parsed_as_float', 'int_feature',
'float_feature', 'str_feature']
schema = text_format.Parse(
"""
feature { name: "int_feature_parsed_as_float" type: FLOAT }
feature { name: "int_feature" type: INT }
feature { name: "float_feature" type: FLOAT }
feature { name: "str_feature" type: BYTES }
""", schema_pb2.Schema())
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[1], [5]], pa.list_(pa.float32())),
pa.array([[1], [5]], pa.list_(pa.int64())),
pa.array([[2.0], [12.34]], pa.list_(pa.float32())),
pa.array([[b'hello'], [b'world']], pa.list_(pa.binary())),
], [
'int_feature_parsed_as_float', 'int_feature', 'float_feature',
'str_feature'
])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False) | csv_decoder.DecodeCSV(
column_names=column_names,
schema=schema,
infer_type_from_schema=True))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_empty_csv(self):
input_lines = []
expected_result = []
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=[]))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_invalid_row(self):
input_lines = ['1,2.0,hello', '5,12.34']
column_names = ['int_feature', 'float_feature', 'str_feature']
with self.assertRaisesRegex( # pylint: disable=g-error-prone-assert-raises
ValueError, '.*Columns do not match specified csv headers.*'):
with beam.Pipeline() as p:
result = (p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, None))
def test_decode_example_with_beam_pipeline(self, example_proto_text,
decoded_record_batch):
example = tf.train.Example()
text_format.Merge(example_proto_text, example)
with beam.Pipeline() as p:
result = (p
| beam.Create([example.SerializeToString()])
| tf_example_decoder.DecodeTFExample())
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(
self, [decoded_record_batch]))
def test_csv_decoder_int64_max(self):
input_lines = ['34', str(sys.maxsize)]
column_names = ['feature']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[34], [sys.maxsize]], pa.list_(pa.int64())),
], ['feature'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_consider_blank_line_single_column(self):
input_lines = ['',
'1']
column_names = ['int_feature']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([None, [1]], pa.list_(pa.int64())),
], ['int_feature'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False) | csv_decoder.DecodeCSV(
column_names=column_names, skip_blank_lines=False))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_large_int_categorical_neg(self):
input_lines = ['34', str(-(sys.maxsize+2))]
column_names = ['feature']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[b'34'], [str(-(sys.maxsize + 2)).encode('utf-8')]],
pa.list_(pa.binary())),
], ['feature'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_with_tab_delimiter(self):
input_lines = ['1\t"this is a \ttext"',
'5\t']
column_names = ['int_feature', 'str_feature']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[1], [5]], pa.list_(pa.int64())),
pa.array([[b'this is a \ttext'], None], pa.list_(pa.binary())),
], ['int_feature', 'str_feature'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names, delimiter='\t'))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_with_space_delimiter(self):
input_lines = ['1 "ab,cd,ef"',
'5 "wx,xy,yz"']
column_names = ['int_feature', 'str_feature']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[1], [5]], pa.list_(pa.int64())),
pa.array([[b'ab,cd,ef'], [b'wx,xy,yz']], pa.list_(pa.binary())),
], ['int_feature', 'str_feature'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names, delimiter=' '))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_with_float_and_string_in_same_column(self):
input_lines = ['2.3,abc',
'abc,2.3']
column_names = ['str_feature1', 'str_feature2']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[b'2.3'], [b'abc']], pa.list_(pa.binary())),
pa.array([[b'abc'], [b'2.3']], pa.list_(pa.binary())),
], ['str_feature1', 'str_feature2'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_with_int_and_float_in_same_column(self):
input_lines = ['2,1.5',
'1.5,2']
column_names = ['float_feature1', 'float_feature2']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[2.0], [1.5]], pa.list_(pa.float32())),
pa.array([[1.5], [2.0]], pa.list_(pa.float32())),
], ['float_feature1', 'float_feature2'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_missing_values(self):
input_lines = ['1,,hello',
',12.34,']
column_names = ['int_feature', 'float_feature', 'str_feature']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[1], None], pa.list_(pa.int64())),
pa.array([None, [12.34]], pa.list_(pa.float32())),
pa.array([[b'hello'], None], pa.list_(pa.binary())),
], ['int_feature', 'float_feature', 'str_feature'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder_with_unicode(self):
input_lines = [u'1,שקרכלשהו,22.34,text field']
column_names = ['int_feature', 'unicode_feature',
'float_feature', 'str_feature']
expected_result = [
pa.RecordBatch.from_arrays([
pa.array([[1]], pa.list_(pa.int64())),
pa.array([[22.34]], pa.list_(pa.float32())),
pa.array([[u'שקרכלשהו'.encode('utf-8')]],
pa.list_(pa.binary())),
pa.array([[b'text field']], pa.list_(pa.binary())),
], ['int_feature', 'float_feature', 'unicode_feature', 'str_feature'])
]
with beam.Pipeline() as p:
result = (
p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(column_names=column_names))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(self, expected_result))
def test_csv_decoder(self,
input_lines,
expected_result,
column_names,
delimiter=',',
skip_blank_lines=True,
schema=None,
multivalent_columns=None,
secondary_delimiter=None):
with beam.Pipeline() as p:
result = (p | beam.Create(input_lines, reshuffle=False)
| csv_decoder.DecodeCSV(
column_names=column_names,
delimiter=delimiter,
skip_blank_lines=skip_blank_lines,
schema=schema,
multivalent_columns=multivalent_columns,
secondary_delimiter=secondary_delimiter))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(
self, expected_result))
def test_batch_examples(self):
examples = [{
'a': np.array([1.0, 2.0], dtype=np.float32),
'b': np.array(['a', 'b', 'c', 'e'])
}, {
'a': np.array([3.0, 4.0, 5.0], dtype=np.float32),
}, {
'b': np.array(['d', 'e', 'f']),
'd': np.array([10, 20, 30], dtype=np.int64),
}, {
'b': np.array(['a', 'b', 'c'])
}, {
'c': np.array(['d', 'e', 'f'])
}]
expected_record_batches = [
pa.RecordBatch.from_arrays([
pa.array([[1.0, 2.0], [3.0, 4.0, 5.0]],
type=pa.list_(pa.float32())),
pa.array([['a', 'b', 'c', 'e'], None])
], ['a', 'b']),
pa.RecordBatch.from_arrays([
pa.array([['d', 'e', 'f'], ['a', 'b', 'c']]),
pa.array([[10, 20, 30], None], type=pa.list_(pa.int64()))
], ['b', 'd']),
pa.RecordBatch.from_arrays([pa.array([['d', 'e', 'f']])], ['c']),
]
with beam.Pipeline() as p:
result = (p
| beam.Create(examples, reshuffle=False)
| batch_util.BatchExamplesToArrowRecordBatches(
desired_batch_size=2))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(
self, expected_record_batches))
def test_batch_serialized_examples(self):
examples = [
"""
features {
feature {
key: "a"
value { float_list { value: [ 1.0, 2.0 ] } }
}
feature {
key: "b"
value { bytes_list { value: [ 'a', 'b', 'c', 'e' ] } }
}
}""",
"""
features {
feature {
key: "a"
value { float_list { value: [ 3.0, 4.0, 5.0 ] } }
}
}""",
"""
features {
feature {
key: "b"
value { bytes_list { value: [ 'd', 'e', 'f' ] } }
}
feature {
key: "d"
value { int64_list { value: [ 10, 20, 30 ] } }
}
}""",
"""
features {
feature {
key: "b"
value { bytes_list { value: [ 'a', 'b', 'c' ] } }
}
}""",
"""
features {
feature {
key: "c"
value { bytes_list { value: [ 'd', 'e', 'f' ] } }
}
}""",
]
serialized_examples = [
text_format.Merge(example_pbtxt,
tf.train.Example()).SerializeToString()
for example_pbtxt in examples
]
expected_record_batches = [
pa.RecordBatch.from_arrays([
pa.array([[1.0, 2.0], [3.0, 4.0, 5.0]],
type=pa.list_(pa.float32())),
pa.array([['a', 'b', 'c', 'e'], None],
type=pa.list_(pa.binary()))
], ['a', 'b']),
pa.RecordBatch.from_arrays([
pa.array([['d', 'e', 'f'], ['a', 'b', 'c']],
type=pa.list_(pa.binary())),
pa.array([[10, 20, 30], None], type=pa.list_(pa.int64()))
], ['b', 'd']),
pa.RecordBatch.from_arrays(
[pa.array([['d', 'e', 'f']], type=pa.list_(pa.binary()))],
['c']),
]
with beam.Pipeline() as p:
result = (p
| beam.Create(serialized_examples, reshuffle=False)
| batch_util.BatchSerializedExamplesToArrowRecordBatches(
desired_batch_size=2))
util.assert_that(
result,
test_util.make_arrow_record_batches_equal_fn(
self, expected_record_batches))
