[NSE-647] Leverage buffered write in shuffle (oap-project#648)

Closes oap-project#647
zhouyuan · Jan 8, 2022 · 2f0ee49 · 2f0ee49
1 parent d5a7f81
commit 2f0ee49
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 2 deletions.
diff --git a/native-sql-engine/cpp/src/jni/jni_wrapper.cc b/native-sql-engine/cpp/src/jni/jni_wrapper.cc
@@ -1051,6 +1051,7 @@ Java_com_intel_oap_vectorized_ShuffleSplitterJniWrapper_nativeMake(
   auto splitOptions = SplitOptions::Defaults();
   splitOptions.write_schema = write_schema;
   splitOptions.prefer_spill = prefer_spill;
+  splitOptions.buffered_write = true;
   if (buffer_size > 0) {
     splitOptions.buffer_size = buffer_size;
   }

diff --git a/native-sql-engine/cpp/src/shuffle/splitter.cc b/native-sql-engine/cpp/src/shuffle/splitter.cc
@@ -433,8 +433,15 @@ arrow::Status Splitter::Split(const arrow::RecordBatch& rb) {
 arrow::Status Splitter::Stop() {
   EVAL_START("write", options_.thread_id)
   // open data file output stream
-  ARROW_ASSIGN_OR_RAISE(data_file_os_,
+  std::shared_ptr<arrow::io::FileOutputStream> fout;
+  ARROW_ASSIGN_OR_RAISE(fout,
                         arrow::io::FileOutputStream::Open(options_.data_file, true));
+  if (options_.buffered_write) {
+    ARROW_ASSIGN_OR_RAISE(data_file_os_, arrow::io::BufferedOutputStream::Create(
+                                             16384, options_.memory_pool, fout));
+  } else {
+    data_file_os_ = fout;
+  }
 
   // stop PartitionWriter and collect metrics
   for (auto pid = 0; pid < num_partitions_; ++pid) {

diff --git a/native-sql-engine/cpp/src/shuffle/splitter.h b/native-sql-engine/cpp/src/shuffle/splitter.h
@@ -230,7 +230,7 @@ class Splitter {
   std::vector<int32_t> sub_dir_selection_;
   std::vector<std::string> configured_dirs_;
 
-  std::shared_ptr<arrow::io::FileOutputStream> data_file_os_;
+  std::shared_ptr<arrow::io::OutputStream> data_file_os_;
 
   // shared by all partition writers
   std::shared_ptr<arrow::ipc::IpcPayload> schema_payload_;

diff --git a/native-sql-engine/cpp/src/shuffle/type.h b/native-sql-engine/cpp/src/shuffle/type.h
@@ -44,6 +44,7 @@ struct SplitOptions {
   arrow::Compression::type compression_type = arrow::Compression::UNCOMPRESSED;
   bool prefer_spill = true;
   bool write_schema = true;
+  bool buffered_write = false;
 
   std::string data_file;