mbutrovich commented on code in PR #6095:
URL: https://github.com/apache/datafusion-comet/pull/6095#discussion_r4075724191
##########
native/core/src/execution/planner.rs:
##########
@@ -3652,15 +3652,26 @@ impl PhysicalPlanner {
}
PartitioningStruct::SinglePartition(_) =>
Ok(CometPartitioning::SinglePartition),
PartitioningStruct::RoundRobinPartition(rr_partition) => {
- // Treat negative max_hash_columns as 0 (no limit)
- let max_hash_columns = if rr_partition.max_hash_columns <= 0 {
- 0
+ let strategy = if rr_partition.positional {
+ // The Spark map partition id, not the DataFusion one:
`jni_api` runs every
+ // native root plan with partition 0 (one Comet execution
per Spark task), so
+ // `ShuffleWriterExec::execute` cannot supply it. See
+ // `RoundRobinStrategy::RowGroups` for why it has to be
this value.
+ RoundRobinStrategy::RowGroups {
+ start_partition: self.partition.max(0) as usize,
Review Comment:
Using the map partition id as the start keeps retries reproducible, but
consecutive tasks start on consecutive partitions. That's the correlation Spark
fixed in SPARK-21782, which is why Spark scrambles the start with `new
XORShiftRandom(partitionId).nextInt(numPartitions)`
([Spark](https://github.com/apache/spark/blob/582e28bf7074598ac85a0b7de9c5916c44cc5661/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala#L428-L442)).
Comet's JVM path does the same
([`CometShuffleExchangeExec.scala`](https://github.com/apache/datafusion-comet/blob/e99307b9df2620a743239abfc032f539f21e7042/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala#L1116-L1129)).
With adjacent starts, every task's run of partitions overlaps its neighbours',
and the partitions past `numMapTasks + groupsPerTask` get nothing. Ten map
tasks of 5,000 rows into 200 partitions at the auto group of 64 leave 112
reducers empty with adjacent start
s and none with scrambled starts.
What do you think about computing the start in
`CometNativeShuffleWriter.buildUnifiedPlan`, which already runs per task with
`context` in scope, and passing it in the proto? `new
XORShiftRandom(context.partitionId()).nextInt(numPartitions)` is still a pure
function of the map partition, so retries stay safe. The planner would also
stop depending on `self.partition`, which removes the `jni_api` partition-0
caveat in the comment above. Spark increments before its first use, so a start
of `nextInt(numPartitions) + 1` would make `groupRows = 1` place rows the way
Spark does for the same row order. A native test that runs several map tasks
with a small number of groups each and asserts the stage-wide spread would
cover this.
##########
native/shuffle/src/comet_partitioning.rs:
##########
@@ -19,6 +19,132 @@ use arrow::row::{OwnedRow, RowConverter};
use datafusion::physical_expr::{LexOrdering, PhysicalExpr};
use std::sync::Arc;
+/// How [`CometPartitioning::RoundRobin`] decides which output partition a row
belongs to.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RoundRobinStrategy {
+ /// Hash each row over its leading `max_hash_columns` columns (`0` meaning
all of them) and
+ /// place it at `pmod(hash, num_partitions)`.
+ ///
+ /// Placement is a pure function of a row's contents, so a re-executed map
task reproduces it
+ /// no matter what its input does. The price is a murmur3 pass per row
that recurses into
+ /// every struct child, plus a per-row gather on flush because adjacent
rows scatter across
+ /// every partition. It is also not really round robin: identical rows
always hash to the same
+ /// partition, so low-cardinality input skews where Spark's round robin
spreads evenly.
+ HashAll { max_hash_columns: usize },
+
+ /// Place rows positionally, in contiguous groups of `group_rows` rows,
counting rows across
+ /// input batch boundaries: the row at task-global ordinal `i` goes to
output partition
+ /// `(start_partition + i / group_rows) % num_partitions`.
+ ///
+ /// This is Spark's own round robin at a coarser granularity — Spark seeds
a counter with
+ /// `XORShiftRandom(partitionId)` and bumps it per row, which is the
`group_rows == 1` case —
+ /// and it inherits Spark's determinism condition exactly: placement is
reproducible when the
+ /// upstream operator replays rows in the same *order*. It deliberately
does not depend on how
+ /// those rows are framed into batches, because no Spark contract covers
framing;
+ /// `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. Keying on a
row ordinal rather
+ /// than a batch ordinal is what lets this strategy rely on the level
Spark already publishes
+ /// instead of an assumption nothing checks.
+ ///
+ /// `start_partition` must be the Spark map partition id. It has to be
distinct across mappers,
+ /// or every task starts at partition 0 and a task emitting fewer groups
than there are output
+ /// partitions leaves the tail empty stage-wide; and it has to be a pure
function of the map
+ /// partition, or a re-executed task does not reproduce its own placement.
Spark seeds
+ /// `XORShiftRandom(partitionId)` for the same two reasons.
+ ///
Review Comment:
The doc comment says a distinct start per mapper prevents an empty tail, and
that Spark seeds `XORShiftRandom(partitionId)` for the same reason. Distinct
isn't enough. Adjacent starts still leave the tail empty whenever the number of
map tasks plus the groups per task is less than the partition count (see the
comment on `planner.rs`). Spark's reason for the random seed is to decorrelate
the starts. Could this say that, whichever start function you end up with? The
same text is in `native_shuffle.md` lines 351-356.
##########
spark/src/main/scala/org/apache/comet/CometConf.scala:
##########
@@ -475,6 +475,42 @@ object CometConf extends ShimCometConf {
"The maximum number of columns to hash for round robin partitioning
must be non-negative.")
.createWithDefault(0)
+ val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_ENABLED:
ConfigEntry[Boolean] =
+
conf("spark.comet.shuffle.native.partitioning.roundrobin.positional.enabled")
+ .category(CATEGORY_SHUFFLE)
+ .doc(
+ "When true, Comet's native round-robin shuffle places rows by position
rather than by " +
+ "hashing their contents, the way Spark's own round robin does: the
row at " +
+ "task-global ordinal i goes to output partition " +
+ "(mapPartitionId + i / groupRows) % numPartitions. This skips a
murmur3 pass over " +
+ "every column of every row and replaces the per-row gather on flush
with a bulk copy " +
+ "per run, which is what dominates the shuffle write on wide nested
schemas. It also " +
+ "spreads duplicate rows evenly, where hashing sends them all to one
partition. " +
+ "Positional placement is only reproducible when the map task replays
rows in the " +
+ "same order, so it is used only where Comet can establish that from
the plan: a " +
+ "native scan under nothing but projections and filters. Any other
plan silently " +
+ "keeps content-hash placement. " +
+ s"Has no effect unless
${COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key} " +
+ "is also true.")
+ .booleanConf
+ .createWithDefault(false)
+
+ val COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_POSITIONAL_GROUP_ROWS:
ConfigEntry[Int] =
+
conf("spark.comet.shuffle.native.partitioning.roundrobin.positional.groupRows")
+ .category(CATEGORY_SHUFFLE)
+ .doc(
+ "Rows per contiguous group under positional round robin. Imbalance
between any two " +
+ "output partitions is bounded by this many rows however the reader
frames its " +
+ "batches, so smaller groups balance better while larger groups
produce fewer, longer " +
Review Comment:
"Imbalance between any two output partitions is bounded by this many rows"
is true for the rows one map task writes, but a reducer sees the sum over all
map tasks. With 50 map tasks and `groupRows = 8192`, the simulation in the
review body gives one reducer 409,600 rows and another 0. This is user-facing
documentation, so could it say "within one map task", and say that stage-wide
balance needs each task to emit many more groups than there are output
partitions? The same sentence is in `comet_partitioning.rs` lines 56-58 and
`native_shuffle.md` lines 358-359.
##########
native/shuffle/src/comet_partitioning.rs:
##########
@@ -19,6 +19,132 @@ use arrow::row::{OwnedRow, RowConverter};
use datafusion::physical_expr::{LexOrdering, PhysicalExpr};
use std::sync::Arc;
+/// How [`CometPartitioning::RoundRobin`] decides which output partition a row
belongs to.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RoundRobinStrategy {
+ /// Hash each row over its leading `max_hash_columns` columns (`0` meaning
all of them) and
+ /// place it at `pmod(hash, num_partitions)`.
+ ///
+ /// Placement is a pure function of a row's contents, so a re-executed map
task reproduces it
+ /// no matter what its input does. The price is a murmur3 pass per row
that recurses into
+ /// every struct child, plus a per-row gather on flush because adjacent
rows scatter across
+ /// every partition. It is also not really round robin: identical rows
always hash to the same
+ /// partition, so low-cardinality input skews where Spark's round robin
spreads evenly.
+ HashAll { max_hash_columns: usize },
+
+ /// Place rows positionally, in contiguous groups of `group_rows` rows,
counting rows across
+ /// input batch boundaries: the row at task-global ordinal `i` goes to
output partition
+ /// `(start_partition + i / group_rows) % num_partitions`.
+ ///
+ /// This is Spark's own round robin at a coarser granularity — Spark seeds
a counter with
+ /// `XORShiftRandom(partitionId)` and bumps it per row, which is the
`group_rows == 1` case —
+ /// and it inherits Spark's determinism condition exactly: placement is
reproducible when the
+ /// upstream operator replays rows in the same *order*. It deliberately
does not depend on how
+ /// those rows are framed into batches, because no Spark contract covers
framing;
+ /// `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. Keying on a
row ordinal rather
+ /// than a batch ordinal is what lets this strategy rely on the level
Spark already publishes
+ /// instead of an assumption nothing checks.
+ ///
+ /// `start_partition` must be the Spark map partition id. It has to be
distinct across mappers,
+ /// or every task starts at partition 0 and a task emitting fewer groups
than there are output
+ /// partitions leaves the tail empty stage-wide; and it has to be a pure
function of the map
+ /// partition, or a re-executed task does not reproduce its own placement.
Spark seeds
+ /// `XORShiftRandom(partitionId)` for the same two reasons.
+ ///
+ /// `group_rows` trades balance against copying. Imbalance between any two
output partitions is
+ /// bounded by `group_rows` rows regardless of how the reader frames
batches, so small groups
+ /// balance better; large groups produce fewer, longer runs to copy on
flush, and a group as
+ /// large as the batch size lets a whole input batch pass through to one
partition untouched.
+ /// [`Self::AUTO_GROUP_ROWS`] picks a value from the batch size and
partition count.
+ RowGroups {
+ start_partition: usize,
+ group_rows: usize,
+ },
+}
+
+impl Default for RoundRobinStrategy {
+ /// Hashing every column, which is what Comet's round robin did before
`RowGroups` existed.
+ fn default() -> Self {
+ Self::HashAll {
+ max_hash_columns: 0,
+ }
+ }
+}
+
+impl RoundRobinStrategy {
+ /// `group_rows` sentinel asking for a value derived from the batch size
and partition count.
+ pub const AUTO_GROUP_ROWS: usize = 0;
+
+ /// Smallest automatically chosen group. A multiple of 8 so that a run
starts on a byte
+ /// boundary of a validity bitmap, which keeps the per-run copy a memcpy
rather than a
+ /// bit-shift for every column.
+ const MIN_AUTO_GROUP_ROWS: usize = 64;
Review Comment:
The multiple-of-8 rationale only holds when every input batch starts on a
group boundary. `row_seq` counts rows across batches, so after a
`CometFilterExec` under the exchange, which the allowlist admits, a batch
starts at an arbitrary ordinal and its first run ends at an arbitrary offset.
Every run after it in that batch starts off a byte boundary. The same happens
with any explicit `groupRows` that isn't a multiple of 8. Is the 64-row floor
there to cap the run count instead? If so, could the comment say that?
--
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]
