xudong963 commented on code in PR #26404: URL: https://github.com/apache/datafusion/pull/26404#discussion_r4319959656
########## datafusion/physical-optimizer/src/eliminate_constant_join.rs: ########## @@ -0,0 +1,202 @@ +// 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. + +//! Replace inner joins against a single row of literals with a filter. + +use std::sync::Arc; + +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{JoinSide, JoinType, internal_datafusion_err}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{Column, LambdaExpr, Literal}; +use datafusion_physical_plan::filter::FilterExec; +use datafusion_physical_plan::joins::NestedLoopJoinExec; +use datafusion_physical_plan::joins::utils::ColumnIndex; +use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +use crate::PhysicalOptimizerRule; + +/// Rewrites an inner [`NestedLoopJoinExec`] whose input provably produces one row +/// of literals into a [`FilterExec`] over the other input, with the literals +/// substituted into the join filter, under a [`ProjectionExec`] that restores the +/// join's output. +/// +/// For example, [`AggregateStatistics`] folds the scalar subquery in +/// +/// ```sql +/// SELECT e.* FROM events e CROSS JOIN (SELECT min(ts) AS first_ts FROM live) l +/// WHERE e.ts < l.first_ts +/// ``` +/// +/// into literals over a [`PlaceholderRowExec`], leaving a join that only filters +/// `events`. A nested-loop join preserves neither input's order, while the filter +/// and projection preserve the other input's ordering and partitioning. This lets +/// later operators, such as an aggregation grouped by a sort key of `events`, +/// use that ordering. +/// +/// [`AggregateStatistics`]: crate::aggregate_statistics::AggregateStatistics +#[derive(Debug, Default)] +pub struct EliminateConstantJoin {} + +impl EliminateConstantJoin { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for EliminateConstantJoin { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + _config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<NestedLoopJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + Ok(match eliminate(join)? { + Some(rewritten) => Transformed::yes(rewritten), + None => Transformed::no(plan), + }) + }) + .data() + } + + fn name(&self) -> &str { + "EliminateConstantJoin" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Rewrites `join` over its non-constant input, if the other input is a constant row. +fn eliminate(join: &NestedLoopJoinExec) -> Result<Option<Arc<dyn ExecutionPlan>>> { + if join.join_type() != &JoinType::Inner { + return Ok(None); + } + if let Some(filter) = join.filter() + // Lambda variables are numbered after the filter's own columns, not the + // other input's. + && filter + .expression() + .exists(|expr| Ok(expr.is::<LambdaExpr>()))? + { + return Ok(None); + } + // Each input's columns as expressions over the non-constant input. + let (left, right, input) = if let Some(row) = constant_row(join.left()) { + (row, columns(join.right()), join.right()) + } else if let Some(row) = constant_row(join.right()) { + (columns(join.left()), row, join.left()) + } else { + return Ok(None); + }; + + let mut input = Arc::clone(input); + if let Some(filter) = join.filter() { + let predicate = Arc::clone(filter.expression()) + .transform_up(|expr| { + let Some(column) = expr.downcast_ref::<Column>() else { + return Ok(Transformed::no(expr)); + }; + // Filter columns index the filter's own schema, which + // `column_indices` maps back to the join inputs. + let ColumnIndex { index, side } = + filter.column_indices().get(column.index()).ok_or_else(|| { + internal_datafusion_err!("Unmapped join filter column") + })?; + let value = match side { + JoinSide::Left => left.get(*index), + JoinSide::Right => right.get(*index), + JoinSide::None => None, + } + .ok_or_else(|| { + internal_datafusion_err!("Join filter column out of range") + })?; + Ok(Transformed::yes(Arc::clone(value))) + }) + .data()?; + input = Arc::new(FilterExec::try_new(predicate, input)?); + } + + let mut outputs = left.into_iter().chain(right).collect::<Vec<_>>(); + if let Some(projection) = join.projection() { + outputs = projection + .iter() + .map(|&index| outputs.get(index).cloned()) + .collect::<Option<_>>() + .ok_or_else(|| internal_datafusion_err!("Join projection out of range"))?; + } + let schema = join.schema(); + let exprs = outputs + .into_iter() + .zip(schema.fields()) + .map(|(expr, field)| ProjectionExpr { + expr, + alias: field.name().clone(), + }); + // Keep the join's metadata, which may come from the constant input. + Ok(Some(Arc::new( + ProjectionExec::try_new_with_schema_metadata(exprs, input, &schema)?, + ))) +} + +/// Returns the literal values of the one row `plan` produces, if it is a stack of +/// projections over a single-partition [`PlaceholderRowExec`] that project only +/// literals and columns of the projection below. +fn constant_row(plan: &Arc<dyn ExecutionPlan>) -> Option<Vec<Arc<dyn PhysicalExpr>>> { + let projection = plan.downcast_ref::<ProjectionExec>()?; + let input = projection.input(); + let input_row = if input.is::<PlaceholderRowExec>() { + // A placeholder emits one row per partition, and its columns are not + // literals. + if input.output_partitioning().partition_count() != 1 { + return None; + } + Vec::new() + } else { + constant_row(input)? + }; + projection + .expr() + .iter() + .map(|ProjectionExpr { expr, .. }| { + if expr.is::<Literal>() { Review Comment: Checked this: constants written in SQL are folded by the logical simplifier, but values folded by `AggregateStatistics` only become literals during physical planning, so `(SELECT min(ts) + 1 ...)` reached the rule as `25 + 1` and was skipped. The rule now substitutes the row below and folds with `PhysicalExprSimplifier`, accepting only a literal result; volatile expressions and ones that fail to evaluate are still rejected. Added a comment and an slt case in b9fa393cd. ########## datafusion/sqllogictest/test_files/eliminate_constant_join.slt: ########## @@ -0,0 +1,126 @@ +# 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. + +# Inner joins against a single row of literals become a filter over the other +# input, which keeps that input's ordering (EliminateConstantJoin). + +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.explain.physical_plan_only = true; + +# Events sorted by account +query I +COPY (SELECT * FROM (VALUES ('a', 10), ('a', 30), ('b', 20), ('b', NULL), ('c', 5)) AS v(acct, t)) +TO 'test_files/scratch/eliminate_constant_join/events.parquet' +STORED AS PARQUET; +---- +5 + +statement ok +CREATE EXTERNAL TABLE events +STORED AS PARQUET +WITH ORDER (acct ASC NULLS LAST) +LOCATION 'test_files/scratch/eliminate_constant_join/events.parquet'; + +query I +COPY (SELECT * FROM (VALUES (25), (40)) AS v(ts)) +TO 'test_files/scratch/eliminate_constant_join/live.parquet' +STORED AS PARQUET; +---- +2 + +statement ok +CREATE EXTERNAL TABLE live +STORED AS PARQUET +LOCATION 'test_files/scratch/eliminate_constant_join/live.parquet'; + +# min(ts) is answered from Parquet statistics, so the join only filters the +# events, and the aggregation keeps their account order +query TT +EXPLAIN SELECT e.acct, count(*) AS n +FROM events e CROSS JOIN (SELECT min(ts) AS first_ts FROM live) l +WHERE e.t < l.first_ts +GROUP BY e.acct; +---- +physical_plan +01)ProjectionExec: expr=[acct@0 as acct, count(Int64(1))@1 as n] +02)--AggregateExec: mode=Single, gby=[acct@0 as acct], aggr=[count(Int64(1))], group_clustering_mode=Full +03)----FilterExec: 25 > t@1, projection=[acct@0] +04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/eliminate_constant_join/events.parquet]]}, projection=[acct, t], output_ordering=[acct@0 ASC NULLS LAST], file_type=parquet, predicate=25 > t@1, pruning_predicate=t_null_count@1 != row_count@2 AND t_min@0 < 25, required_guarantees=[] + +query TI +SELECT e.acct, count(*) AS n +FROM events e CROSS JOIN (SELECT min(ts) AS first_ts FROM live) l +WHERE e.t < l.first_ts +GROUP BY e.acct +ORDER BY e.acct; +---- +a 1 +b 1 +c 1 + +# The constant row's columns stay in the output +query TT +EXPLAIN SELECT e.acct, e.t, c.cutoff +FROM (SELECT 25 AS cutoff) c JOIN events e ON e.t < c.cutoff; +---- +physical_plan +01)ProjectionExec: expr=[acct@0 as acct, t@1 as t, 25 as cutoff] +02)--FilterExec: t@1 < 25 +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/eliminate_constant_join/events.parquet]]}, projection=[acct, t], output_ordering=[acct@0 ASC NULLS LAST], file_type=parquet, predicate=t@1 < 25, pruning_predicate=t_null_count@1 != row_count@2 AND t_min@0 < 25, required_guarantees=[] + +query TII rowsort +SELECT e.acct, e.t, c.cutoff +FROM (SELECT 25 AS cutoff) c JOIN events e ON e.t < c.cutoff; +---- +a 10 25 +b 20 25 +c 5 25 + +# A left join keeps the constant row when nothing matches Review Comment: Added an `EXPLAIN` showing the left join stays a `NestedLoopJoinExec` (b9fa393cd). ########## datafusion/physical-optimizer/src/eliminate_constant_join.rs: ########## @@ -0,0 +1,202 @@ +// 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. + +//! Replace inner joins against a single row of literals with a filter. + +use std::sync::Arc; + +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{JoinSide, JoinType, internal_datafusion_err}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{Column, LambdaExpr, Literal}; +use datafusion_physical_plan::filter::FilterExec; +use datafusion_physical_plan::joins::NestedLoopJoinExec; +use datafusion_physical_plan::joins::utils::ColumnIndex; +use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + +use crate::PhysicalOptimizerRule; + +/// Rewrites an inner [`NestedLoopJoinExec`] whose input provably produces one row +/// of literals into a [`FilterExec`] over the other input, with the literals +/// substituted into the join filter, under a [`ProjectionExec`] that restores the +/// join's output. +/// +/// For example, [`AggregateStatistics`] folds the scalar subquery in +/// +/// ```sql +/// SELECT e.* FROM events e CROSS JOIN (SELECT min(ts) AS first_ts FROM live) l +/// WHERE e.ts < l.first_ts +/// ``` +/// +/// into literals over a [`PlaceholderRowExec`], leaving a join that only filters +/// `events`. A nested-loop join preserves neither input's order, while the filter +/// and projection preserve the other input's ordering and partitioning. This lets +/// later operators, such as an aggregation grouped by a sort key of `events`, +/// use that ordering. +/// +/// [`AggregateStatistics`]: crate::aggregate_statistics::AggregateStatistics +#[derive(Debug, Default)] +pub struct EliminateConstantJoin {} + +impl EliminateConstantJoin { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for EliminateConstantJoin { + fn optimize( + &self, + plan: Arc<dyn ExecutionPlan>, + _config: &ConfigOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + plan.transform_up(|plan| { + let Some(join) = plan.downcast_ref::<NestedLoopJoinExec>() else { + return Ok(Transformed::no(plan)); + }; + Ok(match eliminate(join)? { + Some(rewritten) => Transformed::yes(rewritten), + None => Transformed::no(plan), + }) + }) + .data() + } + + fn name(&self) -> &str { + "EliminateConstantJoin" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Rewrites `join` over its non-constant input, if the other input is a constant row. +fn eliminate(join: &NestedLoopJoinExec) -> Result<Option<Arc<dyn ExecutionPlan>>> { Review Comment: Added in b9fa393cd. `CrossJoinExec` drops input ordering and, when the other input has no statistics, buffers it as the build side, so the rewrite helps there as well. This changes the `j1, LATERAL (SELECT 1)` plans in `cte.slt` and `joins.slt`. -- 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]
