Update ExecutionPlan to know about sortedness and repartitioning optimizer pass respect the invariants (#1776)

* Do not repartition sorted inputs `SortPreservingMerge`

* Add notion of sortedness to `ExecutionPlan`, use to avoid repartitioning when that would result in incorrect behavior

* fix: fix ballitsa

* Update datafusion/src/physical_optimizer/repartition.rs

Co-authored-by: xudong.w <[email protected]>

* Update datafusion/src/physical_optimizer/repartition.rs

Co-authored-by: xudong.w <[email protected]>

* Add more comments

* Remove special `EmptyExec` case

* restore default for benefits_from_input_partitioning

* avoid unecessary check

* default relies_on_input_order to true

* fix test

Co-authored-by: xudong.w <[email protected]>

Author: alamb

ballista/rust/core/src/execution_plans/distributed_query.rs

@@ -33,6 +33,7 @@ use crate::utils::WrappedStream;
 use datafusion::arrow::datatypes::{Schema, SchemaRef};
 use datafusion::error::{DataFusionError, Result};
 use datafusion::logical_plan::LogicalPlan;
+use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, RecordBatchStream,
     SendableRecordBatchStream, Statistics,
@@ -82,6 +83,14 @@ impl ExecutionPlan for DistributedQueryExec {
         Partitioning::UnknownPartitioning(1)
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
+    fn relies_on_input_order(&self) -> bool {
+        false
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![]
     }

ballista/rust/core/src/execution_plans/shuffle_reader.rs

@@ -28,6 +28,7 @@ use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::arrow::error::Result as ArrowResult;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::execution::runtime_env::RuntimeEnv;
+use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::metrics::{
     ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
 };
@@ -85,6 +86,14 @@ impl ExecutionPlan for ShuffleReaderExec {
         Partitioning::UnknownPartitioning(self.partition.len())
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
+    fn relies_on_input_order(&self) -> bool {
+        false
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![]
     }

ballista/rust/core/src/execution_plans/shuffle_writer.rs

@@ -20,6 +20,7 @@
 //! partition is re-partitioned and streamed to disk in Arrow IPC format. Future stages of the query
 //! will use the ShuffleReaderExec to read these results.
 
+use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use parking_lot::Mutex;
 use std::fs::File;
 use std::iter::Iterator;
@@ -334,6 +335,14 @@ impl ExecutionPlan for ShuffleWriterExec {
         self.plan.output_partitioning()
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
+    fn relies_on_input_order(&self) -> bool {
+        false
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![self.plan.clone()]
     }

ballista/rust/core/src/execution_plans/unresolved_shuffle.rs

@@ -23,6 +23,7 @@ use crate::serde::scheduler::PartitionLocation;
 use async_trait::async_trait;
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::execution::runtime_env::RuntimeEnv;
+use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
@@ -85,6 +86,14 @@ impl ExecutionPlan for UnresolvedShuffleExec {
         Partitioning::UnknownPartitioning(self.output_partition_count)
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
+    fn relies_on_input_order(&self) -> bool {
+        false
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![]
     }

ballista/rust/executor/src/collect.rs

@@ -28,6 +28,7 @@ use datafusion::arrow::{
 };
 use datafusion::error::DataFusionError;
 use datafusion::execution::runtime_env::RuntimeEnv;
+use datafusion::physical_plan::expressions::PhysicalSortExpr;
 use datafusion::physical_plan::{
     DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics,
 };
@@ -62,6 +63,10 @@ impl ExecutionPlan for CollectExec {
         Partitioning::UnknownPartitioning(1)
     }
 
+    fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> {
+        None
+    }
+
     fn children(&self) -> Vec<Arc<dyn ExecutionPlan>> {
         vec![self.plan.clone()]
     }