proto wip

Author: adriangb

datafusion/expr/src/logical_plan/tree_node.rs

@@ -882,12 +882,12 @@ impl LogicalPlan {
 /// to perform mutable operations via [`Self::update_plan_from_children`].
 /// After mutating the `LogicalPlanContext.children`, or after creating the `LogicalPlanContext`,
 /// call `update_plan_from_children` to sync.
-/// 
+///
 /// See also:
 /// - [`datafusion_common::tree_node::TreeNode`] trait for tree traversal and mutation utilities.
 /// - [`datafusion_common::tree_node::ConcreteTreeNode`] trait for integrating with the tree node utilities.
 /// - [`datafusion::physical_plan::tree_node::PlanContext`] for a similar context for physical plans.
-/// 
+///
 /// [`datafusion::physical_plan::tree_node::PlanContext`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/tree_node/struct.PlanContext.html
 #[derive(Debug, Clone)]
 pub struct LogicalPlanContext<T: Sized> {

datafusion/optimizer/src/push_down_filter.rs

@@ -41,7 +41,7 @@ use datafusion_expr::{
 use crate::optimizer::ApplyOrder;
 use crate::simplify_expressions::simplify_predicates;
 use crate::utils::{
-    build_schema_remapping, has_all_column_refs, is_restrict_null_predicate
+    build_schema_remapping, has_all_column_refs, is_restrict_null_predicate,
 };
 use crate::{OptimizerConfig, OptimizerRule};
 

datafusion/proto/proto/datafusion.proto

@@ -82,6 +82,10 @@ message SortExprNodeCollection {
   repeated SortExprNode sort_expr_nodes = 1;
 }
 
+message ScanOrderingNode {
+  SortExprNodeCollection preferred_ordering = 1;
+}
+
 message ListingTableScanNode {
   reserved 1; // was string table_name
   TableReference table_name = 14;
@@ -101,6 +105,7 @@ message ListingTableScanNode {
     datafusion_common.ArrowFormat arrow = 16;
   }
   repeated SortExprNodeCollection file_sort_order = 13;
+  ScanOrderingNode ordering = 17;
 }
 
 message ViewTableScanNode {

datafusion/proto/src/generated/pbjson.rs

@@ -57,8 +57,8 @@ use datafusion::{
 };
 use datafusion_common::file_options::file_type::FileType;
 use datafusion_common::{
-    context, internal_datafusion_err, internal_err, not_impl_err, plan_err, Result,
-    TableReference, ToDFSchema,
+    context, internal_datafusion_err, internal_err, not_impl_err, plan_err,
+    DataFusionError, Result, TableReference, ToDFSchema,
 };
 use datafusion_expr::{
     dml,
@@ -495,13 +495,25 @@ impl AsLogicalPlan for LogicalPlanNode {
                     projection = Some(column_indices);
                 }
 
-                LogicalPlanBuilder::scan_with_filters(
+                let ordering = scan.ordering.as_ref().map(|o| {
+                    o.preferred_ordering
+                        .as_ref()
+                        .map(|so| from_proto::parse_sorts(&so.sort_expr_nodes, ctx, extension_codec))
+                        .transpose()
+                }).transpose()?;
+
+                let mut scan = TableScan::try_new(
                     table_name,
                     provider_as_source(Arc::new(provider)),
                     projection,
                     filters,
-                )?
-                .build()
+                    None,
+                )?;
+                if let Some(ordering) = ordering {
+                    scan = scan.with_ordering(ordering);
+                }
+
+                Ok(LogicalPlan::TableScan(scan))
             }
             LogicalPlanType::CustomScan(scan) => {
                 let schema: Schema = convert_required!(scan.schema)?;
@@ -1003,6 +1015,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                 source,
                 filters,
                 projection,
+                ordering,
                 ..
             }) => {
                 let provider = source_as_provider(source)?;
@@ -1118,6 +1131,26 @@ impl AsLogicalPlan for LogicalPlanNode {
                         })
                         .collect::<Result<Vec<_>>>()?;
 
+                    let ordering = ordering
+                        .as_ref()
+                        .map(|ordering| {
+                            Ok::<_, DataFusionError>(protobuf::ScanOrderingNode {
+                                preferred_ordering: ordering
+                                    .preferred_ordering
+                                    .as_ref()
+                                    .map(|order| {
+                                        Ok::<_, DataFusionError>(SortExprNodeCollection {
+                                            sort_expr_nodes: serialize_sorts(
+                                                order,
+                                                extension_codec,
+                                            )?,
+                                        })
+                                    })
+                                    .transpose()?,
+                            })
+                        })
+                        .transpose()?;
+
                     Ok(LogicalPlanNode {
                         logical_plan_type: Some(LogicalPlanType::ListingScan(
                             protobuf::ListingTableScanNode {
@@ -1136,6 +1169,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                                 filters,
                                 target_partitions: options.target_partitions as u32,
                                 file_sort_order: exprs_vec,
+                                ordering,
                             },
                         )),
                     })