def test_remove_field_transform():
one_removed = transform_schema(
TestSchema,
TransformSpec(lambda x: x, edit_fields=None, removed_fields=['int']))
assert set(one_removed.fields.keys()) == {'string', 'double'}
two_removed = transform_schema(
TestSchema,
TransformSpec(lambda x: x,
edit_fields=None,
removed_fields=['int', 'double']))
assert set(two_removed.fields.keys()) == {'string'}
def __init__(
self, dataset: str, batch_size: int = 1, transform_spec: TransformSpec = None
):
"""Initializes a reader depending on the dataset (Petastorm/Parquet).
:param dataset: Path to the dataset.
:param batch_size: How many samples per batch to load (default: ``1``).
:param transform_spec: Petastorm transform spec for data augmentation.
"""
num_workers = int(os.environ["WORLD_SIZE"]) # Is set at lagom startup.
rank = int(os.environ["RANK"])
is_peta_ds = EnvSing.get_instance().exists(
dataset.rstrip("/") + "/_common_metadata"
)
# Make reader only compatible with petastorm dataset.
ds_type = "Petastorm" if is_peta_ds else "Parquet"
print(f"{ds_type} dataset detected in folder {dataset}")
reader_factory = make_reader if is_peta_ds else make_batch_reader
reader = reader_factory(
dataset,
cur_shard=rank,
shard_count=num_workers,
transform_spec=TransformSpec(transform_spec),
)
super().__init__(reader, batch_size=batch_size)
self.iterator = None
def test_select_field_transform():
test_list = [['string', 'double', 'int'], ['int', 'string', 'double'],
['string', 'int'], ['int']]
for selected_fields in test_list:
transformed = transform_schema(
TestSchema, TransformSpec(selected_fields=selected_fields))
assert list(transformed.fields.keys()) == selected_fields
def test_unknown_fields_in_remove_field_transform():
with pytest.warns(UserWarning, match='not part of the schema.*unknown_1'):
one_removed = transform_schema(
TestSchema,
TransformSpec(lambda x: x,
edit_fields=None,
removed_fields=['int', 'unknown_1', 'unknown_2']))
assert set(one_removed.fields.keys()) == {'string', 'double'}
def test_change_field_transform():
one_added = transform_schema(
TestSchema,
TransformSpec(lambda x: x,
edit_fields=[
UnischemaField('double', np.float16, (), None, False)
]))
assert one_added.fields['double'].numpy_dtype == np.float16
def test_add_field_transform():
one_added = transform_schema(
TestSchema,
TransformSpec(lambda x: x,
edit_fields=[
UnischemaField('double2', np.float64, (), None,
False)
]))
assert set(
one_added.fields.keys()) == {'string', 'double', 'double2', 'int'}
def test_transform_spec_support_return_tensor(scalar_dataset, reader_factory):
field1 = UnischemaField(name='abc', shape=(2, 3), numpy_dtype=np.float32)
with pytest.raises(ValueError, match='field abc must be numpy array type'):
ArrowReaderWorker._check_shape_and_ravel('xyz', field1)
with pytest.raises(ValueError, match='field abc must be the shape'):
ArrowReaderWorker._check_shape_and_ravel(np.zeros((2, 5)), field1)
with pytest.raises(
ValueError,
match=
'field abc error: only support row major multi-dimensional array'):
ArrowReaderWorker._check_shape_and_ravel(np.zeros((2, 3), order='F'),
field1)
assert (6, ) == ArrowReaderWorker._check_shape_and_ravel(
np.zeros((2, 3)), field1).shape
for partial_shape in [(2, None), (None, ), (None, None)]:
field_with_unknown_dim = UnischemaField(name='abc',
shape=partial_shape,
numpy_dtype=np.float32)
with pytest.raises(
ValueError,
match='All dimensions of a shape.*must be constant'):
ArrowReaderWorker._check_shape_and_ravel(
np.zeros((2, 3), order='F'), field_with_unknown_dim)
def preproc_fn1(x):
return pd.DataFrame({
'tensor_col_1':
x['id'].map(lambda _: np.random.rand(2, 3)),
'tensor_col_2':
x['id'].map(lambda _: np.random.rand(3, 4, 5)),
})
edit_fields = [
('tensor_col_1', np.float32, (2, 3), False),
('tensor_col_2', np.float32, (3, 4, 5), False),
]
# This spec will remove all input columns and return one new column 'tensor_col_1' with shape (2, 3)
spec1 = TransformSpec(preproc_fn1,
edit_fields=edit_fields,
removed_fields=list(scalar_dataset.data[0].keys()))
with reader_factory(scalar_dataset.url, transform_spec=spec1) as reader:
sample = next(reader)._asdict()
assert len(sample) == 2
assert (2, 3) == sample['tensor_col_1'].shape[1:] and \
(3, 4, 5) == sample['tensor_col_2'].shape[1:]
def __init__(self,
train_path,
val_path,
batch_size=16,
num_minibatchs=8,
num_workers=4,
text_column="",
target_column="target"):
# Tokenised columns replace text based ones
remove_other = lambda rows: rows[[text_column, target_column]]
transform_spec = TransformSpec(
remove_other, selected_fields=[text_column, target_column])
super().__init__(train_path, val_path, batch_size, num_minibatchs,
num_workers, transform_spec)
start = default_timer()
dataset = IterableParquetDataset("Data/Train.parquet", process_rows)
dataloader = DataLoader(dataset, num_workers=4)
list(dataloader)
end = default_timer()
print(end - start)
del dataset, dataloader
del start, end
start = default_timer()
dataset = IterableManualParquetDataset("Data/Train.parquet", process_rows)
dataloader = DataLoader(dataset, num_workers=4)
list(dataloader)
end = default_timer()
print(end - start)
del dataset, dataloader
del start, end
start = default_timer()
dataset = make_batch_reader(Path("Data/Train.parquet").absolute().as_uri(),
workers_count=4,
transform_spec=TransformSpec(
lambda rows: rows[["readme", "target"]],
selected_fields=["readme", "target"]))
dataloader = TransformersDataLoader(dataset, string_column="readme")
list(dataloader)
end = default_timer()
print(end - start)
def test_noop_transform():
transformed_schema = transform_schema(
TestSchema,
TransformSpec(lambda x: x, edit_fields=None, removed_fields=None))
assert set(transformed_schema.fields) == set(TestSchema.fields)