1+ // Licensed to the Apache Software Foundation (ASF) under one
2+ // or more contributor license agreements. See the NOTICE file
3+ // distributed with this work for additional information
4+ // regarding copyright ownership. The ASF licenses this file
5+ // to you under the Apache License, Version 2.0 (the
6+ // "License"); you may not use this file except in compliance
7+ // with the License. You may obtain a copy of the License at
8+ //
9+ // http://www.apache.org/licenses/LICENSE-2.0
10+ //
11+ // Unless required by applicable law or agreed to in writing,
12+ // software distributed under the License is distributed on an
13+ // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+ // KIND, either express or implied. See the License for the
15+ // specific language governing permissions and limitations
16+ // under the License.
17+
18+ //! [`PushDownSort`] pushes sort expressions into table scans to enable
19+ //! sort pushdown optimizations by table providers
20+
21+ use std:: sync:: Arc ;
22+
23+ use crate :: optimizer:: ApplyOrder ;
24+ use crate :: { OptimizerConfig , OptimizerRule } ;
25+
26+ use datafusion_common:: tree_node:: Transformed ;
27+ use datafusion_common:: Result ;
28+ use datafusion_expr:: logical_plan:: { LogicalPlan , TableScan } ;
29+ use datafusion_expr:: { Expr , SortExpr } ;
30+
31+ /// Optimization rule that pushes sort expressions down to table scans
32+ /// when the sort can potentially be optimized by the table provider.
33+ ///
34+ /// This rule looks for `Sort -> TableScan` patterns and moves the sort
35+ /// expressions into the `TableScan.preferred_ordering` field, allowing
36+ /// table providers to potentially optimize the scan based on sort requirements.
37+ #[ derive( Default , Debug ) ]
38+ pub struct PushDownSort { }
39+
40+ impl PushDownSort {
41+ #[ allow( missing_docs) ]
42+ pub fn new ( ) -> Self {
43+ Self { }
44+ }
45+
46+ /// Checks if a sort expression can be pushed down to a table scan.
47+ ///
48+ /// Currently, we only support pushing down simple column references
49+ /// because table providers typically can't optimize complex expressions
50+ /// in sort pushdown.
51+ fn can_pushdown_sort_expr ( expr : & SortExpr ) -> bool {
52+ // Only push down simple column references
53+ matches ! ( expr. expr, Expr :: Column ( _) )
54+ }
55+
56+ /// Checks if all sort expressions in a list can be pushed down.
57+ fn can_pushdown_sort_exprs ( sort_exprs : & [ SortExpr ] ) -> bool {
58+ sort_exprs. iter ( ) . all ( Self :: can_pushdown_sort_expr)
59+ }
60+ }
61+
62+ impl OptimizerRule for PushDownSort {
63+ fn supports_rewrite ( & self ) -> bool {
64+ true
65+ }
66+
67+ fn apply_order ( & self ) -> Option < ApplyOrder > {
68+ Some ( ApplyOrder :: TopDown )
69+ }
70+
71+ fn rewrite (
72+ & self ,
73+ plan : LogicalPlan ,
74+ _config : & dyn OptimizerConfig ,
75+ ) -> Result < Transformed < LogicalPlan > > {
76+ // Look for Sort -> TableScan pattern
77+ let LogicalPlan :: Sort ( sort) = & plan else {
78+ return Ok ( Transformed :: no ( plan) ) ;
79+ } ;
80+
81+ let LogicalPlan :: TableScan ( table_scan) = sort. input . as_ref ( ) else {
82+ return Ok ( Transformed :: no ( plan) ) ;
83+ } ;
84+
85+ // Check if we can push down the sort expressions
86+ if !Self :: can_pushdown_sort_exprs ( & sort. expr ) {
87+ return Ok ( Transformed :: no ( plan) ) ;
88+ }
89+
90+ // If the table scan already has preferred ordering, don't overwrite it
91+ // This preserves any existing sort preferences from other optimizations
92+ if table_scan. preferred_ordering . is_some ( ) {
93+ return Ok ( Transformed :: no ( plan) ) ;
94+ }
95+
96+ // Create new TableScan with preferred ordering
97+ let new_table_scan = TableScan {
98+ table_name : table_scan. table_name . clone ( ) ,
99+ source : Arc :: clone ( & table_scan. source ) ,
100+ projection : table_scan. projection . clone ( ) ,
101+ projected_schema : Arc :: clone ( & table_scan. projected_schema ) ,
102+ filters : table_scan. filters . clone ( ) ,
103+ fetch : table_scan. fetch ,
104+ preferred_ordering : Some ( sort. expr . clone ( ) ) ,
105+ } ;
106+
107+ // The sort can be completely eliminated since we've pushed it down
108+ // The table provider may or may not be able to satisfy the ordering,
109+ // but that's up to the table provider to decide
110+ let new_plan = LogicalPlan :: TableScan ( new_table_scan) ;
111+
112+ Ok ( Transformed :: yes ( new_plan) )
113+ }
114+
115+ fn name ( & self ) -> & str {
116+ "push_down_sort"
117+ }
118+ }
119+
120+ #[ cfg( test) ]
121+ mod tests {
122+ use super :: * ;
123+ use datafusion_expr:: { col, SortExpr } ;
124+
125+ #[ test]
126+ fn test_can_pushdown_sort_expr ( ) {
127+ // Simple column reference should be pushable
128+ let sort_expr = SortExpr :: new ( col ( "a" ) , true , false ) ;
129+ assert ! ( PushDownSort :: can_pushdown_sort_expr( & sort_expr) ) ;
130+
131+ // Complex expression should not be pushable
132+ let sort_expr = SortExpr :: new (
133+ col ( "a" ) + col ( "b" ) ,
134+ true ,
135+ false
136+ ) ;
137+ assert ! ( !PushDownSort :: can_pushdown_sort_expr( & sort_expr) ) ;
138+ }
139+
140+ #[ test]
141+ fn test_name ( ) {
142+ let rule = PushDownSort :: new ( ) ;
143+ assert_eq ! ( rule. name( ) , "push_down_sort" ) ;
144+ }
145+ }
0 commit comments