andygrove commented on PR #5449:
URL:
https://github.com/apache/datafusion-comet/pull/5449#issuecomment-5469861003
Picking up the fixed-size row group idea from @sunchao's comment, because I
think it interacts
with the flush path in a way that isn't obvious and that decides whether
it's worth doing.
The hash isn't the only per-cell cost in the current round-robin path.
Because `HashAll` scatters
adjacent rows across all N partitions, `partition_indices` ends up as a list
of `(batch, row)`
pairs and the flush goes through `interleave_record_batch`
(`partitioners/partitioned_batch_iterator.rs:111`), a per-row gather that
re-walks every column
and every nested child a second time. A good share of what `WholeBatch` buys
us comes from the
new fast path skipping that, not from skipping the hash.
Which is why I'd be careful with the row-group variant as stated. That fast
path fires only when
a flush chunk covers one whole source batch in natural order. With
fixed-size groups of B rows a
group is a strict sub-range of a source batch, and an 8192-row flush chunk
gets stitched from
~`batch_size / B` groups spanning several source batches, so the condition
never holds and we
fall straight back to per-row interleave. We'd keep the hash savings and
hand back the interleave
savings.
Would it work to change `partition_indices` from `Vec<Vec<(u32, u32)>>` to
runs of
`(batch_idx, start, len)` and build the output with `slice` +
`concat_batches`, keeping the
zero-copy return this PR adds for when a single run already fills the chunk?
Then the group size
becomes a knob rather than a cliff: at B = `num_rows` it degenerates to what
this PR does today,
at B = 1 it's Spark's per-row positional assignment, and in between the copy
is a memcpy per run
per buffer instead of a gather per row. It also drops index memory from 8
bytes per row to 12
bytes per run, which is not nothing when we're buffering wide batches.
On sizing, I think the balance question then answers itself analytically
rather than needing to be
measured. With a global row counter the imbalance between any two partitions
is bounded by B rows
regardless of how the reader frames batches, so something like
`B = clamp(batch_size / num_partitions, 64, batch_size)` gives 64-row runs
at the default shuffle
partition count and a worst case of 64 rows of skew. That seems more
defensible than reasoning
about how many batches a task happens to produce.
Two things I'd want to check before committing to it. First, `WholeBatch`
today never produces a
sliced batch — the fast path clones the source batch whole — so nothing here
exercises arrow-rs
IPC serialization of a sliced nested array. If a slice writes the parent's
full buffers rather
than just its window, shuffle write volume regresses and the win evaporates.
Could we add a test
asserting the bytes written for a sliced batch are within noise of the same
rows copied into a
fresh batch?
Second, there's an ownership property we'd be giving up. With whole-batch
assignment each buffered
batch is pinned by exactly one partition, so flushing that partition can
release it. With runs,
every partition holds a slice of every batch again and nothing can be
released until spill or
finish. That's neutral against `HashAll` today, and `count_new_buffers`
dedups by buffer address
so the accounting stays honest either way, but it is a real regression
against `WholeBatch` and
worth being deliberate about rather than discovering later.
Separately, since `HashAll { max_hash_columns }` is being reshaped here
anyway — is
`roundrobin.maxHashColumns` worth keeping? Hashing the first N columns
quietly collapses to N
distinct partitions when the leading columns are low cardinality, so someone
whose leading column
is a date or a partition key gets severe skew with nothing in the plan to
indicate why.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
