Why is dask selecting by index field so slow?

user_1487345 0
Viewed 194

I am reading data stored in parquet files into dask, and then attempting to select the data by a categorical index field. When saving the data in parquet, I partition the data by the same field that I'm using as the index field in dask.

I expect the dask set_index operation to take some time, but it seems like subsequent "select" type operations based on the index field should then be very fast. However, they are quite slow. Here's a minimal example. Warning, this example will create about 9.5G of data on disk, and take a few minutes to run:

import dask
from dask import dataframe as ddf
from dask.distributed import Client, LocalCluster
import pandas as pd
import numpy as np

cluster = LocalCluster(n_workers=16)
client = Client(cluster)

# Create parquet files. Takes about 1.5 mins to run, creates ~9.5G of data
for i in range(0,100):
    category = f'category-{i}'
    data = np.random.random_sample((10000,1000))
    df = pd.DataFrame(data)
    df['category'] = category
    df.columns = df.columns.map(str)
    df.to_parquet('/tmp/dasktest/', partition_cols=['category'], compression=None)

# Load all data, index all data, and then attempt to read data for one category
df = ddf.read_parquet('/tmp/dasktest/', npartitions=16)
df.category = df.category.cat.as_ordered()
df = df.set_index(df.category) #takes only 8s to run, this is fine.

#the below statement takes about 35 seconds to run, even though the data is only 94M on disk, and the parquet store is partitioned by the "category" column
data_df = df.loc[df.category == 'category-1'].compute()

#manually reading the data for one category takes about 200ms
pd_data_df = pd.read_parquet('/tmp/dasktest/category=category-1')

Is there a better way to improve the speed using dask, without manually just reading from the appropriate partitioned folder?

dask pandas dataframe python
2 Answers

Since category already in index try loc without bool select

df.loc['category-1']
user_79645270

Keep in mind that dask operations are lazy by default and are only triggered when needed. So in general, be careful with statements like "I expect line N to be slow and line N + 1 to be fast, but in practice N is fast and N + 1 is slow." - you need to be really sure that the observed execution time is being attributed correctly.

dask.dataframe.set_index has an optional argument compute which is False by default:

compute: bool, default False
Whether or not to trigger an immediate computation. Defaults to False. Note, that even if you set compute=False, an immediate computation will still be triggered if divisions is None.

Because you provide the divisions, the 8 seconds you're seeing are likely just the amount of time required to specify and schedule the operation. A good portion of the execution time of set_index is probably in the 35 seconds you're seeing in the next step.

So you're not seeing the time impact of the set_index operation (which requires reading all partitions) until you're computing the first category.

Have you tried doing this without setting the index and converting df.category to categorical? I would imagine the more operations you do to change the dataframe, the less likely dask will be to be able to use the metadata in the parquet archive to determine which files to read from without loading all of the partitions. The problem you're presenting to dask is "please give me only elements matching this string pattern from a column in a multi-file archive, that I know only exists in the first file".

This is a good test case for whether dask is making good use of the locations of the categorical variable within the parquet archive. But it's a lot to ask of dask, especially since this wasn't even written from a dask dataframe.

user_38887190
Related