andygrove opened a new pull request, #6095: URL: https://github.com/apache/datafusion-comet/pull/6095
## Which issue does this PR close? Alternative to #5449, for #5397. Not intended to land alongside it — this is the same optimization with a different placement key, put up so the two can be compared directly. ## Rationale for this change Comet implements round robin as hash partitioning over every column of every row. On the wide nested schema that motivated #5449 that dominates the shuffle write: `create_murmur3_hashes` recurses into every struct child per row, and the row-level scatter then forces `interleave_record_batch` to walk every column and child again on flush. It is also not round robin. Placement is a pure function of a row's contents, so a column of one repeated value lands entirely on one reducer where Spark's round robin spreads it evenly. There is a test in this PR that pins both halves of that. #5449 replaces the hash with a **batch** counter. My concern with that, laid out in https://github.com/apache/datafusion-comet/pull/5449#issuecomment-5768393056, is that batch framing is the one property no Spark contract covers. `DeterministicLevel.DETERMINATE` promises the same rows in the same order and says nothing about how a downstream operator chunks them, so an operator that spills can reframe under different memory pressure while still honouring it. That gap is what `failOnRetry` in #5449 exists to paper over, and `failOnRetry` and the `INDETERMINATE` declaration cancel each other out: every rollback the declaration asks for goes through a stage resubmission, and `DAGScheduler.submitMissingTasks` bumps the stage attempt on every submission, so every rolled-back task hits the throw. So this PR keys on a **row** ordinal instead. The residual assumption becomes exactly the one Spark's own round robin makes — that the map task replays rows in the same order — which is the assumption the determinism level already describes and the DAGScheduler already knows how to act on. `failOnRetry` is then unnecessary and there is no equivalent here: Spark's normal fault tolerance is preserved. ## What changes are included in this PR? `RoundRobinStrategy` replaces the bare `max_hash_columns` argument of `CometPartitioning::RoundRobin`. `HashAll` is the existing behaviour and stays the default. `RowGroups { start_partition, group_rows }` places the row at task-global ordinal `i` at `(start_partition + i / group_rows) % num_partitions`. The counter runs over rows and carries across batch boundaries, so a group one input batch leaves part-way through is finished by the next and placement is independent of framing. `start_partition` is the Spark map partition id, filled in by `PhysicalPlanner::create_partitioning` from the planner's partition because `jni_api` runs every native root plan with partition 0. `group_rows` defaults to `clamp(batch_size / num_partitions, 64, batch_size)`, resolved on the driver. Two independent gates decide where it is used, and both must hold: - `CometShuffleExchangeExec.replaysRowsInOrder` walks the native subtree fused into the writer. The RDD graph cannot see it, because the whole subtree collapses into one `CometNativeShuffleInputRDD` whose dependencies are its leaves. Deliberately a short allowlist rather than a denylist of known-bad operators: a native scan under nothing but projections and filters. Operators that spill are the interesting exclusion. Anything else keeps `HashAll`. - `CometNativeShuffleInputRDD.getOutputDeterministicLevel` applies Spark's own `isOrderSensitive` rule to everything below that RDD, reporting `INDETERMINATE` over a non-determinate parent. This part is the same as #5449's, and it is the right mechanism. On the flush side, `PartitionIndices` gains a second shape. Positional placement records `(batch, start, len)` runs rather than one `(batch, row)` pair per row, which is smaller against the spill reservation and lets the copy move whole ranges instead of gathering row by row. `RunIterator` builds each output chunk by slicing and concatenating runs, and hands a run covering an entire buffered batch straight through with no copy. A repartitioner picks one shape for its lifetime from the partitioning it was built with, so the scatter paths are byte-for-byte unchanged. One schema-level restriction: positional placement is the only path that hands a sliced array to the IPC writer. arrow-rs truncates a slice's buffers for every type reached here — `ArrayData::slice` pushes the slice into struct children, `get_or_truncate_buffer` handles numeric and temporal, `reencode_offsets` the byte arrays, `get_list_array_buffers` lists and maps, `bit_slice` booleans — except `Utf8View`/`BinaryView`, where it truncates the views buffer but serializes every shared data buffer in full. `create_repartitioner` falls back to `HashAll` for a schema containing one. ## How are these changes tested? Native, in `comet_partitioning.rs` and `multi_partition.rs`: - `positional_placement_is_independent_of_batch_framing` and `positional_placement_survives_reframing` feed the same 1000 rows through five different batch framings (one batch, 1000 single-row batches, ragged, and so on) and assert an identical partitioning each time. This is the property the whole design rests on and the one the batch counter cannot offer. - `positional_placement_is_a_partition_of_the_input` asserts every row written exactly once, each partition in input order. - `positional_placement_is_balanced_within_one_group` asserts the spread stays within one group under deliberately ragged framing. - `positional_placement_starts_at_the_map_partition`, `positional_placement_records_runs_not_rows` (four runs rather than 256 row entries, and no row scratch allocated), plus run-iterator tests for chunking, the zero-copy whole-batch path, and the slice-not-copy path. - `view_types_are_detected_at_any_depth` for the fallback. JVM, in `CometNativePositionalRoundRobinSuite` (new, registered in both PR workflows) and `CometNativeShuffleInputRDDSuite`: - the gating, positive and negative: scan/filter/project takes it, an aggregate or sort under the exchange does not, the config gates it, a hash repartition never takes it; - `checkSparkAnswer` over a repartition, including with a group larger than a batch so the zero-copy path runs; - duplicate rows spread across partitions and partition identically to distinct rows, while `HashAll` collapses them all onto one — the skew described above; - the determinism level over determinate, unordered and indeterminate parents, including through `copyForLocalShuffle`. No benchmark numbers yet. The partitioning-only microbench in #5449 is the right harness and I would rather quote it on the same host for both strategies than repeat numbers taken elsewhere. -- 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]
