This is an automated email from the ASF dual-hosted git repository.
dheres pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new b00723a2a0 Fix and Improve Sort Pushdown for Nested Loop and Hash Join
(#12559)
b00723a2a0 is described below
commit b00723a2a0c0c1a5a30171729f7091407c819110
Author: Berkay Şahin <[email protected]>
AuthorDate: Mon Sep 23 16:31:01 2024 +0300
Fix and Improve Sort Pushdown for Nested Loop and Hash Join (#12559)
* Update sort_pushdown.rs
* Update joins.slt
---
.../core/src/physical_optimizer/sort_pushdown.rs | 128 +++++++++++++++++++--
datafusion/sqllogictest/test_files/joins.slt | 37 ++++++
2 files changed, 153 insertions(+), 12 deletions(-)
diff --git a/datafusion/core/src/physical_optimizer/sort_pushdown.rs
b/datafusion/core/src/physical_optimizer/sort_pushdown.rs
index 41059ef7b1..34410bd6b1 100644
--- a/datafusion/core/src/physical_optimizer/sort_pushdown.rs
+++ b/datafusion/core/src/physical_optimizer/sort_pushdown.rs
@@ -22,22 +22,27 @@ use super::utils::{add_sort_above, is_sort};
use crate::physical_optimizer::utils::{is_sort_preserving_merge, is_union,
is_window};
use crate::physical_plan::filter::FilterExec;
use crate::physical_plan::joins::utils::calculate_join_output_ordering;
-use crate::physical_plan::joins::{HashJoinExec, SortMergeJoinExec};
+use crate::physical_plan::joins::SortMergeJoinExec;
use crate::physical_plan::projection::ProjectionExec;
use crate::physical_plan::repartition::RepartitionExec;
use crate::physical_plan::sorts::sort::SortExec;
use crate::physical_plan::tree_node::PlanContext;
use crate::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
-use datafusion_common::tree_node::{ConcreteTreeNode, Transformed,
TreeNodeRecursion};
+use datafusion_common::tree_node::{
+ ConcreteTreeNode, Transformed, TreeNode, TreeNodeRecursion,
+};
use datafusion_common::{plan_err, JoinSide, Result};
use datafusion_expr::JoinType;
use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr::utils::collect_columns;
use datafusion_physical_expr::{
LexRequirementRef, PhysicalSortExpr, PhysicalSortRequirement,
};
use datafusion_physical_expr_common::sort_expr::LexRequirement;
+use hashbrown::HashSet;
+
/// This is a "data class" we use within the [`EnforceSorting`] rule to push
/// down [`SortExec`] in the plan. In some cases, we can reduce the total
/// computational cost by pushing down `SortExec`s through some executors. The
@@ -253,7 +258,6 @@ fn pushdown_requirement_to_children(
|| plan.as_any().is::<FilterExec>()
// TODO: Add support for Projection push down
|| plan.as_any().is::<ProjectionExec>()
- || plan.as_any().is::<HashJoinExec>()
|| pushdown_would_violate_requirements(parent_required, plan.as_ref())
{
// If the current plan is a leaf node or can not maintain any of the
input ordering, can not pushed down requirements.
@@ -277,15 +281,7 @@ fn pushdown_requirement_to_children(
Ok(Some(vec![req]))
}
} else {
- Ok(Some(
- maintains_input_order
- .into_iter()
- .map(|flag| {
- (flag && !parent_required.is_empty())
- .then(|| parent_required.to_vec())
- })
- .collect(),
- ))
+ handle_custom_pushdown(plan, parent_required, maintains_input_order)
}
// TODO: Add support for Projection push down
}
@@ -483,6 +479,114 @@ fn shift_right_required(
}
}
+/// Handles the custom pushdown of parent-required sorting requirements down to
+/// the child execution plans, considering whether the input order is
maintained.
+///
+/// # Arguments
+///
+/// * `plan` - A reference to an `ExecutionPlan` for which the pushdown will
be applied.
+/// * `parent_required` - The sorting requirements expected by the parent node.
+/// * `maintains_input_order` - A vector of booleans indicating whether each
child
+/// maintains the input order.
+///
+/// # Returns
+///
+/// Returns `Ok(Some(Vec<Option<LexRequirement>>))` if the sorting
requirements can be
+/// pushed down, `Ok(None)` if not. On error, returns a `Result::Err`.
+fn handle_custom_pushdown(
+ plan: &Arc<dyn ExecutionPlan>,
+ parent_required: LexRequirementRef,
+ maintains_input_order: Vec<bool>,
+) -> Result<Option<Vec<Option<LexRequirement>>>> {
+ // If there's no requirement from the parent or the plan has no children,
return early
+ if parent_required.is_empty() || plan.children().is_empty() {
+ return Ok(None);
+ }
+
+ // Collect all unique column indices used in the parent-required sorting
expression
+ let all_indices: HashSet<usize> = parent_required
+ .iter()
+ .flat_map(|order| {
+ collect_columns(&order.expr)
+ .iter()
+ .map(|col| col.index())
+ .collect::<HashSet<_>>()
+ })
+ .collect();
+
+ // Get the number of fields in each child's schema
+ let len_of_child_schemas: Vec<usize> = plan
+ .children()
+ .iter()
+ .map(|c| c.schema().fields().len())
+ .collect();
+
+ // Find the index of the child that maintains input order
+ let Some(maintained_child_idx) = maintains_input_order
+ .iter()
+ .enumerate()
+ .find(|(_, m)| **m)
+ .map(|pair| pair.0)
+ else {
+ return Ok(None);
+ };
+
+ // Check if all required columns come from the child that maintains input
order
+ let start_idx = len_of_child_schemas[..maintained_child_idx]
+ .iter()
+ .sum::<usize>();
+ let end_idx = start_idx + len_of_child_schemas[maintained_child_idx];
+ let all_from_maintained_child =
+ all_indices.iter().all(|i| i >= &start_idx && i < &end_idx);
+
+ // If all columns are from the maintained child, update the parent
requirements
+ if all_from_maintained_child {
+ let sub_offset = len_of_child_schemas
+ .iter()
+ .take(maintained_child_idx)
+ .sum::<usize>();
+ // Transform the parent-required expression for the child schema by
adjusting columns
+ let updated_parent_req = parent_required
+ .iter()
+ .map(|req| {
+ let child_schema =
plan.children()[maintained_child_idx].schema();
+ let updated_columns = req
+ .expr
+ .clone()
+ .transform_up(|expr| {
+ if let Some(col) =
expr.as_any().downcast_ref::<Column>() {
+ let new_index = col.index() - sub_offset;
+ Ok(Transformed::yes(Arc::new(Column::new(
+ child_schema.field(new_index).name(),
+ new_index,
+ ))))
+ } else {
+ Ok(Transformed::no(expr))
+ }
+ })?
+ .data;
+ Ok(PhysicalSortRequirement::new(updated_columns, req.options))
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ // Prepare the result, populating with the updated requirements for
children that maintain order
+ let result = maintains_input_order
+ .iter()
+ .map(|&maintains_order| {
+ if maintains_order {
+ Some(updated_parent_req.clone())
+ } else {
+ None
+ }
+ })
+ .collect();
+
+ Ok(Some(result))
+ } else {
+ Ok(None)
+ }
+}
+
/// Define the Requirements Compatibility
#[derive(Debug)]
enum RequirementsCompatibility {
diff --git a/datafusion/sqllogictest/test_files/joins.slt
b/datafusion/sqllogictest/test_files/joins.slt
index 7d0262952b..64d5f6c7b8 100644
--- a/datafusion/sqllogictest/test_files/joins.slt
+++ b/datafusion/sqllogictest/test_files/joins.slt
@@ -4151,3 +4151,40 @@ DROP TABLE sales_global;
statement ok
DROP TABLE exchange_rates;
+
+# HashJoinExec and NestedLoopJoinExec can propagate SortExec down through its
right child.
+
+statement ok
+CREATE TABLE left_table(a INT, b INT, c INT)
+
+statement ok
+CREATE TABLE right_table(x INT, y INT, z INT)
+
+query TT
+EXPLAIN SELECT * FROM left_table JOIN right_table ON
left_table.a<right_table.x ORDER BY x;
+----
+logical_plan
+01)Sort: right_table.x ASC NULLS LAST
+02)--Inner Join: Filter: left_table.a < right_table.x
+03)----TableScan: left_table projection=[a, b, c]
+04)----TableScan: right_table projection=[x, y, z]
+physical_plan
+01)NestedLoopJoinExec: join_type=Inner, filter=a@0 < x@1
+02)--MemoryExec: partitions=1, partition_sizes=[0]
+03)--SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false]
+04)----MemoryExec: partitions=1, partition_sizes=[0]
+
+query TT
+EXPLAIN SELECT * FROM left_table JOIN right_table ON
left_table.a<right_table.x AND left_table.b=right_table.y ORDER BY x;
+----
+logical_plan
+01)Sort: right_table.x ASC NULLS LAST
+02)--Inner Join: left_table.b = right_table.y Filter: left_table.a <
right_table.x
+03)----TableScan: left_table projection=[a, b, c]
+04)----TableScan: right_table projection=[x, y, z]
+physical_plan
+01)CoalesceBatchesExec: target_batch_size=3
+02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(b@1, y@1)],
filter=a@0 < x@1
+03)----MemoryExec: partitions=1, partition_sizes=[0]
+04)----SortExec: expr=[x@0 ASC NULLS LAST], preserve_partitioning=[false]
+05)------MemoryExec: partitions=1, partition_sizes=[0]
\ No newline at end of file
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
