fix: total_bytes_written is not updated in celeborn partition writer

apache · Aug 20, 2024 · 3a97115 · 3a97115
1 parent 9fcd488
commit 3a97115
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/cpp-ch/local-engine/Shuffle/PartitionWriter.cpp b/cpp-ch/local-engine/Shuffle/PartitionWriter.cpp
@@ -469,6 +469,7 @@ size_t MemorySortCelebornPartitionWriter::evictPartitions()
                 celeborn_client->pushPartitionData(cur_partition_id, data.data(), data.size());
                 shuffle_writer->split_result.total_io_time += push_time_watch.elapsedNanoseconds();
                 shuffle_writer->split_result.partition_lengths[cur_partition_id] += data.size();
+                shuffle_writer->split_result.total_bytes_written += data.size();
             }
             output.restart();
         };
@@ -586,6 +587,7 @@ size_t CelebornPartitionWriter::evictSinglePartition(size_t partition_id)
         shuffle_writer->split_result.total_write_time += push_time_watch.elapsedNanoseconds();
         shuffle_writer->split_result.total_io_time += push_time_watch.elapsedNanoseconds();
         shuffle_writer->split_result.total_serialize_time += serialization_time_watch.elapsedNanoseconds();
+        shuffle_writer->split_result.total_bytes_written += written_bytes;
     };
 
     Stopwatch spill_time_watch;

diff --git a/cpp-ch/local-engine/local_engine_jni.cpp b/cpp-ch/local-engine/local_engine_jni.cpp
@@ -677,6 +677,11 @@ JNIEXPORT jobject Java_org_apache_gluten_vectorized_CHShuffleSplitterJniWrapper_
     const auto * raw_src = reinterpret_cast<const jlong *>(raw_partition_lengths.data());
     env->SetLongArrayRegion(raw_partition_length_arr, 0, raw_partition_lengths.size(), raw_src);
 
+    // AQE has dependency on total_bytes_written, if the data is wrong, it will generate inappropriate plan
+    // add a log here for remining this.
+    if (!result.total_bytes_written)
+        LOG_WARNING(getLogger("_CHShuffleSplitterJniWrapper"), "total_bytes_written is 0, something may be wrong");
+
     jobject split_result = env->NewObject(
         split_result_class,
         split_result_constructor,
@@ -692,6 +697,7 @@ JNIEXPORT jobject Java_org_apache_gluten_vectorized_CHShuffleSplitterJniWrapper_
         result.total_io_time,
         result.total_serialize_time);
 
+
     return split_result;
     LOCAL_ENGINE_JNI_METHOD_END(env, nullptr)
 }