fix tests

Author: buraksenn

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -36,8 +36,7 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 
 use crate::assert_optimized;
 use crate::physical_optimizer::test_util::{
-    csv_exec_ordered, csv_exec_sorted, parquet_exec, parquet_exec_sorted,
-    stream_exec_ordered,
+    csv_exec_ordered, csv_exec_sorted, parquet_exec, parquet_exec_sorted, stream_exec_ordered,
 };
 use datafusion::physical_optimizer::enforce_distribution::EnforceDistribution;
 use datafusion_physical_expr::Partitioning;

datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs

@@ -39,14 +39,14 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
 
 use datafusion_common::tree_node::{TransformedResult, TreeNode};
-use datafusion_common::Result;
-use datafusion_expr::{JoinType, Operator};
-use datafusion_physical_expr::expressions::{self, col, Column};
-use datafusion_physical_expr::PhysicalSortExpr;
-use datafusion_physical_optimizer::test_utils::check_integrity;
-use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext};
+    use datafusion_common::Result;
+    use datafusion_expr::{JoinType, Operator};
+    use datafusion_physical_expr::expressions::{self, col, Column};
+    use datafusion_physical_expr::PhysicalSortExpr;
+    use datafusion_physical_optimizer::test_utils::check_integrity;
+    use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext};
 
-use crate::physical_optimizer::test_util::stream_exec_ordered;
+use crate::physical_optimizer::test_util::stream_exec_ordered_with_projection;
 
 use object_store::memory::InMemory;
 use object_store::ObjectStore;
@@ -190,7 +190,7 @@ async fn test_replace_multiple_input_repartition_1(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -259,7 +259,7 @@ async fn test_with_inter_children_change_only(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr_default("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -362,7 +362,7 @@ async fn test_replace_multiple_input_repartition_2(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -438,7 +438,7 @@ async fn test_replace_multiple_input_repartition_with_extra_steps(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -520,7 +520,7 @@ async fn test_replace_multiple_input_repartition_with_extra_steps_2(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -608,7 +608,7 @@ async fn test_not_replacing_when_no_need_to_preserve_sorting(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -681,7 +681,7 @@ async fn test_with_multiple_replacable_repartitions(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -769,7 +769,7 @@ async fn test_not_replace_with_different_orderings(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -841,7 +841,7 @@ async fn test_with_lost_ordering(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -911,7 +911,7 @@ async fn test_with_lost_and_kept_ordering(
     let schema = create_test_schema()?;
     let sort_exprs = vec![sort_expr("a", &schema)];
     let source = if source_unbounded {
-        stream_exec_ordered(&schema, sort_exprs)
+        stream_exec_ordered_with_projection(&schema, sort_exprs)
     } else {
         memory_exec_sorted(&schema, sort_exprs)
     };
@@ -1017,7 +1017,7 @@ async fn test_with_multiple_child_trees(
 
     let left_sort_exprs = vec![sort_expr("a", &schema)];
     let left_source = if source_unbounded {
-        stream_exec_ordered(&schema, left_sort_exprs)
+        stream_exec_ordered_with_projection(&schema, left_sort_exprs)
     } else {
         memory_exec_sorted(&schema, left_sort_exprs)
     };
@@ -1028,7 +1028,7 @@ async fn test_with_multiple_child_trees(
 
     let right_sort_exprs = vec![sort_expr("a", &schema)];
     let right_source = if source_unbounded {
-        stream_exec_ordered(&schema, right_sort_exprs)
+        stream_exec_ordered_with_projection(&schema, right_sort_exprs)
     } else {
         memory_exec_sorted(&schema, right_sort_exprs)
     };

datafusion/core/tests/physical_optimizer/test_util.rs

@@ -268,6 +268,29 @@ pub fn stream_exec_ordered(
     )
 }
 
+// Creates a stream exec source for the test purposes
+pub fn stream_exec_ordered_with_projection(
+    schema: &SchemaRef,
+    sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+) -> Arc<dyn ExecutionPlan> {
+    let sort_exprs = sort_exprs.into_iter().collect();
+    let projection: Vec<usize> = vec![0, 2, 3];
+
+    Arc::new(
+        StreamingTableExec::try_new(
+            schema.clone(),
+            vec![Arc::new(TestStreamPartition {
+                schema: schema.clone(),
+            }) as _],
+            Some(&projection),
+            vec![sort_exprs],
+            true,
+            None,
+        )
+        .unwrap(),
+    )
+}
+
 /// Create a csv exec for tests
 pub fn csv_exec_ordered(
     schema: &SchemaRef,