andygrove commented on code in PR #6095:
URL: https://github.com/apache/datafusion-comet/pull/6095#discussion_r4076038779
##########
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:
You're right, and I had taken "distinct is enough" on faith rather than
simulating it. Done the way
you suggested: the start is computed in `buildUnifiedPlan` as
`XORShiftRandom(context.partitionId()).nextInt(numPartitions) + 1` and
passed in the proto, so the
planner no longer reads `self.partition` and the `jni_api` partition-0
caveat goes with it. I took
the `+ 1` as well — Spark increments before its first use, so with it
`groupRows = 1` places rows
exactly where Spark's round robin would for the same row order, which makes
the "this is Spark's
round robin at a coarser granularity" claim literally true rather than
nearly true.
Your ten-tasks-of-5,000-into-200-at-64 figure reproduces exactly: 112 empty
with adjacent starts,
zero with scrambled. `CometNativePositionalRoundRobinSuite` now asserts
both, running the placement
formula over the real `positionalStartPartition`, and keeping the adjacent
case in the test so the
hazard stays visible rather than becoming folklore. On the native side the
new test asserts that a
task's groups walk *consecutive* partitions from its start, which is the
property that makes the
stage-wide spread a function of how the starts are chosen rather than of the
data — the thing that
makes your argument load-bearing on this side of the boundary.
##########
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:
Fixed, in the Rust doc and in `native_shuffle.md`. Both now say the starts
have to be decorrelated
rather than distinct, and give your overlap argument as the reason instead
of the empty-tail one.
##########
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:
Capping the run count is the real reason, and the alignment claim was wrong
for exactly the reason
you give — I had it in my head that a batch starts on a group boundary,
which is only true with no
filter and a group that divides the batch size. Reworded to say the floor is
there so that a
partition count far larger than the batch size cannot round `batch_size /
num_partitions` down
towards a handful of rows and turn the flush back into the per-row gather
the strategy exists to
avoid. I kept a sentence saying alignment is explicitly *not* guaranteed,
since 64 looks like an
alignment number and someone will assume it back otherwise.
##########
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:
Changed in all three. The config doc now says the bound holds within one map
task, that a reducer
sees the sum over all of them, and that the stage is only evenly balanced
when each task emits many
more groups than there are output partitions.
--
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]
