Copilot commented on code in PR #2464:
URL: https://github.com/apache/auron/pull/2464#discussion_r3754954817
##########
native-engine/datafusion-ext-plans/src/sort_exec.rs:
##########
@@ -1547,6 +1557,35 @@ mod test {
Ok(())
}
+ #[tokio::test]
+ async fn test_top_k_with_only_sort_column() -> Result<()> {
+ MemManager::init(100);
+ let session_ctx = SessionContext::new();
+ let task_ctx = session_ctx.task_ctx();
+ let input = build_single_column_table("id", (0..10).rev().collect())?;
+ let sort_exprs = vec![PhysicalSortExpr {
+ expr: Arc::new(Column::new("id", 0)),
+ options: SortOptions::default(),
+ }];
+
+ let sort = SortExec::new(input, sort_exprs, Some(6), 0);
+ let output = sort.execute(0, task_ctx)?;
+ let batches = common::collect(output).await?;
+ let expected = r#"+----+
+| id |
++----+
+| 0 |
+| 1 |
+| 2 |
+| 3 |
+| 4 |
+| 5 |
++----+"#;
+ assert_batches_eq!(expected.lines().collect::<Vec<_>>(), &batches);
Review Comment:
The expected pretty-printed output includes leading spaces before the `| ...
|` lines (e.g., `" | id |"`). `assert_batches_eq!` compares lines verbatim, so
these extra spaces are likely to make the test fail even when the batches are
correct. Use the same no-leading-space format as the other tests in this module
(a `Vec<&str>` of lines).
##########
native-engine/datafusion-ext-plans/src/sort_exec.rs:
##########
@@ -675,7 +675,7 @@ impl ExternalSorter {
let sorted_batch = if !self.prune_sort_keys_from_batch.is_all_pruned()
{
take_batch(batch, sorted_indices)?
} else {
- create_zero_column_batch(batch.num_rows())
+ create_zero_column_batch(sorted_indices.len())
};
Review Comment:
In Top-K mode only `sorted_indices` are retained, but the key buffer is
still reserved for the full input (`key_collector.reserve(keys.num_rows(),
keys.size())`). For large inputs with small limits this can over-allocate
memory proportional to the full batch, potentially triggering unnecessary
spills or memory pressure. Consider reserving based on `sorted_indices.len()`
(and an estimated key byte size) instead.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
