lance-format/lance

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

Open

#1,980 创建于 2024年2月19日

在 GitHub 查看
 (8 评论) (0 反应) (0 负责人)Rust (695 fork)github user discovery
enhancementgood first issue

仓库指标

Star
 (6,582 star)
PR 合并指标
 (平均合并 6天 1小时) (30 天内合并 219 个 PR)

描述

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.

贡献者指南