Tabular Datasets¶
Warning
The pyarrow.dataset module is experimental (specifically the classes),
and a stable API is not yet guaranteed.
The pyarrow.dataset module provides functionality to efficiently work with
tabular, potentially larger than memory and multi-file datasets:
A unified interface for different sources: supporting different sources and file formats (Parquet, Feather files) and different file systems (local, cloud).
Discovery of sources (crawling directories, handle directory-based partitioned datasets, basic schema normalization, ..)
Optimized reading with predicate pushdown (filtering rows), projection (selecting columns), parallel reading or fine-grained managing of tasks.
Currently, only Parquet and Feather / Arrow IPC files are supported. The goal is to expand this in the future to other file formats and data sources (e.g. database connections).
For those familiar with the existing pyarrow.parquet.ParquetDataset for
reading Parquet datasets: pyarrow.dataset’s goal is similar but not specific
to the Parquet format and not tied to Python: the same datasets API is exposed
in the R bindings or Arrow. In addition pyarrow.dataset boasts improved
perfomance and new features (e.g. filtering within files rather than only on
partition keys).
Reading Datasets¶
For the examples below, let’s create a small dataset consisting of a directory with two parquet files:
In [1]: import tempfile
In [2]: import pathlib
In [3]: import pyarrow as pa
In [4]: import pyarrow.parquet as pq
In [5]: base = pathlib.Path(tempfile.gettempdir())
In [6]: (base / "parquet_dataset").mkdir(exist_ok=True)
# creating an Arrow Table
In [7]: table = pa.table({'a': range(10), 'b': np.random.randn(10), 'c': [1, 2] * 5})
# writing it into two parquet files
In [8]: pq.write_table(table.slice(0, 5), base / "parquet_dataset/data1.parquet")
In [9]: pq.write_table(table.slice(5, 10), base / "parquet_dataset/data2.parquet")
Dataset discovery¶
A Dataset object can be created with the dataset() function. We
can pass it the path to the directory containing the data files:
In [10]: import pyarrow.dataset as ds
In [11]: dataset = ds.dataset(base / "parquet_dataset", format="parquet")
In [12]: dataset
Out[12]: <pyarrow._dataset.FileSystemDataset at 0x7f3cdab22bb0>
In addition to a base directory path, dataset() accepts a path to a single
file or a list of file paths.
Creating a Dataset object loads nothing into memory, it only crawls the
directory to find all the files:
In [13]: dataset.files
Out[13]: ['/tmp/parquet_dataset/data1.parquet', '/tmp/parquet_dataset/data2.parquet']
… and infers the dataset’s schema (by default from the first file):
In [14]: print(dataset.schema.to_string(show_field_metadata=False))
a: int64
b: double
c: int64
Using the Dataset.to_table() method we can read the dataset (or a portion
of it) into a pyarrow Table (note that depending on the size of your dataset
this can require a lot of memory, see below on filtering / iterative loading):
In [15]: dataset.to_table()
Out[15]:
pyarrow.Table
a: int64
b: double
c: int64
# converting to pandas to see the contents of the scanned table
In [16]: dataset.to_table().to_pandas()
Out[16]:
a b c
0 0 -0.273220 1
1 1 -0.967304 2
2 2 -1.134447 1
3 3 1.325103 2
4 4 0.216083 1
5 5 1.021441 2
6 6 -1.111828 1
7 7 -1.579745 2
8 8 -1.086391 1
9 9 0.220407 2
Reading different file formats¶
The above examples use Parquet files as dataset source but the Dataset API provides a consistent interface across multiple file formats and sources. Currently, Parquet and Feather / Arrow IPC file format are supported; more formats are planned in the future.
If we save the table as a Feather file instead of Parquet files:
In [17]: import pyarrow.feather as feather
In [18]: feather.write_feather(table, base / "data.feather")
then we can read the Feather file using the same functions, but with specifying
format="feather":
In [19]: dataset = ds.dataset(base / "data.feather", format="feather")
In [20]: dataset.to_table().to_pandas().head()
Out[20]:
a b c
0 0 -0.273220 1
1 1 -0.967304 2
2 2 -1.134447 1
3 3 1.325103 2
4 4 0.216083 1
Customizing file formats¶
The format name as a string, like:
ds.dataset(..., format="parquet")
is short hand for a default constructed ParquetFileFormat:
ds.dataset(..., format=ds.ParquetFileForma())
The FileFormat objects can be customized using keywords. For example:
parquet_format = ds.ParquetFileFormat(read_options={'dictionary_columns': ['a']})
ds.dataset(..., format=parquet_format)
Will configure column "a" to be dictionary encoded on scan.
Filtering data¶
To avoid reading all data when only needing a subset, the columns and
filter keywords can be used.
The columns keyword can be used to only read the specified columns:
In [21]: dataset = ds.dataset(base / "parquet_dataset", format="parquet")
In [22]: dataset.to_table(columns=['a', 'b']).to_pandas()
Out[22]:
a b
0 0 -0.273220
1 1 -0.967304
2 2 -1.134447
3 3 1.325103
4 4 0.216083
5 5 1.021441
6 6 -1.111828
7 7 -1.579745
8 8 -1.086391
9 9 0.220407
With the filter keyword, rows which do not match the filter predicate will
not be included in the returned table. The keyword expects a boolean
Expression referencing at least one of the columns:
In [23]: dataset.to_table(filter=ds.field('a') >= 7).to_pandas()
Out[23]:
a b c
0 7 -1.579745 2
1 8 -1.086391 1
2 9 0.220407 2
In [24]: dataset.to_table(filter=ds.field('c') == 2).to_pandas()
Out[24]:
a b c
0 1 -0.967304 2
1 3 1.325103 2
2 5 1.021441 2
3 7 -1.579745 2
4 9 0.220407 2
The easiest way to construct those Expression objects is by using the
field() helper function. Any column - not just partition columns - can be
referenced using the field() function (which creates a
FieldExpression). Operator overloads are provided to compose filters
including the comparisons (equal, larger/less than, etc), set membership
testing, and boolean combinations (and, or, not):
In [25]: ds.field('a') != 3
Out[25]: <pyarrow.dataset.Expression (a != 3:int64)>
In [26]: ds.field('a').isin([1, 2, 3])
Out[26]:
<pyarrow.dataset.Expression (a is in [
1,
2,
3
])>
In [27]: (ds.field('a') > ds.field('b')) & (ds.field('b') > 1)
Out[27]: <pyarrow.dataset.Expression ((a > b) and (b > 1:int64))>
Reading partitioned data¶
Above, a dataset consisting of a flat directory with files was shown. However, a dataset can exploit a nested directory structure defining a partitioned dataset, where the sub-directory names hold information about which subset of the data is stored in that directory.
For example, a dataset partitioned by year and month may look like on disk:
dataset_name/
year=2007/
month=01/
data0.parquet
data1.parquet
...
month=02/
data0.parquet
data1.parquet
...
month=03/
...
year=2008/
month=01/
...
...
The above partitioning scheme is using “/key=value/” directory names, as found in Apache Hive.
Let’s create a small partitioned dataset. The write_to_dataset()
function can write such hive-like partitioned datasets.
In [28]: table = pa.table({'a': range(10), 'b': np.random.randn(10), 'c': [1, 2] * 5,
....: 'part': ['a'] * 5 + ['b'] * 5})
....:
In [29]: pq.write_to_dataset(table, str(base / "parquet_dataset_partitioned"),
....: partition_cols=['part'])
....:
The above created a directory with two subdirectories (“part=a” and “part=b”), and the Parquet files written in those directories no longer include the “part” column.
Reading this dataset with dataset(), we now specify that the dataset
uses a hive-like partitioning scheme with the partitioning keyword:
In [30]: dataset = ds.dataset(str(base / "parquet_dataset_partitioned"), format="parquet",
....: partitioning="hive")
....:
In [31]: dataset.files
Out[31]:
['/tmp/parquet_dataset_partitioned/part=a/c552a15f5c5d41aaa2caf4cc7bf5914b.parquet',
'/tmp/parquet_dataset_partitioned/part=b/0e56b8364a1c4724ab210bb416d3378e.parquet']
Although the partition fields are not included in the actual Parquet files, they will be added back to the resulting table when scanning this dataset:
In [32]: dataset.to_table().to_pandas().head(3)
Out[32]:
a b c part
0 0 -0.173846 1 a
1 1 0.809449 2 a
2 2 1.934439 1 a
We can now filter on the partition keys, which avoids loading files altogether if they do not match the predicate:
In [33]: dataset.to_table(filter=ds.field("part") == "b").to_pandas()
Out[33]:
a b c part
0 5 -0.273986 2 b
1 6 0.403167 1 b
2 7 -0.914403 2 b
3 8 0.662790 1 b
4 9 -0.242684 2 b
Different partitioning schemes¶
The above example uses a hive-like directory scheme, such as “/year=2009/month=11/day=15”.
We specified this passing the partitioning="hive" keyword. In this case,
the types of the partition keys are inferred from the file paths.
It is also possible to explicitly define the schema of the partition keys
using the partitioning() function. For example:
part = ds.partitioning(
pa.schema([("year", pa.int16()), ("month", pa.int8()), ("day", pa.int32())]),
flavor="hive"
)
dataset = ds.dataset(..., partitioning=part)
“Directory partitioning” is also supported, where the segments in the file path represent the values of the partition keys without including the name (the field name are implicit in the segment’s index). For example, given field names “year”, “month”, and “day”, one path might be “/2019/11/15”.
Since the names are not included in the file paths, these must be specified when constructing a directory partitioning:
part = ds.partitioning(field_names=["year", "month", "day"])
Directory partitioning also supports providing a full schema rather than inferring types from file paths.
Reading from cloud storage¶
In addition to local files, pyarrow also supports reading from cloud storage.
Currently, HDFS and
Amazon S3-compatible storage are supported.
When passing a file URI, the file system will be inferred. For example, specifying a S3 path:
dataset = ds.dataset("s3://ursa-labs-taxi-data/", partitioning=["year", "month"])
Typically, you will want to customize the connection parameters, and then
a file system object can be created and passed to the filesystem keyword:
from pyarrow import fs
s3 = fs.S3FileSystem(region="us-east-2")
dataset = ds.dataset("ursa-labs-taxi-data/", filesystem=s3,
partitioning=["year", "month"])
The currently available classes are S3FileSystem and
HadoopFileSystem. See the Filesystem Interface docs for more
details.
Reading from Minio¶
In addition to cloud storage, pyarrow also supports reading from a MinIO object storage instance emulating S3 APIs. Paired with toxiproxy, this is useful for testing or benchmarking.
from pyarrow import fs
# By default, MinIO will listen for unencrypted HTTP traffic.
minio = fs.S3FileSystem(scheme="http", endpoint="localhost:9000")
dataset = ds.dataset("ursa-labs-taxi-data/", filesystem=minio,
partitioning=["year", "month"])
Manual specification of the Dataset¶
The dataset() function allows easy creation of a Dataset viewing a directory,
crawling all subdirectories for files and partitioning information. However
sometimes discovery is not required and the dataset’s files and partitions
are already known (for example, when this information is stored in metadata).
In this case it is possible to create a Dataset explicitly without any
automatic discovery or inference.
For the example here, we are going to use a dataset where the file names contain additional partitioning information:
# creating a dummy dataset: directory with two files
In [34]: table = pa.table({'col1': range(3), 'col2': np.random.randn(3)})
In [35]: (base / "parquet_dataset_manual").mkdir(exist_ok=True)
In [36]: pq.write_table(table, base / "parquet_dataset_manual" / "data_2018.parquet")
In [37]: pq.write_table(table, base / "parquet_dataset_manual" / "data_2019.parquet")
To create a Dataset from a list of files, we need to specify the paths, schema, format, filesystem, and partition expressions manually:
In [38]: from pyarrow import fs
In [39]: schema = pa.schema([("year", pa.int64()), ("col1", pa.int64()), ("col2", pa.float64())])
In [40]: dataset = ds.FileSystemDataset.from_paths(
....: ["data_2018.parquet", "data_2019.parquet"], schema=schema, format=ds.ParquetFileFormat(),
....: filesystem=fs.SubTreeFileSystem(str(base / "parquet_dataset_manual"), fs.LocalFileSystem()),
....: partitions=[ds.field('year') == 2018, ds.field('year') == 2019])
....:
Since we specified the “partition expressions” for our files, this information is materialized as columns when reading the data and can be used for filtering:
In [41]: dataset.to_table().to_pandas()
Out[41]:
year col1 col2
0 2018 0 0.745958
1 2018 1 0.019654
2 2018 2 1.237728
3 2019 0 0.745958
4 2019 1 0.019654
5 2019 2 1.237728
In [42]: dataset.to_table(filter=ds.field('year') == 2019).to_pandas()
Out[42]:
year col1 col2
0 2019 0 0.745958
1 2019 1 0.019654
2 2019 2 1.237728
Manual scheduling¶
The to_table() method loads all selected data into memory
at once resulting in a pyarrow Table. Alternatively, a dataset can also be
scanned one RecordBatch at a time in an iterative manner using the
scan() method:
for scan_task in dataset.scan(columns=[...], filter=...):
for record_batch in scan_task.execute():
# process the record batch