You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/08 06:02:04 UTC
(datafusion) branch main updated: fix: skips projection pruning for whole subtree (#20545)
This is an automated email from the ASF dual-hosted git repository.
github-bot 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 f4e24a54d5 fix: skips projection pruning for whole subtree (#20545)
f4e24a54d5 is described below
commit f4e24a54d52682bd070c1a40a31f55fb5655f014
Author: Acfboy <ac...@outlook.com>
AuthorDate: Wed Apr 8 14:00:17 2026 +0800
fix: skips projection pruning for whole subtree (#20545)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Closes #18816 .
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
In `UserDefinedLogicalNodeCore`, the default implementation of
`necessary_children_exprs ` returns `None`, which signals to the
optimizer that it cannot determine which columns are required from the
child.
The optimizer takes a conservative approach and skips projection pruning
for that node, leading to complex and redundant plans in the subtree.
However, it would make more sense to assume all columns are required and
let the optimizer proceed, rather than giving up on the entire subtree
entirely.
## What changes are included in this PR?
```rust
LogicalPlan::Extension(extension) => {
if let Some(necessary_children_indices) =
extension.node.necessary_children_exprs(indices.indices())
{
...
} else {
// Requirements from parent cannot be routed down to user defined logical plan safely
// Assume it requires all input exprs here
plan.inputs()
.into_iter()
.map(RequiredIndices::new_for_all_exprs)
.collect()
}
}
```
instead of
https://github.com/apache/datafusion/blob/b6d46a63824f003117297848d8d83b659ac2e759/datafusion/optimizer/src/optimize_projections/mod.rs#L331-L337
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
## Are these changes tested?
Yes.
In addition to unit tests, I've also added a complete end-to-end
integration test that reproduces the full scenario in the issue. This
might seem redundant, bloated, or even unnecessary. Please let me know
if I should remove these tests.
An existing test is modified, but I think the newer behavior is
expected.
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
## Are there any user-facing changes?
Yes. But I think the new implementation is the expected behavior.
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
datafusion/core/tests/dataframe/mod.rs | 5 +-
.../optimizer/src/optimize_projections/mod.rs | 117 +++++++++++++---
.../optimizer/tests/optimizer_integration.rs | 155 ++++++++++++++++++++-
3 files changed, 253 insertions(+), 24 deletions(-)
diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs
index 1c661744e0..6e22aeab61 100644
--- a/datafusion/core/tests/dataframe/mod.rs
+++ b/datafusion/core/tests/dataframe/mod.rs
@@ -2458,9 +2458,8 @@ async fn cache_producer_test() -> Result<()> {
@r"
CacheNode
Projection: aggregate_test_100.c2, aggregate_test_100.c3, CAST(CAST(aggregate_test_100.c2 AS Int64) + CAST(aggregate_test_100.c3 AS Int64) AS Int64) AS sum
- Projection: aggregate_test_100.c2, aggregate_test_100.c3
- Limit: skip=0, fetch=1
- TableScan: aggregate_test_100, fetch=1
+ Limit: skip=0, fetch=1
+ TableScan: aggregate_test_100 projection=[c2, c3], fetch=1
"
);
Ok(())
diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs
index 93df300bb5..26fbb71af6 100644
--- a/datafusion/optimizer/src/optimize_projections/mod.rs
+++ b/datafusion/optimizer/src/optimize_projections/mod.rs
@@ -329,29 +329,34 @@ fn optimize_projections(
.collect()
}
LogicalPlan::Extension(extension) => {
- let Some(necessary_children_indices) =
+ if let Some(necessary_children_indices) =
extension.node.necessary_children_exprs(indices.indices())
- else {
- // Requirements from parent cannot be routed down to user defined logical plan safely
- return Ok(Transformed::no(plan));
- };
- let children = extension.node.inputs();
- assert_eq_or_internal_err!(
- children.len(),
- necessary_children_indices.len(),
- "Inconsistent length between children and necessary children indices. \
+ {
+ let children = extension.node.inputs();
+ assert_eq_or_internal_err!(
+ children.len(),
+ necessary_children_indices.len(),
+ "Inconsistent length between children and necessary children indices. \
Make sure `.necessary_children_exprs` implementation of the \
`UserDefinedLogicalNode` is consistent with actual children length \
for the node."
- );
- children
- .into_iter()
- .zip(necessary_children_indices)
- .map(|(child, necessary_indices)| {
- RequiredIndices::new_from_indices(necessary_indices)
- .with_plan_exprs(&plan, child.schema())
- })
- .collect::<Result<Vec<_>>>()?
+ );
+ children
+ .into_iter()
+ .zip(necessary_children_indices)
+ .map(|(child, necessary_indices)| {
+ RequiredIndices::new_from_indices(necessary_indices)
+ .with_plan_exprs(&plan, child.schema())
+ })
+ .collect::<Result<Vec<_>>>()?
+ } else {
+ // Requirements from parent cannot be routed down to user defined logical plan safely
+ // Assume it requires all input exprs here
+ plan.inputs()
+ .into_iter()
+ .map(RequiredIndices::new_for_all_exprs)
+ .collect()
+ }
}
LogicalPlan::EmptyRelation(_)
| LogicalPlan::Values(_)
@@ -1172,6 +1177,57 @@ mod tests {
}
}
+ /// A user-defined node that does NOT implement `necessary_children_exprs`,
+ /// so the optimizer cannot determine which columns are required from its
+ /// children and must assume all columns are needed.
+ #[derive(Debug, Hash, PartialEq, Eq)]
+ struct OpaqueRequirementsUserDefined {
+ input: Arc<LogicalPlan>,
+ schema: DFSchemaRef,
+ }
+
+ // Manual implementation needed because of `schema` field. Comparison excludes this field.
+ impl PartialOrd for OpaqueRequirementsUserDefined {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ self.input
+ .partial_cmp(&other.input)
+ .filter(|cmp| *cmp != Ordering::Equal || self == other)
+ }
+ }
+
+ impl UserDefinedLogicalNodeCore for OpaqueRequirementsUserDefined {
+ fn name(&self) -> &str {
+ "OpaqueRequirementsUserDefined"
+ }
+
+ fn inputs(&self) -> Vec<&LogicalPlan> {
+ vec![&self.input]
+ }
+
+ fn schema(&self) -> &DFSchemaRef {
+ &self.schema
+ }
+
+ fn expressions(&self) -> Vec<Expr> {
+ vec![]
+ }
+
+ fn with_exprs_and_inputs(
+ &self,
+ _exprs: Vec<Expr>,
+ mut inputs: Vec<LogicalPlan>,
+ ) -> Result<Self> {
+ Ok(Self {
+ input: Arc::new(inputs.swap_remove(0)),
+ schema: Arc::clone(&self.schema),
+ })
+ }
+
+ fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
+ write!(f, "OpaqueRequirementsUserDefined")
+ }
+ }
+
#[test]
fn merge_two_projection() -> Result<()> {
let table_scan = test_table_scan()?;
@@ -2204,6 +2260,29 @@ mod tests {
Ok(())
}
+ #[test]
+ fn test_continue_processing_through_extension() -> Result<()> {
+ let table_scan = test_table_scan()?;
+ let plan = LogicalPlanBuilder::from(table_scan.clone())
+ .project(vec![col("a")])?
+ .project(vec![col("a")])?
+ .build()?;
+ let plan = LogicalPlan::Extension(Extension {
+ node: Arc::new(OpaqueRequirementsUserDefined {
+ input: Arc::new(plan),
+ schema: Arc::clone(table_scan.schema()),
+ }),
+ });
+ let plan = optimize(plan).expect("failed to optimize plan");
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ OpaqueRequirementsUserDefined
+ TableScan: test projection=[a]
+ "
+ )
+ }
+
/// tests that it removes an aggregate is never used downstream
#[test]
fn table_unused_aggregate() -> Result<()> {
diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs
index fd4991c244..ed4b0c0410 100644
--- a/datafusion/optimizer/tests/optimizer_integration.rs
+++ b/datafusion/optimizer/tests/optimizer_integration.rs
@@ -16,16 +16,25 @@
// under the License.
use std::any::Any;
+use std::cmp::Ordering;
use std::collections::HashMap;
+use std::fmt::Formatter;
use std::sync::Arc;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
use datafusion_common::config::ConfigOptions;
-use datafusion_common::{Result, TableReference, plan_err};
+use datafusion_common::{
+ DFSchemaRef, Result, ScalarValue, TableReference, ToDFSchema, plan_err,
+};
+use datafusion_expr::expr::Cast;
+use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
use datafusion_expr::planner::ExprPlanner;
use datafusion_expr::test::function_stub::sum_udaf;
-use datafusion_expr::{AggregateUDF, LogicalPlan, ScalarUDF, TableSource, WindowUDF};
+use datafusion_expr::{
+ AggregateUDF, Expr, Extension, LogicalPlan, ScalarUDF, SortExpr,
+ TableProviderFilterPushDown, TableSource, UserDefinedLogicalNodeCore, WindowUDF, col,
+};
use datafusion_functions_aggregate::average::avg_udaf;
use datafusion_functions_aggregate::count::count_udaf;
use datafusion_functions_aggregate::planner::AggregateFunctionPlanner;
@@ -690,6 +699,148 @@ fn test_sql(sql: &str) -> Result<LogicalPlan> {
fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+fn optimize_plan(plan: LogicalPlan) -> Result<LogicalPlan> {
+ let config = OptimizerContext::new().with_skip_failing_rules(false);
+ let optimizer = Optimizer::new();
+ optimizer.optimize(plan, &config, observe)
+}
+
+/// Extension node that does NOT implement `necessary_children_exprs`.
+/// Used to test that the optimizer still processes subtrees below such nodes.
+#[derive(Debug, Hash, PartialEq, Eq)]
+struct OpaqueRequirementsExtension {
+ input: Arc<LogicalPlan>,
+ schema: DFSchemaRef,
+}
+
+impl PartialOrd for OpaqueRequirementsExtension {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ self.input
+ .partial_cmp(&other.input)
+ .filter(|cmp| *cmp != Ordering::Equal || self == other)
+ }
+}
+
+impl UserDefinedLogicalNodeCore for OpaqueRequirementsExtension {
+ fn name(&self) -> &str {
+ "OpaqueRequirementsExtension"
+ }
+
+ fn inputs(&self) -> Vec<&LogicalPlan> {
+ vec![&self.input]
+ }
+
+ fn schema(&self) -> &DFSchemaRef {
+ &self.schema
+ }
+
+ fn expressions(&self) -> Vec<Expr> {
+ vec![]
+ }
+
+ fn with_exprs_and_inputs(
+ &self,
+ _exprs: Vec<Expr>,
+ mut inputs: Vec<LogicalPlan>,
+ ) -> Result<Self> {
+ Ok(Self {
+ input: Arc::new(inputs.swap_remove(0)),
+ schema: Arc::clone(&self.schema),
+ })
+ }
+
+ fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result {
+ write!(f, "OpaqueRequirementsExtension")
+ }
+}
+
+struct InexactFilterTableSource {
+ schema: SchemaRef,
+}
+
+impl TableSource for InexactFilterTableSource {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ fn supports_filters_pushdown(
+ &self,
+ filters: &[&Expr],
+ ) -> Result<Vec<TableProviderFilterPushDown>> {
+ Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
+ }
+}
+
+/// Reproduction of https://github.com/apache/datafusion/issues/18816
+/// Extension nodes without `necessary_children_exprs` should not prevent
+/// the optimizer from pruning unnecessary columns in subtrees.
+#[test]
+fn extension_node_does_not_block_projection_pruning() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ Field::new("ts", DataType::Timestamp(TimeUnit::Nanosecond, None), true),
+ ]));
+
+ let table_source: Arc<dyn TableSource> = Arc::new(InexactFilterTableSource {
+ schema: Arc::clone(&schema),
+ });
+
+ let ts_cast = Expr::Cast(Cast::new(
+ Box::new(col("t.ts")),
+ DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
+ ));
+ let ts_millis_1000 = Expr::Literal(
+ ScalarValue::TimestampMillisecond(Some(1000), Some("UTC".into())),
+ None,
+ );
+ let ts_millis_2000 = Expr::Literal(
+ ScalarValue::TimestampMillisecond(Some(2000), Some("UTC".into())),
+ None,
+ );
+
+ let plan = LogicalPlanBuilder::scan("t", table_source, None)?
+ .project(vec![col("t.a"), ts_cast.alias_qualified(Some("t"), "ts")])?
+ .filter(
+ col("t.ts")
+ .gt(ts_millis_1000)
+ .and(col("t.ts").lt(ts_millis_2000)),
+ )?
+ .sort(vec![
+ SortExpr::new(col("t.a"), true, true),
+ SortExpr::new(col("t.ts"), true, true),
+ ])?
+ .build()?;
+
+ let df_schema = schema.to_dfschema_ref()?;
+ let plan = LogicalPlan::Extension(Extension {
+ node: Arc::new(OpaqueRequirementsExtension {
+ input: Arc::new(plan),
+ schema: df_schema,
+ }),
+ });
+
+ let optimized = optimize_plan(plan)?;
+ assert_snapshot!(
+ format!("{optimized}"),
+ @r#"
+ OpaqueRequirementsExtension
+ Sort: t.a ASC NULLS FIRST, t.ts ASC NULLS FIRST
+ Projection: t.a, CAST(t.ts AS Timestamp(ms, "UTC")) AS ts
+ Projection: t.a, t.ts
+ Filter: __common_expr_3 > TimestampMillisecond(1000, Some("UTC")) AND __common_expr_3 < TimestampMillisecond(2000, Some("UTC"))
+ Projection: CAST(t.ts AS Timestamp(ms, "UTC")) AS __common_expr_3, t.a, t.ts
+ TableScan: t projection=[a, ts], partial_filters=[t.ts > TimestampNanosecond(1000000000, None), t.ts < TimestampNanosecond(2000000000, None), CAST(t.ts AS Timestamp(ms, "UTC")) > TimestampMillisecond(1000, Some("UTC")), CAST(t.ts AS Timestamp(ms, "UTC")) < TimestampMillisecond(2000, Some("UTC"))]
+ "#,
+ );
+
+ Ok(())
+}
+
#[derive(Default)]
struct MyContextProvider {
options: ConfigOptions,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org