lance-format/lance

Add `n_jobs` parameter to `lance.write_dataset` to speed up writing large in-memory tables

Open

#1,980 opened on Feb 19, 2024

View on GitHub
 (8 comments) (0 reactions) (0 assignees)Rust (695 forks)github user discovery
enhancementgood first issue

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.

Contributor guide