diff --git a/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs index d53522f5de7c8..894697595dd3f 100644 --- a/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs +++ b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs @@ -26,7 +26,7 @@ use arrow::datatypes::{DataType, Schema}; use datafusion_common::internal_datafusion_err; use datafusion_common::{ Result, - tree_node::{Transformed, TransformedResult, TreeNode}, + tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion}, }; use datafusion_expr::ColumnarValue; @@ -80,7 +80,7 @@ pub struct DynamicFilterPhysicalExpr { /// when the inner generation hasn't changed (common — updates fire once /// per HashJoin build or once per TopK threshold refresh, but `evaluate` /// is called per batch), the cache serves the remapped expression - /// without re-running the `transform_up` tree walk in + /// without re-running the `transform_down` tree walk in /// [`Self::remap_children`]. Reset on `update()` (by generation bump) /// and populated with `None` on `with_new_children` (each derived /// filter owns its own cache). @@ -234,7 +234,7 @@ impl DynamicFilterPhysicalExpr { if let Some(remapped_children) = remapped_children { // Remap the children to the new children // of the expression. - expr.transform_up(|child| { + expr.transform_down(|child| { // Check if this is any of our original children if let Some(pos) = children.iter().position(|c| c.as_ref() == child.as_ref()) @@ -242,7 +242,9 @@ impl DynamicFilterPhysicalExpr { // If so, remap it to the current children // of the expression. let new_child = Arc::clone(&remapped_children[pos]); - Ok(Transformed::yes(new_child)) + // The replacement is already in the target schema. Do not + // traverse it or remap any of its descendants again. + Ok(Transformed::new(new_child, true, TreeNodeRecursion::Jump)) } else { // Otherwise, just return the expression Ok(Transformed::no(child)) diff --git a/datafusion/sqllogictest/test_files/topk_dynamic_filter_remap.slt b/datafusion/sqllogictest/test_files/topk_dynamic_filter_remap.slt new file mode 100644 index 0000000000000..309457b85ce76 --- /dev/null +++ b/datafusion/sqllogictest/test_files/topk_dynamic_filter_remap.slt @@ -0,0 +1,51 @@ +# 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. + +statement ok +SET datafusion.execution.target_partitions = 1; + +statement ok +SET datafusion.execution.parquet.pushdown_filters = true; + +statement ok +COPY ( + SELECT value AS b FROM generate_series(0, 31999) +) TO 'test_files/scratch/topk_dynamic_filter_remap.parquet' +STORED AS PARQUET OPTIONS ('format.max_row_group_size' '500'); + +statement ok +CREATE EXTERNAL TABLE t STORED AS PARQUET +LOCATION 'test_files/scratch/topk_dynamic_filter_remap.parquet'; + +query I +WITH s AS (SELECT b, 31999 - b AS p, 31999 - b AS q FROM t) +SELECT p FROM s +ORDER BY CASE WHEN b % 2 = 0 THEN p ELSE q END, b +LIMIT 3; +---- +0 +1 +2 + +statement ok +DROP TABLE t; + +statement ok +SET datafusion.execution.target_partitions = 4; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters;