(auron) branch master updated: [AURON #2417] log rss write size metrics (#2418)

[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 a1e48f3f [AURON #2417] log rss write size metrics (#2418)
a1e48f3f is described below

commit a1e48f3fa9b86c406d643944f2ca3fb16ae2721b
Author: cxzl25 <3898450+cxz...@users.noreply.github.com>
AuthorDate: Thu Jul 23 10:56:20 2026 +0800

    [AURON #2417] log rss write size metrics (#2418)
    
    # Which issue does this PR close?
    
    Closes #2417
    
    # Rationale for this change
    
    Now when shuffle is written, there is a shuffle size after log output
    compression, but there is no output compression size when using rss.
    
    
    
https://github.com/apache/auron/blob/2889b563139e6495495b61d0902e3f93f0ebe8bb/native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs#L156
    
    # What changes are included in this PR?
    Output the compressed size when writing RSS.
    
    # Are there any user-facing changes?
    No
    
    # How was this patch tested?
    Production environment verification
    
    ```
    2026-07-22 20:10:07.090 (+80.707s) [INFO] 
[datafusion_ext_plans::shuffle::buffered_data:196] (stage: 0, partition: 1694, 
tid: 0) - all buffered data drained to rss, pushed_compressed=74.2 MiB, 
uncompressed_mem=528.1 MiB
    2026-07-22 20:10:07.110 (+80.727s) [INFO] 
[datafusion_ext_plans::shuffle::buffered_data:196] (stage: 0, partition: 385, 
tid: 2) - all buffered data drained to rss, pushed_compressed=81.3 MiB, 
uncompressed_mem=531.6 MiB
    2026-07-22 20:10:09.505 (+83.122s) [INFO] 
[datafusion_ext_plans::shuffle::buffered_data:196] (stage: 0, partition: 3824, 
tid: 1) - all buffered data drained to rss, pushed_compressed=80.0 MiB, 
uncompressed_mem=528.4 MiB
    ```
    
    
    # Was this patch authored or co-authored using generative AI tooling?
    - [ ] Yes
    - [ ] No
---
 native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs | 8 +++++++-
 native-engine/datafusion-ext-plans/src/shuffle/rss.rs           | 7 +++++++
 2 files changed, 14 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 40524296..32f1e296 100644
--- a/native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs
+++ b/native-engine/datafusion-ext-plans/src/shuffle/buffered_data.rs
@@ -173,6 +173,7 @@ impl BufferedData {
         let output_io_time = self.output_io_time.clone();
         let mut iter = self.into_sorted_batches()?;
         let mut writer = 
IpcCompressionWriter::new(RssWriter::new(rss_partition_writer.clone(), 0));
+        let mut total_pushed: u64 = 0;
 
         while let Some((partition_id, batch_iter)) = 
iter.next_partition_chunk() {
             if !is_task_running() {
@@ -186,12 +187,17 @@ impl BufferedData {
                     .with_timer(|| writer.write_batch(batch.num_rows(), 
batch.columns()))?;
             }
             output_io_time.with_timer(|| writer.finish_current_buf())?;
+            total_pushed += writer.inner().total_written();
         }
 
         output_io_time.with_timer(
             || 
jni_call!(AuronRssPartitionWriterBase(rss_partition_writer.as_obj()).flush() -> 
()),
         )?;
-        log::info!("all buffered data drained to rss");
+        log::info!(
+            "all buffered data drained to rss, pushed_compressed={}, 
uncompressed_mem={}",
+            ByteSize(total_pushed),
+            mem_used
+        );
         Ok(())
     }
 
diff --git a/native-engine/datafusion-ext-plans/src/shuffle/rss.rs 
b/native-engine/datafusion-ext-plans/src/shuffle/rss.rs
index dfdff5f3..18a21167 100644
--- a/native-engine/datafusion-ext-plans/src/shuffle/rss.rs
+++ b/native-engine/datafusion-ext-plans/src/shuffle/rss.rs
@@ -21,6 +21,7 @@ use jni::objects::GlobalRef;
 pub struct RssWriter {
     rss_partition_writer: GlobalRef,
     partition_id: usize,
+    total_written: u64,
 }
 
 impl RssWriter {
@@ -28,8 +29,13 @@ impl RssWriter {
         Self {
             rss_partition_writer,
             partition_id,
+            total_written: 0,
         }
     }
+
+    pub fn total_written(&self) -> u64 {
+        self.total_written
+    }
 }
 
 impl Write for RssWriter {
@@ -40,6 +46,7 @@ impl Write for RssWriter {
             AuronRssPartitionWriterBase(self.rss_partition_writer.as_obj())
                 .write(self.partition_id as i32, buf.as_obj()) -> ()
         )?;
+        self.total_written += buf_len as u64;
         Ok(buf_len)
     }