Update ExecutionPlan to know about sortedness and repartitioning optimizer pass respect the invariants

datafusion/src/physical_plan/analyze.rs

@@ -30,6 +30,7 @@ use crate::{
 use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch};
 use futures::StreamExt;
 
+use super::expressions::PhysicalSortExpr;
 use super::{stream::RecordBatchReceiverStream, Distribution, SendableRecordBatchStream};
 use crate::execution::runtime_env::RuntimeEnv;
 use async_trait::async_trait;
@@ -82,6 +83,10 @@ impl ExecutionPlan for AnalyzeExec {
         Partitioning::UnknownPartitioning(1)
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn with_new_children(
         &self,
         mut children: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/coalesce_batches.rs

@@ -38,6 +38,7 @@ use async_trait::async_trait;
 use futures::stream::{Stream, StreamExt};
 use log::debug;
 
+use super::expressions::PhysicalSortExpr;
 use super::metrics::{BaselineMetrics, MetricsSet};
 use super::{metrics::ExecutionPlanMetricsSet, Statistics};
 
@@ -97,6 +98,10 @@ impl ExecutionPlan for CoalesceBatchesExec {
         self.input.output_partitioning()
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/coalesce_partitions.rs

@@ -31,6 +31,7 @@ use arrow::record_batch::RecordBatch;
 use arrow::{datatypes::SchemaRef, error::Result as ArrowResult};
 
 use super::common::AbortOnDropMany;
+use super::expressions::PhysicalSortExpr;
 use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
 use super::{RecordBatchStream, Statistics};
 use crate::error::{DataFusionError, Result};
@@ -86,6 +87,10 @@ impl ExecutionPlan for CoalescePartitionsExec {
         Partitioning::UnknownPartitioning(1)
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/cross_join.rs

@@ -27,6 +27,7 @@ use arrow::record_batch::RecordBatch;
 
 use futures::{Stream, TryStreamExt};
 
+use super::expressions::PhysicalSortExpr;
 use super::{
     coalesce_partitions::CoalescePartitionsExec, join_utils::check_join_is_valid,
     ColumnStatistics, Statistics,
@@ -137,6 +138,10 @@ impl ExecutionPlan for CrossJoinExec {
         self.right.output_partitioning()
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     async fn execute(
         &self,
         partition: usize,

datafusion/src/physical_plan/empty.rs

@@ -28,6 +28,7 @@ use arrow::array::NullArray;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
 
+use super::expressions::PhysicalSortExpr;
 use super::{common, SendableRecordBatchStream, Statistics};
 
 use crate::execution::runtime_env::RuntimeEnv;
@@ -98,6 +99,10 @@ impl ExecutionPlan for EmptyExec {
         Partitioning::UnknownPartitioning(1)
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/explain.rs

@@ -30,7 +30,7 @@ use crate::{
 };
 use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch};
 
-use super::SendableRecordBatchStream;
+use super::{expressions::PhysicalSortExpr, SendableRecordBatchStream};
 use crate::execution::runtime_env::RuntimeEnv;
 use crate::physical_plan::metrics::{ExecutionPlanMetricsSet, MemTrackingMetrics};
 use async_trait::async_trait;
@@ -89,6 +89,10 @@ impl ExecutionPlan for ExplainExec {
         Partitioning::UnknownPartitioning(1)
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/file_format/avro.rs

@@ -19,6 +19,7 @@
 #[cfg(feature = "avro")]
 use crate::avro_to_arrow;
 use crate::error::{DataFusionError, Result};
+use crate::physical_plan::expressions::PhysicalSortExpr;
 use crate::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
@@ -74,6 +75,10 @@ impl ExecutionPlan for AvroExec {
         Partitioning::UnknownPartitioning(self.base_config.file_groups.len())
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         Vec::new()
     }

datafusion/src/physical_plan/file_format/csv.rs

@@ -18,6 +18,7 @@
 //! Execution plan for reading CSV files
 
 use crate::error::{DataFusionError, Result};
+use crate::physical_plan::expressions::PhysicalSortExpr;
 use crate::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
@@ -88,6 +89,10 @@ impl ExecutionPlan for CsvExec {
         Partitioning::UnknownPartitioning(self.base_config.file_groups.len())
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         // this is a leaf node and has no children
         vec![]

datafusion/src/physical_plan/file_format/json.rs

@@ -20,6 +20,7 @@ use async_trait::async_trait;
 
 use crate::error::{DataFusionError, Result};
 use crate::execution::runtime_env::RuntimeEnv;
+use crate::physical_plan::expressions::PhysicalSortExpr;
 use crate::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
@@ -65,6 +66,10 @@ impl ExecutionPlan for NdJsonExec {
         Partitioning::UnknownPartitioning(self.base_config.file_groups.len())
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         Vec::new()
     }

datafusion/src/physical_plan/file_format/parquet.rs

@@ -24,6 +24,7 @@ use std::{any::Any, convert::TryInto};
 use crate::datasource::file_format::parquet::ChunkObjectReader;
 use crate::datasource::object_store::ObjectStore;
 use crate::datasource::PartitionedFile;
+use crate::physical_plan::expressions::PhysicalSortExpr;
 use crate::{
     error::{DataFusionError, Result},
     logical_plan::{Column, Expr},
@@ -174,6 +175,10 @@ impl ExecutionPlan for ParquetExec {
         Partitioning::UnknownPartitioning(self.base_config.file_groups.len())
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/filter.rs

@@ -23,6 +23,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
+use super::expressions::PhysicalSortExpr;
 use super::{RecordBatchStream, SendableRecordBatchStream, Statistics};
 use crate::error::{DataFusionError, Result};
 use crate::physical_plan::{
@@ -104,6 +105,15 @@ impl ExecutionPlan for FilterExec {
         self.input.output_partitioning()
     }
 
+    fn maintains_input_order(&self) -> bool {
+        // tell optimizer this operator doesn't reorder its input
+        true
+    }
+
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        self.input.output_ordering()
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/hash_aggregate.rs

@@ -52,6 +52,7 @@ use crate::execution::runtime_env::RuntimeEnv;
 use async_trait::async_trait;
 
 use super::common::AbortOnDropSingle;
+use super::expressions::PhysicalSortExpr;
 use super::metrics::{
     self, BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput,
 };
@@ -208,6 +209,10 @@ impl ExecutionPlan for HashAggregateExec {
         self.input.output_partitioning()
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     async fn execute(
         &self,
         partition: usize,
@@ -1148,6 +1153,10 @@ mod tests {
             Partitioning::UnknownPartitioning(1)
         }
 
+        fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+            None
+        }
+
         fn with_new_children(
             &self,
             _: Vec<Arc<dyn ExecutionPlan>>,

datafusion/src/physical_plan/hash_join.rs

@@ -54,6 +54,7 @@ use hashbrown::raw::RawTable;
 
 use super::{
     coalesce_partitions::CoalescePartitionsExec,
+    expressions::PhysicalSortExpr,
     join_utils::{build_join_schema, check_join_is_valid, ColumnIndex, JoinOn, JoinSide},
 };
 use super::{
@@ -278,6 +279,10 @@ impl ExecutionPlan for HashJoinExec {
         self.right.output_partitioning()
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     async fn execute(
         &self,
         partition: usize,

datafusion/src/physical_plan/limit.rs

@@ -35,6 +35,7 @@ use arrow::datatypes::SchemaRef;
 use arrow::error::Result as ArrowResult;
 use arrow::record_batch::RecordBatch;
 
+use super::expressions::PhysicalSortExpr;
 use super::{
     metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet},
     RecordBatchStream, SendableRecordBatchStream, Statistics,
@@ -99,6 +100,22 @@ impl ExecutionPlan for GlobalLimitExec {
         Partitioning::UnknownPartitioning(1)
     }
 
+    fn relies_on_input_order(&self) -> bool {
+        self.input.output_ordering().is_some()
+    }
+
+    fn maintains_input_order(&self) -> bool {
+        true
+    }
+
+    fn benefits_from_input_partitioning(&self) -> bool {
+        false
+    }
+
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        self.input.output_ordering()
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,
@@ -232,6 +249,24 @@ impl ExecutionPlan for LocalLimitExec {
         self.input.output_partitioning()
     }
 
+    fn relies_on_input_order(&self) -> bool {
+        self.input.output_ordering().is_some()
+    }
+
+    fn benefits_from_input_partitioning(&self) -> bool {
+        false
+    }
+
+    // Local limit does not make any attempt to maintain the input
+    // sortedness (if there is more than one partition)
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        if self.output_partitioning().partition_count() == 1 {
+            self.input.output_ordering()
+        } else {
+            None
+        }
+    }
+
     fn with_new_children(
         &self,
         children: Vec<Arc<dyn ExecutionPlan>>,
@@ -300,11 +335,6 @@ impl ExecutionPlan for LocalLimitExec {
             _ => Statistics::default(),
         }
     }
-
-    fn should_repartition_children(&self) -> bool {
-        // No reason to repartition children as this node is just limiting each input partition.
-        false
-    }
 }
 
 /// Truncate a RecordBatch to maximum of n rows

datafusion/src/physical_plan/memory.rs

@@ -22,6 +22,7 @@ use std::any::Any;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
+use super::expressions::PhysicalSortExpr;
 use super::{
     common, project_schema, DisplayFormatType, ExecutionPlan, Partitioning,
     RecordBatchStream, SendableRecordBatchStream, Statistics,
@@ -77,6 +78,10 @@ impl ExecutionPlan for MemoryExec {
         Partitioning::UnknownPartitioning(self.partitions.len())
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn with_new_children(
         &self,
         _: Vec<Arc<dyn ExecutionPlan>>,