This is an automated email from the ASF dual-hosted git repository.
alamb 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 9ab597b251 Simplify Join Partition Rules (#10911)
9ab597b251 is described below
commit 9ab597b2511b23dc951c9f35d00d38c2071e9fea
Author: Berkay Şahin <[email protected]>
AuthorDate: Fri Jun 14 23:33:19 2024 +0300
Simplify Join Partition Rules (#10911)
* clean-up
* add asym case
* fix import errors
* Update utils.rs
* Update hash_join.rs
---
datafusion/core/src/dataframe/mod.rs | 11 ++--
datafusion/physical-plan/src/joins/hash_join.rs | 38 ++++--------
.../physical-plan/src/joins/nested_loop_join.rs | 32 +++--------
.../physical-plan/src/joins/sort_merge_join.rs | 37 ++++++------
.../physical-plan/src/joins/symmetric_hash_join.rs | 12 +---
datafusion/physical-plan/src/joins/utils.rs | 67 +++++++++++++++-------
6 files changed, 90 insertions(+), 107 deletions(-)
diff --git a/datafusion/core/src/dataframe/mod.rs
b/datafusion/core/src/dataframe/mod.rs
index 950cb7ddb2..b5c58eff57 100644
--- a/datafusion/core/src/dataframe/mod.rs
+++ b/datafusion/core/src/dataframe/mod.rs
@@ -3100,10 +3100,7 @@ mod tests {
let join_schema = physical_plan.schema();
match join_type {
- JoinType::Inner
- | JoinType::Left
- | JoinType::LeftSemi
- | JoinType::LeftAnti => {
+ JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti => {
let left_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
Arc::new(Column::new_with_schema("c1", &join_schema)?),
Arc::new(Column::new_with_schema("c2", &join_schema)?),
@@ -3113,7 +3110,10 @@ mod tests {
&Partitioning::Hash(left_exprs,
default_partition_count)
);
}
- JoinType::Right | JoinType::RightSemi | JoinType::RightAnti =>
{
+ JoinType::Inner
+ | JoinType::Right
+ | JoinType::RightSemi
+ | JoinType::RightAnti => {
let right_exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
Arc::new(Column::new_with_schema("c2_c1",
&join_schema)?),
Arc::new(Column::new_with_schema("c2_c2",
&join_schema)?),
@@ -3133,6 +3133,7 @@ mod tests {
Ok(())
}
+
#[tokio::test]
async fn nested_explain_should_fail() -> Result<()> {
let ctx = SessionContext::new();
diff --git a/datafusion/physical-plan/src/joins/hash_join.rs
b/datafusion/physical-plan/src/joins/hash_join.rs
index cd66ab093f..5353092d5c 100644
--- a/datafusion/physical-plan/src/joins/hash_join.rs
+++ b/datafusion/physical-plan/src/joins/hash_join.rs
@@ -23,6 +23,7 @@ use std::sync::Arc;
use std::task::Poll;
use std::{any::Any, vec};
+use super::utils::asymmetric_join_output_partitioning;
use super::{
utils::{OnceAsync, OnceFut},
PartitionMode,
@@ -34,10 +35,10 @@ use crate::{
execution_mode_from_children, handle_state,
hash_utils::create_hashes,
joins::utils::{
- adjust_indices_by_join_type, adjust_right_output_partitioning,
- apply_join_filter_to_indices, build_batch_from_indices,
build_join_schema,
- check_join_is_valid, estimate_join_statistics,
get_final_indices_from_bit_map,
- need_produce_result_in_final, partitioned_join_output_partitioning,
+ adjust_indices_by_join_type, apply_join_filter_to_indices,
+ build_batch_from_indices, build_join_schema, check_join_is_valid,
+ estimate_join_statistics, get_final_indices_from_bit_map,
+ need_produce_result_in_final, symmetric_join_output_partitioning,
BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMap,
JoinHashMapOffset,
JoinHashMapType, JoinOn, JoinOnRef, StatefulStreamResult,
},
@@ -490,33 +491,16 @@ impl HashJoinExec {
on,
);
- // Get output partitioning:
- let left_columns_len = left.schema().fields.len();
let mut output_partitioning = match mode {
- PartitionMode::CollectLeft => match join_type {
- JoinType::Inner | JoinType::Right =>
adjust_right_output_partitioning(
- right.output_partitioning(),
- left_columns_len,
- ),
- JoinType::RightSemi | JoinType::RightAnti => {
- right.output_partitioning().clone()
- }
- JoinType::Left
- | JoinType::LeftSemi
- | JoinType::LeftAnti
- | JoinType::Full => Partitioning::UnknownPartitioning(
- right.output_partitioning().partition_count(),
- ),
- },
- PartitionMode::Partitioned => partitioned_join_output_partitioning(
- join_type,
- left.output_partitioning(),
- right.output_partitioning(),
- left_columns_len,
- ),
+ PartitionMode::CollectLeft => {
+ asymmetric_join_output_partitioning(left, right, &join_type)
+ }
PartitionMode::Auto => Partitioning::UnknownPartitioning(
right.output_partitioning().partition_count(),
),
+ PartitionMode::Partitioned => {
+ symmetric_join_output_partitioning(left, right, &join_type)
+ }
};
// Determine execution mode by checking whether this join is pipeline
diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs
b/datafusion/physical-plan/src/joins/nested_loop_join.rs
index 18518600ef..6be124cce0 100644
--- a/datafusion/physical-plan/src/joins/nested_loop_join.rs
+++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs
@@ -25,18 +25,19 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::task::Poll;
+use super::utils::{asymmetric_join_output_partitioning,
need_produce_result_in_final};
use crate::coalesce_batches::concat_batches;
use crate::coalesce_partitions::CoalescePartitionsExec;
use crate::joins::utils::{
- adjust_indices_by_join_type, adjust_right_output_partitioning,
- apply_join_filter_to_indices, build_batch_from_indices, build_join_schema,
- check_join_is_valid, estimate_join_statistics,
get_final_indices_from_bit_map,
- BuildProbeJoinMetrics, ColumnIndex, JoinFilter, OnceAsync, OnceFut,
+ adjust_indices_by_join_type, apply_join_filter_to_indices,
build_batch_from_indices,
+ build_join_schema, check_join_is_valid, estimate_join_statistics,
+ get_final_indices_from_bit_map, BuildProbeJoinMetrics, ColumnIndex,
JoinFilter,
+ OnceAsync, OnceFut,
};
use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
use crate::{
execution_mode_from_children, DisplayAs, DisplayFormatType, Distribution,
- ExecutionMode, ExecutionPlan, ExecutionPlanProperties, Partitioning,
PlanProperties,
+ ExecutionMode, ExecutionPlan, ExecutionPlanProperties, PlanProperties,
RecordBatchStream, SendableRecordBatchStream,
};
@@ -55,8 +56,6 @@ use
datafusion_physical_expr::equivalence::join_equivalence_properties;
use futures::{ready, Stream, StreamExt, TryStreamExt};
use parking_lot::Mutex;
-use super::utils::need_produce_result_in_final;
-
/// Shared bitmap for visited left-side indices
type SharedBitmapBuilder = Mutex<BooleanBufferBuilder>;
/// Left (build-side) data
@@ -228,21 +227,8 @@ impl NestedLoopJoinExec {
&[],
);
- // Get output partitioning,
- let output_partitioning = match join_type {
- JoinType::Inner | JoinType::Right =>
adjust_right_output_partitioning(
- right.output_partitioning(),
- left.schema().fields().len(),
- ),
- JoinType::RightSemi | JoinType::RightAnti => {
- right.output_partitioning().clone()
- }
- JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti |
JoinType::Full => {
- Partitioning::UnknownPartitioning(
- right.output_partitioning().partition_count(),
- )
- }
- };
+ let output_partitioning =
+ asymmetric_join_output_partitioning(left, right, &join_type);
// Determine execution mode:
let mut mode = execution_mode_from_children([left, right]);
@@ -673,7 +659,7 @@ mod tests {
use datafusion_execution::runtime_env::{RuntimeConfig, RuntimeEnv};
use datafusion_expr::Operator;
use datafusion_physical_expr::expressions::{BinaryExpr, Literal};
- use datafusion_physical_expr::PhysicalExpr;
+ use datafusion_physical_expr::{Partitioning, PhysicalExpr};
fn build_table(
a: (&str, &Vec<i32>),
diff --git a/datafusion/physical-plan/src/joins/sort_merge_join.rs
b/datafusion/physical-plan/src/joins/sort_merge_join.rs
index 8da345cdfc..01abb30181 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join.rs
@@ -30,12 +30,22 @@ use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
+use crate::expressions::PhysicalSortExpr;
+use crate::joins::utils::{
+ build_join_schema, check_join_is_valid, estimate_join_statistics,
+ symmetric_join_output_partitioning, JoinFilter, JoinOn, JoinOnRef,
+};
+use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet};
+use crate::{
+ execution_mode_from_children, metrics, DisplayAs, DisplayFormatType,
Distribution,
+ ExecutionPlan, ExecutionPlanProperties, PhysicalExpr, PlanProperties,
+ RecordBatchStream, SendableRecordBatchStream, Statistics,
+};
+
use arrow::array::*;
use arrow::compute::{self, concat_batches, take, SortOptions};
use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
use arrow::error::ArrowError;
-use futures::{Stream, StreamExt};
-use hashbrown::HashSet;
use datafusion_common::{
internal_err, not_impl_err, plan_err, DataFusionError, JoinSide, JoinType,
Result,
@@ -45,17 +55,8 @@ use datafusion_execution::TaskContext;
use datafusion_physical_expr::equivalence::join_equivalence_properties;
use datafusion_physical_expr::{PhysicalExprRef, PhysicalSortRequirement};
-use crate::expressions::PhysicalSortExpr;
-use crate::joins::utils::{
- build_join_schema, check_join_is_valid, estimate_join_statistics,
- partitioned_join_output_partitioning, JoinFilter, JoinOn, JoinOnRef,
-};
-use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet};
-use crate::{
- execution_mode_from_children, metrics, DisplayAs, DisplayFormatType,
Distribution,
- ExecutionPlan, ExecutionPlanProperties, PhysicalExpr, PlanProperties,
- RecordBatchStream, SendableRecordBatchStream, Statistics,
-};
+use futures::{Stream, StreamExt};
+use hashbrown::HashSet;
/// join execution plan executes partitions in parallel and combines them into
a set of
/// partitions.
@@ -220,14 +221,8 @@ impl SortMergeJoinExec {
join_on,
);
- // Get output partitioning:
- let left_columns_len = left.schema().fields.len();
- let output_partitioning = partitioned_join_output_partitioning(
- join_type,
- left.output_partitioning(),
- right.output_partitioning(),
- left_columns_len,
- );
+ let output_partitioning =
+ symmetric_join_output_partitioning(left, right, &join_type);
// Determine execution mode:
let mode = execution_mode_from_children([left, right]);
diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
index e11e6dd2f6..813f670147 100644
--- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
+++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs
@@ -42,7 +42,7 @@ use crate::joins::stream_join_utils::{
};
use crate::joins::utils::{
apply_join_filter_to_indices, build_batch_from_indices, build_join_schema,
- check_join_is_valid, partitioned_join_output_partitioning, ColumnIndex,
JoinFilter,
+ check_join_is_valid, symmetric_join_output_partitioning, ColumnIndex,
JoinFilter,
JoinHashMapType, JoinOn, JoinOnRef, StatefulStreamResult,
};
use crate::{
@@ -271,14 +271,8 @@ impl SymmetricHashJoinExec {
join_on,
);
- // Get output partitioning:
- let left_columns_len = left.schema().fields.len();
- let output_partitioning = partitioned_join_output_partitioning(
- join_type,
- left.output_partitioning(),
- right.output_partitioning(),
- left_columns_len,
- );
+ let output_partitioning =
+ symmetric_join_output_partitioning(left, right, &join_type);
// Determine execution mode:
let mode = execution_mode_from_children([left, right]);
diff --git a/datafusion/physical-plan/src/joins/utils.rs
b/datafusion/physical-plan/src/joins/utils.rs
index c08b0e3d09..dfa1fd4763 100644
--- a/datafusion/physical-plan/src/joins/utils.rs
+++ b/datafusion/physical-plan/src/joins/utils.rs
@@ -25,7 +25,9 @@ use std::sync::Arc;
use std::task::{Context, Poll};
use crate::metrics::{self, ExecutionPlanMetricsSet, MetricBuilder};
-use crate::{ColumnStatistics, ExecutionPlan, Partitioning, Statistics};
+use crate::{
+ ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning,
Statistics,
+};
use arrow::array::{
downcast_array, new_null_array, Array, BooleanBufferBuilder, UInt32Array,
@@ -428,27 +430,6 @@ fn check_join_set_is_valid(
Ok(())
}
-/// Calculate the OutputPartitioning for Partitioned Join
-pub fn partitioned_join_output_partitioning(
- join_type: JoinType,
- left_partitioning: &Partitioning,
- right_partitioning: &Partitioning,
- left_columns_len: usize,
-) -> Partitioning {
- match join_type {
- JoinType::Inner | JoinType::Left | JoinType::LeftSemi |
JoinType::LeftAnti => {
- left_partitioning.clone()
- }
- JoinType::RightSemi | JoinType::RightAnti =>
right_partitioning.clone(),
- JoinType::Right => {
- adjust_right_output_partitioning(right_partitioning,
left_columns_len)
- }
- JoinType::Full => {
-
Partitioning::UnknownPartitioning(right_partitioning.partition_count())
- }
- }
-}
-
/// Adjust the right out partitioning to new Column Index
pub fn adjust_right_output_partitioning(
right_partitioning: &Partitioning,
@@ -1539,6 +1520,48 @@ pub enum StatefulStreamResult<T> {
Continue,
}
+pub(crate) fn symmetric_join_output_partitioning(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ join_type: &JoinType,
+) -> Partitioning {
+ let left_columns_len = left.schema().fields.len();
+ let left_partitioning = left.output_partitioning();
+ let right_partitioning = right.output_partitioning();
+ match join_type {
+ JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti => {
+ left_partitioning.clone()
+ }
+ JoinType::RightSemi | JoinType::RightAnti =>
right_partitioning.clone(),
+ JoinType::Inner | JoinType::Right => {
+ adjust_right_output_partitioning(right_partitioning,
left_columns_len)
+ }
+ JoinType::Full => {
+ // We could also use left partition count as they are necessarily
equal.
+
Partitioning::UnknownPartitioning(right_partitioning.partition_count())
+ }
+ }
+}
+
+pub(crate) fn asymmetric_join_output_partitioning(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ join_type: &JoinType,
+) -> Partitioning {
+ match join_type {
+ JoinType::Inner | JoinType::Right => adjust_right_output_partitioning(
+ right.output_partitioning(),
+ left.schema().fields().len(),
+ ),
+ JoinType::RightSemi | JoinType::RightAnti =>
right.output_partitioning().clone(),
+ JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti |
JoinType::Full => {
+ Partitioning::UnknownPartitioning(
+ right.output_partitioning().partition_count(),
+ )
+ }
+ }
+}
+
#[cfg(test)]
mod tests {
use std::pin::Pin;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
