SubhamSinghal commented on code in PR #21621:
URL: https://github.com/apache/datafusion/pull/21621#discussion_r3986864957


##########
datafusion/optimizer/src/push_down_limit.rs:
##########
@@ -29,6 +29,9 @@ use datafusion_common::utils::combine_limit;
 use datafusion_expr::logical_plan::{Join, JoinType, Limit, LogicalPlan};
 use datafusion_expr::{FetchType, SkipType, lit};
 
+mod topk_through_join;
+use topk_through_join::push_topk_through_join;
+
 /// Optimization rule that tries to push down `LIMIT`.

Review Comment:
   Addressed in 8dbc1dd53034ec22e39c20a515f54f1addf60cc1



##########
datafusion/optimizer/src/push_down_limit.rs:
##########
@@ -47,146 +50,159 @@ impl OptimizerRule for PushDownLimit {
         true
     }
 
-    #[expect(clippy::only_used_in_recursion)]
     fn rewrite(
         &self,
         plan: LogicalPlan,
         config: &dyn OptimizerConfig,
     ) -> Result<Transformed<LogicalPlan>> {
-        let LogicalPlan::Limit(mut limit) = plan else {
-            return Ok(Transformed::no(plan));
-        };
+        match plan {
+            LogicalPlan::Limit(limit) => rewrite_limit(limit, config),
+            LogicalPlan::Sort(s) if s.fetch.is_some() => {
+                push_topk_through_join(LogicalPlan::Sort(s))
+            }
+            other => Ok(Transformed::no(other)),
+        }
+    }
 
-        // Currently only rewrite if skip and fetch are both literals
-        let SkipType::Literal(skip) = limit.get_skip_type()? else {
+    fn name(&self) -> &str {
+        "push_down_limit"
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(ApplyOrder::TopDown)
+    }
+}
+
+/// Limit-side dispatch (split out from `rewrite` so that the top-level
+/// match in `OptimizerRule::rewrite` reads as a parallel branch alongside
+/// the Sort-with-fetch handler).
+#[expect(clippy::only_used_in_recursion)]

Review Comment:
   Addressed in 
https://github.com/apache/datafusion/commit/8dbc1dd53034ec22e39c20a515f54f1addf60cc1



-- 
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]

Reply via email to