lance-format / lance-format/lance

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

Open
#1,980 8 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement good first issue
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.