def _iter_impl(self):
"""
The Data Loader iterator stops the for-loop when reader runs out of samples.
"""
# As we iterate over incoming samples, we are going to store them in `self._batch_acc`, until we have a batch of
# the requested batch_size ready.
keys = None
if self.shuffling_queue_capacity > 0:
# We can not know what is the reasonable number to use for the extra capacity, so we set a huge number
# and give up on the unbound growth protection mechanism.
min_after_dequeue = self.shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(
self.shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=100000000)
else:
self._shuffling_buffer = NoopShufflingBuffer()
for row in self.reader:
# Default collate does not work nicely on namedtuples and treat them as lists
# Using dict will result in the yielded structures being dicts as well
row_as_dict = row._asdict()
keys = row_as_dict.keys()
# Promote some types that are incompatible with pytorch to be pytorch friendly.
_sanitize_pytorch_types(row_as_dict)
# Add rows to shuffling buffer
if not self.reader.is_batched_reader:
self._shuffling_buffer.add_many([row_as_dict])
else:
# Transposition:
# row_as_dict: {'a': [1,2,3], 'b':[4,5,6]}
# row_group_as_tuple: [(1, 4), (2, 5), (3, 6)]
# The order within a tuple is defined by key order in 'keys'
row_group_as_tuple = list(zip(*(row_as_dict[k] for k in keys)))
# Adding data as 'row-by-row' into a shuffling buffer. This is a pretty
# slow implementation though. Probably can comeup with a faster way to shuffle,
# perhaps at the expense of a larger memory consumption...
self._shuffling_buffer.add_many(row_group_as_tuple)
# _yield_batches will emit as much batches as are allowed by the shuffling_buffer (RandomShufflingBuffer
# will avoid underflowing below a certain number of samples to guarantee some samples decorrelation)
for batch in self._yield_batches(keys):
yield batch
# Once reader can not read new rows, we might still have a bunch of rows waiting in the shuffling buffer.
# Telling shuffling buffer that we are finished allows to deplete the buffer completely, regardless its
# min_after_dequeue setting.
self._shuffling_buffer.finish()
for batch in self._yield_batches(keys):
yield batch
# Yield the last and partial batch
if self._batch_acc:
yield self.collate_fn(self._batch_acc)
def __init__(self,
reader,
batch_size=1,
collate_fn=decimal_friendly_collate,
shuffling_queue_capacity=0):
"""
Initializes a data loader object, with a default collate.
Number of epochs is defined by the configuration of the reader argument.
An optional shuffling queue is created if shuffling_queue_capacity is greater than 0. No samples will be
returned to a user by the ``DataLoader`` until the queue is full. After that, batches of `batch_size`
will be created by uniformly sampling the shuffling queue. Once no more samples are available from the data
reader, the shuffling queue is allowed to be consumed till no further samples are available.
Note that the last returned batch could have less then ``batch_size`` samples.
NOTE: if you are using ``make_batch_reader``, this shuffling queue will be randomizing the order of the
entire batches and not changing the order of elements within a batch. This is likely not what you intend to do.
:param reader: petastorm Reader instance
:param batch_size: the number of items to return per batch; factored into the len() of this reader
:param collate_fn: an optional callable to merge a list of samples to form a mini-batch.
:param shuffling_queue_capacity: Queue capacity is passed to the underlying :class:`tf.RandomShuffleQueue`
instance. If set to 0, no suffling will be done.
"""
self.reader = reader
self.batch_size = batch_size
self.collate_fn = collate_fn
# _batch_acc accumulates samples for a single batch.
self._batch_acc = []
if shuffling_queue_capacity > 0:
# We can not know what is the reasonable number to use for the extra capacity, so we set a huge number
# and give up on the unbound growth protection mechanism.
min_after_dequeue = shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(
shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=100000000)
else:
self._shuffling_buffer = NoopShufflingBuffer()
def make_reader(dataset_url,
schema_fields=None,
reader_pool_type='thread',
workers_count=10,
pyarrow_serialize=False,
results_queue_size=50,
shuffle_row_groups=True,
shuffle_row_drop_partitions=1,
predicate=None,
rowgroup_selector=None,
num_epochs=1,
cur_shard=None,
shard_count=None,
cache_type='null',
cache_location=None,
cache_size_limit=None,
cache_row_size_estimate=None,
cache_extra_settings=None,
hdfs_driver='libhdfs3',
reader_engine='reader_v1',
reader_engine_params=None):
"""
Creates an instance of Reader for reading Petastorm datasets. A Petastorm dataset is a dataset generated using
:func:`~petastorm.etl.dataset_metadata.materialize_dataset` context manager as explained
`here <https://petastorm.readthedocs.io/en/latest/readme_include.html#generating-a-dataset>`_.
See :func:`~petastorm.make_batch_reader` to read from a Parquet store that was not generated using
:func:`~petastorm.etl.dataset_metadata.materialize_dataset`.
:param dataset_url: an filepath or a url to a parquet directory,
e.g. ``'hdfs://some_hdfs_cluster/user/yevgeni/parquet8'``, or ``'file:///tmp/mydataset'``
or ``'s3://bucket/mydataset'``.
:param schema_fields: Can be: a list of unischema fields and/or regex pattern strings; ``None`` to read all fields;
an NGram object, then it will return an NGram of the specified fields.
:param reader_pool_type: A string denoting the reader pool type. Should be one of ['thread', 'process', 'dummy']
denoting a thread pool, process pool, or running everything in the master thread. Defaults to 'thread'
:param workers_count: An int for the number of workers to use in the reader pool. This only is used for the
thread or process pool. Defaults to 10
:param pyarrow_serialize: Whether to use pyarrow for serialization. Currently only applicable to process pool.
Defaults to False.
:param results_queue_size: Size of the results queue to store prefetched rows. Currently only applicable to
thread reader pool type.
:param shuffle_row_groups: Whether to shuffle row groups (the order in which full row groups are read)
:param shuffle_row_drop_partitions: This is is a positive integer which determines how many partitions to
break up a row group into for increased shuffling in exchange for worse performance (extra reads).
For example if you specify 2 each row group read will drop half of the rows within every row group and
read the remaining rows in separate reads. It is recommended to keep this number below the regular row
group size in order to not waste reads which drop all rows.
:param predicate: instance of :class:`.PredicateBase` object to filter rows to be returned by reader. The predicate
will be passed a single row and must return a boolean value indicating whether to include it in the results.
:param rowgroup_selector: instance of row group selector object to select row groups to be read
:param num_epochs: An epoch is a single pass over all rows in the dataset. Setting ``num_epochs`` to
``None`` will result in an infinite number of epochs.
:param cur_shard: An int denoting the current shard number. Each node reading a shard should
pass in a unique shard number in the range [0, shard_count). shard_count must be supplied as well.
Defaults to None
:param shard_count: An int denoting the number of shards to break this dataset into. Defaults to None
:param cache_type: A string denoting the cache type, if desired. Options are [None, 'null', 'local-disk'] to
either have a null/noop cache or a cache implemented using diskcache. Caching is useful when communication
to the main data store is either slow or expensive and the local machine has large enough storage
to store entire dataset (or a partition of a dataset if shard_count is used). By default will be a null cache.
:param cache_location: A string denoting the location or path of the cache.
:param cache_size_limit: An int specifying the size limit of the cache in bytes
:param cache_row_size_estimate: An int specifying the estimated size of a row in the dataset
:param cache_extra_settings: A dictionary of extra settings to pass to the cache implementation,
:param hdfs_driver: A string denoting the hdfs driver to use (if using a dataset on hdfs). Current choices are
libhdfs (java through JNI) or libhdfs3 (C++)
:param reader_engine: Multiple engine implementations exist ('reader_v1' and 'experimental_reader_v2'). 'reader_v1'
(the default value) selects a stable reader implementation.
:param reader_engine_params: For advanced usage: a dictionary with arguments passed directly to a reader
implementation constructor chosen by ``reader_engine`` argument. You should not use this parameter, unless you
fine-tuning of a reader.
:return: A :class:`Reader` object
"""
if dataset_url is None or not isinstance(dataset_url, six.string_types):
raise ValueError("""dataset_url must be a string""")
dataset_url = dataset_url[:-1] if dataset_url[-1] == '/' else dataset_url
logger.debug('dataset_url: %s', dataset_url)
resolver = FilesystemResolver(dataset_url, hdfs_driver=hdfs_driver)
filesystem = resolver.filesystem()
dataset_path = resolver.get_dataset_path()
if cache_type is None or cache_type == 'null':
cache = NullCache()
elif cache_type == 'local-disk':
cache = LocalDiskCache(cache_location, cache_size_limit,
cache_row_size_estimate, **cache_extra_settings
or {})
else:
raise ValueError('Unknown cache_type: {}'.format(cache_type))
# Fail if this is a non-petastorm dataset. Typically, a Parquet store will have hundred thousands rows in a single
# rowgroup. Using PyDictReaderWorker or ReaderV2 implementation is very inefficient as it processes data on a
# row by row basis. ArrowReaderWorker (used by make_batch_reader) is much more efficient in these cases.
try:
dataset_metadata.get_schema_from_dataset_url(dataset_url)
except PetastormMetadataError:
raise RuntimeError(
'Currently make_reader supports reading only Petastorm datasets. '
'To read from a non-Petastorm Parquet store use make_batch_reader')
if reader_engine == 'reader_v1':
if reader_pool_type == 'thread':
reader_pool = ThreadPool(workers_count, results_queue_size)
elif reader_pool_type == 'process':
if pyarrow_serialize:
serializer = PyArrowSerializer()
else:
serializer = PickleSerializer()
reader_pool = ProcessPool(workers_count, serializer)
elif reader_pool_type == 'dummy':
reader_pool = DummyPool()
else:
raise ValueError(
'Unknown reader_pool_type: {}'.format(reader_pool_type))
# Create a dictionary with all ReaderV2 parameters, so we can merge with reader_engine_params if specified
kwargs = {
'schema_fields': schema_fields,
'reader_pool': reader_pool,
'shuffle_row_groups': shuffle_row_groups,
'shuffle_row_drop_partitions': shuffle_row_drop_partitions,
'predicate': predicate,
'rowgroup_selector': rowgroup_selector,
'num_epochs': num_epochs,
'cur_shard': cur_shard,
'shard_count': shard_count,
'cache': cache,
}
if reader_engine_params:
kwargs.update(reader_engine_params)
try:
return Reader(filesystem,
dataset_path,
worker_class=PyDictReaderWorker,
**kwargs)
except PetastormMetadataError as e:
logger.error('Unexpected exception: %s', str(e))
raise RuntimeError(
'make_reader has failed. If you were trying to open a Parquet store that was not '
'created using Petastorm materialize_dataset and it contains only scalar columns, '
'you may use make_batch_reader to read it.\n'
'Inner exception: %s', str(e))
elif reader_engine == 'experimental_reader_v2':
if reader_pool_type == 'thread':
decoder_pool = ThreadPoolExecutor(workers_count)
elif reader_pool_type == 'process':
decoder_pool = ProcessPoolExecutor(workers_count)
elif reader_pool_type == 'dummy':
decoder_pool = SameThreadExecutor()
else:
raise ValueError(
'Unknown reader_pool_type: {}'.format(reader_pool_type))
# TODO(yevgeni): once ReaderV2 is ready to be out of experimental status, we should extend
# the make_reader interfaces to take shuffling buffer parameters explicitly
shuffling_queue = RandomShufflingBuffer(
1000, 800) if shuffle_row_groups else NoopShufflingBuffer()
# Create a dictionary with all ReaderV2 parameters, so we can merge with reader_engine_params if specified
kwargs = {
'schema_fields': schema_fields,
'predicate': predicate,
'rowgroup_selector': rowgroup_selector,
'num_epochs': num_epochs,
'cur_shard': cur_shard,
'shard_count': shard_count,
'cache': cache,
'decoder_pool': decoder_pool,
'shuffling_queue': shuffling_queue,
'shuffle_row_groups': shuffle_row_groups,
'shuffle_row_drop_partitions': shuffle_row_drop_partitions,
}
if reader_engine_params:
kwargs.update(reader_engine_params)
return ReaderV2(dataset_url, **kwargs)
else:
raise ValueError(
'Unexpected value of reader_engine argument \'%s\'. '
'Supported reader_engine values are \'reader_v1\' and \'experimental_reader_v2\'',
reader_engine)
class DataLoader(object):
"""
A data loader adaptor for ``torch.utils.data.DataLoader``.
This class iterates and returns items from the Reader in batches.
This loader can be used as an iterator and will terminate when the reader used in the construction of the class
runs out of samples.
"""
def __init__(self, reader, batch_size=1, collate_fn=decimal_friendly_collate,
shuffling_queue_capacity=0):
"""
Initializes a data loader object, with a default collate.
Number of epochs is defined by the configuration of the reader argument.
An optional shuffling queue is created if shuffling_queue_capacity is greater than 0. No samples will be
returned to a user by the ``DataLoader`` until the queue is full. After that, batches of `batch_size`
will be created by uniformly sampling the shuffling queue. Once no more samples are available from the data
reader, the shuffling queue is allowed to be consumed till no further samples are available.
Note that the last returned batch could have less then ``batch_size`` samples.
NOTE: if you are using ``make_batch_reader``, this shuffling queue will be randomizing the order of the
entire batches and not changing the order of elements within a batch. This is likely not what you intend to do.
:param reader: petastorm Reader instance
:param batch_size: the number of items to return per batch; factored into the len() of this reader
:param collate_fn: an optional callable to merge a list of samples to form a mini-batch.
:param shuffling_queue_capacity: Queue capacity is passed to the underlying :class:`tf.RandomShuffleQueue`
instance. If set to 0, no suffling will be done.
"""
self.reader = reader
self.batch_size = batch_size
self.collate_fn = collate_fn
# _batch_acc accumulates samples for a single batch.
self._batch_acc = []
if shuffling_queue_capacity > 0:
# We can not know what is the reasonable number to use for the extra capacity, so we set a huge number
# and give up on the unbound growth protection mechanism.
min_after_dequeue = shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=100000000)
else:
self._shuffling_buffer = NoopShufflingBuffer()
def __iter__(self):
"""
The Data Loader iterator stops the for-loop when reader runs out of samples.
"""
# As we iterate over incoming samples, we are going to store them in `self._batch_acc`, until we have a batch of
# the requested batch_size ready.
keys = None
for row in self.reader:
# Default collate does not work nicely on namedtuples and treat them as lists
# Using dict will result in the yielded structures being dicts as well
row_as_dict = row._asdict()
keys = row_as_dict.keys()
# Promote some types that are incompatible with pytorch to be pytorch friendly.
_sanitize_pytorch_types(row_as_dict)
# Add rows to shuffling buffer
if not self.reader.is_batched_reader:
self._shuffling_buffer.add_many([row_as_dict])
else:
# Transposition:
# row_as_dict: {'a': [1,2,3], 'b':[4,5,6]}
# row_group_as_tuple: [(1, 4), (2, 5), (3, 6)]
# The order within a tuple is defined by key order in 'keys'
row_group_as_tuple = list(zip(*(row_as_dict[k] for k in keys)))
# Adding data as 'row-by-row' into a shuffling buffer. This is a pretty
# slow implementation though. Probably can comeup with a faster way to shuffle,
# perhaps at the expense of a larger memory consumption...
self._shuffling_buffer.add_many(row_group_as_tuple)
# _yield_batches will emit as much batches as are allowed by the shuffling_buffer (RandomShufflingBuffer
# will avoid underflowing below a certain number of samples to guarantee some samples decorrelation)
for batch in self._yield_batches(keys):
yield batch
# Once reader can not read new rows, we might still have a bunch of rows waiting in the shuffling buffer.
# Telling shuffling buffer that we are finished allows to deplete the buffer completely, regardless its
# min_after_dequeue setting.
self._shuffling_buffer.finish()
for batch in self._yield_batches(keys):
yield batch
# Yield the last and partial batch
if self._batch_acc:
yield self.collate_fn(self._batch_acc)
def _yield_batches(self, keys):
while self._shuffling_buffer.can_retrieve():
post_shuffled_row = self._shuffling_buffer.retrieve()
if not isinstance(post_shuffled_row, dict):
# This is for the case of batched reads. Here we restore back the
# dictionary format of records
post_shuffled_row = dict(zip(keys, post_shuffled_row))
self._batch_acc.append(post_shuffled_row)
# Batch is ready? Collate and emmit
if len(self._batch_acc) == self.batch_size:
yield self.collate_fn(self._batch_acc)
self._batch_acc = []
# Functions needed to treat data loader as a context manager
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.reader.stop()
self.reader.join()
def make_reader(dataset_url,
schema_fields=None,
reader_pool_type='thread', workers_count=10, pyarrow_serialize=False,
shuffle_row_groups=True, shuffle_row_drop_partitions=1,
predicate=None,
rowgroup_selector=None,
num_epochs=1,
cur_shard=None, shard_count=None,
cache_type='null', cache_location=None, cache_size_limit=None,
cache_row_size_estimate=None, cache_extra_settings=None,
hdfs_driver='libhdfs3',
infer_schema=False,
reader_engine='reader_v1', reader_engine_params=None):
"""
Factory convenience method for :class:`Reader`.
:param dataset_url: an filepath or a url to a parquet directory,
e.g. ``'hdfs://some_hdfs_cluster/user/yevgeni/parquet8'``, or ``'file:///tmp/mydataset'``
or ``'s3://bucket/mydataset'``.
:param schema_fields: Either list of unischema fields to subset, or ``None`` to read all fields.
OR an NGram object, then it will return an NGram of the specified properties.
:param reader_pool_type: A string denoting the reader pool type. Should be one of ['thread', 'process', 'dummy']
denoting a thread pool, process pool, or running everything in the master thread. Defaults to 'thread'
:param workers_count: An int for the number of workers to use in the reader pool. This only is used for the
thread or process pool. Defaults to 10
:param pyarrow_serialize: Whether to use pyarrow for serialization. Currently only applicable to process pool.
Defaults to False.
:param shuffle_row_groups: Whether to shuffle row groups (the order in which full row groups are read)
:param shuffle_row_drop_partitions: This is is a positive integer which determines how many partitions to
break up a row group into for increased shuffling in exchange for worse performance (extra reads).
For example if you specify 2 each row group read will drop half of the rows within every row group and
read the remaining rows in separate reads. It is recommended to keep this number below the regular row
group size in order to not waste reads which drop all rows.
:param predicate: instance of :class:`.PredicateBase` object to filter rows to be returned by reader.
:param rowgroup_selector: instance of row group selector object to select row groups to be read
:param num_epochs: An epoch is a single pass over all rows in the dataset. Setting ``num_epochs`` to
``None`` will result in an infinite number of epochs.
:param cur_shard: An int denoting the current shard number. Each node reading a shard should
pass in a unique shard number in the range [0, shard_count). shard_count must be supplied as well.
Defaults to None
:param shard_count: An int denoting the number of shards to break this dataset into. Defaults to None
:param cache_type: A string denoting the cache type, if desired. Options are [None, 'null', 'local-disk'] to
either have a null/noop cache or a cache implemented using diskcache. Caching is useful when communication
to the main data store is either slow or expensive and the local machine has large enough storage
to store entire dataset (or a partition of a dataset if shard_count is used). By default will be a null cache.
:param cache_location: A string denoting the location or path of the cache.
:param cache_size_limit: An int specifying the size limit of the cache in bytes
:param cache_row_size_estimate: An int specifying the estimated size of a row in the dataset
:param cache_extra_settings: A dictionary of extra settings to pass to the cache implementation,
:param hdfs_driver: A string denoting the hdfs driver to use (if using a dataset on hdfs). Current choices are
libhdfs (java through JNI) or libhdfs3 (C++)
:param infer_schema: Whether to infer the unischema object from the parquet schema.
Only works for schemas containing certain scalar type. This option allows getting around explicitly
generating petastorm metadata using :func:`petastorm.etl.dataset_metadata.materialize_dataset` or
petastorm-generate-metadata.py
:param reader_engine: Multiple engine implementations exist ('reader_v1' and 'experimental_reader_v2'). 'reader_v1'
(the default value) selects a stable reader implementation.
:param reader_engine_params: For advanced usage: a dictionary with arguments passed directly to a reader
implementation constructor chosen by ``reader_engine`` argument. You should not use this parameter, unless you
fine-tuning of a reader.
:return: A :class:`Reader` object
"""
if dataset_url is None or not isinstance(dataset_url, six.string_types):
raise ValueError("""dataset_url must be a string""")
dataset_url = dataset_url[:-1] if dataset_url[-1] == '/' else dataset_url
logger.debug('dataset_url: %s', dataset_url)
resolver = FilesystemResolver(dataset_url, hdfs_driver=hdfs_driver)
filesystem = resolver.filesystem()
dataset_path = resolver.get_dataset_path()
if cache_type is None or cache_type == 'null':
cache = NullCache()
elif cache_type == 'local-disk':
cache = LocalDiskCache(cache_location, cache_size_limit, cache_row_size_estimate, **cache_extra_settings or {})
else:
raise ValueError('Unknown cache_type: {}'.format(cache_type))
if reader_engine == 'reader_v1':
if reader_pool_type == 'thread':
reader_pool = ThreadPool(workers_count)
elif reader_pool_type == 'process':
reader_pool = ProcessPool(workers_count, pyarrow_serialize=pyarrow_serialize)
elif reader_pool_type == 'dummy':
reader_pool = DummyPool()
else:
raise ValueError('Unknown reader_pool_type: {}'.format(reader_pool_type))
# Create a dictionary with all ReaderV2 parameters, so we can merge with reader_engine_params if specified
kwargs = {
'schema_fields': schema_fields,
'reader_pool': reader_pool,
'shuffle_row_groups': shuffle_row_groups,
'shuffle_row_drop_partitions': shuffle_row_drop_partitions,
'predicate': predicate,
'rowgroup_selector': rowgroup_selector,
'num_epochs': num_epochs,
'cur_shard': cur_shard,
'shard_count': shard_count,
'cache': cache,
'infer_schema': infer_schema,
}
if reader_engine_params:
kwargs.update(reader_engine_params)
return Reader(filesystem, dataset_path, **kwargs)
elif reader_engine == 'experimental_reader_v2':
if reader_pool_type == 'thread':
decoder_pool = ThreadPoolExecutor(workers_count)
elif reader_pool_type == 'process':
decoder_pool = ProcessPoolExecutor(workers_count)
elif reader_pool_type == 'dummy':
decoder_pool = SameThreadExecutor()
else:
raise ValueError('Unknown reader_pool_type: {}'.format(reader_pool_type))
# TODO(yevgeni): once ReaderV2 is ready to be out of experimental status, we should extend
# the make_reader interfaces to take shuffling buffer parameters explicitly
shuffling_queue = RandomShufflingBuffer(1000, 800) if shuffle_row_groups else NoopShufflingBuffer()
# Create a dictionary with all ReaderV2 parameters, so we can merge with reader_engine_params if specified
kwargs = {
'schema_fields': schema_fields,
'predicate': predicate,
'rowgroup_selector': rowgroup_selector,
'num_epochs': num_epochs,
'cur_shard': cur_shard,
'shard_count': shard_count,
'cache': cache,
'decoder_pool': decoder_pool,
'shuffling_queue': shuffling_queue,
'shuffle_row_groups': shuffle_row_groups,
'shuffle_row_drop_partitions': shuffle_row_drop_partitions,
'infer_schema': infer_schema,
}
if reader_engine_params:
kwargs.update(reader_engine_params)
return ReaderV2(dataset_url, **kwargs)
else:
raise ValueError('Unexpected value of reader_engine argument \'%s\'. '
'Supported reader_engine values are \'reader_v1\' and \'experimental_reader_v2\'',
reader_engine)
def test_noop_shuffling_buffer():
"""Noop should not do any shuffling. Add/retrieve some items while checking correcness of can_retrieve"""
q = NoopShufflingBuffer()
# Empty buffer. Can add, can not retrieve with zero size
assert q.size == 0
assert q.can_add()
assert not q.can_retrieve()
# Try adding some items. Check queue size and can_retrieve indicator
q.add_many([1])
assert q.can_add()
assert q.size == 1
q.add_many([2, 3])
assert q.size == 3
assert 1 == q.retrieve()
assert q.can_retrieve()
assert 2 == q.retrieve()
assert 3 == q.retrieve()
assert not q.can_retrieve()
# No effect is expected in noop implementation
q.finish()
def test_noop_shuffling_buffer_stream_through():
"""Feed a 0:99 sequence through a NoopShufflingBuffer. Check that the order has not changed."""
expected = list(range(100))
actual = _feed_a_sequence_through_the_queue(NoopShufflingBuffer(), expected)
assert expected == actual
def __init__(self,
dataset_url,
schema_fields=None,
predicate=None,
rowgroup_selector=None,
num_epochs=1,
sequence=None,
cur_shard=None,
shard_count=None,
read_timeout_s=None,
cache=None,
loader_pool=None,
decoder_pool=None,
shuffling_queue=None,
shuffle_row_groups=True,
shuffle_row_drop_partitions=1,
pyarrow_filesystem=None,
hdfs_driver='libhdfs3'):
"""Initializes a reader object.
:param dataset_url: an filepath or a url to a parquet directory,
e.g. 'hdfs://some_hdfs_cluster/user/yevgeni/parquet8', or '/tmp/mydataset'
or ``'s3://bucket/mydataset'``.
:param schema_fields:
Either list of unischema fields to subset, or None to read all fields.
OR an NGram object, then it will return an NGram of the specified properties.
:param predicate: instance of predicate object to filter rows to be returned by reader.
:param rowgroup_selector: instance of row group selector object to select row groups to be read
:param reader_pool: parallelization pool. ThreadPool(10) (10 threads) is used by default.
This pool is a custom implementation used to parallelize reading data from the dataset.
Any object from workers_pool package can be used (e.g. ProcessPool)
:param num_epochs: An epoch is a single pass over all samples in the dataset. Setting num_epochs to 'None' will
result in an infinite number of epochs.
:param sequence: This is deprecated. To use sequence/ngram, please supply the argument in schema_fields instead.
:param cur_shard: An int denoting the current shard number. Each node reading a shard should
pass in a unique shard number in the range [0, shard_count).
shard count must be supplied as well.
:param shard_count An int denoting the number of shards to break this dataset into.
:param read_timeout_s: A numeric with the amount of time in seconds you would like to give a read before it
times out and raises an EmptyResultError. Pass in None for an infinite timeout
:param cache: An object conforming to `cache.CacheBase` interface. Before loading row groups from a parquet file
the Reader will attempt to load these values from cache. Caching is useful when communication
to the main data store is either slow or expensive and the local machine has large enough storage
to store entire dataset (or a partition of a dataset if num_training_partitions is used).
:param decoder_pool: An instance of a concurrent.futures pool executor used for decoding. If None,
a default ThreadPoolExecutor(5) will be used.
:param loader_pool: An instance of a concurrent.futures pool executor used for decoding. If None,
a default ThreadPoolExecutor(5) will be used.
:param hdfs_driver: A string denoting the hdfs driver to use (if using a dataset on hdfs). Current choices are
libhdfs (java through JNI) or libhdfs3 (C++)
By default, `NullCache` implementation
"""
# 1. Resolve dataset path (hdfs://, file://) and open the parquet storage (dataset)
# 2. Get a list of all groups
# 3. Filter rowgroups
# a. predicates
# b. row-group selector (our indexing mechanism)
# c. partition: used to get a subset of data for distributed training
# 4. Launch a new thread running `worker_loop` function.
if dataset_url is None or not isinstance(dataset_url,
six.string_types):
raise ValueError("""dataset_url must be a string""")
if not (isinstance(schema_fields, collections.Iterable)
or isinstance(schema_fields, NGram) or schema_fields is None):
raise ValueError(
"""Fields must be either None, an iterable collection of Unischema fields or an NGram
object.""")
if sequence is not None:
raise ValueError(
"""'sequence' argument of Reader object is deprecated. Please pass an NGram instance to
'schema_fields' argument instead.""")
# Can not rely on a check in epochs.py since it runs on a separate thread. Inform user earlier about invalid
# argument value.
if num_epochs is not None and (not isinstance(num_epochs, int)
or num_epochs < 1):
raise ValueError('iterations must be positive integer or None')
self.ngram = schema_fields if isinstance(schema_fields,
NGram) else None
if self.ngram and not self.ngram.timestamp_overlap and shuffle_row_drop_partitions > 1:
raise NotImplementedError(
'Using timestamp_overlap=False is not implemented with'
' shuffle_options.shuffle_row_drop_partitions > 1')
cache = cache or NullCache()
dataset_url = dataset_url[:-1] if dataset_url[
-1] == '/' else dataset_url
# 1. Resolve dataset path (hdfs://, file://) and open the parquet storage (dataset)
logger.debug('dataset_url: %s', dataset_url)
if pyarrow_filesystem is not None:
filesystem = pyarrow_filesystem
dataset_path = urlparse(dataset_url).path
else:
resolver = FilesystemResolver(dataset_url)
filesystem = resolver.filesystem()
dataset_path = resolver.get_dataset_path()
self._dataset = pq.ParquetDataset(dataset_path,
filesystem=filesystem,
validate_schema=False)
shuffle_row_drop_partitions = self._normalize_shuffle_options(
shuffle_row_drop_partitions, self._dataset)
stored_schema = infer_or_load_unischema(self._dataset)
# Make a schema view (a view is a Unischema containing only a subset of fields
# Will raise an exception if invalid schema fields are in schema_fields
fields = schema_fields if isinstance(schema_fields,
collections.Iterable) else None
self.schema = stored_schema.create_schema_view(
fields) if fields else stored_schema
# 2. Get a list of all groups
row_groups = dataset_metadata.load_row_groups(self._dataset)
# 3. Filter rowgroups
filtered_row_groups, worker_predicate = self._filter_row_groups(
self._dataset, row_groups, predicate, rowgroup_selector, cur_shard,
shard_count)
epoch_items = self._apply_row_drop_partition(
filtered_row_groups, shuffle_row_drop_partitions)
# 4. Launch a new thread running `worker_loop` function.
def epochs_iterator():
return epoch_generator(epoch_items, num_epochs, shuffle_row_groups)
self._results_queue = Queue(_OUTPUT_QUEUE_SIZE)
loader = RowGroupLoader(dataset_url,
self.schema,
self.ngram,
cache,
worker_predicate,
hdfs_driver=hdfs_driver)
decoder = RowDecoder(self.schema, self.ngram)
self._loader_pool = loader_pool or ThreadPoolExecutor(5)
self._decoder_pool = decoder_pool or ThreadPoolExecutor(5)
self._stop_flow_manager_event = threading.Event()
self._diags = Counter()
if not shuffling_queue:
shuffling_queue = NoopShufflingBuffer()
self._flow_manager_thread = threading.Thread(
target=worker_loop,
args=(epochs_iterator, self._loader_pool, loader,
self._decoder_pool, decoder, shuffling_queue,
self._results_queue, self._stop_flow_manager_event,
self._diags))
self._flow_manager_thread.daemon = True
self._flow_manager_thread.start()
self._read_timeout_s = read_timeout_s
self.batched_output = False
class DataLoader(object):
"""
A data loader adaptor for ``torch.utils.data.DataLoader``.
This class iterates and returns items from the Reader in batches.
This loader can be used as an iterator and will terminate when the reader used in the construction of the class
runs out of samples.
"""
def __init__(self,
reader,
batch_size=1,
collate_fn=decimal_friendly_collate,
shuffling_queue_capacity=0):
"""
Initializes a data loader object, with a default collate.
Number of epochs is defined by the configuration of the reader argument.
An optional shuffling queue is created if shuffling_queue_capacity is greater than 0. No samples will be
returned to a user by the ``DataLoader`` until the queue is full. After that, batches of `batch_size`
will be created by uniformly sampling the shuffling queue. Once no more samples are available from the data
reader, the shuffling queue is allowed to be consumed till no further samples are available.
Note that the last returned batch could have less then ``batch_size`` samples.
NOTE: if you are using ``make_batch_reader``, this shuffling queue will be randomizing the order of the
entire batches and not changing the order of elements within a batch. This is likely not what you intend to do.
:param reader: petastorm Reader instance
:param batch_size: the number of items to return per batch; factored into the len() of this reader
:param collate_fn: an optional callable to merge a list of samples to form a mini-batch.
:param shuffling_queue_capacity: Queue capacity is passed to the underlying :class:`tf.RandomShuffleQueue`
instance. If set to 0, no suffling will be done.
"""
self.reader = reader
self.batch_size = batch_size
self.collate_fn = collate_fn
if shuffling_queue_capacity > 0:
# We always add and remove one sample from the shuffling buffer in the steady state. min_after_dequeue
# is important when enqueue and dequeue rates are not the same.
# By setting `min_after_dequeue = shuffling_queue_capacity - 1` we will just keep the buffer full.
min_after_dequeue = shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(
shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=0)
else:
self._shuffling_buffer = NoopShufflingBuffer()
def __iter__(self):
"""
The Data Loader iterator stops the for-loop when reader runs out of samples.
"""
batch = []
for row in self.reader:
# Default collate does not work nicely on namedtuples and treat them as lists
# Using dict will result in the yielded structures being dicts as well
row_as_dict = row._asdict()
# Promote some types that are incompatible with pytorch to be pytorch friendly.
_sanitize_pytorch_types(row_as_dict)
# Add rows to shuffling buffer
self._shuffling_buffer.add_many([row_as_dict])
# If the shuffling buffer is ready, pull the resampled rows out and try add to our accumulated batch
while self._shuffling_buffer.can_retrieve():
post_shuffled_row = self._shuffling_buffer.retrieve()
batch.append(post_shuffled_row)
# Batch is ready? Collate and emmit
if len(batch) == self.batch_size:
yield self.collate_fn(batch)
batch = []
# Once reader can not emit new rows, we might still have a bunch of rows waiting in the shuffling buffer.
# Telling shuffling buffer thatt we are finished allows to deplete the buffer completely, regardless its
# min_after_dequeue setting.
self._shuffling_buffer.finish()
# Some code duplication with the main reading loop. Not sure how to reorganize to avoid this without
# overcomplicating the code too much.
while self._shuffling_buffer.can_retrieve():
post_shuffled_row = self._shuffling_buffer.retrieve()
batch.append(post_shuffled_row)
# Batch is ready? Collate and emmit
if len(batch) == self.batch_size:
yield self.collate_fn(batch)
batch = []
if batch:
yield self.collate_fn(batch)
# Functions needed to treat data loader as a context manager
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.reader.stop()
self.reader.join()
class DataLoader(object):
"""
A data loader adaptor for ``torch.utils.data.DataLoader``.
This class iterates and returns items from the Reader in batches.
This loader can be used as an iterator and will terminate when the reader used in the construction of the class
runs out of samples.
"""
def __init__(self,
reader,
batch_size=1,
collate_fn=torch.as_tensor,
shuffling_queue_capacity=0):
"""
Initializes a data loader object, with a default collate.
Number of epochs is defined by the configuration of the reader argument.
An optional shuffling queue is created if shuffling_queue_capacity is greater than 0. No samples will be
returned to a user by the ``DataLoader`` until the queue is full. After that, batches of `batch_size`
will be created by uniformly sampling the shuffling queue. Once no more samples are available from the data
reader, the shuffling queue is allowed to be consumed till no further samples are available.
Note that the last returned batch could have less then ``batch_size`` samples.
NOTE: if you are using ``make_batch_reader``, this shuffling queue will be randomizing the order of the
entire batches and not changing the order of elements within a batch. This is likely not what you intend to do.
:param reader: petastorm Reader instance
:param batch_size: the number of items to return per batch; factored into the len() of this reader
:param collate_fn: an optional callable to merge a list of samples to form a mini-batch.
:param shuffling_queue_capacity: Queue capacity is passed to the underlying :class:`tf.RandomShuffleQueue`
instance. If set to 0, no suffling will be done.
"""
self.reader = reader
self.batch_size = batch_size
self.collate_fn = collate_fn
# _batch_acc accumulates samples for a single batch.
self._batch_acc = []
if shuffling_queue_capacity > 0:
# We can not know what is the reasonable number to use for the extra capacity, so we set a huge number
# and give up on the unbound growth protection mechanism.
min_after_dequeue = shuffling_queue_capacity - 1
self._shuffling_buffer = RandomShufflingBuffer(
shuffling_queue_capacity,
min_after_retrieve=min_after_dequeue,
extra_capacity=100000000)
else:
self._shuffling_buffer = NoopShufflingBuffer()
def __iter__(self):
"""
The Data Loader iterator stops the for-loop when reader runs out of samples.
"""
# As we iterate over incoming samples, we are going to store them in `self._batch_acc`, until we have a batch of
# the requested batch_size ready.
keys = None
for row in self.reader:
# Default collate does not work nicely on namedtuples and treat them as lists
# Using dict will result in the yielded structures being dicts as well
row_as_dict = row._asdict()
keys = row_as_dict.keys()
# Promote some types that are incompatible with pytorch to be pytorch friendly.
_sanitize_pytorch_types(row_as_dict)
# Add rows to shuffling buffer
for k, v in row_as_dict.items():
if not self.reader.is_batched_reader:
row_as_dict[k] = self.collate_fn([v])
else:
row_as_dict[k] = self.collate_fn(v)
self._shuffling_buffer.add_many(row_as_dict)
# _yield_batches will emit as much batches as are allowed by the shuffling_buffer (RandomShufflingBuffer
# will avoid underflowing below a certain number of samples to guarantee some samples decorrelation)
for batch in self._yield_batches(keys):
yield batch
# Once reader can not read new rows, we might still have a bunch of rows waiting in the shuffling buffer.
# Telling shuffling buffer that we are finished allows to deplete the buffer completely, regardless its
# min_after_dequeue setting.
self._shuffling_buffer.finish()
for batch in self._yield_batches(keys):
yield batch
def _yield_batches(self, keys):
while self._shuffling_buffer.can_retrieve(self.batch_size):
batch = self._shuffling_buffer.retrieve(self.batch_size)
yield batch
# Functions needed to treat data loader as a context manager
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.reader.stop()
self.reader.join()
