From 3a97115ec31bcb0639e8b6fad38c473751c4a508 Mon Sep 17 00:00:00 2001 From: lgbo-ustc Date: Tue, 20 Aug 2024 16:44:34 +0800 Subject: [PATCH] fix: total_bytes_written is not updated in celeborn partition writer --- cpp-ch/local-engine/Shuffle/PartitionWriter.cpp | 2 ++ cpp-ch/local-engine/local_engine_jni.cpp | 6 ++++++ 2 files changed, 8 insertions(+) diff --git a/cpp-ch/local-engine/Shuffle/PartitionWriter.cpp b/cpp-ch/local-engine/Shuffle/PartitionWriter.cpp index 2f22d0e24139e..79d640d3b2bca 100644 --- 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 index 828556b4abf66..4118f95768e82 100644 --- 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(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) }