Repository: apache/datafusion-comet
Pull request #1730: feat: Improve performance tracing feature
Author: login: andygrove
Created: 2025-05-10T15:44:22Z | Merged: 2025-05-13T21:44:19Z

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->

N/A

## Rationale for this change

<!--
 Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
 Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
-->

Various improvements to the performance tracing feature.

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->

- Log Comet and Spark task memory usage per thread
- Refactor code to use `withTracing` helper method
- Fixes bug where `TracingGuard` was dropped too soon in native code
- Show stage ID in event trace
- Change timings from nanos to micros (this is the expected unit in the viewing tool)
- Record memory usage in MB rather than bytes
- Improve documentation
 
![2025-05-12_16-49_1](https://github.com/user-attachments/assets/e5fee389-c681-4209-9eb5-593a3fee5e8d)

![2025-05-12_16-49](https://github.com/user-attachments/assets/8d9c4d98-fa8b-40be-ad33-48dc7c7d3baa)

## How are these changes tested?

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->


Linked issue #123: Add sqllogictest in ci to verify execution results against vanilla spark sql.
### What is the problem the feature request solves?

[sqllogictest](https://github.com/hydromatic/sql-logic-test) is a commonly used data driven approach to verify query engine.  We can use this to verify the results of comet are same as  vanilla spark sql.

### Describe the potential solution

_No response_

### Additional context

_No response_
Linked issue comments: user: login: liurenjie1024
created_at: 2024-02-27T07:12:07Z
body: I'm happy to take this is we have reached consensus on this.
author_association: NONE
user: login: viirya
created_at: 2024-02-27T07:16:06Z
body: Hmm, I guess the question is, how do you compare with Spark result in sqllogictest?

For example, in current unit test or e2e test like CometTPCDSQuerySuite, we compare Comet query result with Spark result generated in runtime by disabling Comet or golden filed generated by Spark.

But for sqllogictest, as I saw its usage in DataFusion, it works like comparing actual output with previous output (in sqllogictest file) to check output consistency. So it cannot be used to run Spark in runtime to get expected output.
author_association: MEMBER
user: login: liurenjie1024
created_at: 2024-02-27T07:21:57Z
body: >But for sqllogictest, as I saw its usage in DataFusion, it works like comparing actual output with expected output (in sqllogictest file). So it cannot be used to run Spark in runtime to get expected output.

Why? We can use spark connect.
author_association: NONE
user: login: viirya
created_at: 2024-02-27T07:24:10Z
body: > Why? We can use spark connect.

Can sqllogictest run Spark connect? And how do you compare it in sqllogictest?
author_association: MEMBER
user: login: liurenjie1024
created_at: 2024-02-27T07:25:23Z
body: > > Why? We can use spark connect.
> 
> Can sqllogictest run Spark connect?

There are two ways to do this:

1. We can use sqllogictest as a library.
2. We develop a test executor for sqllogictest.

Please note that the link I provided is java's implementation of sqllogictest.
author_association: NONE
user: login: liurenjie1024
created_at: 2024-02-27T07:26:05Z
body: >  And how do you compare it in sqllogictest?

We just run the slt files against both vanilla spark and comet, and they both should be ok.
author_association: NONE
user: login: viirya
created_at: 2024-02-27T07:57:21Z
body: I'm not sure if this makes sense. I tend to think this doesn't make too much sense at this moment.

* It is not something ready to use. As you said, we need to implement a sqllogictest test executor for Spark to talk to Spark connect. I think this makes thing complicated. For the test, we need Spark server, Spark connect client, sqllogictest text executor, sqllogictest... too many components which increase maintenance effort. 
* The slt files need to be run with both Spark and Comet. Although we have unit tests also doing like that, the number is limited.
* We will soon have Spark test pipelines which cover Spark compatibility which leverage existing Spark tests. I think Spark test itself has most correct and highest coverage to make sure Comet's Spark compatibility. If there are very corner cases which are missing in Spark tests, I think these cases should be less. We probably can add the tests directly into Spark or just in Comet as unit test.

author_association: MEMBER
user: login: liurenjie1024
created_at: 2024-02-27T08:07:23Z
body: Sounds reasonable, if we can port spark tests in this repo, it would be enough. Thanks for sharing!
author_association: NONE

Diff:
diff --git a/docs/source/contributor-guide/tracing.md b/docs/source/contributor-guide/tracing.md
index 47c0cc2f2c..b9b4fe0dcc 100644
--- a/docs/source/contributor-guide/tracing.md
+++ b/docs/source/contributor-guide/tracing.md
@@ -19,18 +19,20 @@ under the License.
 
 # Tracing
 
-Tracing can be enabled by setting `spark.comet.tracing.enabled=true`. Additionally, enabling the `jemalloc` 
-feature will enable tracing of native memory allocations.
+Tracing can be enabled by setting `spark.comet.tracing.enabled=true`.
+
+With this feature enabled, each Spark executor will write a JSON event log file in
+Chrome's [Trace Event Format]. The file will be written to the executor's current working
+directory with the filename `comet-event-trace.json`.
+
+[Trace Event Format]: https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview?tab=t.0#heading=h.yr4qxyxotyw
+
+Additionally, enabling the `jemalloc` feature will enable tracing of native memory allocations.
 
 ```shell
 make release COMET_FEATURES="jemalloc"
 ```
 
-With this feature enabled, each Spark executor will write a JSON event log file in 
-Chrome's [Trace Event Format].
-
-[Trace Event Format]: https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview?tab=t.0#heading=h.yr4qxyxotyw
-
 Example output:
 
 ```json
@@ -43,10 +45,20 @@ Example output:
 { "name": "CometExecIterator_getNextBatch", "cat": "PERF", "ph": "B", "pid": 1, "tid": 5, "ts": 10109479156 },
 ```
 
-Traces can be viewed with [Trace Viewer]. 
+Traces can be viewed with [Trace Viewer].
 
 [Trace Viewer]: https://github.com/catapult-project/catapult/blob/main/tracing/README.md
 
 Example trace visualization:
 
 ![tracing](../_static/images/tracing.png)
+
+## Definition of Labels
+
+| Label                 | Meaning                                                        |
+| --------------------- | -------------------------------------------------------------- |
+| jvm_heapUsed          | JVM heap memory usage of live objects for the executor process |
+| jemalloc_allocated    | Native memory usage for the executor process                   |
+| task_memory_comet_NNN | Off-heap memory allocated by Comet for query execution         |
+| task_memory_spark_NNN | On-heap & Off-heap memory allocated by Spark                   |
+| comet_shuffle_NNN     | Off-heap memory allocated by Comet for columnar shuffle        |
diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
index b7c62d4d92..abff7043a6 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -67,9 +67,9 @@ use crate::execution::operators::ScanExec;
 use crate::execution::shuffle::{read_ipc_compressed, CompressionCodec};
 use crate::execution::spark_plan::SparkPlan;
 
-use crate::execution::tracing::TraceGuard;
-use crate::execution::tracing::{log_counter, trace_begin, trace_end};
+use crate::execution::tracing::{log_memory_usage, trace_begin, trace_end, with_trace};
 
+use datafusion_comet_proto::spark_operator::operator::OpStruct;
 use log::info;
 use once_cell::sync::Lazy;
 #[cfg(feature = "jemalloc")]
@@ -162,90 +162,90 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
     tracing_enabled: jboolean,
 ) -> jlong {
     try_unwrap_or_throw(&e, |mut env| {
-        let _ = TraceGuard::new("createPlan", tracing_enabled != JNI_FALSE);
+        with_trace("createPlan", tracing_enabled != JNI_FALSE, || {
+            // Init JVM classes
+            JVMClasses::init(&mut env);
 
-        // Init JVM classes
-        JVMClasses::init(&mut env);
-
-        let start = Instant::now();
-
-        let array = unsafe { JPrimitiveArray::from_raw(serialized_query) };
-        let bytes = env.convert_byte_array(array)?;
-
-        // Deserialize query plan
-        let spark_plan = serde::deserialize_op(bytes.as_slice())?;
+            let start = Instant::now();
 
-        let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?);
+            let array = unsafe { JPrimitiveArray::from_raw(serialized_query) };
+            let bytes = env.convert_byte_array(array)?;
 
-        // Get the global references of input sources
-        let mut input_sources = vec![];
-        let iter_array = JObjectArray::from_raw(iterators);
-        let num_inputs = env.get_array_length(&iter_array)?;
-        for i in 0..num_inputs {
-            let input_source = env.get_object_array_element(&iter_array, i)?;
-            let input_source = Arc::new(jni_new_global_ref!(env, input_source)?);
-            input_sources.push(input_source);
-        }
+            // Deserialize query plan
+            let spark_plan = serde::deserialize_op(bytes.as_slice())?;
 
-        // Create DataFusion memory pool
-        let task_memory_manager =
-            Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?);
-
-        let memory_pool_type = env.get_string(&JString::from_raw(memory_pool_type))?.into();
-        let memory_pool_config = parse_memory_pool_config(
-            off_heap_mode != JNI_FALSE,
-            memory_pool_type,
-            memory_limit,
-            memory_limit_per_task,
-        )?;
-        let memory_pool =
-            create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id);
-
-        // Get local directories for storing spill files
-        let local_dirs_array = JObjectArray::from_raw(local_dirs);
-        let num_local_dirs = env.get_array_length(&local_dirs_array)?;
-        let mut local_dirs = vec![];
-        for i in 0..num_local_dirs {
-            let local_dir: JString = env.get_object_array_element(&local_dirs_array, i)?.into();
-            let local_dir = env.get_string(&local_dir)?;
-            local_dirs.push(local_dir.into());
-        }
+            let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?);
 
-        // We need to keep the session context alive. Some session state like temporary
-        // dictionaries are stored in session context. If it is dropped, the temporary
-        // dictionaries will be dropped as well.
-        let session =
-            prepare_datafusion_session_context(batch_size as usize, memory_pool, local_dirs)?;
+            // Get the global references of input sources
+            let mut input_sources = vec![];
+            let iter_array = JObjectArray::from_raw(iterators);
+            let num_inputs = env.get_array_length(&iter_array)?;
+            for i in 0..num_inputs {
+                let input_source = env.get_object_array_element(&iter_array, i)?;
+                let input_source = Arc::new(jni_new_global_ref!(env, input_source)?);
+                input_sources.push(input_source);
+            }
 
-        let plan_creation_time = start.elapsed();
+            // Create DataFusion memory pool
+            let task_memory_manager =
+                Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?);
 
-        let metrics_update_interval = if metrics_update_interval > 0 {
-            Some(Duration::from_millis(metrics_update_interval as u64))
-        } else {
-            None
-        };
+            let memory_pool_type = env.get_string(&JString::from_raw(memory_pool_type))?.into();
+            let memory_pool_config = parse_memory_pool_config(
+                off_heap_mode != JNI_FALSE,
+                memory_pool_type,
+                memory_limit,
+                memory_limit_per_task,
+            )?;
+            let memory_pool =
+                create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id);
+
+            // Get local directories for storing spill files
+            let local_dirs_array = JObjectArray::from_raw(local_dirs);
+            let num_local_dirs = env.get_array_length(&local_dirs_array)?;
+            let mut local_dirs = vec![];
+            for i in 0..num_local_dirs {
+                let local_dir: JString = env.get_object_array_element(&local_dirs_array, i)?.into();
+                let local_dir = env.get_string(&local_dir)?;
+                local_dirs.push(local_dir.into());
+            }
 
-        let exec_context = Box::new(ExecutionContext {
-            id,
-            task_attempt_id,
-            spark_plan,
-            partition_count: partition_count as usize,
-            root_op: None,
-            scans: vec![],
-            input_sources,
-            stream: None,
-            metrics,
-            metrics_update_interval,
-            metrics_last_update_time: Instant::now(),
-            plan_creation_time,
-            session_ctx: Arc::new(session),
-            debug_native: debug_native == 1,
-            explain_native: explain_native == 1,
-            memory_pool_config,
-            tracing_enabled: tracing_enabled != JNI_FALSE,
-        });
-
-        Ok(Box::into_raw(exec_context) as i64)
+            // We need to keep the session context alive. Some session state like temporary
+            // dictionaries are stored in session context. If it is dropped, the temporary
+            // dictionaries will be dropped as well.
+            let session =
+                prepare_datafusion_session_context(batch_size as usize, memory_pool, local_dirs)?;
+
+            let plan_creation_time = start.elapsed();
+
+            let metrics_update_interval = if metrics_update_interval > 0 {
+                Some(Duration::from_millis(metrics_update_interval as u64))
+            } else {
+                None
+            };
+
+            let exec_context = Box::new(ExecutionContext {
+                id,
+                task_attempt_id,
+                spark_plan,
+                partition_count: partition_count as usize,
+                root_op: None,
+                scans: vec![],
+                input_sources,
+                stream: None,
+                metrics,
+                metrics_update_interval,
+                metrics_last_update_time: Instant::now(),
+                plan_creation_time,
+                session_ctx: Arc::new(session),
+                debug_native: debug_native == 1,
+                explain_native: explain_native == 1,
+                memory_pool_config,
+                tracing_enabled: tracing_enabled != JNI_FALSE,
+            });
+
+            Ok(Box::into_raw(exec_context) as i64)
+        })
     })
 }
 
@@ -366,117 +366,122 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
     exec_context: jlong,
     array_addrs: jlongArray,
     schema_addrs: jlongArray,
-    tracing_enabled: jboolean,
 ) -> jlong {
     try_unwrap_or_throw(&e, |mut env| {
-        let _ = TraceGuard::new("executePlan", tracing_enabled != JNI_FALSE);
-
         // Retrieve the query
         let exec_context = get_execution_context(exec_context);
 
+        let tracing_event_name = match &exec_context.spark_plan.op_struct {
+            Some(OpStruct::ShuffleWriter(_)) => "executePlan(ShuffleWriter)",
+            _ => "executePlan",
+        };
+
         if exec_context.tracing_enabled {
             #[cfg(feature = "jemalloc")]
             {
                 let e = epoch::mib().unwrap();
                 let allocated = stats::allocated::mib().unwrap();
                 e.advance().unwrap();
-                use crate::execution::tracing::log_counter;
-                log_counter("jemalloc_allocated", allocated.read().unwrap() as u64);
+                use crate::execution::tracing::log_memory_usage;
+                log_memory_usage("jemalloc_allocated", allocated.read().unwrap() as u64);
             }
         }
 
-        let exec_context_id = exec_context.id;
-
-        // Initialize the execution stream.
-        // Because we don't know if input arrays are dictionary-encoded when we create
-        // query plan, we need to defer stream initialization to first time execution.
-        if exec_context.root_op.is_none() {
-            let start = Instant::now();
-            let planner = PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx))
-                .with_exec_id(exec_context_id);
-            let (scans, root_op) = planner.create_plan(
-                &exec_context.spark_plan,
-                &mut exec_context.input_sources.clone(),
-                exec_context.partition_count,
-            )?;
-            let physical_plan_time = start.elapsed();
-
-            exec_context.plan_creation_time += physical_plan_time;
-            exec_context.root_op = Some(Arc::clone(&root_op));
-            exec_context.scans = scans;
+        with_trace(tracing_event_name, exec_context.tracing_enabled, || {
+            let exec_context_id = exec_context.id;
+
+            // Initialize the execution stream.
+            // Because we don't know if input arrays are dictionary-encoded when we create
+            // query plan, we need to defer stream initialization to first time execution.
+            if exec_context.root_op.is_none() {
+                let start = Instant::now();
+                let planner = PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx))
+                    .with_exec_id(exec_context_id);
+                let (scans, root_op) = planner.create_plan(
+                    &exec_context.spark_plan,
+                    &mut exec_context.input_sources.clone(),
+                    exec_context.partition_count,
+                )?;
+                let physical_plan_time = start.elapsed();
+
+                exec_context.plan_creation_time += physical_plan_time;
+                exec_context.root_op = Some(Arc::clone(&root_op));
+                exec_context.scans = scans;
+
+                if exec_context.explain_native {
+                    let formatted_plan_str =
+                        DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true);
+                    info!("Comet native query plan:\n{formatted_plan_str:}");
+                }
 
-            if exec_context.explain_native {
-                let formatted_plan_str =
-                    DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true);
-                info!("Comet native query plan:\n{formatted_plan_str:}");
+                let task_ctx = exec_context.session_ctx.task_ctx();
+                let stream = exec_context
+                    .root_op
+                    .as_ref()
+                    .unwrap()
+                    .native_plan
+                    .execute(partition as usize, task_ctx)?;
+                exec_context.stream = Some(stream);
+            } else {
+                // Pull input batches
+                pull_input_batches(exec_context)?;
             }
 
-            let task_ctx = exec_context.session_ctx.task_ctx();
-            let stream = exec_context
-                .root_op
-                .as_ref()
-                .unwrap()
-                .native_plan
-                .execute(partition as usize, task_ctx)?;
-            exec_context.stream = Some(stream);
-        } else {
-            // Pull input batches
-            pull_input_batches(exec_context)?;
-        }
-
-        loop {
-            // Polling the stream.
-            let next_item = exec_context.stream.as_mut().unwrap().next();
-            let poll_output = get_runtime().block_on(async { poll!(next_item) });
-
-            // update metrics at interval
-            if let Some(interval) = exec_context.metrics_update_interval {
-                let now = Instant::now();
-                if now - exec_context.metrics_last_update_time >= interval {
-                    update_metrics(&mut env, exec_context)?;
-                    exec_context.metrics_last_update_time = now;
+            loop {
+                // Polling the stream.
+                let next_item = exec_context.stream.as_mut().unwrap().next();
+                let poll_output = get_runtime().block_on(async { poll!(next_item) });
+
+                // update metrics at interval
+                if let Some(interval) = exec_context.metrics_update_interval {
+                    let now = Instant::now();
+                    if now - exec_context.metrics_last_update_time >= interval {
+                        update_metrics(&mut env, exec_context)?;
+                        exec_context.metrics_last_update_time = now;
+                    }
                 }
-            }
 
-            match poll_output {
-                Poll::Ready(Some(output)) => {
-                    // prepare output for FFI transfer
-                    return prepare_output(
-                        &mut env,
-                        array_addrs,
-                        schema_addrs,
-                        output?,
-                        exec_context.debug_native,
-                    );
-                }
-                Poll::Ready(None) => {
-                    // Reaches EOF of output.
-                    if exec_context.explain_native {
-                        if let Some(plan) = &exec_context.root_op {
-                            let formatted_plan_str =
-                                DisplayableExecutionPlan::with_metrics(plan.native_plan.as_ref())
-                                    .indent(true);
-                            info!(
-                                "Comet native query plan with metrics (Plan #{} Stage {} Partition {}):\
-                            \n plan creation (including CometScans fetching first batches) took {:?}:\
-                            \n{formatted_plan_str:}",
-                                plan.plan_id, stage_id, partition, exec_context.plan_creation_time
-                            );
+                match poll_output {
+                    Poll::Ready(Some(output)) => {
+                        // prepare output for FFI transfer
+                        return prepare_output(
+                            &mut env,
+                            array_addrs,
+                            schema_addrs,
+                            output?,
+                            exec_context.debug_native,
+                        );
+                    }
+                    Poll::Ready(None) => {
+                        // Reaches EOF of output.
+                        if exec_context.explain_native {
+                            if let Some(plan) = &exec_context.root_op {
+                                let formatted_plan_str = DisplayableExecutionPlan::with_metrics(
+                                    plan.native_plan.as_ref(),
+                                )
+                                .indent(true);
+                                info!(
+                                    "Comet native query plan with metrics (Plan #{} Stage {} Partition {}):\
+                                \n plan creation (including CometScans fetching first batches) took {:?}:\
+                                \n{formatted_plan_str:}",
+                                    plan.plan_id, stage_id, partition, exec_context.plan_creation_time
+                                );
+                            }
                         }
+                        return Ok(-1);
+                    }
+                    // A poll pending means there are more than one blocking operators,
+                    // we don't need go back-forth between JVM/Native. Just keeping polling.
+                    Poll::Pending => {
+                        // Pull input batches
+                        pull_input_batches(exec_context)?;
+
+                        // Output not ready yet
+                        continue;
                     }
-                    return Ok(-1);
-                }
-                // A poll pending means there are more than one blocking operators,
-                // we don't need go back-forth between JVM/Native. Just keeping polling.
-                Poll::Pending => {
-                    // Pull input batches
-                    pull_input_batches(exec_context)?;
-
-                    // Output not ready yet
-                    continue;
                 }
             }
-        }
+        })
     })
 }
 
@@ -565,71 +570,76 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_writeSortedFileNative
     tracing_enabled: jboolean,
 ) -> jlongArray {
     try_unwrap_or_throw(&e, |mut env| unsafe {
-        let _ = TraceGuard::new("writeSortedFileNative", tracing_enabled != JNI_FALSE);
-
-        let data_types = convert_datatype_arrays(&mut env, serialized_datatypes)?;
-
-        let row_address_array = JLongArray::from_raw(row_addresses);
-        let row_num = env.get_array_length(&row_address_array)? as usize;
-        let row_addresses = env.get_array_elements(&row_address_array, ReleaseMode::NoCopyBack)?;
-
-        let row_size_array = JIntArray::from_raw(row_sizes);
-        let row_sizes = env.get_array_elements(&row_size_array, ReleaseMode::NoCopyBack)?;
-
-        let row_addresses_ptr = row_addresses.as_ptr();
-        let row_sizes_ptr = row_sizes.as_ptr();
-
-        let output_path: String = env
-            .get_string(&JString::from_raw(file_path))
-            .unwrap()
-            .into();
-
-        let checksum_enabled = checksum_enabled == 1;
-        let current_checksum = if current_checksum == i64::MIN {
-            // Initial checksum is not available.
-            None
-        } else {
-            Some(current_checksum as u32)
-        };
-
-        let compression_codec: String = env
-            .get_string(&JString::from_raw(compression_codec))
-            .unwrap()
-            .into();
-
-        let compression_codec = match compression_codec.as_str() {
-            "zstd" => CompressionCodec::Zstd(compression_level),
-            "lz4" => CompressionCodec::Lz4Frame,
-            "snappy" => CompressionCodec::Snappy,
-            _ => CompressionCodec::Lz4Frame,
-        };
-
-        let (written_bytes, checksum) = process_sorted_row_partition(
-            row_num,
-            batch_size as usize,
-            row_addresses_ptr,
-            row_sizes_ptr,
-            &data_types,
-            output_path,
-            prefer_dictionary_ratio,
-            checksum_enabled,
-            checksum_algo,
-            current_checksum,
-            &compression_codec,
-        )?;
-
-        let checksum = if let Some(checksum) = checksum {
-            checksum as i64
-        } else {
-            // Spark checksums (CRC32 or Adler32) are both u32, so we use i64::MIN to indicate
-            // checksum is not available.
-            i64::MIN
-        };
-
-        let long_array = env.new_long_array(2)?;
-        env.set_long_array_region(&long_array, 0, &[written_bytes, checksum])?;
-
-        Ok(long_array.into_raw())
+        with_trace(
+            "writeSortedFileNative",
+            tracing_enabled != JNI_FALSE,
+            || {
+                let data_types = convert_datatype_arrays(&mut env, serialized_datatypes)?;
+
+                let row_address_array = JLongArray::from_raw(row_addresses);
+                let row_num = env.get_array_length(&row_address_array)? as usize;
+                let row_addresses =
+                    env.get_array_elements(&row_address_array, ReleaseMode::NoCopyBack)?;
+
+                let row_size_array = JIntArray::from_raw(row_sizes);
+                let row_sizes = env.get_array_elements(&row_size_array, ReleaseMode::NoCopyBack)?;
+
+                let row_addresses_ptr = row_addresses.as_ptr();
+                let row_sizes_ptr = row_sizes.as_ptr();
+
+                let output_path: String = env
+                    .get_string(&JString::from_raw(file_path))
+                    .unwrap()
+                    .into();
+
+                let checksum_enabled = checksum_enabled == 1;
+                let current_checksum = if current_checksum == i64::MIN {
+                    // Initial checksum is not available.
+                    None
+                } else {
+                    Some(current_checksum as u32)
+                };
+
+                let compression_codec: String = env
+                    .get_string(&JString::from_raw(compression_codec))
+                    .unwrap()
+                    .into();
+
+                let compression_codec = match compression_codec.as_str() {
+                    "zstd" => CompressionCodec::Zstd(compression_level),
+                    "lz4" => CompressionCodec::Lz4Frame,
+                    "snappy" => CompressionCodec::Snappy,
+                    _ => CompressionCodec::Lz4Frame,
+                };
+
+                let (written_bytes, checksum) = process_sorted_row_partition(
+                    row_num,
+                    batch_size as usize,
+                    row_addresses_ptr,
+                    row_sizes_ptr,
+                    &data_types,
+                    output_path,
+                    prefer_dictionary_ratio,
+                    checksum_enabled,
+                    checksum_algo,
+                    current_checksum,
+                    &compression_codec,
+                )?;
+
+                let checksum = if let Some(checksum) = checksum {
+                    checksum as i64
+                } else {
+                    // Spark checksums (CRC32 or Adler32) are both u32, so we use i64::MIN to indicate
+                    // checksum is not available.
+                    i64::MIN
+                };
+
+                let long_array = env.new_long_array(2)?;
+                env.set_long_array_region(&long_array, 0, &[written_bytes, checksum])?;
+
+                Ok(long_array.into_raw())
+            },
+        )
     })
 }
 
@@ -643,11 +653,17 @@ pub extern "system" fn Java_org_apache_comet_Native_sortRowPartitionsNative(
     tracing_enabled: jboolean,
 ) {
     try_unwrap_or_throw(&e, |_| {
-        let _ = TraceGuard::new("sortRowPartitionsNative", tracing_enabled != JNI_FALSE);
-        // SAFETY: JVM unsafe memory allocation is aligned with long.
-        let array = unsafe { std::slice::from_raw_parts_mut(address as *mut i64, size as usize) };
-        array.rdxsort();
-        Ok(())
+        with_trace(
+            "sortRowPartitionsNative",
+            tracing_enabled != JNI_FALSE,
+            || {
+                // SAFETY: JVM unsafe memory allocation is aligned with long.
+                let array =
+                    unsafe { std::slice::from_raw_parts_mut(address as *mut i64, size as usize) };
+                array.rdxsort();
+                Ok(())
+            },
+        )
     })
 }
 
@@ -665,12 +681,13 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock(
     tracing_enabled: jboolean,
 ) -> jlong {
     try_unwrap_or_throw(&e, |mut env| {
-        let _ = TraceGuard::new("decodeShuffleBlock", tracing_enabled != JNI_FALSE);
-        let raw_pointer = env.get_direct_buffer_address(&byte_buffer)?;
-        let length = length as usize;
-        let slice: &[u8] = unsafe { std::slice::from_raw_parts(raw_pointer, length) };
-        let batch = read_ipc_compressed(slice)?;
-        prepare_output(&mut env, array_addrs, schema_addrs, batch, false)
+        with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || {
+            let raw_pointer = env.get_direct_buffer_address(&byte_buffer)?;
+            let length = length as usize;
+            let slice: &[u8] = unsafe { std::slice::from_raw_parts(raw_pointer, length) };
+            let batch = read_ipc_compressed(slice)?;
+            prepare_output(&mut env, array_addrs, schema_addrs, batch, false)
+        })
     })
 }
 
@@ -707,7 +724,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_traceEnd(
 #[no_mangle]
 /// # Safety
 /// This function is inherently unsafe since it deals with raw pointers passed from JNI.
-pub unsafe extern "system" fn Java_org_apache_comet_Native_logCounter(
+pub unsafe extern "system" fn Java_org_apache_comet_Native_logMemoryUsage(
     e: JNIEnv,
     _class: JClass,
     name: jstring,
@@ -715,7 +732,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_logCounter(
 ) {
     try_unwrap_or_throw(&e, |mut env| {
         let name: String = env.get_string(&JString::from_raw(name)).unwrap().into();
-        log_counter(&name, value as u64);
+        log_memory_usage(&name, value as u64);
         Ok(())
     })
 }
diff --git a/native/core/src/execution/shuffle/shuffle_writer.rs b/native/core/src/execution/shuffle/shuffle_writer.rs
index 24f4d48104..e2619eed70 100644
--- a/native/core/src/execution/shuffle/shuffle_writer.rs
+++ b/native/core/src/execution/shuffle/shuffle_writer.rs
@@ -18,7 +18,7 @@
 //! Defines the External shuffle repartition plan.
 
 use crate::execution::shuffle::{CompressionCodec, ShuffleBlockWriter};
-use crate::execution::tracing::TraceGuard;
+use crate::execution::tracing::{with_trace, with_trace_async};
 use arrow::compute::interleave_record_batch;
 use async_trait::async_trait;
 use datafusion::common::utils::proxy::VecAllocExt;
@@ -215,45 +215,48 @@ async fn external_shuffle(
     codec: CompressionCodec,
     tracing_enabled: bool,
 ) -> Result<SendableRecordBatchStream> {
-    let _ = TraceGuard::new("external_shuffle", tracing_enabled);
-
-    let schema = input.schema();
+    with_trace_async("external_shuffle", tracing_enabled, || async {
+        let schema = input.schema();
+
+        let mut repartitioner: Box<dyn ShufflePartitioner> = match &partitioning {
+            any if any.partition_count() == 1 => {
+                Box::new(SinglePartitionShufflePartitioner::try_new(
+                    output_data_file,
+                    output_index_file,
+                    Arc::clone(&schema),
+                    metrics,
+                    context.session_config().batch_size(),
+                    codec,
+                )?)
+            }
+            _ => Box::new(MultiPartitionShuffleRepartitioner::try_new(
+                partition,
+                output_data_file,
+                output_index_file,
+                Arc::clone(&schema),
+                partitioning,
+                metrics,
+                context.runtime_env(),
+                context.session_config().batch_size(),
+                codec,
+                tracing_enabled,
+            )?),
+        };
 
-    let mut repartitioner: Box<dyn ShufflePartitioner> = match &partitioning {
-        any if any.partition_count() == 1 => Box::new(SinglePartitionShufflePartitioner::try_new(
-            output_data_file,
-            output_index_file,
-            Arc::clone(&schema),
-            metrics,
-            context.session_config().batch_size(),
-            codec,
-        )?),
-        _ => Box::new(MultiPartitionShuffleRepartitioner::try_new(
-            partition,
-            output_data_file,
-            output_index_file,
-            Arc::clone(&schema),
-            partitioning,
-            metrics,
-            context.runtime_env(),
-            context.session_config().batch_size(),
-            codec,
-            tracing_enabled,
-        )?),
-    };
-
-    while let Some(batch) = input.next().await {
-        // Block on the repartitioner to insert the batch and shuffle the rows
-        // into the corresponding partition buffer.
-        // Otherwise, pull the next batch from the input stream might overwrite the
-        // current batch in the repartitioner.
-        block_on(repartitioner.insert_batch(batch?))?;
-    }
+        while let Some(batch) = input.next().await {
+            // Block on the repartitioner to insert the batch and shuffle the rows
+            // into the corresponding partition buffer.
+            // Otherwise, pull the next batch from the input stream might overwrite the
+            // current batch in the repartitioner.
+            block_on(repartitioner.insert_batch(batch?))?;
+        }
 
-    repartitioner.shuffle_write().await?;
+        repartitioner.shuffle_write()?;
 
-    // shuffle writer always has empty output
-    Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))))
+        // shuffle writer always has empty output
+        Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))) as SendableRecordBatchStream)
+    })
+    .await
 }
 
 struct ShuffleRepartitionerMetrics {
@@ -306,7 +309,7 @@ trait ShufflePartitioner: Send + Sync {
     /// Insert a batch into the partitioner
     async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()>;
     /// Write shuffle data and shuffle index file to disk
-    async fn shuffle_write(&mut self) -> Result<()>;
+    fn shuffle_write(&mut self) -> Result<()>;
 }
 
 /// A partitioner that uses a hash function to partition data into multiple partitions
@@ -559,7 +562,7 @@ impl MultiPartitionShuffleRepartitioner {
             result
         };
         if grow_result.is_err() {
-            self.spill().await?;
+            self.spill()?;
         }
 
         Ok(())
@@ -611,9 +614,7 @@ impl MultiPartitionShuffleRepartitioner {
         PartitionedBatchesProducer::new(buffered_batches, indices, self.batch_size)
     }
 
-    async fn spill(&mut self) -> Result<()> {
-        let _ = TraceGuard::new("spill", self.tracing_enabled);
-
+    fn spill(&mut self) -> Result<()> {
         log::debug!(
             "ShuffleRepartitioner spilling shuffle data of {} to disk while inserting ({} time(s) so far)",
             self.used(),
@@ -625,22 +626,24 @@ impl MultiPartitionShuffleRepartitioner {
             return Ok(());
         }
 
-        let num_output_partitions = self.partition_writers.len();
-        let mut partitioned_batches = self.partitioned_batches();
-        let mut spilled_bytes = 0;
+        with_trace("shuffle_spill", self.tracing_enabled, || {
+            let num_output_partitions = self.partition_writers.len();
+            let mut partitioned_batches = self.partitioned_batches();
+            let mut spilled_bytes = 0;
 
-        for partition_id in 0..num_output_partitions {
-            let partition_writer = &mut self.partition_writers[partition_id];
-            let mut iter = partitioned_batches.produce(partition_id);
-            spilled_bytes += partition_writer.spill(&mut iter, &self.runtime, &self.metrics)?;
-        }
+            for partition_id in 0..num_output_partitions {
+                let partition_writer = &mut self.partition_writers[partition_id];
+                let mut iter = partitioned_batches.produce(partition_id);
+                spilled_bytes += partition_writer.spill(&mut iter, &self.runtime, &self.metrics)?;
+            }
 
-        let mut timer = self.metrics.mempool_time.timer();
-        self.reservation.free();
-        timer.stop();
-        self.metrics.spill_count.add(1);
-        self.metrics.spilled_bytes.add(spilled_bytes);
-        Ok(())
+            let mut timer = self.metrics.mempool_time.timer();
+            self.reservation.free();
+            timer.stop();
+            self.metrics.spill_count.add(1);
+            self.metrics.spilled_bytes.add(spilled_bytes);
+            Ok(())
+        })
     }
 }
 
@@ -650,96 +653,97 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartitioner {
     /// This function will slice input batch according to configured batch size and then
     /// shuffle rows into corresponding partition buffer.
     async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()> {
-        let _ = TraceGuard::new("insert_batch", self.tracing_enabled);
-        let start_time = Instant::now();
-        let mut start = 0;
-        while start < batch.num_rows() {
-            let end = (start + self.batch_size).min(batch.num_rows());
-            let batch = batch.slice(start, end - start);
-            self.partitioning_batch(batch).await?;
-            start = end;
-        }
-        self.metrics.input_batches.add(1);
-        self.metrics
-            .baseline
-            .elapsed_compute()
-            .add_duration(start_time.elapsed());
-        Ok(())
+        with_trace_async("shuffle_insert_batch", self.tracing_enabled, || async {
+            let start_time = Instant::now();
+            let mut start = 0;
+            while start < batch.num_rows() {
+                let end = (start + self.batch_size).min(batch.num_rows());
+                let batch = batch.slice(start, end - start);
+                self.partitioning_batch(batch).await?;
+                start = end;
+            }
+            self.metrics.input_batches.add(1);
+            self.metrics
+                .baseline
+                .elapsed_compute()
+                .add_duration(start_time.elapsed());
+            Ok(())
+        })
+        .await
     }
 
     /// Writes buffered shuffled record batches into Arrow IPC bytes.
-    async fn shuffle_write(&mut self) -> Result<()> {
-        let _ = TraceGuard::new("shuffle_write", self.tracing_enabled);
+    fn shuffle_write(&mut self) -> Result<()> {
+        with_trace("shuffle_write", self.tracing_enabled, || {
+            let start_time = Instant::now();
 
-        let start_time = Instant::now();
-
-        let mut partitioned_batches = self.partitioned_batches();
-        let num_output_partitions = self.partition_indices.len();
-        let mut offsets = vec![0; num_output_partitions + 1];
-
-        let data_file = self.output_data_file.clone();
-        let index_file = self.output_index_file.clone();
-
-        let output_data = OpenOptions::new()
-            .write(true)
-            .create(true)
-            .truncate(true)
-            .open(data_file)
-            .map_err(|e| DataFusionError::Execution(format!("shuffle write error: {:?}", e)))?;
+            let mut partitioned_batches = self.partitioned_batches();
+            let num_output_partitions = self.partition_indices.len();
+            let mut offsets = vec![0; num_output_partitions + 1];
 
-        let mut output_data = BufWriter::new(output_data);
+            let data_file = self.output_data_file.clone();
+            let index_file = self.output_index_file.clone();
 
-        #[allow(clippy::needless_range_loop)]
-        for i in 0..num_output_partitions {
-            offsets[i] = output_data.stream_position()?;
+            let output_data = OpenOptions::new()
+                .write(true)
+                .create(true)
+                .truncate(true)
+                .open(data_file)
+                .map_err(|e| DataFusionError::Execution(format!("shuffle write error: {:?}", e)))?;
+
+            let mut output_data = BufWriter::new(output_data);
+
+            #[allow(clippy::needless_range_loop)]
+            for i in 0..num_output_partitions {
+                offsets[i] = output_data.stream_position()?;
+
+                // if we wrote a spill file for this partition then copy the
+                // contents into the shuffle file
+                if let Some(spill_data) = self.partition_writers[i].spill_file.as_ref() {
+                    let mut spill_file =
+                        BufReader::new(File::open(spill_data.temp_file.path()).map_err(to_df_err)?);
+                    let mut write_timer = self.metrics.write_time.timer();
+                    std::io::copy(&mut spill_file, &mut output_data).map_err(to_df_err)?;
+                    write_timer.stop();
+                }
 
-            // if we wrote a spill file for this partition then copy the
-            // contents into the shuffle file
-            if let Some(spill_data) = self.partition_writers[i].spill_file.as_ref() {
-                let mut spill_file =
-                    BufReader::new(File::open(spill_data.temp_file.path()).map_err(to_df_err)?);
-                let mut write_timer = self.metrics.write_time.timer();
-                std::io::copy(&mut spill_file, &mut output_data).map_err(to_df_err)?;
-                write_timer.stop();
+                // Write in memory batches to output data file
+                let mut partition_iter = partitioned_batches.produce(i);
+                Self::shuffle_write_partition(
+                    &mut partition_iter,
+                    &mut self.shuffle_block_writer,
+                    &mut output_data,
+                    &self.metrics.encode_time,
+                    &self.metrics.write_time,
+                )?;
             }
 
-            // Write in memory batches to output data file
-            let mut partition_iter = partitioned_batches.produce(i);
-            Self::shuffle_write_partition(
-                &mut partition_iter,
-                &mut self.shuffle_block_writer,
-                &mut output_data,
-                &self.metrics.encode_time,
-                &self.metrics.write_time,
-            )?;
-        }
-
-        let mut write_timer = self.metrics.write_time.timer();
-        output_data.flush()?;
-        write_timer.stop();
+            let mut write_timer = self.metrics.write_time.timer();
+            output_data.flush()?;
+            write_timer.stop();
 
-        // add one extra offset at last to ease partition length computation
-        offsets[num_output_partitions] = output_data.stream_position().map_err(to_df_err)?;
+            // add one extra offset at last to ease partition length computation
+            offsets[num_output_partitions] = output_data.stream_position().map_err(to_df_err)?;
 
-        let mut write_timer = self.metrics.write_time.timer();
-        let mut output_index =
-            BufWriter::new(File::create(index_file).map_err(|e| {
+            let mut write_timer = self.metrics.write_time.timer();
+            let mut output_index = BufWriter::new(File::create(index_file).map_err(|e| {
                 DataFusionError::Execution(format!("shuffle write error: {:?}", e))
             })?);
-        for offset in offsets {
-            output_index
-                .write_all(&(offset as i64).to_le_bytes()[..])
-                .map_err(to_df_err)?;
-        }
-        output_index.flush()?;
-        write_timer.stop();
+            for offset in offsets {
+                output_index
+                    .write_all(&(offset as i64).to_le_bytes()[..])
+                    .map_err(to_df_err)?;
+            }
+            output_index.flush()?;
+            write_timer.stop();
 
-        self.metrics
-            .baseline
-            .elapsed_compute()
-            .add_duration(start_time.elapsed());
+            self.metrics
+                .baseline
+                .elapsed_compute()
+                .add_duration(start_time.elapsed());
 
-        Ok(())
+            Ok(())
+        })
     }
 }
 
@@ -883,7 +887,7 @@ impl ShufflePartitioner for SinglePartitionShufflePartitioner {
         Ok(())
     }
 
-    async fn shuffle_write(&mut self) -> Result<()> {
+    fn shuffle_write(&mut self) -> Result<()> {
         let start_time = Instant::now();
         let concatenated_batch = self.concat_buffered_batches()?;
 
@@ -1265,7 +1269,7 @@ mod test {
         assert!(repartitioner.partition_writers[0].spill_file.is_none());
         assert!(repartitioner.partition_writers[1].spill_file.is_none());
 
-        repartitioner.spill().await.unwrap();
+        repartitioner.spill().unwrap();
 
         // after spill, there should be spill files
         assert!(repartitioner.partition_writers[0].spill_file.is_some());
diff --git a/native/core/src/execution/tracing.rs b/native/core/src/execution/tracing.rs
index 8925f429dd..0555e6b147 100644
--- a/native/core/src/execution/tracing.rs
+++ b/native/core/src/execution/tracing.rs
@@ -36,13 +36,15 @@ impl Recorder {
             .create(true)
             .append(true)
             .open("comet-event-trace.json")
-            .unwrap();
+            .expect("Error writing tracing");
 
         let mut writer = BufWriter::new(file);
 
         // Write start of JSON array. Note that there is no requirement to write
         // the closing ']'.
-        writer.write_all("[ ".as_bytes()).unwrap();
+        writer
+            .write_all("[ ".as_bytes())
+            .expect("Error writing tracing");
         Self {
             now: Instant::now(),
             writer: Arc::new(Mutex::new(writer)),
@@ -56,14 +58,17 @@ impl Recorder {
         self.log_event(name, "E")
     }
 
-    pub fn log_counter(&self, name: &str, value: u64) {
+    pub fn log_memory_usage(&self, name: &str, usage_bytes: u64) {
+        let usage_mb = (usage_bytes as f64 / 1024.0 / 1024.0) as usize;
         let json = format!(
-            "{{ \"name\": \"{name}\", \"cat\": \"PERF\", \"ph\": \"C\", \"pid\": 1, \"tid\": {}, \"ts\": {}, \"args\": {{ \"{name}\": {value} }} }},\n",
+            "{{ \"name\": \"{name}\", \"cat\": \"PERF\", \"ph\": \"C\", \"pid\": 1, \"tid\": {}, \"ts\": {}, \"args\": {{ \"{name}\": {usage_mb} }} }},\n",
             Self::get_thread_id(),
-            self.now.elapsed().as_nanos()
+            self.now.elapsed().as_micros()
         );
         let mut writer = self.writer.lock().unwrap();
-        writer.write_all(json.as_bytes()).unwrap();
+        writer
+            .write_all(json.as_bytes())
+            .expect("Error writing tracing");
     }
 
     fn log_event(&self, name: &str, ph: &str) {
@@ -71,10 +76,12 @@ impl Recorder {
             "{{ \"name\": \"{}\", \"cat\": \"PERF\", \"ph\": \"{ph}\", \"pid\": 1, \"tid\": {}, \"ts\": {} }},\n",
             name,
             Self::get_thread_id(),
-            self.now.elapsed().as_nanos()
+            self.now.elapsed().as_micros()
         );
         let mut writer = self.writer.lock().unwrap();
-        writer.write_all(json.as_bytes()).unwrap();
+        writer
+            .write_all(json.as_bytes())
+            .expect("Error writing tracing");
     }
 
     fn get_thread_id() -> u64 {
@@ -83,7 +90,7 @@ impl Recorder {
             .trim_start_matches("ThreadId(")
             .trim_end_matches(")")
             .parse()
-            .unwrap()
+            .expect("Error parsing thread id")
     }
 }
 
@@ -95,30 +102,41 @@ pub(crate) fn trace_end(name: &str) {
     RECORDER.end_task(name);
 }
 
-pub(crate) fn log_counter(name: &str, value: u64) {
-    RECORDER.log_counter(name, value);
-}
-pub(crate) struct TraceGuard<'a> {
-    label: &'a str,
-    tracing_enabled: bool,
+pub(crate) fn log_memory_usage(name: &str, value: u64) {
+    RECORDER.log_memory_usage(name, value);
 }
 
-impl<'a> TraceGuard<'a> {
-    pub fn new(label: &'a str, tracing_enabled: bool) -> Self {
-        if tracing_enabled {
-            trace_begin(label);
-        }
-        Self {
-            label,
-            tracing_enabled,
-        }
+pub(crate) fn with_trace<T, F>(label: &str, tracing_enabled: bool, f: F) -> T
+where
+    F: FnOnce() -> T,
+{
+    if tracing_enabled {
+        trace_begin(label);
+    }
+
+    let result = f();
+
+    if tracing_enabled {
+        trace_end(label);
     }
+
+    result
 }
 
-impl Drop for TraceGuard<'_> {
-    fn drop(&mut self) {
-        if self.tracing_enabled {
-            trace_end(self.label);
-        }
+pub(crate) async fn with_trace_async<F, Fut, T>(label: &str, tracing_enabled: bool, f: F) -> T
+where
+    F: FnOnce() -> Fut,
+    Fut: std::future::Future<Output = T>,
+{
+    if tracing_enabled {
+        trace_begin(label);
+    }
+
+    let result = f().await;
+
+    if tracing_enabled {
+        trace_end(label);
     }
+
+    result
 }
diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
index 219bc67c24..23c753dd5e 100644
--- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
+++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
@@ -40,7 +40,7 @@ public class CometTaskMemoryManager {
   /** The id uniquely identifies the native plan this memory manager is associated to */
   private final long id;
 
-  private final TaskMemoryManager internal;
+  public final TaskMemoryManager internal;
   private final NativeMemoryConsumer nativeMemoryConsumer;
   private final AtomicLong used = new AtomicLong();
 
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
index c939e7428e..a58ec7851b 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java
@@ -234,7 +234,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
 
       Native _native = new Native();
       if (tracingEnabled) {
-        _native.logCounter("BypassMergeShortShuffle", allocator.getUsed());
+        _native.logMemoryUsage("comet_shuffle_", allocator.getUsed());
       }
 
       long spillRecords = 0;
@@ -247,7 +247,7 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
       }
 
       if (tracingEnabled) {
-        _native.logCounter("BypassMergeShortShuffle", allocator.getUsed());
+        _native.logMemoryUsage("comet_shuffle_", allocator.getUsed());
       }
 
       if (outputRows != spillRecords) {
diff --git a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
index 686633e7f7..552a189ab4 100644
--- a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
+++ b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometUnsafeShuffleWriter.java
@@ -213,13 +213,13 @@ public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOEx
     if (tracingEnabled) {
       nativeLib.traceBegin("CometUnsafeShuffleWriter");
     }
-    String offheapMemKey = "offheap_shuffle_" + Thread.currentThread().getId();
+    String offheapMemKey = "comet_shuffle_" + Thread.currentThread().getId();
     try {
       while (records.hasNext()) {
         insertRecordIntoSorter(records.next());
       }
       if (tracingEnabled) {
-        nativeLib.logCounter(offheapMemKey, this.allocator.getUsed());
+        nativeLib.logMemoryUsage(offheapMemKey, this.allocator.getUsed());
       }
       closeAndWriteOutput();
       success = true;
@@ -242,7 +242,7 @@ public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOEx
         }
       }
       if (tracingEnabled) {
-        nativeLib.logCounter(offheapMemKey, this.allocator.getUsed());
+        nativeLib.logMemoryUsage(offheapMemKey, this.allocator.getUsed());
       }
     }
   }
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index 7dc65687ab..349c9cd581 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -28,6 +28,7 @@ import org.apache.spark.sql.comet.CometMetricNode
 import org.apache.spark.sql.vectorized._
 
 import org.apache.comet.CometConf.{COMET_BATCH_SIZE, COMET_DEBUG_ENABLED, COMET_EXEC_MEMORY_POOL_TYPE, COMET_EXPLAIN_NATIVE_ENABLED, COMET_METRICS_UPDATE_INTERVAL}
+import org.apache.comet.Tracing.withTrace
 import org.apache.comet.vector.NativeUtil
 
 /**
@@ -60,6 +61,7 @@ class CometExecIterator(
     with Logging {
 
   private val tracingEnabled = CometConf.COMET_TRACING_ENABLED.get()
+  private val memoryMXBean = ManagementFactory.getMemoryMXBean
   private val nativeLib = new Native()
   private val nativeUtil = new NativeUtil()
   private val cometTaskMemoryManager = new CometTaskMemoryManager(id)
@@ -134,33 +136,20 @@ class CometExecIterator(
   def getNextBatch(): Option[ColumnarBatch] = {
     assert(partitionIndex >= 0 && partitionIndex < numParts)
 
-    // TODO introduce withTracing method to avoid the try/finish pattern here
-    try {
-      if (tracingEnabled) {
-        nativeLib.traceBegin("CometExecIterator_getNextBatch")
-        val memoryMXBean = ManagementFactory.getMemoryMXBean
-        val heap = memoryMXBean.getHeapMemoryUsage
-        nativeLib.logCounter("jvm_heapUsed", heap.getUsed)
-      }
-
-      nativeUtil.getNextBatch(
-        numOutputCols,
-        (arrayAddrs, schemaAddrs) => {
-          val ctx = TaskContext.get()
-          nativeLib.executePlan(
-            ctx.stageId(),
-            partitionIndex,
-            plan,
-            arrayAddrs,
-            schemaAddrs,
-            tracingEnabled)
-        })
-    } finally {
-      if (tracingEnabled) {
-        nativeLib.traceEnd("CometExecIterator_getNextBatch")
-        nativeLib.logCounter("CometTaskMemoryManager", cometTaskMemoryManager.getUsed)
-      }
+    if (tracingEnabled) {
+      traceMemoryUsage()
     }
+
+    val ctx = TaskContext.get()
+    withTrace(
+      s"getNextBatch[JVM] stage=${ctx.stageId()}",
+      tracingEnabled, {
+        nativeUtil.getNextBatch(
+          numOutputCols,
+          (arrayAddrs, schemaAddrs) => {
+            nativeLib.executePlan(ctx.stageId(), partitionIndex, plan, arrayAddrs, schemaAddrs)
+          })
+      })
   }
 
   override def hasNext: Boolean = {
@@ -213,8 +202,9 @@ class CometExecIterator(
       }
       nativeUtil.close()
       nativeLib.releasePlan(plan)
+
       if (tracingEnabled) {
-        nativeLib.logCounter("CometTaskMemoryManager", cometTaskMemoryManager.getUsed)
+        traceMemoryUsage()
       }
 
       // The allocator thoughts the exported ArrowArray and ArrowSchema structs are not released,
@@ -239,4 +229,14 @@ class CometExecIterator(
       closed = true
     }
   }
+
+  private def traceMemoryUsage(): Unit = {
+    nativeLib.logMemoryUsage("jvm_heapUsed", memoryMXBean.getHeapMemoryUsage.getUsed)
+    val totalTaskMemory = cometTaskMemoryManager.internal.getMemoryConsumptionForThisTask
+    val cometTaskMemory = cometTaskMemoryManager.getUsed
+    val sparkTaskMemory = totalTaskMemory - cometTaskMemory
+    val threadId = Thread.currentThread().getId
+    nativeLib.logMemoryUsage(s"task_memory_comet_$threadId", cometTaskMemory)
+    nativeLib.logMemoryUsage(s"task_memory_spark_$threadId", sparkTaskMemory)
+  }
 }
diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala
index 0a08e29eeb..afaa3d17db 100644
--- a/spark/src/main/scala/org/apache/comet/Native.scala
+++ b/spark/src/main/scala/org/apache/comet/Native.scala
@@ -91,8 +91,7 @@ class Native extends NativeBase {
       partition: Int,
       plan: Long,
       arrayAddrs: Array[Long],
-      schemaAddrs: Array[Long],
-      tracingEnabled: Boolean): Long
+      schemaAddrs: Array[Long]): Long
 
   /**
    * Release and drop the native query plan object and context object.
@@ -178,10 +177,28 @@ class Native extends NativeBase {
       schemaAddrs: Array[Long],
       tracingEnabled: Boolean): Long
 
+  /**
+   * Log the beginning of an event.
+   * @param name
+   *   The name of the event.
+   */
   @native def traceBegin(name: String): Unit
 
+  /**
+   * Log the end of an event.
+   * @param name
+   *   The name of the event.
+   */
   @native def traceEnd(name: String): Unit
 
-  @native def logCounter(name: String, value: Long): Unit
+  /**
+   * Log the amount of memory currently in use.
+   *
+   * @param name
+   *   Type of memory e.g. jvm, native, off-heap pool
+   * @param memoryUsageBytes
+   *   Number of bytes in use
+   */
+  @native def logMemoryUsage(name: String, memoryUsageBytes: Long): Unit
 
 }
diff --git a/spark/src/main/scala/org/apache/comet/Tracing.scala b/spark/src/main/scala/org/apache/comet/Tracing.scala
new file mode 100644
index 0000000000..13f44ce270
--- /dev/null
+++ b/spark/src/main/scala/org/apache/comet/Tracing.scala
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet
+
+object Tracing {
+
+  private val nativeLib = new Native
+
+  def withTrace[T](label: String, tracingEnabled: Boolean, fun: => T): T = {
+    try {
+      if (tracingEnabled) {
+        nativeLib.traceBegin(label)
+      }
+      fun
+    } finally {
+      if (tracingEnabled) {
+        nativeLib.traceEnd(label)
+      }
+    }
+  }
+
+}


Reviews:
user: login: andygrove
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T01:06:41Z
user: login: comphead
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T01:40:41Z
user: login: comphead
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T01:43:17Z
user: login: andygrove
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T12:29:02Z
user: login: andygrove
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T12:36:18Z
user: login: comphead
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T21:23:05Z
user: login: comphead
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T21:25:45Z
user: login: comphead
body: lgtm thanks @andygrove 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T21:27:23Z
user: login: andygrove
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T21:27:47Z
user: login: comphead
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-13T21:36:40Z
user: login: andygrove
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-13T21:43:59Z

Review comments:
diff_hunk: @@ -95,30 +102,41 @@ pub(crate) fn trace_end(name: &str) {
     RECORDER.end_task(name);
 }
 
-pub(crate) fn log_counter(name: &str, value: u64) {
-    RECORDER.log_counter(name, value);
-}
-pub(crate) struct TraceGuard<'a> {
-    label: &'a str,
-    tracing_enabled: bool,
+pub(crate) fn log_memory_usage(name: &str, value: u64) {
+    RECORDER.log_memory_usage(name, value);
 }
 
-impl<'a> TraceGuard<'a> {
-    pub fn new(label: &'a str, tracing_enabled: bool) -> Self {
-        if tracing_enabled {
-            trace_begin(label);
-        }
-        Self {
-            label,
-            tracing_enabled,
-        }
+pub(crate) fn with_trace<T, F>(label: &str, tracing_enabled: bool, f: F) -> T
path: native/core/src/execution/tracing.rs
user: login: andygrove
body: @comphead as you requested on the original PR, I had now added `with_trace`.
created_at: 2025-05-13T01:06:40Z
author_association: MEMBER
diff_hunk: @@ -650,96 +653,97 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartitioner {
     /// This function will slice input batch according to configured batch size and then
     /// shuffle rows into corresponding partition buffer.
     async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()> {
-        let _ = TraceGuard::new("insert_batch", self.tracing_enabled);
-        let start_time = Instant::now();
-        let mut start = 0;
-        while start < batch.num_rows() {
-            let end = (start + self.batch_size).min(batch.num_rows());
-            let batch = batch.slice(start, end - start);
-            self.partitioning_batch(batch).await?;
-            start = end;
-        }
-        self.metrics.input_batches.add(1);
-        self.metrics
-            .baseline
-            .elapsed_compute()
-            .add_duration(start_time.elapsed());
-        Ok(())
+        with_trace_async("shuffle_insert_batch", self.tracing_enabled, || async {
+            let start_time = Instant::now();
+            let mut start = 0;
+            while start < batch.num_rows() {
+                let end = (start + self.batch_size).min(batch.num_rows());
+                let batch = batch.slice(start, end - start);
+                self.partitioning_batch(batch).await?;
+                start = end;
+            }
+            self.metrics.input_batches.add(1);
+            self.metrics
+                .baseline
+                .elapsed_compute()
+                .add_duration(start_time.elapsed());
+            Ok(())
+        })
+        .await
     }
 
     /// Writes buffered shuffled record batches into Arrow IPC bytes.
     async fn shuffle_write(&mut self) -> Result<()> {
-        let _ = TraceGuard::new("shuffle_write", self.tracing_enabled);
+        with_trace("shuffle_write", self.tracing_enabled, || {
path: native/core/src/execution/shuffle/shuffle_writer.rs
user: login: comphead
body: ```suggestion
        with_trace_async("shuffle_write", self.tracing_enabled, || {
```

?
created_at: 2025-05-13T01:40:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -625,22 +626,24 @@ impl MultiPartitionShuffleRepartitioner {
             return Ok(());
         }
 
-        let num_output_partitions = self.partition_writers.len();
-        let mut partitioned_batches = self.partitioned_batches();
-        let mut spilled_bytes = 0;
+        with_trace("shuffle_spill", self.tracing_enabled, || {
path: native/core/src/execution/shuffle/shuffle_writer.rs
user: login: comphead
body: ```suggestion
        with_trace_async("shuffle_spill", self.tracing_enabled, || {
```

I just assumed if method is async then the `trace_async` needs to be called like in other similar methods?
created_at: 2025-05-13T01:43:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -625,22 +626,24 @@ impl MultiPartitionShuffleRepartitioner {
             return Ok(());
         }
 
-        let num_output_partitions = self.partition_writers.len();
-        let mut partitioned_batches = self.partitioned_batches();
-        let mut spilled_bytes = 0;
+        with_trace("shuffle_spill", self.tracing_enabled, || {
path: native/core/src/execution/shuffle/shuffle_writer.rs
user: login: andygrove
body: The closure here does not call any async code. The `spill` method was marked as `async` but didn't need to be so I have now updated that.
created_at: 2025-05-13T12:29:01Z
author_association: MEMBER
diff_hunk: @@ -650,96 +653,97 @@ impl ShufflePartitioner for MultiPartitionShuffleRepartitioner {
     /// This function will slice input batch according to configured batch size and then
     /// shuffle rows into corresponding partition buffer.
     async fn insert_batch(&mut self, batch: RecordBatch) -> Result<()> {
-        let _ = TraceGuard::new("insert_batch", self.tracing_enabled);
-        let start_time = Instant::now();
-        let mut start = 0;
-        while start < batch.num_rows() {
-            let end = (start + self.batch_size).min(batch.num_rows());
-            let batch = batch.slice(start, end - start);
-            self.partitioning_batch(batch).await?;
-            start = end;
-        }
-        self.metrics.input_batches.add(1);
-        self.metrics
-            .baseline
-            .elapsed_compute()
-            .add_duration(start_time.elapsed());
-        Ok(())
+        with_trace_async("shuffle_insert_batch", self.tracing_enabled, || async {
+            let start_time = Instant::now();
+            let mut start = 0;
+            while start < batch.num_rows() {
+                let end = (start + self.batch_size).min(batch.num_rows());
+                let batch = batch.slice(start, end - start);
+                self.partitioning_batch(batch).await?;
+                start = end;
+            }
+            self.metrics.input_batches.add(1);
+            self.metrics
+                .baseline
+                .elapsed_compute()
+                .add_duration(start_time.elapsed());
+            Ok(())
+        })
+        .await
     }
 
     /// Writes buffered shuffled record batches into Arrow IPC bytes.
     async fn shuffle_write(&mut self) -> Result<()> {
-        let _ = TraceGuard::new("shuffle_write", self.tracing_enabled);
+        with_trace("shuffle_write", self.tracing_enabled, || {
path: native/core/src/execution/shuffle/shuffle_writer.rs
user: login: andygrove
body: There is no async code in `shuffle_write`, so I removed the `async` keyword.
created_at: 2025-05-13T12:36:18Z
author_association: MEMBER
diff_hunk: @@ -43,10 +45,20 @@ Example output:
 { "name": "CometExecIterator_getNextBatch", "cat": "PERF", "ph": "B", "pid": 1, "tid": 5, "ts": 10109479156 },
 ```
 
-Traces can be viewed with [Trace Viewer]. 
+Traces can be viewed with [Trace Viewer].
 
 [Trace Viewer]: https://github.com/catapult-project/catapult/blob/main/tracing/README.md
 
 Example trace visualization:
 
 ![tracing](../_static/images/tracing.png)
+
+## Definition of Labels
+
+| Label                 | Meaning                                                        |
+| --------------------- | -------------------------------------------------------------- |
+| jvm_heapUsed          | JVM heap memory usage of live objects for the executor process |
+| jemalloc_allocated    | Native memory usage for the executor process                   |
+| task_memory_comet_NNN | Off-heap memory allocated by Comet for query execution         |
path: docs/source/contributor-guide/tracing.md
user: login: comphead
body: what is NNN is it a task number or thread number or query number?
created_at: 2025-05-13T21:23:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -40,7 +40,7 @@ public class CometTaskMemoryManager {
   /** The id uniquely identifies the native plan this memory manager is associated to */
   private final long id;
 
-  private final TaskMemoryManager internal;
+  public final TaskMemoryManager internal;
path: spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
user: login: comphead
body: we prob should have another name for that? 
created_at: 2025-05-13T21:25:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -43,10 +45,20 @@ Example output:
 { "name": "CometExecIterator_getNextBatch", "cat": "PERF", "ph": "B", "pid": 1, "tid": 5, "ts": 10109479156 },
 ```
 
-Traces can be viewed with [Trace Viewer]. 
+Traces can be viewed with [Trace Viewer].
 
 [Trace Viewer]: https://github.com/catapult-project/catapult/blob/main/tracing/README.md
 
 Example trace visualization:
 
 ![tracing](../_static/images/tracing.png)
+
+## Definition of Labels
+
+| Label                 | Meaning                                                        |
+| --------------------- | -------------------------------------------------------------- |
+| jvm_heapUsed          | JVM heap memory usage of live objects for the executor process |
+| jemalloc_allocated    | Native memory usage for the executor process                   |
+| task_memory_comet_NNN | Off-heap memory allocated by Comet for query execution         |
path: docs/source/contributor-guide/tracing.md
user: login: andygrove
body: Thread number
created_at: 2025-05-13T21:27:47Z
author_association: MEMBER
diff_hunk: @@ -43,10 +45,20 @@ Example output:
 { "name": "CometExecIterator_getNextBatch", "cat": "PERF", "ph": "B", "pid": 1, "tid": 5, "ts": 10109479156 },
 ```
 
-Traces can be viewed with [Trace Viewer]. 
+Traces can be viewed with [Trace Viewer].
 
 [Trace Viewer]: https://github.com/catapult-project/catapult/blob/main/tracing/README.md
 
 Example trace visualization:
 
 ![tracing](../_static/images/tracing.png)
+
+## Definition of Labels
+
+| Label                 | Meaning                                                        |
+| --------------------- | -------------------------------------------------------------- |
+| jvm_heapUsed          | JVM heap memory usage of live objects for the executor process |
+| jemalloc_allocated    | Native memory usage for the executor process                   |
+| task_memory_comet_NNN | Off-heap memory allocated by Comet for query execution         |
path: docs/source/contributor-guide/tracing.md
user: login: comphead
body: > Thread number

it might be beneficial for the user to mention it

created_at: 2025-05-13T21:36:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -40,7 +40,7 @@ public class CometTaskMemoryManager {
   /** The id uniquely identifies the native plan this memory manager is associated to */
   private final long id;
 
-  private final TaskMemoryManager internal;
+  public final TaskMemoryManager internal;
path: spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
user: login: andygrove
body: yes, probably
created_at: 2025-05-13T21:43:59Z
author_association: MEMBER

PR comments:
user: login: codecov-commenter
created_at: 2025-05-10T16:48:24Z
body: ## [Codecov](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `46.87500%` with `17 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 58.65%. Comparing base ([`f09f8af`](https://app.codecov.io/gh/apache/datafusion-comet/commit/f09f8af64c6599255e116a376f4f008f2fd63b43?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`d1c77c4`](https://app.codecov.io/gh/apache/datafusion-comet/commit/d1c77c45072226dc0ddb92f14665ade4a71baaec?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 1150 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...ain/scala/org/apache/comet/CometExecIterator.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fcomet%2FCometExecIterator.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9jb21ldC9Db21ldEV4ZWNJdGVyYXRvci5zY2FsYQ==) | 47.36% | [9 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...park/src/main/scala/org/apache/comet/Tracing.scala](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fcomet%2FTracing.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9jb21ldC9UcmFjaW5nLnNjYWxh) | 62.50% | [1 Missing and 2 partials :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...ion/shuffle/CometBypassMergeSortShuffleWriter.java](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fspark%2Fsql%2Fcomet%2Fexecution%2Fshuffle%2FCometBypassMergeSortShuffleWriter.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL3NwYXJrL3NxbC9jb21ldC9leGVjdXRpb24vc2h1ZmZsZS9Db21ldEJ5cGFzc01lcmdlU29ydFNodWZmbGVXcml0ZXIuamF2YQ==) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...et/execution/shuffle/CometUnsafeShuffleWriter.java](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&filepath=spark%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fspark%2Fsql%2Fcomet%2Fexecution%2Fshuffle%2FCometUnsafeShuffleWriter.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-c3Bhcmsvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL3NwYXJrL3NxbC9jb21ldC9leGVjdXRpb24vc2h1ZmZsZS9Db21ldFVuc2FmZVNodWZmbGVXcml0ZXIuamF2YQ==) | 33.33% | [2 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |

<details><summary>Additional details and impacted files</summary>



```diff
@@             Coverage Diff              @@
##               main    #1730      +/-   ##
============================================
+ Coverage     56.12%   58.65%   +2.52%     
- Complexity      976     1139     +163     
============================================
  Files           119      130      +11     
  Lines         11743    12699     +956     
  Branches       2251     2376     +125     
============================================
+ Hits           6591     7448     +857     
- Misses         4012     4062      +50     
- Partials       1140     1189      +49     
```
</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/apache/datafusion-comet/pull/1730?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: andygrove
created_at: 2025-05-13T21:44:10Z
body: Thanks for the reviews @comphead 
author_association: MEMBER