refactor: make join projection pushdown schema-aware via ColumnIndex/… #23185
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -706,6 +706,80 @@ pub fn try_pushdown_through_join( | |
| })) | ||
| } | ||
|
|
||
| pub(crate) fn try_pushdown_through_join_with_column_indices( | ||
|
Contributor
There was a problem hiding this comment. this looks pretty similar to
Contributor
Author
There was a problem hiding this comment.
Contributor
There was a problem hiding this comment.
I suggest keep as a compatibily wrapper and delegate and then mark it as "deprecate" per https://datafusion.apache.org/contributor-guide/api-health.html#deprecation-guidelines |
||
| projection: &ProjectionExec, | ||
| join_left: &Arc<dyn ExecutionPlan>, | ||
| join_right: &Arc<dyn ExecutionPlan>, | ||
| join_on: JoinOnRef, | ||
| schema: &SchemaRef, | ||
| filter: Option<&JoinFilter>, | ||
| column_indices: &[ColumnIndex], | ||
| ) -> Result<Option<JoinData>> { | ||
| // Convert projected expressions to columns. We can not proceed if this is not possible. | ||
| let Some(projection_as_columns) = physical_to_column_exprs(projection.expr()) else { | ||
| return Ok(None); | ||
| }; | ||
|
|
||
| if projection_as_columns.len() >= schema.fields().len() { | ||
| return Ok(None); | ||
| } | ||
| let mut left_proj: Vec<(Column, String)> = Vec::new(); | ||
| let mut right_proj: Vec<(Column, String)> = Vec::new(); | ||
| let mut seen_right = false; | ||
| for (col, alias) in &projection_as_columns { | ||
| let origin = &column_indices[col.index()]; | ||
| match origin.side { | ||
| // Keep the "left block before right block" contiguity the current | ||
| // pushdown supports; a left column after a right one is "mixed". | ||
| JoinSide::Left => { | ||
| if seen_right { | ||
| return Ok(None); | ||
| } | ||
| left_proj.push((Column::new(col.name(), origin.index), alias.clone())); | ||
| } | ||
| JoinSide::Right => { | ||
| seen_right = true; | ||
| right_proj.push((Column::new(col.name(), origin.index), alias.clone())); | ||
| } | ||
| // Synthetic column (e.g. mark): belongs to neither child. | ||
| // Phase 2 declines; Phase 3 keeps it at the join output instead. | ||
| JoinSide::None => return Ok(None), | ||
| } | ||
| } | ||
|
|
||
| // Parity: neither side fully dropped. | ||
| if left_proj.is_empty() || right_proj.is_empty() { | ||
| return Ok(None); | ||
| } | ||
|
|
||
| // `left_proj` / `right_proj` carry *child* indices (from `column_indices`), | ||
| // so the shared `update_join_*` helpers must use a 0 column-index offset for | ||
| // both sides (the offset bridges child -> join-output index, which is the | ||
| // identity here). | ||
| let new_filter = if let Some(filter) = filter { | ||
| match update_join_filter(&left_proj, &right_proj, filter, 0) { | ||
| Some(updated) => Some(updated), | ||
| None => return Ok(None), | ||
| } | ||
| } else { | ||
| None | ||
| }; | ||
|
|
||
| let Some(new_on) = update_join_on(&left_proj, &right_proj, join_on, 0) else { | ||
| return Ok(None); | ||
| }; | ||
|
|
||
| let (new_left, new_right) = | ||
| new_join_children_from_groups(&left_proj, &right_proj, join_left, join_right)?; | ||
|
|
||
| Ok(Some(JoinData { | ||
| projected_left_child: new_left, | ||
| projected_right_child: new_right, | ||
| join_filter: new_filter, | ||
| join_on: new_on, | ||
| })) | ||
| } | ||
|
|
||
| /// This function checks if `plan` is a [`ProjectionExec`], and inspects its | ||
| /// input(s) to test whether it can push `plan` under its input(s). This function | ||
| /// will operate on the entire tree and may ultimately remove `plan` entirely | ||
|
|
@@ -880,6 +954,34 @@ pub fn new_join_children( | |
| Ok((new_left, new_right)) | ||
| } | ||
|
|
||
| /// Build the projected left and right children from side-grouped projection | ||
| /// columns whose indices are already *child*-relative (e.g. derived from a | ||
| /// join's `ColumnIndex`). Unlike [`new_join_children`], this does not infer | ||
| /// child ownership from output position, so it is safe for join schemas whose | ||
| /// output is not a plain `left ++ right` (used by the schema-aware | ||
| /// `try_pushdown_through_join`). | ||
| fn new_join_children_from_groups( | ||
| left_proj: &[(Column, String)], | ||
| right_proj: &[(Column, String)], | ||
| left_child: &Arc<dyn ExecutionPlan>, | ||
| right_child: &Arc<dyn ExecutionPlan>, | ||
| ) -> Result<(ProjectionExec, ProjectionExec)> { | ||
| let build = |cols: &[(Column, String)], child: &Arc<dyn ExecutionPlan>| { | ||
| ProjectionExec::try_new( | ||
| cols.iter().map(|(col, alias)| ProjectionExpr { | ||
| expr: Arc::new(Column::new(col.name(), col.index())) as _, | ||
| alias: alias.clone(), | ||
| }), | ||
| Arc::clone(child), | ||
| ) | ||
| }; | ||
|
|
||
| Ok(( | ||
| build(left_proj, left_child)?, | ||
| build(right_proj, right_child)?, | ||
| )) | ||
| } | ||
|
|
||
| /// Checks three conditions for pushing a projection down through a join: | ||
| /// - Projection must narrow the join output schema. | ||
| /// - Columns coming from left/right tables must be collected at the left/right | ||
|
|
@@ -946,14 +1048,10 @@ pub fn update_join_on( | |
| .map(|(left, right)| (left, right)) | ||
| .unzip(); | ||
|
|
||
| let new_left = new_columns_for_join_on(&left_idx, proj_left_exprs, 0)?; | ||
| let new_right = | ||
| new_columns_for_join_on(&right_idx, proj_right_exprs, left_field_size)?; | ||
| Some(new_left.into_iter().zip(new_right).collect()) | ||
| } | ||
|
|
||
| /// Tries to update the column indices of a [`JoinFilter`] as if the input of | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we still need try_pushdown_through_join after this PR? Perhaps we should remove it (or migrate other sites over that still use it)?