Skip to content

perf: the native multi-partition shuffle writer copies each partition's tail batch and flushes the output file once per partition #7071

Closed

Description

@andygrove

What is the problem the feature request solves?

The native multi-partition shuffle writer still pays per-partition fixed costs when it writes its final output, and copies each partition's tail batch a second time. With 82k rows of a mixed schema and codec None, going from 200 to 8000 partitions takes the same bytes from 3.6 ms to 26.7 ms (measured when #5905 was filed; #5565 and #5916 have since removed some of it).

Still on main:

  • Tail chunks are copied twice. PartitionedBatchIterator emits full batch_size chunks (which bypass the coalescer) plus one tail chunk. In LocalPartitionWriter::finish_partition and PartitionedSpill::write each partition gets a fresh BufBatchWriter with its own BatchCoalescer, so the tail has nothing to coalesce with and is copy_rows'd and re-emitted unchanged. With 200 partitions and 8192-row inputs a partition sees about 41 rows per input batch, so the whole payload is gathered by interleave and then copied again.
  • Every partition flushes and seeks the output file. LocalPartitionWriter::finish_partition calls output_writer.stream_position() on a BufWriter, which flushes and seeks, so each partition leaves as its own write(2) plus lseek.

Describe the potential solution

  • Give multi-partition writers a passthrough mode (BufBatchWriter::new_for_partition) that writes batches as they arrive and keeps coalescing only for the single-partition output. Do not use passthrough for schemas with view types: interleave_views re-points views at the buffered batches' data buffers and the IPC writer serializes every referenced buffer, while arrow's coalescer compacts sparse views.
  • Track partition offsets from a running byte total and check it against stream_position() once in finish_all.

Draft #5908 is the work in progress for both. It currently conflicts with main, and a rebuild onto main is pending, so please coordinate there before starting.

Additional context

Findings W2 and W3 in #5905; also items in #5002. The IPC-context reuse part of W3 landed in #5565 and the spill-file flush in #5916.

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

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions