fix(query): avoid spill at Top-N limit boundary (#20186)

Author: dantengsky

src/query/pipeline/transforms/src/processors/transforms/sorts/sort_merge_limit.rs

@@ -57,7 +57,8 @@ impl<R: Rows> MergeSort<R> for TransformSortMergeLimit<R> {
         self.next_index += 1;
 
         let mut cursor = Cursor::new(input_index, init_rows);
-        self.num_bytes += block.memory_size() as u64;
+        let block_num_bytes = block.memory_size() as u64;
+        self.num_bytes += block_num_bytes;
         self.num_rows += block.num_rows();
         let cur_index = input_index;
         self.buffer.insert(cur_index, block);
@@ -81,6 +82,15 @@ impl<R: Rows> MergeSort<R> for TransformSortMergeLimit<R> {
             cursor.advance();
         }
 
+        // String views may keep source buffers alive after filtering or slicing. Compact only
+        // blocks that remain in the Top-N candidate set so discarded blocks avoid the copy.
+        if let Some(block) = self.buffer.remove(&cur_index) {
+            self.num_bytes -= block_num_bytes;
+            let block = block.maybe_gc();
+            self.num_bytes += block.memory_size() as u64;
+            self.buffer.insert(cur_index, block);
+        }
+
         Ok(())
     }
 
@@ -116,6 +126,57 @@ impl<R: Rows> MergeSort<R> for TransformSortMergeLimit<R> {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use databend_common_exception::Result;
+    use databend_common_expression::Column;
+    use databend_common_expression::DataBlock;
+    use databend_common_expression::FromData;
+    use databend_common_expression::types::Int32Type;
+    use databend_common_expression::types::StringType;
+
+    use super::MergeSort;
+    use super::TransformSortMergeLimit;
+    use crate::sorts::core::Rows;
+    use crate::sorts::core::SimpleRowsAsc;
+
+    #[test]
+    fn test_top_n_compacts_retained_string_views() -> Result<()> {
+        const SOURCE_ROWS: i32 = 2_000;
+        const LIMIT: usize = 10;
+
+        let payload_suffix = "x".repeat(256);
+        let keys = (0..SOURCE_ROWS).collect::<Vec<_>>();
+        let payloads = keys
+            .iter()
+            .map(|key| format!("{key:08}-{payload_suffix}"))
+            .collect::<Vec<_>>();
+        let block = DataBlock::new_from_columns(vec![
+            Int32Type::from_data(keys),
+            StringType::from_data(payloads),
+        ])
+        .slice(0..LIMIT);
+        let rows = SimpleRowsAsc::<Int32Type>::from_column(&block.get_by_offset(0).to_column())?;
+
+        let mut sort = TransformSortMergeLimit::create(4_096, LIMIT);
+        sort.add_block(block, rows)?;
+
+        let retained = sort.buffer.values().next().unwrap();
+        let Column::String(payloads) = retained.get_by_offset(1).to_column() else {
+            unreachable!("expected string payload column")
+        };
+        assert_eq!(payloads.total_bytes_len(), LIMIT * (8 + 1 + 256));
+        assert!(
+            payloads.total_buffer_len() < 16 * 1024,
+            "Top-N retained {} bytes of source string buffers",
+            payloads.total_buffer_len(),
+        );
+        assert_eq!(sort.num_bytes().0, retained.memory_size() as u64);
+
+        Ok(())
+    }
+}
+
 #[derive(Clone, Copy)]
 struct LocalCursorOrder;
 

src/query/service/src/pipelines/processors/transforms/sort/sort_builder.rs

@@ -212,7 +212,7 @@ impl<S: SortSpiller> TransformSortBuilder<S> {
     }
 
     fn should_use_sort_limit(&self) -> bool {
-        self.limit.map(|limit| limit < 10000).unwrap_or_default()
+        self.limit.map(|limit| limit <= 10000).unwrap_or_default()
     }
 
     fn new_base(&self) -> Base<S> {

tests/sqllogictests/suites/mode/standalone/explain/sort.test

@@ -138,6 +138,59 @@ digraph {
     17 -> 18 [ label = "" ]
 }
 
+# LIMIT 10000 is the inclusive boundary for the Top-N merge path.
+# Disable lazy read so this only verifies sort implementation selection.
+query T
+settings(lazy_read_threshold = 0) explain pipeline select a, b from t1 order by a limit 10000;
+----
+digraph {
+    0 [ label = "PartitionStreamSource" ]
+    1 [ label = "AsyncReadDataTransform" ]
+    2 [ label = "DeserializeDataTransform" ]
+    3 [ label = "Resize" ]
+    4 [ label = "SortPartialTransform" ]
+    5 [ label = "SortPartialTransform" ]
+    6 [ label = "SortPartialTransform" ]
+    7 [ label = "SortPartialTransform" ]
+    8 [ label = "TransformSortMergeLimit" ]
+    9 [ label = "TransformSortMergeLimit" ]
+    10 [ label = "TransformSortMergeLimit" ]
+    11 [ label = "TransformSortMergeLimit" ]
+    12 [ label = "KWayMergePartitioner" ]
+    13 [ label = "KWayMergeWorker" ]
+    14 [ label = "KWayMergeWorker" ]
+    15 [ label = "KWayMergeWorker" ]
+    16 [ label = "KWayMergeWorker" ]
+    17 [ label = "KWayMergeCombiner" ]
+    18 [ label = "LimitTransform" ]
+    19 [ label = "CompoundBlockOperator(Project)" ]
+    0 -> 1 [ label = "" ]
+    1 -> 2 [ label = "" ]
+    2 -> 3 [ label = "" ]
+    3 -> 4 [ label = "from: 0, to: 0" ]
+    3 -> 5 [ label = "from: 1, to: 0" ]
+    3 -> 6 [ label = "from: 2, to: 0" ]
+    3 -> 7 [ label = "from: 3, to: 0" ]
+    4 -> 8 [ label = "" ]
+    5 -> 9 [ label = "" ]
+    6 -> 10 [ label = "" ]
+    7 -> 11 [ label = "" ]
+    8 -> 12 [ label = "from: 0, to: 0" ]
+    9 -> 12 [ label = "from: 0, to: 1" ]
+    10 -> 12 [ label = "from: 0, to: 2" ]
+    11 -> 12 [ label = "from: 0, to: 3" ]
+    12 -> 13 [ label = "from: 0, to: 0" ]
+    12 -> 14 [ label = "from: 1, to: 0" ]
+    12 -> 15 [ label = "from: 2, to: 0" ]
+    12 -> 16 [ label = "from: 3, to: 0" ]
+    13 -> 17 [ label = "from: 0, to: 0" ]
+    14 -> 17 [ label = "from: 0, to: 1" ]
+    15 -> 17 [ label = "from: 0, to: 2" ]
+    16 -> 17 [ label = "from: 0, to: 3" ]
+    17 -> 18 [ label = "" ]
+    18 -> 19 [ label = "" ]
+}
+
 
 # Sort spilling
 statement ok