Add LimitPushdown optimization rule and CoalesceBatchesExec fetch

datafusion/core/src/dataframe/mod.rs

@@ -2132,7 +2132,7 @@ mod tests {
         assert_physical_plan(
             &df,
             vec![
-            "CoalesceBatchesExec: target_batch_size=8192",
+            "CoalesceBatchesExec: target_batch_size=8192, fetch=None",
             "  FilterExec: id@0 = 1 AND name@1 = a",
             "    AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[]",
             "      MemoryExec: partitions=1, partition_sizes=[1]",
@@ -2181,7 +2181,7 @@ mod tests {
         assert_physical_plan(
             &df,
             vec![
-            "CoalesceBatchesExec: target_batch_size=8192",
+            "CoalesceBatchesExec: target_batch_size=8192, fetch=None",
             "  FilterExec: id@0 = 1",
             "    AggregateExec: mode=Single, gby=[id@0 as id, name@1 as name], aggr=[]",
             "      MemoryExec: partitions=1, partition_sizes=[1]",
@@ -2232,7 +2232,7 @@ mod tests {
         assert_physical_plan(
             &df,
             vec![
-                "CoalesceBatchesExec: target_batch_size=8192",
+                "CoalesceBatchesExec: target_batch_size=8192, fetch=None",
                 "  FilterExec: id@0 = 1",
                 "    AggregateExec: mode=Single, gby=[id@0 as id], aggr=[]",
                 "      MemoryExec: partitions=1, partition_sizes=[1]",

datafusion/core/src/physical_optimizer/enforce_distribution.rs

@@ -3024,7 +3024,7 @@ pub(crate) mod tests {
         // data is already sorted
         let expected = &[
             "SortPreservingMergeExec: [a@0 ASC]",
-            "CoalesceBatchesExec: target_batch_size=4096",
+            "CoalesceBatchesExec: target_batch_size=4096, fetch=None",
             "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC]",
         ];
         assert_optimized!(expected, exec, true);
@@ -3036,7 +3036,7 @@ pub(crate) mod tests {
         let expected = &[
             "SortExec: expr=[a@0 ASC], preserve_partitioning=[false]",
             "CoalescePartitionsExec",
-            "CoalesceBatchesExec: target_batch_size=4096",
+            "CoalesceBatchesExec: target_batch_size=4096, fetch=None",
             "ParquetExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC]",
         ];
         assert_optimized!(expected, exec, false);

datafusion/core/src/physical_optimizer/enforce_sorting.rs

@@ -816,14 +816,14 @@ mod tests {
             "  FilterExec: NOT non_nullable_col@1",
             "    SortExec: expr=[non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[false]",
             "      BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
-            "        CoalesceBatchesExec: target_batch_size=128",
+            "        CoalesceBatchesExec: target_batch_size=128, fetch=None",
             "          SortExec: expr=[non_nullable_col@1 DESC], preserve_partitioning=[false]",
             "            MemoryExec: partitions=1, partition_sizes=[0]"];
 
         let expected_optimized = ["WindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: CurrentRow, end_bound: Following(NULL), is_causal: false }]",
             "  FilterExec: NOT non_nullable_col@1",
             "    BoundedWindowAggExec: wdw=[count: Ok(Field { name: \"count\", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Range, start_bound: Preceding(NULL), end_bound: CurrentRow, is_causal: false }], mode=[Sorted]",
-            "      CoalesceBatchesExec: target_batch_size=128",
+            "      CoalesceBatchesExec: target_batch_size=128, fetch=None",
             "        SortExec: expr=[non_nullable_col@1 DESC], preserve_partitioning=[false]",
             "          MemoryExec: partitions=1, partition_sizes=[0]"];
         assert_optimized!(expected_input, expected_optimized, physical_plan, true);