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 0x7f06a34aeea0>

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.140763  1
1  1 -0.238656  2
2  2 -1.604585  1
3  3 -0.484210  2
4  4 -0.040160  1
5  5 -0.892570  2
6  6  1.494783  1
7  7  0.273504  2
8  8 -0.054818  1
9  9 -1.140701  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.140763  1
1  1 -0.238656  2
2  2 -1.604585  1
3  3 -0.484210  2
4  4 -0.040160  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.140763
1  1 -0.238656
2  2 -1.604585
3  3 -0.484210
4  4 -0.040160
5  5 -0.892570
6  6  1.494783
7  7  0.273504
8  8 -0.054818
9  9 -1.140701

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  0.273504  2
1  8 -0.054818  1
2  9 -1.140701  2

In [24]: dataset.to_table(filter=ds.field('c') == 2).to_pandas()
Out[24]: 
   a         b  c
0  1 -0.238656  2
1  3 -0.484210  2
2  5 -0.892570  2
3  7  0.273504  2
4  9 -1.140701  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/a841830a8bc14986bd1238641721d9ce.parquet',
 '/tmp/parquet_dataset_partitioned/part=b/c18dac16e14b43a78c229af08f69f9eb.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.417472  1    a
1  1 -0.178685  2    a
2  2  0.614755  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.932564  2    b
1  6  1.030880  1    b
2  7  0.669259  2    b
3  8 -1.101198  1    b
4  9 -0.501647  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.421649
1  2018     1 -0.944103
2  2018     2 -1.042005
3  2019     0  0.421649
4  2019     1 -0.944103
5  2019     2 -1.042005

In [42]: dataset.to_table(filter=ds.field('year') == 2019).to_pandas()
Out[42]: 
   year  col1      col2
0  2019     0  0.421649
1  2019     1 -0.944103
2  2019     2 -1.042005

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