Dask DataFrame.to_parquet fails on read - repartition - write operation

user_2105924 0
Viewed 506

I have the following workflow.

def read_file(path, indx):
    df = pd.read_parquet(path)
    df.index = [indx] * len(df)
    return df

files_list = get_all_files()  # list of 10k parquet files, each about 1MB
df = dask.dataframe.from_delayed([dask.delayed(read_file)(x, indx) for (indx, x) in enumerate(files_list)])
df.divisions = list(range(10000)) + [9999] # each divisions include 1 file
new_divisions = [0, 10, 23, 45, ...., 9999] # new_divisions that reduces number of partitions by putting a bunch of files into same partitions. 
df = df.repartition(divisions = new_divisions)
df.to_parquet("fewer_files") # This causes dask to essentially freeze and no files get written

The new divisions are chosen so that the total memory of the files in each partition doesn't exceed 1000 MB. However, the final to_parquet call hangs forever. On the dask dashboard, there is no activity. The memory consumed by all workers remains very small (55MB), at least in the dashboard; but I suspect it might just be not updating since everything becomes super slow. The python process running the code keeps increasing the memory consumption (the virtual memory in Mac keeps increasing; I let it go upto 30GB).

If there are only about 200 files in the files_list, the code works just fine. Here is what the df.visualize() looks like when there are 236 files in files_list which gets repartitioned into 41 partitions: Task graph for df

Any idea on what might be causing the df.to_parquet to freeze when there are 10k files? When I print df before computation it shows the following:

npartitions=65, Dask Name: repartition-merge, 26417 tasks

Also, I can get the df.get_partition(0).to_parquet or other partition to work fairly quickly. However, df.to_parquet on the whole dataset fails. Is the 26K tasks simply too much to handle for 4 workers in my laptop?

dask-distributed dask parquet pandas python
3 Answers

Use dask.dataframe.read_parquet or other dask I/O implementations, not dask.delayed wrapping pandas I/O operations, whenever possible. Giving dask direct access to the file object or filepath allows the scheduler to quickly assess the steps in the job and accurately estimate the job size & requirements without executing the full workflow.

Explanation

By using dask.delayed with the pandas read_parquet reader, you're essentially robbing dask of the ability to peek into the file structure in order to help schedule the job, and also to open and close the files multiple times when running the full job (a problem you haven't even gotten to yet).

When everything fits neatly into memory, using dask.dataframe.read_parquet and the delayed method you use are very similar. The difference comes when the optimal strategy is not simply "read in all the data and then figure out what to do with it". Specifically, you are performing many reindexing and sorting operations, all of which require dask to know a lot about the contents of the files before the index-manipulation operations can even be scheduled.

Essentially, wrapping something in dask.delayed tells dask "here's this unknown block of code. Run it as a pure-python black box lots of times. The dask.dataframe and dask.array interfaces have smaller APIs and less interoperability compared with their pandas and numpy counterparts, but what you get for this is dask actually knows what's going on under the hood and can optimize it for you. When you use dask.delayed, you're gaining flexibility at the expense of dask's ability to tune the operation for you.

Example

As an exmaple, I'll create a large number of tiny files:

In [9]: tinydf = pd.DataFrame({"col1": [11, 21], "col2": [12, 22]})
   ...: for i in range(1000):
   ...:     tinydf.to_parquet(f"myfile_{i}.parquet")

dask.dataframe.read_parquet

Now, let's read this in with dask.dataframe.read_parquet:

In [10]: df = dask.dataframe.read_parquet([f"myfile_{i}.parquet" for i in range(1000)])

Note that this is lightning fast. We can take a peek at the high-level task graph by inspecting the dask attribute:

In [13]: df.dask
Out[13]:
HighLevelGraph with 1 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x15f79e2f0>
 0. read-parquet-e38709bfe39c7f8dfb5c4abf2fd08b50

Note that dask.dataframe.read_parquet is a single concept to dask. It can tune and optimize however it needs within this task. That includes "peeking" at the files to understand their column structure, look at the metadata file/attributes, etc., without reading in all the data.

In [30]: df.divisions = list(range(0, 2001, 2))

In [31]: df = df.repartition(divisions=list(range(0, 2001, 500)))

In [33]: df.dask
Out[33]:
HighLevelGraph with 2 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x168b5fcd0>
 0. read-parquet-e38709bfe39c7f8dfb5c4abf2fd08b50
 1. repartition-merge-bc42fb2f09234f7656901995bf3b29fa

The high level graph for the full workflow has two steps! Dask understands the operation in terms of file I/O and repartitions. It can decide how to split up these tasks in order to stay within memory limits and spread workload across workers, all without bogging down the scheduler.

dask.delayed(pd.read_parquet)

On the other hand, what happens if we do this with dask.delayed?

In [14]: def read_file(path, indx):
    ...:     df = pd.read_parquet(path)
    ...:     df.index = [indx] * len(df)
    ...:     return df
    ...:
    ...:
    ...: files_list = [f"myfile_{i}.parquet" for i in range(1000)]
    ...: df = dask.dataframe.from_delayed(
    ...:     [dask.delayed(read_file)(x, indx) for (indx, x) in enumerate(files_list)]
    ...: )

The dataframe preview ends up looking similar, but if we peek under the hood at the high level task graph, we can see that dask needs to read in all of the data before it even knows what the index looks like!

In [16]: df.dask
Out[16]:
HighLevelGraph with 1001 layers.
<dask.highlevelgraph.HighLevelGraph object at 0x168bf6230>
 0. read_file-b7aed020-1dc7-4872-a37d-d514b407a7d8
 1. read_file-a0462606-999b-4af1-9977-acb562edab67
 2. read_file-286df439-df34-4a5a-baf9-75dd0a5ae09b
 3. read_file-4db8c178-a67e-4775-b117-228ac607f02f
 4. read_file-a19d6144-5560-4da7-a1f5-8dc92b3ccf1c
 
 # yeah... really there are 1000 of these...

 998. read_file-d0cbd4a4-c255-4a77-a905-199bc289a0b5
 999. read_file-45a80080-426a-48fd-8dcb-9ba7565307f1
 1000. from-delayed-833eff6e232da1e10ca7221b961c21c1

To make matters worse, each pd.read_parquet uses the default pandas read behavior, which is to assume the data can fit into memory and just read the whole file in at once. Pandas does NOT return a file object - it loads all the data and returns a DataFrame before dask even sees it.

Because of this, dask is essentially prevented from getting to the scheduling bit until all of the read has already been done, and it has very little to work with in terms of workload balancing, memory management, etc. It can try to get a sneak-peek at the workload by executing the first task, but this is still a read of the whole first file.

This only gets worse when we start trying to shuffle the index. I won't go into it here, but you get the idea...

user_38887190

The new divisions are chosen so that the total memory of the files in each partition doesn't exceed 1000 MB.

If the main consideration for repartitioning is memory, it might be a good idea to use .repartition(partition_size='1000MB'). The script looks like:

df = dask.dataframe.from_delayed([dask.delayed(read_file)(x, indx) for (indx, x) in enumerate(files_list)])
df = df.repartition(partition_size='1000MB')
df.to_parquet("fewer_files")
user_106935960

When Dask reads a Parquet file, it opens the file and might attempt to read one partition in memory to infer the data types the file contains. This depends on the engine used (pyarrow or fastparquet) and on if there's a _metadata file present). Both opening a file and inferring memory types are slow operations.

I suspect that is the reason why you can run your computation on 200 files, while it stalls on 10k files.

A good practice is to think about how you'll want to read your data when you're writing it.

For your specific case I'd recommend reading in the parquet data 200 files at a time, and then write your to a much smaller number of parquet files when you have it all in memory. typically a number of files equal to the number of workers you have available is sensible, assuming data from a single file fits in worker memory.

Asking your OS for 10.000 file handles may also be causing some unexpected behaviour, and could also be a cause.

Running this command on my mac system tells me my OS is limiting me to 256 open file descriptors by default:

launchctl limit maxfiles

maxfiles    256            unlimited
user_11857740
Related