andygrove commented on code in PR #6430:
URL: https://github.com/apache/datafusion-comet/pull/6430#discussion_r4139273069
##########
native/shuffle/src/partitioners/partitioned_batch_iterator.rs:
##########
@@ -239,6 +262,51 @@ impl Iterator for RowIterator<'_> {
}
}
+/// Gather scattered Boolean state without rebuilding Arrow's MutableArrayData
+/// descriptors for every reducer. Clustered selections retain Arrow's range
copies.
+/// The producer supplies matching immutable batches and valid row indices,
just as
+/// for the ordinary Arrow interleave path.
Review Comment:
Arrow's `interleave` has specialized paths for primitives, bytes, views,
dictionaries, structs, lists, maps and REE. Boolean still falls back in 59.3,
in 60.0 and on arrow-rs main, and I couldn't find an issue or PR for it. This
kernel would fit there naturally. Doing it upstream would also cover Booleans
nested inside structs, lists and maps, which the top-level check here can't
see. Could you open an arrow-rs issue or PR and link it from the doc comment
here? Then the Comet copy can go away with the Arrow upgrade that includes it.
##########
native/shuffle/src/partitioners/partitioned_batch_iterator.rs:
##########
@@ -239,6 +262,51 @@ impl Iterator for RowIterator<'_> {
}
}
+/// Gather scattered Boolean state without rebuilding Arrow's MutableArrayData
+/// descriptors for every reducer. Clustered selections retain Arrow's range
copies.
+/// The producer supplies matching immutable batches and valid row indices,
just as
+/// for the ordinary Arrow interleave path.
+fn interleave_shuffle_batches(
+ batches: &[&RecordBatch],
+ indices: &[(usize, usize)],
+) -> Result<RecordBatch, ArrowError> {
+ let schema = batches[0].schema_ref();
+ // Arrow already copies clustered bitmap ranges efficiently. A false
positive only
+ // selects that existing path; the probe is not used to copy or validate
indices.
+ // A four-row stride detects eight-row runs regardless of their starting
alignment.
+ let clustered = (0..indices.len().saturating_sub(4)).step_by(4).any(|i| {
+ indices[i].0 == indices[i + 4].0 && indices[i].1.checked_add(4) ==
Some(indices[i + 4].1)
+ });
Review Comment:
Thanks for the careful numbers in the gist. One gap: on `runs8`, `prefix1`
and `prefix4` the head falls back to Arrow, so those rows compare Arrow with
Arrow. I timed the `collect_bool` gather against Arrow `interleave` on
8,192-row chunks (x86_64, Arrow 59.3.0, release). Negative means the direct
gather is faster:
| Selection | 4 batches | 4 batches, nullable | 64 batches | 64 batches,
nullable |
|---|---|---|---|---|
| hash, 4 partitions (probe falls back) | -85% | -81% | -87% | -84% |
| 8-row runs (probe falls back) | -50% | -38% | -64% | -52% |
| 16-row runs | -21% | +12% | -50% | -24% |
| 32-row runs | +12% | +59% | -38% | -4% |
| 64-row runs | +58% | +180% | -27% | +29% |
So the direct gather still wins on 8-row runs. Break-even falls between 16
and 64 rows, depending on nulls and the number of buffered batches. With plain
hash partitioning the probe sends 100% of full chunks back to Arrow at 2 or 4
partitions, 78% at 6 and 38% at 8. At the writer level, 4 partitions gets no
gain at all. A probe that only catches 64-row runs (the same check with a
stride and distance of 32) gets it back and still falls back on key-clustered
input:
| Writer run (ms) | Base | This PR | 64-row probe | No probe |
|---|---|---|---|---|
| `SUM` decimal state, 4 partitions | 43.5 | 43.5 | 36.5 | 36.2 |
| 8 nullable Booleans, 4 partitions | 107.0 | 105.7 | 35.7 | 35.4 |
| 8 nullable Booleans, 200 partitions, 64 rows per key | 40.4 | 40.4 | 41.3
| 30.2 |
Would you be open to raising the threshold so the probe only catches long
runs? It may be worth dropping the probe entirely: with no probe, my
clustered-key runs were 2 to 25% faster than this PR. Either way, could you add
a unit test for the probe itself? The differential test still passes when
`clustered` is forced to `true`.
--
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]
