Use LogicalPlanStats

peter-toth · Nov 29, 2024 · 7711964 · 7711964
1 parent 13eca9f
commit 7711964
Show file tree

Hide file tree

Showing 37 changed files with 2,127 additions and 1,606 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -122,6 +122,7 @@ datafusion-proto = { path = "datafusion/proto", version = "43.0.0" }
 datafusion-proto-common = { path = "datafusion/proto-common", version = "43.0.0" }
 datafusion-sql = { path = "datafusion/sql", version = "43.0.0" }
 doc-comment = "0.3"
+enumset = "1.1.5"
 env_logger = "0.11"
 futures = "0.3"
 half = { version = "2.2.1", default-features = false }

diff --git a/datafusion-cli/Cargo.lock b/datafusion-cli/Cargo.lock
diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml
@@ -110,6 +110,7 @@ datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-optimizer = { workspace = true }
 datafusion-physical-plan = { workspace = true }
 datafusion-sql = { workspace = true }
+enumset = { workspace = true }
 flate2 = { version = "1.0.24", optional = true }
 futures = { workspace = true }
 glob = "0.3.0"

diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs
@@ -69,6 +69,7 @@ use arrow::{
     util::pretty::pretty_format_batches,
 };
 use async_trait::async_trait;
+use enumset::enum_set;
 use futures::{Stream, StreamExt};
 
 use datafusion::execution::session_state::SessionStateBuilder;
@@ -97,8 +98,8 @@ use datafusion::{
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::ScalarValue;
+use datafusion_expr::logical_plan::tree_node::LogicalPlanPattern;
 use datafusion_expr::{FetchType, Projection, SortExpr};
-use datafusion_optimizer::optimizer::ApplyOrder;
 use datafusion_optimizer::AnalyzerRule;
 
 /// Execute the specified sql and return the resulting record batches
@@ -343,10 +344,6 @@ impl OptimizerRule for TopKOptimizerRule {
         "topk"
     }
 
-    fn apply_order(&self) -> Option<ApplyOrder> {
-        Some(ApplyOrder::TopDown)
-    }
-
     fn supports_rewrite(&self) -> bool {
         true
     }
@@ -357,38 +354,47 @@ impl OptimizerRule for TopKOptimizerRule {
         plan: LogicalPlan,
         _config: &dyn OptimizerConfig,
     ) -> Result<Transformed<LogicalPlan>, DataFusionError> {
-        // Note: this code simply looks for the pattern of a Limit followed by a
-        // Sort and replaces it by a TopK node. It does not handle many
-        // edge cases (e.g multiple sort columns, sort ASC / DESC), etc.
-        let LogicalPlan::Limit(ref limit, _) = plan else {
-            return Ok(Transformed::no(plan));
-        };
-        let FetchType::Literal(Some(fetch)) = limit.get_fetch_type()? else {
-            return Ok(Transformed::no(plan));
-        };
+        plan.transform_down_with_subqueries(|plan| {
+            if !plan.stats().contains_all_patterns(enum_set!(
+                LogicalPlanPattern::LogicalPlanLimit
+                    | LogicalPlanPattern::LogicalPlanSort
+            )) {
+                return Ok(Transformed::jump(plan));
+            }
 
-        if let LogicalPlan::Sort(
-            Sort {
-                ref expr,
-                ref input,
-                ..
-            },
-            _,
-        ) = limit.input.as_ref()
-        {
-            if expr.len() == 1 {
-                // we found a sort with a single sort expr, replace with a a TopK
-                return Ok(Transformed::yes(LogicalPlan::extension(Extension {
-                    node: Arc::new(TopKPlanNode {
-                        k: fetch,
-                        input: input.as_ref().clone(),
-                        expr: expr[0].clone(),
-                    }),
-                })));
+            // Note: this code simply looks for the pattern of a Limit followed by a
+            // Sort and replaces it by a TopK node. It does not handle many
+            // edge cases (e.g multiple sort columns, sort ASC / DESC), etc.
+            let LogicalPlan::Limit(ref limit, _) = plan else {
+                return Ok(Transformed::no(plan));
+            };
+            let FetchType::Literal(Some(fetch)) = limit.get_fetch_type()? else {
+                return Ok(Transformed::no(plan));
+            };
+
+            if let LogicalPlan::Sort(
+                Sort {
+                    ref expr,
+                    ref input,
+                    ..
+                },
+                _,
+            ) = limit.input.as_ref()
+            {
+                if expr.len() == 1 {
+                    // we found a sort with a single sort expr, replace with a a TopK
+                    return Ok(Transformed::yes(LogicalPlan::extension(Extension {
+                        node: Arc::new(TopKPlanNode {
+                            k: fetch,
+                            input: input.as_ref().clone(),
+                            expr: expr[0].clone(),
+                        }),
+                    })));
+                }
             }
-        }
 
-        Ok(Transformed::no(plan))
+            Ok(Transformed::no(plan))
+        })
     }
 }
 

diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml
@@ -49,6 +49,7 @@ datafusion-expr-common = { workspace = true }
 datafusion-functions-aggregate-common = { workspace = true }
 datafusion-functions-window-common = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
+enumset = { workspace = true }
 indexmap = { workspace = true }
 paste = "^1.0"
 recursive = { workspace = true }