perf: rm coalesce batch when target_batch_size > fetch limit (#5658)

* fix: rm coalesce > limit

* fix: only rm one&test: sqlness

src/query/src/optimizer/windowed_sort.rs

@@ -80,6 +80,9 @@ impl WindowedSortPhysicalRule {
                     let preserve_partitioning = sort_exec.preserve_partitioning();
 
                     let sort_input = remove_repartition(sort_exec.input().clone())?.data;
+                    let sort_input =
+                        remove_coalesce_batches_exec(sort_input, sort_exec.fetch())?.data;
+
                     // Gets scanner info from the input without repartition before filter.
                     let Some(scanner_info) = fetch_partition_range(sort_input.clone())? else {
                         return Ok(Transformed::no(plan));
@@ -237,6 +240,33 @@ fn remove_repartition(
     })
 }
 
+/// Remove `CoalesceBatchesExec` if the limit is less than the batch size.
+///
+/// so that if limit is too small we can avoid need to scan for more rows than necessary
+fn remove_coalesce_batches_exec(
+    plan: Arc<dyn ExecutionPlan>,
+    fetch: Option<usize>,
+) -> DataFusionResult<Transformed<Arc<dyn ExecutionPlan>>> {
+    let Some(fetch) = fetch else {
+        return Ok(Transformed::no(plan));
+    };
+
+    // Avoid removing multiple coalesce batches
+    let mut is_done = false;
+
+    plan.transform_down(|plan| {
+        if let Some(coalesce_batches_exec) = plan.as_any().downcast_ref::<CoalesceBatchesExec>() {
+            let target_batch_size = coalesce_batches_exec.target_batch_size();
+            if fetch < target_batch_size && !is_done {
+                is_done = true;
+                return Ok(Transformed::yes(coalesce_batches_exec.input().clone()));
+            }
+        }
+
+        Ok(Transformed::no(plan))
+    })
+}
+
 /// Resolves alias of the time index column.
 ///
 /// i.e if a is time index, alias= {a:b, b:c}, then result should be {a, b}(not {a, c}) because projection is not transitive