(auron) branch master updated: [AURON #2446] Flush buffered data based on row count (#2447)

[slfan1989]

This is an automated email from the ASF dual-hosted git repository.

slfan1989 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git


The following commit(s) were added to refs/heads/master by this push:
     new 23fc4f531 [AURON #2446] Flush buffered data based on row count (#2447)
23fc4f531 is described below

commit 23fc4f53107afb5c50706072319c6fa0a4e6445d
Author: Junfan Zhang <zus...@apache.org>
AuthorDate: Fri Jul 31 08:44:58 2026 +0800

    [AURON #2446] Flush buffered data based on row count (#2447)
    
    # Which issue does this PR close?
    
    Closes #2446
    
    this PR is to fix the incorrect flushing logic, that should be based on
    the row count rather than the bytes
    
    # Rationale for this change
    
    Without this fix, the flushing will be triggered too frequently.
    
    # What changes are included in this PR?
    
    # Are there any user-facing changes?
    
    # How was this patch tested?
    
    # Was this patch authored or co-authored using generative AI tooling?
    - [x] Yes
    - [ ] No
    
    Generated-by: OpenAI Codex (GPT-5)
---
 .../src/shuffle/buffered_data.rs                   | 33 +++++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs 
b/native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs
index 32f1e2969..07954f641 100644
--- a/native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs
+++ b/native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs
@@ -94,7 +94,7 @@ impl BufferedData {
 
         let suggested_batch_size =
             compute_suggested_batch_size_for_output(self.staging_mem_used, 
self.staging_num_rows);
-        if self.staging_mem_used > suggested_batch_size {
+        if self.staging_num_rows >= suggested_batch_size {
             self.flush_staging()?;
         }
         Ok(())
@@ -399,6 +399,37 @@ mod test {
         Ok(batch)
     }
 
+    #[test]
+    fn test_add_batch_flushes_by_row_count() -> Result<()> {
+        let values = (0..1000).collect::<Vec<_>>();
+        let record_batch = build_table_i32(("a", &values), ("b", &values), 
("c", &values))?;
+        let batch_mem_used = record_batch.get_batch_mem_size() * 2;
+        let suggested_batch_size =
+            compute_suggested_batch_size_for_output(batch_mem_used, 
record_batch.num_rows());
+
+        // Ensure this fixture distinguishes byte count from row count.
+        assert!(batch_mem_used > suggested_batch_size);
+        assert!(record_batch.num_rows() < suggested_batch_size);
+
+        let num_batches = 
suggested_batch_size.div_ceil(record_batch.num_rows());
+        let mut buffered_data =
+            BufferedData::new(Partitioning::RoundRobinPartitioning(4), 0, 
Time::new());
+
+        for _ in 1..num_batches {
+            buffered_data.add_batch(record_batch.clone())?;
+        }
+        assert!(buffered_data.sorted_batches.is_empty());
+
+        buffered_data.add_batch(record_batch)?;
+        assert_eq!(buffered_data.sorted_batches.len(), 1);
+        assert_eq!(
+            buffered_data.sorted_batches[0].num_rows(),
+            num_batches * values.len()
+        );
+        assert!(buffered_data.staging_batches.is_empty());
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_round_robin() -> Result<()> {
         let record_batch = build_table_i32(