Skip to content

Send map workers only their shard's rows instead of the whole table - #8611

Open
behroozazarkhalili wants to merge 1 commit into
huggingface:mainfrom
behroozazarkhalili:fix/1992-map-num-proc-overhead
Open

Send map workers only their shard's rows instead of the whole table#8611
behroozazarkhalili wants to merge 1 commit into
huggingface:mainfrom
behroozazarkhalili:fix/1992-map-num-proc-overhead

Conversation

@behroozazarkhalili

Copy link
Copy Markdown
Contributor

Resolves #1992

Root cause

Dataset.map(num_proc>1) shards the table with Arrow slices (table.slice), and Arrow slices share their parent buffers. When a shard is pickled for a worker, InMemoryTable.__getstate__ serializes the sliced arrays' full buffers, so every worker receives the whole table. A 2-process map therefore moved 2x the dataset through the pool queues, which is why it came out slower than a single process on in-memory datasets.

Fix

Right before each job is submitted to the pool, compact its shard to only its rows:

  • Take the shard's indices when it carries an indices mapping (after select, filter, shuffle, train_test_split), so index-mapped shards no longer ship the full table either.
  • Copy each chunk individually so the chunk layout is preserved (chunks above the 32-bit offset limit stay separate).
  • Rebuild string_view / binary_view columns (top-level or nested) from their values, because concatenating a view array keeps the parent value buffer.
  • In-memory ConcatenationTable inputs (the result of concatenate_datasets) are compacted through their combined table; memory-mapped tables are left alone because they pickle by path.
  • Compaction happens per job, lazily, on a shallow copy of the shard, so fingerprints, features, format and cache-file names are unchanged, and the source shards are not mutated.
  • Any Arrow error during compaction falls back to the uncompacted shard, so unsupported types still map exactly as before.

Measurements

100 unique 140-byte rows, num_proc=2, bytes of each worker's pickled shard and how many of the 100 source rows appear in it:

dataset before after
Dataset.from_dict 15,910 B, 100 rows 8,507 B, 50 rows
after non-contiguous select 16,967 B, 100 rows 8,507 B, 50 rows
concatenate_datasets (70 + 30 rows) 100 rows in one worker 50 rows
string_view column 100 rows 50 rows

Map results, order, features and fingerprints match upstream for all of these, including formatted (numpy/torch) datasets, Array2D extension columns, large types, and num_proc > num_shards.

Tests

New tests/test_map_multiprocessing.py: captures the exact bytes each worker receives and asserts other shards' rows are absent for plain, indexed, concatenated and view-typed inputs; checks results and fingerprints across memory/indexed/disk storage; preserves chunk boundaries above the offset limit and nested chunked features; and exercises the fallback path by injecting Arrow errors.

…table

With num_proc>1, Dataset.map shards the table with Arrow slices, and
Arrow slices share their parent buffers. Pickling a shard for a worker
therefore serialized the entire table once per shard, so a 2-process
map moved 2x the dataset through the pool queues and was slower than a
single process on in-memory datasets (huggingface#1992).

Before a job is submitted, compact its shard to only its rows: take the
shard's indices when it carries an indices mapping (select, filter,
shuffle, train_test_split), copy each chunk so the chunk layout, and
therefore 32-bit offset limits, are preserved, and rebuild string_view
and binary_view columns from their values since concatenating a view
array keeps the parent buffer. In-memory ConcatenationTables are
compacted through their combined table; memory-mapped tables are left
alone because they pickle by path. Compaction is per job, lazily, on a
shallow copy of the shard so fingerprints, features and format are
unchanged, and any Arrow error falls back to the uncompacted shard so
unsupported types still map.

Measured with 100 unique rows and 2 workers: each worker payload now
carries 50 rows instead of 100 for plain, indexed and concatenated
in-memory datasets.

Fixes huggingface#1992
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

datasets.map multi processing much slower than single processing

1 participant