Remove unused fetch inside ExternalSorter (#15525)

Author: 2010YOUY01

datafusion/physical-plan/src/sorts/sort.rs

@@ -207,8 +207,6 @@ struct ExternalSorter {
     expr: Arc<[PhysicalSortExpr]>,
     /// RowConverter corresponding to the sort expressions
     sort_keys_row_converter: Arc<RowConverter>,
-    /// If Some, the maximum number of output rows that will be produced
-    fetch: Option<usize>,
     /// The target number of rows for output batches
     batch_size: usize,
     /// If the in size of buffered memory batches is below this size,
@@ -262,7 +260,6 @@ impl ExternalSorter {
         schema: SchemaRef,
         expr: LexOrdering,
         batch_size: usize,
-        fetch: Option<usize>,
         sort_spill_reservation_bytes: usize,
         sort_in_place_threshold_bytes: usize,
         metrics: &ExecutionPlanMetricsSet,
@@ -307,7 +304,6 @@ impl ExternalSorter {
             expr: expr.into(),
             sort_keys_row_converter: Arc::new(converter),
             metrics,
-            fetch,
             reservation,
             spill_manager,
             merge_reservation,
@@ -330,10 +326,8 @@ impl ExternalSorter {
 
         let size = get_reserved_byte_for_record_batch(&input);
         if self.reservation.try_grow(size).is_err() {
-            self.sort_or_spill_in_mem_batches(false).await?;
-            // We've already freed more than half of reserved memory,
-            // so we can grow the reservation again. There's nothing we can do
-            // if this try_grow fails.
+            self.sort_and_spill_in_mem_batches().await?;
+            // After spilling all in-memory batches, the retry should succeed
             self.reservation.try_grow(size)?;
         }
 
@@ -367,7 +361,7 @@ impl ExternalSorter {
             // `in_mem_batches` and the memory limit is almost reached, merging
             // them with the spilled files at the same time might cause OOM.
             if !self.in_mem_batches.is_empty() {
-                self.sort_or_spill_in_mem_batches(true).await?;
+                self.sort_and_spill_in_mem_batches().await?;
             }
 
             for spill in self.finished_spill_files.drain(..) {
@@ -386,7 +380,7 @@ impl ExternalSorter {
                 .with_expressions(expressions.as_ref())
                 .with_metrics(self.metrics.baseline.clone())
                 .with_batch_size(self.batch_size)
-                .with_fetch(self.fetch)
+                .with_fetch(None)
                 .with_reservation(self.merge_reservation.new_empty())
                 .build()
         } else {
@@ -532,28 +526,15 @@ impl ExternalSorter {
         Ok(())
     }
 
-    /// Sorts the in_mem_batches and potentially spill the sorted batches.
-    ///
-    /// If the memory usage has dropped by a factor of 2, it might be a sort with
-    /// fetch (e.g. sorting 1M rows but only keep the top 100), so we keep the
-    /// sorted entries inside `in_mem_batches` to be sorted in the next iteration.
-    /// Otherwise, we spill the sorted run to free up memory for inserting more batches.
-    ///
-    /// # Arguments
-    ///
-    /// * `force_spill` - If true, the method will spill the in-memory batches
-    ///   even if the memory usage has not dropped by a factor of 2. Otherwise it will
-    ///   only spill when the memory usage has dropped by the pre-defined factor.
-    ///
-    async fn sort_or_spill_in_mem_batches(&mut self, force_spill: bool) -> Result<()> {
+    /// Sorts the in-memory batches and merges them into a single sorted run, then writes
+    /// the result to spill files.
+    async fn sort_and_spill_in_mem_batches(&mut self) -> Result<()> {
         // Release the memory reserved for merge back to the pool so
         // there is some left when `in_mem_sort_stream` requests an
         // allocation. At the end of this function, memory will be
         // reserved again for the next spill.
         self.merge_reservation.free();
 
-        let before = self.reservation.size();
-
         let mut sorted_stream =
             self.in_mem_sort_stream(self.metrics.baseline.intermediate())?;
         // After `in_mem_sort_stream()` is constructed, all `in_mem_batches` is taken
@@ -568,7 +549,6 @@ impl ExternalSorter {
         // sort-preserving merge and incrementally append to spill files.
         let mut globally_sorted_batches: Vec<RecordBatch> = vec![];
 
-        let mut spilled = false;
         while let Some(batch) = sorted_stream.next().await {
             let batch = batch?;
             let sorted_size = get_reserved_byte_for_record_batch(&batch);
@@ -579,7 +559,6 @@ impl ExternalSorter {
                 globally_sorted_batches.push(batch);
                 self.consume_and_spill_append(&mut globally_sorted_batches)
                     .await?; // reservation is freed in spill()
-                spilled = true;
             } else {
                 globally_sorted_batches.push(batch);
             }
@@ -589,33 +568,17 @@ impl ExternalSorter {
         // upcoming `self.reserve_memory_for_merge()` may fail due to insufficient memory.
         drop(sorted_stream);
 
-        // Sorting may free up some memory especially when fetch is `Some`. If we have
-        // not freed more than 50% of the memory, then we have to spill to free up more
-        // memory for inserting more batches.
-        if (self.reservation.size() > before / 2) || force_spill {
-            // We have not freed more than 50% of the memory, so we have to spill to
-            // free up more memory
-            self.consume_and_spill_append(&mut globally_sorted_batches)
-                .await?;
-            spilled = true;
-        }
-
-        if spilled {
-            // There might be some buffered batches that haven't trigger a spill yet.
-            self.consume_and_spill_append(&mut globally_sorted_batches)
-                .await?;
-            self.spill_finish().await?;
-        } else {
-            // If the memory limit has reached before calling this function, and it
-            // didn't spill anything, it means this is a sorting with fetch top K
-            // element: after sorting only the top K elements will be kept in memory.
-            // For simplicity, those sorted top K entries are put back to unsorted
-            // `in_mem_batches` to be consumed by the next sort/merge.
-            if !self.in_mem_batches.is_empty() {
-                return internal_err!("in_mem_batches should be cleared before");
-            }
+        self.consume_and_spill_append(&mut globally_sorted_batches)
+            .await?;
+        self.spill_finish().await?;
 
-            self.in_mem_batches = std::mem::take(&mut globally_sorted_batches);
+        // Sanity check after spilling
+        let buffers_cleared_property =
+            self.in_mem_batches.is_empty() && globally_sorted_batches.is_empty();
+        if !buffers_cleared_property {
+            return internal_err!(
+                "in_mem_batches and globally_sorted_batches should be cleared before"
+            );
         }
 
         // Reserve headroom for next sort/merge
@@ -740,7 +703,7 @@ impl ExternalSorter {
             .with_expressions(expressions.as_ref())
             .with_metrics(metrics)
             .with_batch_size(self.batch_size)
-            .with_fetch(self.fetch)
+            .with_fetch(None)
             .with_reservation(self.merge_reservation.new_empty())
             .build()
     }
@@ -761,7 +724,6 @@ impl ExternalSorter {
         );
         let schema = batch.schema();
 
-        let fetch = self.fetch;
         let expressions: LexOrdering = self.expr.iter().cloned().collect();
         let row_converter = Arc::clone(&self.sort_keys_row_converter);
         let stream = futures::stream::once(async move {
@@ -775,9 +737,9 @@ impl ExternalSorter {
             let sorted = if is_multi_column_with_lists(&sort_columns) {
                 // lex_sort_to_indices doesn't support List with more than one column
                 // https://github.com/apache/arrow-rs/issues/5454
-                sort_batch_row_based(&batch, &expressions, row_converter, fetch)?
+                sort_batch_row_based(&batch, &expressions, row_converter, None)?
             } else {
-                sort_batch(&batch, &expressions, fetch)?
+                sort_batch(&batch, &expressions, None)?
             };
 
             metrics.record_output(sorted.num_rows());
@@ -1244,7 +1206,6 @@ impl ExecutionPlan for SortExec {
                     input.schema(),
                     self.expr.clone(),
                     context.session_config().batch_size(),
-                    self.fetch,
                     execution_options.sort_spill_reservation_bytes,
                     execution_options.sort_in_place_threshold_bytes,
                     &self.metrics_set,