initial stab

Author: adriangb

datafusion/catalog/src/default_table_source.rs

@@ -73,6 +73,7 @@ impl TableSource for DefaultTableSource {
         &self,
         filter: &[&Expr],
     ) -> datafusion_common::Result<Vec<TableProviderFilterPushDown>> {
+        #[allow(deprecated)]
         self.table_provider.supports_filters_pushdown(filter)
     }
 

datafusion/catalog/src/table.rs

@@ -25,13 +25,14 @@ use arrow::datatypes::SchemaRef;
 use async_trait::async_trait;
 use datafusion_common::Result;
 use datafusion_common::{not_impl_err, Constraints, Statistics};
-use datafusion_expr::Expr;
+use datafusion_expr::{Expr, SortExpr};
 
 use datafusion_expr::dml::InsertOp;
 use datafusion_expr::{
     CreateExternalTable, LogicalPlan, TableProviderFilterPushDown, TableType,
 };
 use datafusion_physical_plan::ExecutionPlan;
+use itertools::Itertools;
 
 /// A table which can be queried and modified.
 ///
@@ -163,6 +164,7 @@ pub trait TableProvider: Debug + Sync + Send {
     /// because inexact filters do not guarantee that every filtered row is
     /// removed, so applying the limit could lead to too few rows being available
     /// to return as a final result.
+    #[deprecated(since = "50.0.0", note = "Use `scan_with_options` instead")]
     async fn scan(
         &self,
         state: &dyn Session,
@@ -171,6 +173,25 @@ pub trait TableProvider: Debug + Sync + Send {
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>>;
 
+    async fn scan_with_options(
+        &self,
+        state: &dyn Session,
+        options: ScanArgs,
+    ) -> Result<ScanResult> {
+        let ScanArgs {
+            preferred_ordering: _,
+            filters,
+            projection,
+            limit,
+        } = options;
+        let filters = filters.unwrap_or_default();
+        #[allow(deprecated)]
+        let plan = self
+            .scan(state, projection.as_ref(), &filters, limit)
+            .await?;
+        Ok(ScanResult::new(plan, vec![]))
+    }
+
     /// Specify if DataFusion should provide filter expressions to the
     /// TableProvider to apply *during* the scan.
     ///
@@ -251,6 +272,7 @@ pub trait TableProvider: Debug + Sync + Send {
     ///     }
     /// }
     /// ```
+    #[deprecated(since = "50.0.0", note = "Use `scan_with_options` instead")]
     fn supports_filters_pushdown(
         &self,
         filters: &[&Expr],
@@ -299,6 +321,74 @@ pub trait TableProvider: Debug + Sync + Send {
     }
 }
 
+#[derive(Debug, Clone, Default)]
+pub struct ScanArgs {
+    preferred_ordering: Option<Vec<SortExpr>>,
+    filters: Option<Vec<Expr>>,
+    projection: Option<Vec<usize>>,
+    limit: Option<usize>,
+}
+
+impl ScanArgs {
+    pub fn with_projection(mut self, projection: Option<Vec<usize>>) -> Self {
+        self.projection = projection;
+        self
+    }
+
+    pub fn projection(&self) -> Option<Vec<usize>> {
+        self.projection.clone()
+    }
+
+    pub fn with_filters(mut self, filters: Option<Vec<Expr>>) -> Self {
+        self.filters = filters;
+        self
+    }
+
+    pub fn filters(&self) -> Option<&[Expr]> {
+        self.filters.as_deref()
+    }
+
+    pub fn with_limit(mut self, limit: Option<usize>) -> Self {
+        self.limit = limit;
+        self
+    }
+
+    pub fn limit(&self) -> Option<usize> {
+        self.limit
+    }
+
+    pub fn with_preferred_ordering(mut self, ordering: Option<Vec<SortExpr>>) -> Self {
+        self.preferred_ordering = ordering;
+        self
+    }
+
+    pub fn preferred_ordering(&self) -> Option<&[SortExpr]> {
+        self.preferred_ordering.as_deref()
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct ScanResult {
+    /// The ExecutionPlan to run.
+    plan: Arc<dyn ExecutionPlan>,
+    // Remaining filters that were not completely evaluated during `scan_with_options()`.
+    filters: Vec<Expr>,
+}
+
+impl ScanResult {
+    pub fn new(plan: Arc<dyn ExecutionPlan>, filters: Vec<Expr>) -> Self {
+        Self { plan, filters }
+    }
+
+    pub fn plan(&self) -> Arc<dyn ExecutionPlan> {
+        Arc::clone(&self.plan)
+    }
+
+    pub fn filters(&self) -> &[Expr] {
+        &self.filters
+    }
+}
+
 /// A factory which creates [`TableProvider`]s at runtime given a URL.
 ///
 /// For example, this can be used to create a table "on the fly"

datafusion/core/src/datasource/listing/table.rs

@@ -29,7 +29,7 @@ use crate::{
 use arrow::datatypes::{DataType, Field, SchemaBuilder, SchemaRef};
 use arrow_schema::Schema;
 use async_trait::async_trait;
-use datafusion_catalog::{Session, TableProvider};
+use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider};
 use datafusion_common::{
     config_datafusion_err, config_err, internal_err, plan_err, project_schema,
     stats::Precision, Constraints, DataFusionError, Result, SchemaExt,
@@ -1166,6 +1166,22 @@ impl TableProvider for ListingTable {
         filters: &[Expr],
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        let options = ScanArgs::default()
+            .with_projection(projection.cloned())
+            .with_filters(Some(filters.to_vec()))
+            .with_limit(limit);
+        Ok(self.scan_with_options(state, options).await?.plan())
+    }
+
+    async fn scan_with_options(
+        &self,
+        state: &dyn Session,
+        options: ScanArgs,
+    ) -> Result<ScanResult> {
+        let projection = options.projection();
+        let filters = options.filters().map(|f| f.to_vec()).unwrap_or_default();
+        let limit = options.limit();
+
         // extract types of partition columns
         let table_partition_cols = self
             .options
@@ -1195,21 +1211,36 @@ impl TableProvider for ListingTable {
 
         // if no files need to be read, return an `EmptyExec`
         if partitioned_file_lists.is_empty() {
-            let projected_schema = project_schema(&self.schema(), projection)?;
-            return Ok(Arc::new(EmptyExec::new(projected_schema)));
+            let projected_schema = project_schema(&self.schema(), projection.as_ref())?;
+            return Ok(ScanResult::new(
+                Arc::new(EmptyExec::new(projected_schema)),
+                filters.clone(),
+            ));
         }
 
-        let output_ordering = self.try_create_output_ordering()?;
+        let known_file_ordering = self.try_create_output_ordering()?;
+        let desired_file_ordering = match options.preferred_ordering() {
+            Some(ordering) if !ordering.is_empty() => {
+                // Prefer the ordering requested by the query to any inherint file ordering
+                create_ordering(&self.table_schema, &[ordering.to_vec()])?
+                    .first()
+                    .cloned()
+            }
+            Some(_) | None => {
+                // If the query did not request a specific ordering, fall back to any inherent file ordering
+                known_file_ordering.first().cloned()
+            }
+        };
         match state
             .config_options()
             .execution
             .split_file_groups_by_statistics
             .then(|| {
-                output_ordering.first().map(|output_ordering| {
+                desired_file_ordering.map(|ordering| {
                     FileScanConfig::split_groups_by_statistics_with_target_partitions(
                         &self.table_schema,
                         &partitioned_file_lists,
-                        output_ordering,
+                        &ordering,
                         self.options.target_partitions,
                     )
                 })
@@ -1230,13 +1261,17 @@ impl TableProvider for ListingTable {
         let Some(object_store_url) =
             self.table_paths.first().map(ListingTableUrl::object_store)
         else {
-            return Ok(Arc::new(EmptyExec::new(Arc::new(Schema::empty()))));
+            return Ok(ScanResult::new(
+                Arc::new(EmptyExec::new(Arc::new(Schema::empty()))),
+                filters.clone(),
+            ));
         };
 
         let file_source = self.create_file_source_with_schema_adapter()?;
 
         // create the execution plan
-        self.options
+        let plan = self
+            .options
             .format
             .create_physical_plan(
                 state,
@@ -1248,14 +1283,16 @@ impl TableProvider for ListingTable {
                 .with_file_groups(partitioned_file_lists)
                 .with_constraints(self.constraints.clone())
                 .with_statistics(statistics)
-                .with_projection(projection.cloned())
+                .with_projection(projection)
                 .with_limit(limit)
-                .with_output_ordering(output_ordering)
+                .with_output_ordering(known_file_ordering)
                 .with_table_partition_cols(table_partition_cols)
                 .with_expr_adapter(self.expr_adapter_factory.clone())
                 .build(),
             )
-            .await
+            .await?;
+
+        Ok(ScanResult::new(plan, filters.clone()))
     }
 
     fn supports_filters_pushdown(

datafusion/core/src/physical_planner.rs

@@ -60,6 +60,7 @@ use crate::schema_equivalence::schema_satisfied_by;
 use arrow::array::{builder::StringBuilder, RecordBatch};
 use arrow::compute::SortOptions;
 use arrow::datatypes::{Schema, SchemaRef};
+use datafusion_catalog::ScanArgs;
 use datafusion_common::display::ToStringifiedPlan;
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
@@ -85,7 +86,7 @@ use datafusion_expr::{
 use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr};
 use datafusion_physical_expr::expressions::{Column, Literal};
 use datafusion_physical_expr::{
-    create_physical_sort_exprs, LexOrdering, PhysicalSortExpr,
+    conjunction, create_physical_sort_exprs, LexOrdering, PhysicalSortExpr,
 };
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::empty::EmptyExec;
@@ -452,16 +453,22 @@ impl DefaultPhysicalPlanner {
                 projection,
                 filters,
                 fetch,
+                preferred_ordering,
                 ..
             }) => {
                 let source = source_as_provider(source)?;
                 // Remove all qualifiers from the scan as the provider
                 // doesn't know (nor should care) how the relation was
                 // referred to in the query
                 let filters = unnormalize_cols(filters.iter().cloned());
-                source
-                    .scan(session_state, projection.as_ref(), &filters, *fetch)
-                    .await?
+                let opts = ScanArgs::default()
+                    .with_projection(projection.clone())
+                    .with_filters(Some(filters).clone())
+                    .with_preferred_ordering(preferred_ordering.clone())
+                    .with_limit(*fetch);
+                let res = source.scan_with_options(session_state, opts).await?;
+                // TODO: move FilterExec wrapping logic from filter pushdown rule to here?
+                res.plan()
             }
             LogicalPlan::Values(Values { values, schema }) => {
                 let exec_schema = schema.as_ref().to_owned().into();

datafusion/expr/src/logical_plan/plan.rs

@@ -2522,6 +2522,8 @@ pub struct TableScan {
     pub filters: Vec<Expr>,
     /// Optional number of rows to read
     pub fetch: Option<usize>,
+    /// Preferred read ordering of the table
+    pub preferred_ordering: Option<Vec<SortExpr>>,
 }
 
 impl Debug for TableScan {
@@ -2640,8 +2642,14 @@ impl TableScan {
             projected_schema,
             filters,
             fetch,
+            preferred_ordering: None,
         })
     }
+
+    pub fn with_preferred_ordering(mut self, ordering: Option<Vec<SortExpr>>) -> Self {
+        self.preferred_ordering = ordering;
+        self
+    }
 }
 
 // Repartition the plan based on a partitioning scheme.
@@ -4823,6 +4831,7 @@ mod tests {
             projected_schema: Arc::clone(&schema),
             filters: vec![],
             fetch: None,
+            preferred_ordering: None,
         }));
         let col = schema.field_names()[0].clone();
 
@@ -4853,6 +4862,7 @@ mod tests {
             projected_schema: Arc::clone(&unique_schema),
             filters: vec![],
             fetch: None,
+            preferred_ordering: None,
         }));
         let col = schema.field_names()[0].clone();
 

datafusion/expr/src/logical_plan/tree_node.rs

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

datafusion/optimizer/src/lib.rs

@@ -58,6 +58,7 @@ pub mod optimizer;
 pub mod propagate_empty_relation;
 pub mod push_down_filter;
 pub mod push_down_limit;
+pub mod push_down_sort;
 pub mod replace_distinct_aggregate;
 pub mod scalar_subquery_to_join;
 pub mod simplify_expressions;

datafusion/optimizer/src/optimize_projections/mod.rs

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

datafusion/optimizer/src/optimizer.rs

@@ -51,6 +51,7 @@ use crate::plan_signature::LogicalPlanSignature;
 use crate::propagate_empty_relation::PropagateEmptyRelation;
 use crate::push_down_filter::PushDownFilter;
 use crate::push_down_limit::PushDownLimit;
+use crate::push_down_sort::PushDownSort;
 use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate;
 use crate::scalar_subquery_to_join::ScalarSubqueryToJoin;
 use crate::simplify_expressions::SimplifyExpressions;
@@ -243,6 +244,8 @@ impl Optimizer {
             // Filters can't be pushed down past Limits, we should do PushDownFilter after PushDownLimit
             Arc::new(PushDownLimit::new()),
             Arc::new(PushDownFilter::new()),
+            // Push down sort requirements to TableScan preferred_ordering
+            Arc::new(PushDownSort::new()),
             Arc::new(SingleDistinctToGroupBy::new()),
             // The previous optimizations added expressions and projections,
             // that might benefit from the following rules