This tutorial will cover the following topics:
TileDB can model dataframes either as dense or sparse arrays. Storing a dataframe as a (1D) dense array allows for rapid slicing on row indices. On the other hand, storing the dataframe as a ND sparse array, specifying any subset of the columns to act as the dimensions, allows for rapid slicing on range predicates on those column dimensions.
In either case and in addition to the slicing predicate, TileDB allows for very fast subselection of columns. This is because it implements a "columnar" format and, therefore, it fetches from persistent storage only data from the requested columns.
This notebook was run on a 2.3 GHz Intel Core i9, 8 cores, 16GB RAM, running MacOS Mojave.
We will use the NYC Taxi Trip dataset and specifically the yellow taxi trip records which has this schema.
We will focus on ingesting the data from January 2020, namely file yellow_tripdata_2020-01.csv. The file is about 560MB.
!wget https://s3.amazonaws.com/nyc-tlc/trip+data/yellow_tripdata_2020-01.csv
!ls -alh yellow_tripdata_2020-01.csv
You need to install TileDB-Py, the Python wrapper of TileDB Embedded, as follows:
# Pip:
$ pip install tiledb
# Or Conda:
$ conda install -c conda-forge tiledb-py
The notebook was run using Pandas 1.1.0.
Note that the TileDB core is a C++ library. To boost performance when integrating with pandas, we use Apache Arrow to achieve zero-copy when returning results from TileDB into pandas dataframes. You need to install pyarrow to take advantage of this optimization:
# Pip:
$ pip install pyarrow
# Or Conda:
$ conda install -c conda-forge pyarrow
Import TileDB and check the versions of the C++ core and TileDB-Py respectively.
import tiledb, numpy as np
# Version of TileDB core (C++ library)
tiledb.libtiledb.version()
# Version of TileDB-Py (Python wrapper)
tiledb.__version__
Before we start, we create the TileDB context passing a configuration parameter around memory allocation during read queries that will be explained in a later tutorial. That needs to be set at the very beginning of the code and before any other TileDB function is called.
cfg = tiledb.Ctx().config()
cfg.update(
{
'py.init_buffer_bytes': 1024**2 * 50
}
)
tiledb.default_ctx(cfg)
We also enable the TileDB stats so that we can get some insight into performance.
tiledb.stats_enable()
We ingest the yellow_tripdata_2020-01.csv CSV file into a TileDB dense array as shown below. The command takes the taxi CSV file and ingests it into a 1D dense array called taxi_dense_array. It sets the tile extent to 100K, which means that groups of 100K rows each across every column will comprise the atomic unit of compression and IO (i.e., a data tile). Two of the columns (tpep_dropoff_datetime and tpep_pickup_datetime) are dates, so we make sure to parse them as such. Finally, one of the columns (store_and_fwd_flag) may have nulls, so we explicitly set some null value.
%%time
tiledb.stats_reset()
tiledb.from_csv("taxi_dense_array", "yellow_tripdata_2020-01.csv",
tile = 100000,
parse_dates=['tpep_dropoff_datetime', 'tpep_pickup_datetime'],
fillna={'store_and_fwd_flag': ''})
tiledb.stats_dump()
From the stats, the actual write time in TileDB took under 1 second (the rest was mostly parsing the CSV in Pandas). The raw uncompressed CSV data was about 870 MB in binary format, and those got compressed down to about 131 MB in TileDB. There are 18 columns written as attributes, one of which is var-sized (of string type, as we will see in the schema below).
Next, let's open the written array and inspect the TileDB schema.
A = tiledb.open("taxi_dense_array")
print(A.schema)
That shows the 18 columns being stored as attributes, along with their types and filters (e.g., zstd compression, which is the default). There is a single dimension __tiledb_rows, which corresponds to the row indices. This essentially means that you will be able to slice fast across the row indices.
In order to see the number of rows ingested into this array, you can use the non-empty domain. The range below is inclusive and states that there are 6,405,008 rows in the array.
print(A.nonempty_domain())
Let's reset the stats and perform a full read of the array (all rows and all columns). The result is stored directly in a pandas dataframe. Note that ranges with df are always inclusive.
%%time
tiledb.stats_reset()
df = A.df[0:6405007]
df
tiledb.stats_dump()
This operation fetches the entire array / dataframe from the disk, decompresses all tiles and creates a pandas dataframe with the result. The whole process takes 1.2 seconds in TileDB core (C++) and about 0.7 seconds on the Python wrapper side for buffer conversion.
The stats are quite informative. They break down how long it took to read from storage and unfilter (i.e., decompress), how many cells were fetched, what is the percentage of useful results, etc.
However, note that you do not need to read the entire dataframe in main memory in order to process it. You can efficiently slice any subset of rows directly from storage as follows. TileDB makes very lightweight use of main memory to process the result. Note that df[] works with mulit-index semantics and thus can take multi-range subarrays as well.
%%time
df = A.df[0:999]
df
Notice how much faster that operation was, taking only a few milliseconds.
Finally, you can slice any subset of columns, without fetching all the columns first in a pandas dataframe.
%%time
df = A.query(attrs=['tpep_dropoff_datetime', 'fare_amount']).df[0:6405007]
df
Once again, that operation was much faster than fetching the entire dataframe in main memory. The stats also inform you about how many attributes (i.e., columns) were retrieved, which is two in this example.
Remember to close the array when you are done.
A.close()
Storing the dataframe as a 1D dense array allowed us to rapidly slice on row indexes. But what if we wished to slice fast on predicates applied to column values, such as dropoff time and fare amount? For such scenarios and if you know for a fact that the majority of your workloads involve applying a range (or equality) predicate on a specific subset of columns, you can create a sparse array with those columns set as the dimensions.
This can be done as follows. Instead of the tile argument we used in dense arrays, we use capacity to determine how many rows to group in a data tile (read about the difference between dense and sparse data tiles). Also index_col determines which columns will act as dimensions.
%%time
tiledb.stats_reset()
tiledb.from_csv("taxi_sparse_array", "yellow_tripdata_2020-01.csv",
capacity=100000,
sparse=True,
index_col=['tpep_dropoff_datetime', 'fare_amount'],
parse_dates=['tpep_dropoff_datetime', 'tpep_pickup_datetime'],
fillna={'store_and_fwd_flag': ''})
tiledb.stats_dump()
Once again, most of the total ingestion time is spent on parsing on the pandas side. Notice that the R-tree (which is 2D) this time is slightly larger, as this is the main indexing method is sparse arrays. It is still tiny though relative to the entire array size, which is ~100MB.
Note that you can choose any subset of columns as the dimensions (any number with different types, even strings).
Let's open the array and print the schema.
A = tiledb.open("taxi_sparse_array")
print(A.schema)
Observe that now the array is sparse, having 16 attributes and 2 dimensions. Also notice that, by default, the array allows duplicates. This can be turned off by passing allows_duplicates=False in from_csv, which will return an error if the CSV contains rows with identical coordinates along the array dimensions.
Let's print the non-empty domain for the sparse array.
A.nonempty_domain()
The first range corresponds to tpep_dropoff_datetime and the second to fare_amount.
Now let's slice the whole array into a pandas dataframe.
%%time
tiledb.stats_reset()
df = A.query().df[:]
df
tiledb.stats_dump()
Notice that this takes longer than the dense case. This is because the sparse case involves more advanced indexing and copying operations than dense. However, the real benefit of sparse dataframe modeling is the ability to slice rapidly with range conditions on the indexed dimensions, without having to fetch the entire dataframe in main memory.
%%time
df = A.df[np.datetime64("2020-07-01"):np.datetime64("2020-10-01"), 5.5:12.5]
df
This is truly rapid. In the dense case, you would have to load the whole dataframe in main memory and then slice using pandas.
You can subset on attributes as follows.
%%time
df = A.query(attrs=['trip_distance']).df[:]
df
By default, TileDB fetches also the coordinate values and sets them as pandas indices. To disable them, you can run:
%%time
df = A.query(dims=False, attrs=['trip_distance']).df[:]
df
Wer can also subselect on dimensions:
%%time
df = A.query(dims=['tpep_dropoff_datetime'], attrs=['trip_distance']).df[:]
df
Finally, you can choose even attributes to act as dataframe indices using the index_col argument.
%%time
df = A.query(index_col=['trip_distance'], attrs=['passenger_count', 'trip_distance']).df[:]
df
For convenience, TileDB can also return dataframe results as an Arrow Table as follows:
%%time
df = A.query(return_arrow=True, index_col=['trip_distance'], attrs=['passenger_count', 'trip_distance']).df[:]
df
Since we are done, we can close the array.
A.close()
You can also store a pandas dataframe you already created in main memory into a TileDB array. The following will create a new TileDB array and write the contents of a pandas dataframe.
# First read some data into a pandas dataframe
A = tiledb.open("taxi_sparse_array")
df = A.query(attrs=['passenger_count', 'trip_distance']).df[:]
df
# Create and write into a TileDB array
tiledb.from_pandas("sliced_taxi_sparse_array", df)
Let's inspect the schema.
A2 = tiledb.open("sliced_taxi_sparse_array")
A2.schema
Reading the array back:
A2.df[:]
Lastly, we close the opened arrays.
A.close()
A2.close()
One of the cool things about TileDB is that it offers a powerful integration with embedded MariaDB. This allows for execution of arbitrary SQL queries directly on TileDB arrays (both dense and sparse). We took appropriate care to push the fast slicing and attribute subsetting portions of the query down to TileDB, leaving the rest of the SQL execution to MariaDB. In other words, we made MariaDB take advantage of the multi-dimensional indexing and columnar format of TileDB!
To install this capability, run:
conda install -c conda-forge libtiledb-sql-py
The usage is very simple and intuitive. All results are retunred directly as pandas dataframes.
import tiledb.sql, pandas as pd
db = tiledb.sql.connect()
%%time
pd.read_sql(sql="SELECT AVG(trip_distance) FROM taxi_dense_array WHERE __tiledb_rows >= 0 AND __tiledb_rows <1000", con=db)
%%time
pd.read_sql(sql="SELECT AVG(trip_distance) FROM taxi_sparse_array WHERE tpep_dropoff_datetime <= '2019-07-31' AND fare_amount < 5.5", con=db)
So far we have explained how to store TileDB arrays to the local disk. TileDB is optimized for numerous storage backends, including AWS S3, Azure Blob Storage and more. The entire functionality shown above (including SQL queries with embedded MariaDB) "just works" by replacing the array names taxi_dense_array and taxi_sparse_array with a URI that points to another backend, e.g., s3://<my_bucket>/<path>/array_name. The TileDB data format is cloud-native (based on immutable objects for fast updates and time traveling, to be covered in later tutorials) and the storage engine takes it to the extreme to implement parallel IO while minimizing the communication with the backend wherever possible.
In order to be able to support numerous storage backends, we abstracted all IO (e.g., read, write, remove, move, list, etc.) behind a Virtual Filesystem (VFS) class, which we exposed in our APIs as it is useful beyond the array internals. Everything we describe below in this section "just works" for any other storage backend URI.
For example, you can use the VFS functionality to list the contents of an array folder:
vfs = tiledb.VFS()
vfs.ls("taxi_sparse_array")
Or remove the arrays we created.
vfs.remove_dir("taxi_dense_array")
vfs.remove_dir("taxi_sparse_array")
vfs.remove_dir("sliced_taxi_sparse_array")
Also you can remove the CSV file as follows.
vfs.remove_file('yellow_tripdata_2020-01.csv')