address some feedback

Author: adriangb

datafusion/expr/src/logical_plan/mod.rs

@@ -41,8 +41,9 @@ pub use plan::{
     projection_schema, Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct,
     DistinctOn, EmptyRelation, Explain, ExplainFormat, ExplainOption, Extension,
     FetchType, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning,
-    PlanType, Projection, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan,
-    Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window,
+    PlanType, Projection, RecursiveQuery, Repartition, ScanOrdering, SkipType, Sort,
+    StringifiedPlan, Subquery, SubqueryAlias, TableScan, ToStringifiedPlan, Union,
+    Unnest, Values, Window,
 };
 pub use statement::{
     Deallocate, Execute, Prepare, SetVariable, Statement, TransactionAccessMode,

datafusion/expr/src/logical_plan/plan.rs

@@ -2537,6 +2537,43 @@ impl PartialOrd for Window {
     }
 }
 
+#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Default)]
+pub struct ScanOrdering {
+    /// Optional preferred ordering for the scan that matches the output order of upstream query nodes.
+    /// It is optional / best effort for the scan to produce this ordering.
+    /// If the scan produces this exact ordering and sets it's properties to reflect this upstream sorts may be optimized away.
+    /// Otherwise the sorts may remain in place but partial ordering may be exploited e.g. to do early stopping or reduce complexity of the sort.
+    /// Thus it is recommended for the scan to also do a best effort to produce partially sorted data if possible.
+    pub preferred_ordering: Option<Vec<SortExpr>>,
+}
+
+impl ScanOrdering {
+    /// Create a new ScanOrdering
+    pub fn with_preferred_ordering(mut self, preferred_ordering: Vec<SortExpr>) -> Self {
+        self.preferred_ordering = Some(preferred_ordering);
+        self
+    }
+}
+
+impl Debug for ScanOrdering {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        let ordering_display = self
+            .preferred_ordering
+            .as_ref()
+            .map(|ordering| {
+                ordering
+                    .iter()
+                    .map(|e| e.to_string())
+                    .collect::<Vec<String>>()
+                    .join(", ")
+            })
+            .unwrap_or_else(|| "None".to_string());
+        f.debug_struct("ScanOrdering")
+            .field("preferred_ordering", &ordering_display)
+            .finish_non_exhaustive()
+    }
+}
+
 /// Produces rows from a table provider by reference or from the context
 #[derive(Clone)]
 pub struct TableScan {
@@ -2552,8 +2589,8 @@ pub struct TableScan {
     pub filters: Vec<Expr>,
     /// Optional number of rows to read
     pub fetch: Option<usize>,
-    /// Optional preferred ordering for the scan
-    pub preferred_ordering: Option<Vec<SortExpr>>,
+    /// Ordering for the scan
+    pub ordering: Option<ScanOrdering>,
 }
 
 impl Debug for TableScan {
@@ -2565,7 +2602,7 @@ impl Debug for TableScan {
             .field("projected_schema", &self.projected_schema)
             .field("filters", &self.filters)
             .field("fetch", &self.fetch)
-            .field("preferred_ordering", &self.preferred_ordering)
+            .field("ordering", &self.ordering)
             .finish_non_exhaustive()
     }
 }
@@ -2577,7 +2614,7 @@ impl PartialEq for TableScan {
             && self.projected_schema == other.projected_schema
             && self.filters == other.filters
             && self.fetch == other.fetch
-            && self.preferred_ordering == other.preferred_ordering
+            && self.ordering == other.ordering
     }
 }
 
@@ -2598,21 +2635,21 @@ impl PartialOrd for TableScan {
             /// Optional number of rows to read
             pub fetch: &'a Option<usize>,
             /// Optional preferred ordering for the scan
-            pub preferred_ordering: &'a Option<Vec<SortExpr>>,
+            pub ordering: &'a Option<ScanOrdering>,
         }
         let comparable_self = ComparableTableScan {
             table_name: &self.table_name,
             projection: &self.projection,
             filters: &self.filters,
             fetch: &self.fetch,
-            preferred_ordering: &self.preferred_ordering,
+            ordering: &self.ordering,
         };
         let comparable_other = ComparableTableScan {
             table_name: &other.table_name,
             projection: &other.projection,
             filters: &other.filters,
             fetch: &other.fetch,
-            preferred_ordering: &other.preferred_ordering,
+            ordering: &other.ordering,
         };
         comparable_self.partial_cmp(&comparable_other)
     }
@@ -2625,7 +2662,7 @@ impl Hash for TableScan {
         self.projected_schema.hash(state);
         self.filters.hash(state);
         self.fetch.hash(state);
-        self.preferred_ordering.hash(state);
+        self.ordering.hash(state);
     }
 }
 
@@ -2679,7 +2716,7 @@ impl TableScan {
             projected_schema,
             filters,
             fetch,
-            preferred_ordering: None,
+            ordering: None,
         })
     }
 
@@ -2734,11 +2771,8 @@ impl TableScan {
     /// This is purely an optimization hint. The table provider may choose to ignore
     /// the preferred ordering if it cannot be efficiently satisfied, and the query
     /// execution engine should not rely on the data being returned in this order.
-    pub fn with_preferred_ordering(
-        mut self,
-        preferred_ordering: Option<Vec<SortExpr>>,
-    ) -> Self {
-        self.preferred_ordering = preferred_ordering;
+    pub fn with_ordering(mut self, ordering: ScanOrdering) -> Self {
+        self.ordering = Some(ordering);
         self
     }
 }
@@ -4965,7 +4999,7 @@ mod tests {
             projected_schema: Arc::clone(&schema),
             filters: vec![],
             fetch: None,
-            preferred_ordering: None,
+            ordering: None,
         }));
         let col = schema.field_names()[0].clone();
 
@@ -4996,7 +5030,7 @@ mod tests {
             projected_schema: Arc::clone(&unique_schema),
             filters: vec![],
             fetch: None,
-            preferred_ordering: None,
+            ordering: None,
         }));
         let col = schema.field_names()[0].clone();
 

datafusion/expr/src/logical_plan/tree_node.rs

@@ -599,7 +599,7 @@ impl LogicalPlan {
                 projected_schema,
                 filters,
                 fetch,
-                preferred_ordering,
+                ordering,
             }) => filters.map_elements(f)?.update_data(|filters| {
                 LogicalPlan::TableScan(TableScan {
                     table_name,
@@ -608,7 +608,7 @@ impl LogicalPlan {
                     projected_schema,
                     filters,
                     fetch,
-                    preferred_ordering,
+                    ordering,
                 })
             }),
             LogicalPlan::Distinct(Distinct::On(DistinctOn {

datafusion/optimizer/src/optimize_projections/mod.rs

@@ -242,7 +242,7 @@ fn optimize_projections(
                 filters,
                 fetch,
                 projected_schema: _,
-                preferred_ordering,
+                ordering,
             } = table_scan;
 
             // Get indices referred to in the original (schema with all fields)
@@ -258,7 +258,10 @@ fn optimize_projections(
                 filters,
                 fetch,
             )
-            .map(|s| s.with_preferred_ordering(preferred_ordering))
+            .map(|s| match ordering {
+                Some(ordering) => s.with_ordering(ordering),
+                None => s,
+            })
             .map(LogicalPlan::TableScan)
             .map(Transformed::yes);
         }

datafusion/optimizer/src/push_down_filter.rs

@@ -3055,7 +3055,7 @@ mod tests {
             projection,
             source: Arc::new(test_provider),
             fetch: None,
-            preferred_ordering: None,
+            ordering: None,
         });
 
         Ok(LogicalPlanBuilder::from(table_scan))

datafusion/optimizer/src/push_down_sort.rs

@@ -25,8 +25,8 @@ use crate::{OptimizerConfig, OptimizerRule};
 
 use datafusion_common::tree_node::Transformed;
 use datafusion_common::Result;
-use datafusion_expr::logical_plan::{LogicalPlan, TableScan};
-use datafusion_expr::{Expr, SortExpr};
+use datafusion_expr::logical_plan::LogicalPlan;
+use datafusion_expr::{Expr, ScanOrdering, SortExpr};
 
 /// Optimization rule that pushes sort expressions down to table scans
 /// when the sort can potentially be optimized by the table provider.
@@ -123,22 +123,10 @@ impl OptimizerRule for PushDownSort {
             return Ok(Transformed::no(plan));
         }
 
-        // If the table scan already has preferred ordering, don't overwrite it
-        // This preserves any existing sort preferences from other optimizations
-        if table_scan.preferred_ordering.is_some() {
-            return Ok(Transformed::no(plan));
-        }
-
         // Create new TableScan with preferred ordering
-        let new_table_scan = TableScan {
-            table_name: table_scan.table_name.clone(),
-            source: Arc::clone(&table_scan.source),
-            projection: table_scan.projection.clone(),
-            projected_schema: Arc::clone(&table_scan.projected_schema),
-            filters: table_scan.filters.clone(),
-            fetch: table_scan.fetch,
-            preferred_ordering: Some(sort.expr.clone()),
-        };
+        let new_table_scan = table_scan.clone().with_ordering(
+            ScanOrdering::default().with_preferred_ordering(sort.expr.clone()),
+        );
 
         // Preserve the Sort node as a fallback while passing the ordering
         // preference to the TableScan as an optimization hint
@@ -275,7 +263,7 @@ mod tests {
             LogicalPlan::Sort(sort) => {
                 // Check that TableScan has preferred_ordering
                 if let LogicalPlan::TableScan(ts) = sort.input.as_ref() {
-                    assert!(ts.preferred_ordering.is_some());
+                    assert!(ts.ordering.is_some());
                 } else {
                     panic!("Expected TableScan input");
                 }

datafusion/proto/src/logical_plan/mod.rs

@@ -271,7 +271,7 @@ fn from_table_source(
         projected_schema,
         filters: vec![],
         fetch: None,
-        preferred_ordering: None,
+        ordering: None,
     });
 
     LogicalPlanNode::try_from_logical_plan(&r, extension_codec)