andygrove commented on code in PR #2501:
URL:
https://github.com/apache/datafusion-ballista/pull/2501#discussion_r4161011854
##########
ballista/scheduler/src/state/aqe/planner.rs:
##########
@@ -609,10 +603,13 @@ impl AdaptivePlanner {
// without restricting by join type -- safe in one process, but
// Ballista runs one task per probe partition. Demote the
unsafe
// ones straight after, while `EnsureRequirements` can still
add
- // the repartitions a `Partitioned` join needs.
- "join_selection" => {
- vec![r, Arc::new(DemoteUnsafeBroadcastJoinRule::default())]
- }
+ // the repartitions a `Partitioned` join needs. A swap can also
+ // move a broadcast onto the probe side; read it partitioned.
+ "join_selection" => vec![
+ r,
+ Arc::new(DemoteUnsafeBroadcastJoinRule::default()),
+ Arc::new(PartitionProbeSideBroadcastRule::default()),
Review Comment:
I'm curious about the size of the new build side. After the swap it's
whatever used to be the probe side, which can be well over
`broadcast_join_threshold_bytes`. In
`swapped_broadcast_is_read_partitioned_on_the_probe_side` it's `mid` at 2 GB,
against the 128 MB default. Before this change one task built that hash table.
Now every probe task builds its own copy. With the default
`max_partitions_per_task = 0` that's roughly one copy per executor at a time,
but with `max_partitions_per_task = 1` it's one per concurrent task. The 55.0.0
upgrade guide tells users peak memory is about the threshold times concurrent
tasks, and this path isn't bounded by the threshold. Would it make sense to
only do this when the new build side fits the threshold, and otherwise read the
probe side partitioned but demote the join to `Partitioned`, the way
`DemoteUnsafeBroadcastJoinRule` does? Happy for that to be a follow-up if the
benchmark numbers look good.
##########
ballista/scheduler/src/state/aqe/execution_plan/exchange.rs:
##########
@@ -175,6 +175,32 @@ impl ExchangeExec {
)
}
+ /// The inverse of [`Self::to_broadcast`]: the same stage read partitioned,
+ /// keeping this exchange's plan id so it replaces it in place.
+ ///
+ /// `None` when the stage was written with a partition count other than its
+ /// input's, as when a stage that ran as a hash shuffle is later broadcast:
+ /// reading it partitioned would drop or misplace locations.
+ pub fn to_partitioned(&self) -> Option<Self> {
+ let input_partitions =
self.input.output_partitioning().partition_count();
+ if let Some(written) = self.shuffle_partitions.lock().as_ref()
+ && written.len() != input_partitions
Review Comment:
I think the count check might not fully cover the case the doc comment
describes. When a stage written as a hash shuffle gets broadcast via
`to_broadcast`, it usually has `target_partitions` buckets, and its input often
has `target_partitions` partitions too, so the check passes. The partitioned
read then advertises the input's partitioning, but the data is bucketed on the
shuffle keys. I tried an input partitioned on `Hash([x], 4)` under a stage
written as `Hash([y], 4)`. After the rule the exchange reports `Hash([x@0], 4)`
and the join reports `Hash([x@1], 4)`, so a parent grouping or joining on `x`
could skip a repartition it needs. I couldn't find a default-config path where
`join_selection` swaps an exchange built by `to_broadcast`, since it compares
the same stats `SelectJoinRule` already used. Setting
`datafusion.optimizer.use_statistics_registry = true` can make the two disagree
though. Would it work to advertise `UnknownPartitioning(n)` once the stage has
been written? Or
`to_broadcast` could keep the original partitioning so `to_partitioned` can
restore `Hash(y, K)`, which would also make the count check unnecessary.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/partition_probe_side.rs:
##########
@@ -0,0 +1,185 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use datafusion::common::Result;
+use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
+use datafusion::physical_plan::{
+ ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions,
+};
+use std::sync::Arc;
+
+/// Reads a broadcast exchange on the probe side of a `CollectLeft` join
+/// partitioned instead.
+///
+/// When a broadcast build side measures larger than expected, `join_selection`
+/// swaps it onto the probe side, where its single partition would pin the join
+/// stage to one task. Runs straight after `join_selection`.
+///
+/// Null-aware joins track probe-side NULLs in-process, so they keep a single
+/// task. Broadcasts over ordered inputs are left to the k-way merge reader.
+#[derive(Debug, Default)]
+pub struct PartitionProbeSideBroadcastRule {}
+
+/// `node`'s children with its probe-side broadcast read partitioned, if it has
+/// one to replace.
+fn with_partitioned_probe(
+ node: &Arc<dyn ExecutionPlan>,
+) -> Option<Vec<Arc<dyn ExecutionPlan>>> {
+ let join = node.downcast_ref::<HashJoinExec>()?;
+ if join.null_aware || *join.partition_mode() != PartitionMode::CollectLeft
{
+ return None;
+ }
+ let probe = join.right().downcast_ref::<ExchangeExec>()?;
+ if !probe.broadcast ||
probe.input().properties().output_ordering().is_some() {
Review Comment:
Could you say a bit more about skipping ordered inputs? In
`BallistaAdapter::build_reader` a broadcast exchange always becomes
`try_new_broadcast`, and it's the partitioned read that would go through the
k-way merge reader. So as written, an ordered broadcast on the probe side stays
single-task. Is there a case where `RangeShuffleReaderExec` would get this
wrong? If it's just being conservative that's fine too, but the doc comment on
line 37 reads as if the merge reader handles these.
##########
ballista/scheduler/src/state/execution_stage.rs:
##########
@@ -1402,14 +1402,6 @@ impl StageOutput {
partition_locations
}
-
- /// returns vector of partition locations
- /// which is compatible with ShuffleReader vector format
- /// supporting broadcast shuffle read.
- /// All partitions are merged into one
- pub fn partition_locations_broadcast(self) -> Vec<Vec<PartitionLocation>> {
Review Comment:
`StageOutput` is public through
`ballista_scheduler::state::execution_stage`, so removing
`partition_locations_broadcast` is a small breaking change. Could we keep it
with a `#[deprecated]` note pointing at `partition_locations(n)`? Then this
doesn't need the `api-change` label or an upgrade guide entry.
##########
ballista/scheduler/src/state/aqe/optimizer_rule/partition_probe_side.rs:
##########
@@ -0,0 +1,185 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::state::aqe::execution_plan::ExchangeExec;
+use datafusion::common::Result;
+use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion::config::ConfigOptions;
+use datafusion::physical_optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode};
+use datafusion::physical_plan::{
+ ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions,
+};
+use std::sync::Arc;
+
+/// Reads a broadcast exchange on the probe side of a `CollectLeft` join
+/// partitioned instead.
+///
+/// When a broadcast build side measures larger than expected, `join_selection`
+/// swaps it onto the probe side, where its single partition would pin the join
+/// stage to one task. Runs straight after `join_selection`.
+///
+/// Null-aware joins track probe-side NULLs in-process, so they keep a single
+/// task. Broadcasts over ordered inputs are left to the k-way merge reader.
+#[derive(Debug, Default)]
+pub struct PartitionProbeSideBroadcastRule {}
+
+/// `node`'s children with its probe-side broadcast read partitioned, if it has
+/// one to replace.
+fn with_partitioned_probe(
+ node: &Arc<dyn ExecutionPlan>,
+) -> Option<Vec<Arc<dyn ExecutionPlan>>> {
+ let join = node.downcast_ref::<HashJoinExec>()?;
+ if join.null_aware || *join.partition_mode() != PartitionMode::CollectLeft
{
Review Comment:
Should this check the join type too? Right now the rule is only safe because
`DemoteUnsafeBroadcastJoinRule` runs just before it and has already demoted
`Left`, `Full`, `LeftSemi`, `LeftAnti` and `LeftMark` joins. The unit tests
actually assert that a non-null-aware `LeftAnti` join gets rewritten, and in
Ballista that rewrite would give wrong answers, since each probe task would
emit its own unmatched build rows. Checking
`!collect_left_broadcast_safe(*join.join_type())` would make the rule safe on
its own. It covers the null-aware case too, since DataFusion only allows
`null_aware` on `LeftAnti`. I tried it locally, and once the test helper uses
`Inner` for the non-null-aware cases the only change is the `join_type=` line
in three snapshots. Small related thing: the doc comment says the rule runs
straight after `join_selection`, but it really runs after the demote rule,
which is the ordering it depends on.
--
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]
