zhztheplayer commented on code in PR #9944:
URL: https://github.com/apache/incubator-gluten/pull/9944#discussion_r2151661018
##########
cpp/core/jni/JniWrapper.cc:
##########
@@ -780,141 +780,223 @@ JNIEXPORT jlong JNICALL
Java_org_apache_gluten_columnarbatch_ColumnarBatchJniWra
JNI_METHOD_END(kInvalidObjectHandle)
}
-// Shuffle
-JNIEXPORT jlong JNICALL
Java_org_apache_gluten_vectorized_ShuffleWriterJniWrapper_nativeMake( // NOLINT
+JNIEXPORT jlong JNICALL
+Java_org_apache_gluten_vectorized_LocalPartitionWriterJniWrapper_createPartitionWriter(
// NOLINT
JNIEnv* env,
jobject wrapper,
- jstring partitioningNameJstr,
jint numPartitions,
- jint bufferSize,
- jint mergeBufferSize,
- jdouble mergeThreshold,
jstring codecJstr,
jstring codecBackendJstr,
jint compressionLevel,
jint compressionBufferSize,
- jint diskWriteBufferSize,
jint compressionThreshold,
- jstring compressionModeJstr,
- jint initialSortBufferSize,
- jboolean useRadixSort,
- jstring dataFileJstr,
+ jint mergeBufferSize,
+ jdouble mergeThreshold,
jint numSubDirs,
- jstring localDirsJstr,
- jdouble reallocThreshold,
- jlong firstBatchHandle,
- jlong taskAttemptId,
- jint startPartitionId,
+ jint shuffleFileBufferSize,
+ jstring dataFileJstr,
+ jstring localDirsJstr) {
+ JNI_METHOD_START
+
+ const auto ctx = getRuntime(env, wrapper);
+
+ auto dataFile = jStringToCString(env, dataFileJstr);
+ auto localDirs = splitPaths(jStringToCString(env, localDirsJstr));
+
+ auto partitionWriterOptions = std::make_shared<LocalPartitionWriterOptions>(
+ shuffleFileBufferSize, compressionBufferSize, compressionThreshold,
mergeBufferSize, mergeThreshold, numSubDirs);
+
+ auto partitionWriter = std::make_shared<LocalPartitionWriter>(
+ numPartitions,
+ createArrowIpcCodec(getCompressionType(env, codecJstr),
getCodecBackend(env, codecBackendJstr), compressionLevel),
+ ctx->memoryManager(),
+ partitionWriterOptions,
+ dataFile,
+ std::move(localDirs));
+
+ return ctx->saveObject(partitionWriter);
+ JNI_METHOD_END(kInvalidObjectHandle)
+}
+
+JNIEXPORT jlong JNICALL
+Java_org_apache_gluten_vectorized_CelebornPartitionWriterJniWrapper_createPartitionWriter(
// NOLINT
Review Comment:
(minor)
The function could be moved to VeloxJniWrapper.cc given the relevant Java
JNI wrapper is moved, ditto for the one for Uniffle.
Once https://github.com/apache/incubator-gluten/issues/9639 is done we will
no longer need to manage these module mappings.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
