avantgardnerio commented on code in PR #2685: URL: https://github.com/apache/datafusion-ballista/pull/2685#discussion_r4407262437
########## ballista/scheduler/src/state/pipelined_join_selection.rs: ########## @@ -0,0 +1,521 @@ +// 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. + +//! Size-aware join selection for pipelined jobs. +//! +//! A pipelined job is planned with the static planner, so AQE never chooses +//! its join strategies. The scheduler plans it with hash joins, and +//! [`PipelinedJoinSelection`] then turns every `Partitioned` hash join whose +//! build side is not known to fit +//! `ballista.optimizer.hash_join_max_build_partition_bytes` per partition back +//! into a `SortMergeJoinExec`, which can spill. Small build sides stay hash +//! joins, and the distributed planner promotes those under the broadcast +//! threshold to `CollectLeft` as usual. +//! +//! The rule runs among the physical optimizer rules, right after +//! DataFusion's `JoinSelection`, so that the later rules sort and partition +//! the inputs of the sort-merge joins it creates. Rewriting a finished plan +//! instead would lose orderings that a parent relied on: a hash join keeps +//! its probe side's order, so the planner may have dropped a sort above it. +//! +//! This is stricter than AQE, which keeps a hash join when the size is +//! unknown: a pipelined job holds every hash table of the job at once, a hash +//! build cannot spill, and a failed pipelined task fails the job, so a build of +//! unknown size is treated as too large. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use datafusion::arrow::compute::SortOptions; +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::StatisticsContext; +use datafusion::physical_plan::{ExecutionPlan, StatisticsArgs}; +use log::{debug, warn}; + +use crate::state::aqe::execution_plan::estimate_output_byte_size; + +/// Replaces `Partitioned` hash joins whose build side may not fit in memory +/// with sort-merge joins. It must run before the rules that enforce +/// distribution and sorting; see the module docs. +#[derive(Debug, Default)] +pub struct PipelinedJoinSelection; + +impl PipelinedJoinSelection { + /// Creates a new `PipelinedJoinSelection` rule. + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for PipelinedJoinSelection { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let max_build_bytes = config + .extensions + .get::<BallistaConfig>() + .cloned() + .unwrap_or_default() + .hash_join_max_build_partition_bytes(); + // 0 disables the check, as it does under AQE. + if max_build_bytes == 0 { + return Ok(plan); + } + + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<HashJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + // `CollectLeft` joins are left to the distributed planner's + // broadcast check, and a null-aware anti join has no sort-merge + // equivalent. + if *join.partition_mode() != PartitionMode::Partitioned || join.null_aware { + return Ok(Transformed::no(plan)); + } + + // The inputs are not repartitioned yet. A `Partitioned` join's + // will be, into `target_partitions` partitions. + let build_partition_bytes = estimated_bytes(join.left().as_ref()) + .map(|bytes| bytes.div_ceil(config.execution.target_partitions.max(1))); + if build_partition_bytes.is_some_and(|bytes| bytes <= max_build_bytes) { + return Ok(Transformed::no(plan)); + } + + match to_sort_merge_join(join) { + Ok(sort_merge_join) => { + debug!( + "pipelined join selection: demoting {:?} hash join to sort-merge join, \ + estimated build bytes per partition={build_partition_bytes:?}, \ + hash_join_max_build_partition_bytes={max_build_bytes}", + join.join_type(), + ); + Ok(Transformed::yes(sort_merge_join)) + } + // A join type the sort-merge join does not support keeps its + // hash join. + Err(e) => { + debug!( + "pipelined join selection: keeping {:?} hash join, no sort-merge equivalent: {e}", + join.join_type(), + ); + Ok(Transformed::no(plan)) + } + } + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "PipelinedJoinSelection" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Plans a pipelined job's logical plan with size-aware join selection. +/// +/// Without AQE there are no measured sizes to pick join strategies from, so +/// every join is planned as a hash join, as AQE does whatever +/// `datafusion.optimizer.prefer_hash_join` says, and [`PipelinedJoinSelection`] +/// falls back to a sort-merge join wherever the build side may not fit. +pub(crate) async fn create_pipelined_physical_plan( + state: &SessionState, + logical_plan: &LogicalPlan, +) -> Result<Arc<dyn ExecutionPlan>> { + let mut rules = state.physical_optimizers().to_vec(); + // Without `JoinSelection` the partition modes stay `Auto`, so the rule + // finds no `Partitioned` join and leaves the plan as it is. + let position = match rules + .iter() + .position(|rule| rule.name() == "join_selection") + { + Some(index) => index + 1, + None => { Review Comment: If `join_selection` is missing, the rule lands at position 0 where every join is still `Auto`, so it never fires. But `prefer_hash_join = true` is still forced below, so the result is all hash joins with no size check, which is less safe than plain static planning. Would `plan_err!` fit better here than `warn!`? ########## ballista/scheduler/src/state/pipelined_join_selection.rs: ########## @@ -0,0 +1,521 @@ +// 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. + +//! Size-aware join selection for pipelined jobs. +//! +//! A pipelined job is planned with the static planner, so AQE never chooses +//! its join strategies. The scheduler plans it with hash joins, and +//! [`PipelinedJoinSelection`] then turns every `Partitioned` hash join whose +//! build side is not known to fit +//! `ballista.optimizer.hash_join_max_build_partition_bytes` per partition back +//! into a `SortMergeJoinExec`, which can spill. Small build sides stay hash +//! joins, and the distributed planner promotes those under the broadcast +//! threshold to `CollectLeft` as usual. +//! +//! The rule runs among the physical optimizer rules, right after +//! DataFusion's `JoinSelection`, so that the later rules sort and partition +//! the inputs of the sort-merge joins it creates. Rewriting a finished plan +//! instead would lose orderings that a parent relied on: a hash join keeps +//! its probe side's order, so the planner may have dropped a sort above it. +//! +//! This is stricter than AQE, which keeps a hash join when the size is +//! unknown: a pipelined job holds every hash table of the job at once, a hash Review Comment: > holds every hash table of the job at once The check is per join, per partition, so N joins each just under the limit all stay hash joins and are all resident together. Is a job-wide or per-executor budget planned for #2657, or is per-join the intended bound? ########## ballista/scheduler/src/state/pipelined_join_selection.rs: ########## @@ -0,0 +1,521 @@ +// 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. + +//! Size-aware join selection for pipelined jobs. +//! +//! A pipelined job is planned with the static planner, so AQE never chooses +//! its join strategies. The scheduler plans it with hash joins, and +//! [`PipelinedJoinSelection`] then turns every `Partitioned` hash join whose +//! build side is not known to fit +//! `ballista.optimizer.hash_join_max_build_partition_bytes` per partition back +//! into a `SortMergeJoinExec`, which can spill. Small build sides stay hash +//! joins, and the distributed planner promotes those under the broadcast +//! threshold to `CollectLeft` as usual. +//! +//! The rule runs among the physical optimizer rules, right after +//! DataFusion's `JoinSelection`, so that the later rules sort and partition +//! the inputs of the sort-merge joins it creates. Rewriting a finished plan +//! instead would lose orderings that a parent relied on: a hash join keeps +//! its probe side's order, so the planner may have dropped a sort above it. +//! +//! This is stricter than AQE, which keeps a hash join when the size is +//! unknown: a pipelined job holds every hash table of the job at once, a hash +//! build cannot spill, and a failed pipelined task fails the job, so a build of +//! unknown size is treated as too large. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use datafusion::arrow::compute::SortOptions; +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::StatisticsContext; +use datafusion::physical_plan::{ExecutionPlan, StatisticsArgs}; +use log::{debug, warn}; + +use crate::state::aqe::execution_plan::estimate_output_byte_size; + +/// Replaces `Partitioned` hash joins whose build side may not fit in memory +/// with sort-merge joins. It must run before the rules that enforce +/// distribution and sorting; see the module docs. +#[derive(Debug, Default)] +pub struct PipelinedJoinSelection; + +impl PipelinedJoinSelection { + /// Creates a new `PipelinedJoinSelection` rule. + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for PipelinedJoinSelection { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let max_build_bytes = config + .extensions + .get::<BallistaConfig>() + .cloned() + .unwrap_or_default() + .hash_join_max_build_partition_bytes(); + // 0 disables the check, as it does under AQE. + if max_build_bytes == 0 { + return Ok(plan); + } + + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<HashJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + // `CollectLeft` joins are left to the distributed planner's + // broadcast check, and a null-aware anti join has no sort-merge + // equivalent. + if *join.partition_mode() != PartitionMode::Partitioned || join.null_aware { + return Ok(Transformed::no(plan)); + } + + // The inputs are not repartitioned yet. A `Partitioned` join's + // will be, into `target_partitions` partitions. + let build_partition_bytes = estimated_bytes(join.left().as_ref()) + .map(|bytes| bytes.div_ceil(config.execution.target_partitions.max(1))); + if build_partition_bytes.is_some_and(|bytes| bytes <= max_build_bytes) { + return Ok(Transformed::no(plan)); + } + + match to_sort_merge_join(join) { + Ok(sort_merge_join) => { + debug!( + "pipelined join selection: demoting {:?} hash join to sort-merge join, \ + estimated build bytes per partition={build_partition_bytes:?}, \ + hash_join_max_build_partition_bytes={max_build_bytes}", + join.join_type(), + ); + Ok(Transformed::yes(sort_merge_join)) + } + // A join type the sort-merge join does not support keeps its Review Comment: > A join type the sort-merge join does not support In DataFusion 55.2 `SortMergeJoinExec::try_new` accepts every `JoinType`, Mark included; it only errors on key or sort-option validation. So this arm is really "construction failed", and it keeps a possibly oversized hash join. Should it propagate the error instead? ########## ballista/scheduler/src/state/pipelined_join_selection.rs: ########## @@ -0,0 +1,521 @@ +// 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. + +//! Size-aware join selection for pipelined jobs. +//! +//! A pipelined job is planned with the static planner, so AQE never chooses +//! its join strategies. The scheduler plans it with hash joins, and +//! [`PipelinedJoinSelection`] then turns every `Partitioned` hash join whose +//! build side is not known to fit +//! `ballista.optimizer.hash_join_max_build_partition_bytes` per partition back +//! into a `SortMergeJoinExec`, which can spill. Small build sides stay hash +//! joins, and the distributed planner promotes those under the broadcast +//! threshold to `CollectLeft` as usual. +//! +//! The rule runs among the physical optimizer rules, right after +//! DataFusion's `JoinSelection`, so that the later rules sort and partition +//! the inputs of the sort-merge joins it creates. Rewriting a finished plan +//! instead would lose orderings that a parent relied on: a hash join keeps +//! its probe side's order, so the planner may have dropped a sort above it. +//! +//! This is stricter than AQE, which keeps a hash join when the size is +//! unknown: a pipelined job holds every hash table of the job at once, a hash +//! build cannot spill, and a failed pipelined task fails the job, so a build of +//! unknown size is treated as too large. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use datafusion::arrow::compute::SortOptions; +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::StatisticsContext; +use datafusion::physical_plan::{ExecutionPlan, StatisticsArgs}; +use log::{debug, warn}; + +use crate::state::aqe::execution_plan::estimate_output_byte_size; + +/// Replaces `Partitioned` hash joins whose build side may not fit in memory +/// with sort-merge joins. It must run before the rules that enforce +/// distribution and sorting; see the module docs. +#[derive(Debug, Default)] +pub struct PipelinedJoinSelection; + +impl PipelinedJoinSelection { + /// Creates a new `PipelinedJoinSelection` rule. + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for PipelinedJoinSelection { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let max_build_bytes = config + .extensions + .get::<BallistaConfig>() + .cloned() + .unwrap_or_default() + .hash_join_max_build_partition_bytes(); + // 0 disables the check, as it does under AQE. + if max_build_bytes == 0 { + return Ok(plan); + } + + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<HashJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + // `CollectLeft` joins are left to the distributed planner's + // broadcast check, and a null-aware anti join has no sort-merge + // equivalent. + if *join.partition_mode() != PartitionMode::Partitioned || join.null_aware { + return Ok(Transformed::no(plan)); + } + + // The inputs are not repartitioned yet. A `Partitioned` join's + // will be, into `target_partitions` partitions. + let build_partition_bytes = estimated_bytes(join.left().as_ref()) + .map(|bytes| bytes.div_ceil(config.execution.target_partitions.max(1))); Review Comment: Nit: with `repartition_joins` off, or when `EnforceDistribution` reuses an existing hash partitioning with a different count, the join runs with fewer input partitions than `target_partitions`, so each one is larger than this estimate. ########## ballista/scheduler/src/state/pipelined_join_selection.rs: ########## @@ -0,0 +1,521 @@ +// 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. + +//! Size-aware join selection for pipelined jobs. +//! +//! A pipelined job is planned with the static planner, so AQE never chooses +//! its join strategies. The scheduler plans it with hash joins, and +//! [`PipelinedJoinSelection`] then turns every `Partitioned` hash join whose +//! build side is not known to fit +//! `ballista.optimizer.hash_join_max_build_partition_bytes` per partition back +//! into a `SortMergeJoinExec`, which can spill. Small build sides stay hash +//! joins, and the distributed planner promotes those under the broadcast +//! threshold to `CollectLeft` as usual. +//! +//! The rule runs among the physical optimizer rules, right after +//! DataFusion's `JoinSelection`, so that the later rules sort and partition +//! the inputs of the sort-merge joins it creates. Rewriting a finished plan +//! instead would lose orderings that a parent relied on: a hash join keeps +//! its probe side's order, so the planner may have dropped a sort above it. +//! +//! This is stricter than AQE, which keeps a hash join when the size is +//! unknown: a pipelined job holds every hash table of the job at once, a hash +//! build cannot spill, and a failed pipelined task fails the job, so a build of +//! unknown size is treated as too large. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use datafusion::arrow::compute::SortOptions; +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::StatisticsContext; +use datafusion::physical_plan::{ExecutionPlan, StatisticsArgs}; +use log::{debug, warn}; + +use crate::state::aqe::execution_plan::estimate_output_byte_size; + +/// Replaces `Partitioned` hash joins whose build side may not fit in memory +/// with sort-merge joins. It must run before the rules that enforce +/// distribution and sorting; see the module docs. +#[derive(Debug, Default)] +pub struct PipelinedJoinSelection; + +impl PipelinedJoinSelection { + /// Creates a new `PipelinedJoinSelection` rule. + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for PipelinedJoinSelection { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let max_build_bytes = config + .extensions + .get::<BallistaConfig>() + .cloned() + .unwrap_or_default() + .hash_join_max_build_partition_bytes(); + // 0 disables the check, as it does under AQE. + if max_build_bytes == 0 { + return Ok(plan); + } + + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<HashJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + // `CollectLeft` joins are left to the distributed planner's Review Comment: The distributed planner can turn a `CollectLeft` join back into `Partitioned` after this rule has run (`planner.rs`, unsafe join type or over `broadcast_join_threshold_bytes`), so those joins never get size-checked. In practice they are under DataFusion's `hash_join_single_partition_threshold`, so it only matters if that threshold is raised a lot. Maybe worth a sentence in the module docs? ########## ballista/scheduler/src/state/pipelined_join_selection.rs: ########## @@ -0,0 +1,521 @@ +// 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. + +//! Size-aware join selection for pipelined jobs. +//! +//! A pipelined job is planned with the static planner, so AQE never chooses +//! its join strategies. The scheduler plans it with hash joins, and +//! [`PipelinedJoinSelection`] then turns every `Partitioned` hash join whose +//! build side is not known to fit +//! `ballista.optimizer.hash_join_max_build_partition_bytes` per partition back +//! into a `SortMergeJoinExec`, which can spill. Small build sides stay hash +//! joins, and the distributed planner promotes those under the broadcast +//! threshold to `CollectLeft` as usual. +//! +//! The rule runs among the physical optimizer rules, right after +//! DataFusion's `JoinSelection`, so that the later rules sort and partition +//! the inputs of the sort-merge joins it creates. Rewriting a finished plan +//! instead would lose orderings that a parent relied on: a hash join keeps +//! its probe side's order, so the planner may have dropped a sort above it. +//! +//! This is stricter than AQE, which keeps a hash join when the size is +//! unknown: a pipelined job holds every hash table of the job at once, a hash +//! build cannot spill, and a failed pipelined task fails the job, so a build of +//! unknown size is treated as too large. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use datafusion::arrow::compute::SortOptions; +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::StatisticsContext; +use datafusion::physical_plan::{ExecutionPlan, StatisticsArgs}; +use log::{debug, warn}; + +use crate::state::aqe::execution_plan::estimate_output_byte_size; + +/// Replaces `Partitioned` hash joins whose build side may not fit in memory +/// with sort-merge joins. It must run before the rules that enforce +/// distribution and sorting; see the module docs. +#[derive(Debug, Default)] +pub struct PipelinedJoinSelection; + +impl PipelinedJoinSelection { + /// Creates a new `PipelinedJoinSelection` rule. + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for PipelinedJoinSelection { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let max_build_bytes = config + .extensions + .get::<BallistaConfig>() + .cloned() + .unwrap_or_default() + .hash_join_max_build_partition_bytes(); + // 0 disables the check, as it does under AQE. + if max_build_bytes == 0 { + return Ok(plan); + } + + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<HashJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + // `CollectLeft` joins are left to the distributed planner's + // broadcast check, and a null-aware anti join has no sort-merge + // equivalent. + if *join.partition_mode() != PartitionMode::Partitioned || join.null_aware { + return Ok(Transformed::no(plan)); + } + + // The inputs are not repartitioned yet. A `Partitioned` join's + // will be, into `target_partitions` partitions. + let build_partition_bytes = estimated_bytes(join.left().as_ref()) + .map(|bytes| bytes.div_ceil(config.execution.target_partitions.max(1))); + if build_partition_bytes.is_some_and(|bytes| bytes <= max_build_bytes) { + return Ok(Transformed::no(plan)); + } + + match to_sort_merge_join(join) { + Ok(sort_merge_join) => { + debug!( + "pipelined join selection: demoting {:?} hash join to sort-merge join, \ + estimated build bytes per partition={build_partition_bytes:?}, \ + hash_join_max_build_partition_bytes={max_build_bytes}", + join.join_type(), + ); + Ok(Transformed::yes(sort_merge_join)) + } + // A join type the sort-merge join does not support keeps its + // hash join. + Err(e) => { + debug!( + "pipelined join selection: keeping {:?} hash join, no sort-merge equivalent: {e}", + join.join_type(), + ); + Ok(Transformed::no(plan)) + } + } + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "PipelinedJoinSelection" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Plans a pipelined job's logical plan with size-aware join selection. +/// +/// Without AQE there are no measured sizes to pick join strategies from, so +/// every join is planned as a hash join, as AQE does whatever +/// `datafusion.optimizer.prefer_hash_join` says, and [`PipelinedJoinSelection`] +/// falls back to a sort-merge join wherever the build side may not fit. +pub(crate) async fn create_pipelined_physical_plan( + state: &SessionState, + logical_plan: &LogicalPlan, +) -> Result<Arc<dyn ExecutionPlan>> { + let mut rules = state.physical_optimizers().to_vec(); + // Without `JoinSelection` the partition modes stay `Auto`, so the rule + // finds no `Partitioned` join and leaves the plan as it is. + let position = match rules + .iter() + .position(|rule| rule.name() == "join_selection") + { + Some(index) => index + 1, + None => { + warn!( + "pipelined join selection: no join_selection rule in the session, \ + so hash joins are not demoted to sort-merge joins" + ); + 0 + } + }; + rules.insert(position, Arc::new(PipelinedJoinSelection::new())); + let mut planning_state = SessionStateBuilder::new_from_existing(state.clone()) + .with_physical_optimizer_rules(rules) + .build(); + planning_state + .config_mut() + .options_mut() + .optimizer + .prefer_hash_join = true; + planning_state.create_physical_plan(logical_plan).await +} + +/// Estimated output size of `plan` in bytes, or `None` if unknown. +/// +/// Uses `total_byte_size` when it is known, otherwise estimates it from the +/// row count as AQE's broadcast check does. A zero is treated as unknown, since +/// statistics can report zero for inputs they did not measure. +fn estimated_bytes(plan: &dyn ExecutionPlan) -> Option<usize> { + let stats = StatisticsContext::new() + .compute(plan, &StatisticsArgs::new()) + .ok()?; + if let Some(&bytes) = stats.total_byte_size.get_value() + && bytes != 0 Review Comment: The reasoning for treating zero as unknown fits `Inexact(0)`, but an `Exact(0)` build side is known empty and gets demoted to a sort-merge join too. Only costs a needless sort, so minor. ########## ballista/scheduler/src/state/pipelined_join_selection.rs: ########## @@ -0,0 +1,521 @@ +// 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. + +//! Size-aware join selection for pipelined jobs. +//! +//! A pipelined job is planned with the static planner, so AQE never chooses +//! its join strategies. The scheduler plans it with hash joins, and +//! [`PipelinedJoinSelection`] then turns every `Partitioned` hash join whose +//! build side is not known to fit +//! `ballista.optimizer.hash_join_max_build_partition_bytes` per partition back +//! into a `SortMergeJoinExec`, which can spill. Small build sides stay hash +//! joins, and the distributed planner promotes those under the broadcast +//! threshold to `CollectLeft` as usual. +//! +//! The rule runs among the physical optimizer rules, right after +//! DataFusion's `JoinSelection`, so that the later rules sort and partition +//! the inputs of the sort-merge joins it creates. Rewriting a finished plan +//! instead would lose orderings that a parent relied on: a hash join keeps +//! its probe side's order, so the planner may have dropped a sort above it. +//! +//! This is stricter than AQE, which keeps a hash join when the size is +//! unknown: a pipelined job holds every hash table of the job at once, a hash +//! build cannot spill, and a failed pipelined task fails the job, so a build of +//! unknown size is treated as too large. + +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use datafusion::arrow::compute::SortOptions; +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_expr::PhysicalExprRef; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::StatisticsContext; +use datafusion::physical_plan::{ExecutionPlan, StatisticsArgs}; +use log::{debug, warn}; + +use crate::state::aqe::execution_plan::estimate_output_byte_size; + +/// Replaces `Partitioned` hash joins whose build side may not fit in memory +/// with sort-merge joins. It must run before the rules that enforce +/// distribution and sorting; see the module docs. +#[derive(Debug, Default)] +pub struct PipelinedJoinSelection; + +impl PipelinedJoinSelection { + /// Creates a new `PipelinedJoinSelection` rule. + pub fn new() -> Self { + Self + } +} + +impl PhysicalOptimizerRule for PipelinedJoinSelection { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let max_build_bytes = config + .extensions + .get::<BallistaConfig>() + .cloned() + .unwrap_or_default() + .hash_join_max_build_partition_bytes(); + // 0 disables the check, as it does under AQE. + if max_build_bytes == 0 { + return Ok(plan); + } + + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<HashJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + // `CollectLeft` joins are left to the distributed planner's + // broadcast check, and a null-aware anti join has no sort-merge + // equivalent. + if *join.partition_mode() != PartitionMode::Partitioned || join.null_aware { + return Ok(Transformed::no(plan)); + } + + // The inputs are not repartitioned yet. A `Partitioned` join's + // will be, into `target_partitions` partitions. + let build_partition_bytes = estimated_bytes(join.left().as_ref()) + .map(|bytes| bytes.div_ceil(config.execution.target_partitions.max(1))); + if build_partition_bytes.is_some_and(|bytes| bytes <= max_build_bytes) { + return Ok(Transformed::no(plan)); + } + + match to_sort_merge_join(join) { + Ok(sort_merge_join) => { + debug!( + "pipelined join selection: demoting {:?} hash join to sort-merge join, \ + estimated build bytes per partition={build_partition_bytes:?}, \ + hash_join_max_build_partition_bytes={max_build_bytes}", + join.join_type(), + ); + Ok(Transformed::yes(sort_merge_join)) + } + // A join type the sort-merge join does not support keeps its + // hash join. + Err(e) => { + debug!( + "pipelined join selection: keeping {:?} hash join, no sort-merge equivalent: {e}", + join.join_type(), + ); + Ok(Transformed::no(plan)) + } + } + }) + .map(|t| t.data) + } + + fn name(&self) -> &str { + "PipelinedJoinSelection" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Plans a pipelined job's logical plan with size-aware join selection. +/// +/// Without AQE there are no measured sizes to pick join strategies from, so +/// every join is planned as a hash join, as AQE does whatever +/// `datafusion.optimizer.prefer_hash_join` says, and [`PipelinedJoinSelection`] +/// falls back to a sort-merge join wherever the build side may not fit. +pub(crate) async fn create_pipelined_physical_plan( + state: &SessionState, + logical_plan: &LogicalPlan, +) -> Result<Arc<dyn ExecutionPlan>> { + let mut rules = state.physical_optimizers().to_vec(); + // Without `JoinSelection` the partition modes stay `Auto`, so the rule + // finds no `Partitioned` join and leaves the plan as it is. + let position = match rules + .iter() + .position(|rule| rule.name() == "join_selection") + { + Some(index) => index + 1, + None => { + warn!( + "pipelined join selection: no join_selection rule in the session, \ + so hash joins are not demoted to sort-merge joins" + ); + 0 + } + }; + rules.insert(position, Arc::new(PipelinedJoinSelection::new())); + let mut planning_state = SessionStateBuilder::new_from_existing(state.clone()) + .with_physical_optimizer_rules(rules) + .build(); + planning_state + .config_mut() + .options_mut() + .optimizer + .prefer_hash_join = true; + planning_state.create_physical_plan(logical_plan).await +} + +/// Estimated output size of `plan` in bytes, or `None` if unknown. +/// +/// Uses `total_byte_size` when it is known, otherwise estimates it from the +/// row count as AQE's broadcast check does. A zero is treated as unknown, since +/// statistics can report zero for inputs they did not measure. +fn estimated_bytes(plan: &dyn ExecutionPlan) -> Option<usize> { + let stats = StatisticsContext::new() + .compute(plan, &StatisticsArgs::new()) + .ok()?; + if let Some(&bytes) = stats.total_byte_size.get_value() + && bytes != 0 + { + return Some(bytes); + } + let num_rows = *stats.num_rows.get_value()?; + if num_rows == 0 { + return None; + } + estimate_output_byte_size(&plan.schema(), num_rows, &stats.column_statistics) +} + +/// The sort-merge join equivalent of a `Partitioned` hash join, over the same +/// inputs. The later optimizer rules sort and partition them on the join keys. +/// The hash join's projection, if any, is kept as a `ProjectionExec` on top. +fn to_sort_merge_join(join: &HashJoinExec) -> Result<Arc<dyn ExecutionPlan>> { + let sort_merge_join = Arc::new(SortMergeJoinExec::try_new( + Arc::clone(join.left()), + Arc::clone(join.right()), + join.on().to_vec(), + join.filter().cloned(), + *join.join_type(), + vec![SortOptions::default(); join.on().len()], + join.null_equality, + )?) as Arc<dyn ExecutionPlan>; + + let Some(projection) = join.projection.as_deref() else { + return Ok(sort_merge_join); + }; + let schema = sort_merge_join.schema(); + let exprs = projection.iter().map(|&index| { + let name = schema.field(index).name(); + ( + Arc::new(Column::new(name, index)) as PhysicalExprRef, + name.clone(), + ) + }); + Ok(Arc::new(ProjectionExec::try_new(exprs, sort_merge_join)?)) +} + +#[cfg(test)] +mod tests { + use super::*; + use ballista_core::config::BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES; + use ballista_core::extension::SessionConfigExt; + use datafusion::arrow::array::{Int64Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::stats::Precision; + use datafusion::common::tree_node::TreeNodeRecursion; + use datafusion::common::{ColumnStatistics, JoinType, NullEquality, Statistics}; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::test::exec::StatisticsExec; + use datafusion::prelude::{SessionConfig, SessionContext, col}; + + const PARTITIONS: usize = 16; + + /// An input with the given size statistics. The rule runs before the + /// inputs are repartitioned, so it sees them bare. + fn input( + name: &str, + num_rows: Precision<usize>, + bytes: Precision<usize>, + ) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![ + Field::new(format!("{name}_key"), DataType::Int64, false), + Field::new(format!("{name}_value"), DataType::Int64, false), + ]); + let stats = Statistics { + num_rows, + total_byte_size: bytes, + column_statistics: vec![ColumnStatistics::new_unknown(); 2], + }; + Arc::new(StatisticsExec::new(stats, schema)) + } + + fn hash_join( + build: Arc<dyn ExecutionPlan>, + mode: PartitionMode, + projection: Option<Vec<usize>>, + ) -> Arc<dyn ExecutionPlan> { + let probe = input( + "probe", + Precision::Exact(1 << 30), + Precision::Exact(16 << 30), + ); + let on = vec![( + Arc::new(Column::new("build_key", 0)) as PhysicalExprRef, + Arc::new(Column::new("probe_key", 0)) as PhysicalExprRef, + )]; + Arc::new( + HashJoinExec::try_new( + build, + probe, + on, + None, + &JoinType::Inner, Review Comment: Nit: only `Inner` is tested. A `Left` or `Full` join with a filter would also cover the sort-merge join's filter path after demotion. -- 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]
