def test_arrow_table_caching(tmpdir):
cache = LocalDiskArrowTableCache(tmpdir.strpath, 10 * MB, 4)
df = pd.DataFrame(np.random.randn(50, 4), columns=list('ABCD'))
dummy_table = pa.Table.from_pandas(df)
table_from_cache = cache.get('my_key', lambda: dummy_table)
assert table_from_cache == dummy_table
cache.get('my_key', _should_never_be_called)
def make_batch_reader(dataset_url,
schema_fields=None,
reader_pool_type='thread',
workers_count=10,
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'):
"""
Creates an instance of Reader for reading batches out of a non-Petastorm Parquet store.
Currently, only stores having native scalar parquet data types are supported.
Use :func:`~petastorm.make_reader` to read Petastorm Parquet stores generated with
:func:`~petastorm.etl.dataset_metadata.materialize_dataset`.
NOTE: only scalar columns are currently supported.
: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: A list of regex pattern strings. Only columns matching at least one of the
patterns in the list will be loaded.
: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 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 pandas DataFrame object and must return a pandas Series with boolean values of matching
dimensions.
: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++)
: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.parsed_dataset_url().path
if cache_type is None or cache_type == 'null':
cache = NullCache()
elif cache_type == 'local-disk':
cache = LocalDiskArrowTableCache(cache_location, cache_size_limit,
cache_row_size_estimate,
**cache_extra_settings or {})
else:
raise ValueError('Unknown cache_type: {}'.format(cache_type))
if reader_pool_type == 'thread':
reader_pool = ThreadPool(workers_count)
elif reader_pool_type == 'process':
serializer = ArrowTableSerializer()
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))
return Reader(filesystem,
dataset_path,
schema_fields=schema_fields,
worker_class=ArrowReaderWorker,
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)
def make_batch_reader(dataset_url_or_urls,
schema_fields=None,
reader_pool_type='thread',
workers_count=10,
shuffle_row_groups=True,
shuffle_row_drop_partitions=1,
predicate=None,
rowgroup_selector=None,
num_epochs=1,
cur_shard=None,
shard_count=None,
shard_seed=None,
cache_type='null',
cache_location=None,
cache_size_limit=None,
cache_row_size_estimate=None,
cache_extra_settings=None,
hdfs_driver='libhdfs3',
transform_spec=None,
filters=None,
storage_options=None,
zmq_copy_buffers=True,
filesystem=None):
"""
Creates an instance of Reader for reading batches out of a non-Petastorm Parquet store.
Currently, only stores having native scalar parquet data types are supported.
Use :func:`~petastorm.make_reader` to read Petastorm Parquet stores generated with
:func:`~petastorm.etl.dataset_metadata.materialize_dataset`.
NOTE: only scalar columns or array type (of primitive type element) columns are currently supported.
NOTE: If without `schema_fields` specified, the reader schema will be inferred from parquet dataset. then the
reader schema fields order will preserve parqeut dataset fields order (partition column come first), but if
setting `transform_spec` and specified `TransformSpec.selected_fields`, then the reader schema fields order
will be the order of 'selected_fields'.
:param dataset_url_or_urls: a url to a parquet directory or a url list (with the same scheme) to parquet files.
e.g. ``'hdfs://some_hdfs_cluster/user/yevgeni/parquet8'``, or ``'file:///tmp/mydataset'``,
or ``'s3://bucket/mydataset'``, or ``'gs://bucket/mydataset'``,
or ``[file:///tmp/mydataset/00000.parquet, file:///tmp/mydataset/00001.parquet]``.
:param schema_fields: A list of regex pattern strings. Only columns matching at least one of the
patterns in the list will be loaded.
: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 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 pandas DataFrame object and must return a pandas Series with boolean values of matching
dimensions.
: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 shard_seed: Random seed to shuffle row groups for data sharding. 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 transform_spec: An instance of :class:`~petastorm.transform.TransformSpec` object defining how a record
is transformed after it is loaded and decoded. The transformation occurs on a worker thread/process (depends
on the ``reader_pool_type`` value).
:param filters: (List[Tuple] or List[List[Tuple]]): Standard PyArrow filters.
These will be applied when loading the parquet file with PyArrow. More information
here: https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetDataset.html
:param storage_options: Dict of kwargs forwarded to ``fsspec`` to initialize the filesystem.
:param zmq_copy_buffers: A bool indicating whether to use 0mq copy buffers with ProcessPool.
:param filesystem: An instance of ``pyarrow.FileSystem`` to use. Will ignore storage_options and
other filesystem configs if it's provided.
:return: A :class:`Reader` object
"""
dataset_url_or_urls = normalize_dataset_url_or_urls(dataset_url_or_urls)
filesystem, dataset_path_or_paths = get_filesystem_and_path_or_paths(
dataset_url_or_urls,
hdfs_driver,
storage_options=storage_options,
filesystem=filesystem)
try:
dataset_metadata.get_schema_from_dataset_url(
dataset_url_or_urls,
hdfs_driver=hdfs_driver,
storage_options=storage_options,
filesystem=filesystem)
warnings.warn(
'Please use make_reader (instead of \'make_batch_dataset\' function to read this dataset. '
'You may get unexpected results. '
'Currently make_batch_reader supports reading only Parquet stores that contain '
'standard Parquet data types and do not require petastorm decoding.'
)
except PetastormMetadataError:
pass
if cache_type is None or cache_type == NULL_CACHE:
cache = NullCache()
elif cache_type == LOCAL_DISK_CACHE:
cache = LocalDiskArrowTableCache(cache_location, cache_size_limit,
cache_row_size_estimate,
**cache_extra_settings or {})
else:
raise ValueError('Unknown cache_type: {}'.format(cache_type))
if reader_pool_type == 'thread':
reader_pool = ThreadPool(workers_count)
elif reader_pool_type == 'process':
serializer = ArrowTableSerializer()
reader_pool = ProcessPool(workers_count,
serializer,
zmq_copy_buffers=zmq_copy_buffers)
elif reader_pool_type == 'dummy':
reader_pool = DummyPool()
else:
raise ValueError(
'Unknown reader_pool_type: {}'.format(reader_pool_type))
return Reader(filesystem,
dataset_path_or_paths,
schema_fields=schema_fields,
worker_class=ArrowReaderWorker,
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,
shard_seed=shard_seed,
cache=cache,
transform_spec=transform_spec,
is_batched_reader=True,
filters=filters)