Add `n_jobs` parameter to `lance.write_dataset` to speed up writing large in-memory tables
#1,980 opened on Feb 19, 2024
Repository metrics
- Stars
- (6,582 stars)
- PR merge metrics
- (Avg merge 6d 1h) (219 merged PRs in 30d)
Description
I was just testing the lance dataset writer, and to my surprise there is a lot of headroom when using multi-processing.
This is what I did:
I have many identical tables (polars dataframes), and I'm writing them to Azure Blob Storage.
I repeated each 3 times and these are the median values
Writing 1 table: 3 minutes (3min/table) Writing 2 tables, each in its own thread (2 threads): 3 minutes (1.5min/table) Writing 4 tables, each in its own thread (4 threads): 4.5 minutes (1.125 min/table) Writing 12 tables, each in its own thread (12 threads, and maxed out RAM and other system resources): 15 minutes (1.25min/table)
Seems to me the writer is not using all the cores/network that are available - I see that when writing a single table I am getting < 10% CPU usage, and only 20% network bandwidth usage.
Btw, this is what I used to write the tables in parallel:
from joblib import Parallel, delayed, parallel_backend
def exec_in_parallel(callables:list[(Callable, 'args', 'kwargs')], n_jobs=1):
'''Use joblib to execute the callables via c() for c in callables'''
with parallel_backend('threading', n_jobs=n_jobs):
ans = Parallel(verbose=False)(delayed(c)(*args, **kwargs) for c, args, kwargs in callables)
return ans
def test_punish_write_lance(dfs, n_files, n_cores):
files = [f'tf{i}.lance' for i in range(n_files)]
exec_in_parallel([
(lance.write_dataset, [], {'uri':f'az://lance/{f}', 'data_obj':df}) for df, f in zip(dfs,files)
], n_jobs=n_cores)
At least in my set up I will benefit from writing my tables in parallel, but it'd be much better to have the writer write each table at maximum speed - that's usually what most people want.
I'll repeat my testing when v0.2 comes out.