lance-format / lance-format/lance
Add `n_jobs` parameter to `lance.write_dataset` to speed up writing large in-memory tables
Nobody has claimed this yet.
- Dominant language
- Rust
- Stars
- 7.1k
- Forks
- 852
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 272
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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at the lance.write_dataset entry point and trace how in-memory tables are written to Azure Blob Storage. Check the existing writer tests and measure the current single-table behavior; done means the new n_jobs option enables parallel work for one call while preserving the existing writing behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- azure, python, rust
- Domain
- api, cloud, data-engineering, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100