cache multiple inputs and split

apache · Jul 3, 2024 · 26dc7b6 · 26dc7b6
1 parent a0f4fd2
commit 26dc7b6
Show file tree

Hide file tree

Showing 13 changed files with 343 additions and 288 deletions.
diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
@@ -953,7 +953,8 @@ JNIEXPORT jlong JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrappe
 JNIEXPORT jobject JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrapper_stop( // NOLINT
     JNIEnv* env,
     jobject wrapper,
-    jlong shuffleWriterHandle) {
+    jlong shuffleWriterHandle,
+    jlong memLimit) {
   JNI_METHOD_START
   auto ctx = gluten::getRuntime(env, wrapper);
 
@@ -963,7 +964,7 @@ JNIEXPORT jobject JNICALL Java_org_apache_gluten_vectorized_ShuffleWriterJniWrap
     throw gluten::GlutenException(errorMessage);
   }
 
-  gluten::arrowAssertOkOrThrow(shuffleWriter->stop(), "Native shuffle write: ShuffleWriter stop failed");
+  gluten::arrowAssertOkOrThrow(shuffleWriter->stop(memLimit), "Native shuffle write: ShuffleWriter stop failed");
 
   const auto& partitionLengths = shuffleWriter->partitionLengths();
   auto partitionLengthArr = env->NewLongArray(partitionLengths.size());

diff --git a/cpp/core/shuffle/ShuffleWriter.h b/cpp/core/shuffle/ShuffleWriter.h
@@ -39,7 +39,7 @@ class ShuffleWriter : public Reclaimable {
 
   virtual arrow::Status write(std::shared_ptr<ColumnarBatch> cb, int64_t memLimit) = 0;
 
-  virtual arrow::Status stop() = 0;
+  virtual arrow::Status stop(int64_t memLimit) = 0;
 
   int32_t numPartitions() const {
     return numPartitions_;