[VL] Fix shuffle ignores uncompressed compression configuration (#3672)

apache · Nov 13, 2023 · 139c8a1 · 139c8a1
1 parent 1f6c6c3
commit 139c8a1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 12 deletions.
diff --git a/cpp/core/jni/JniCommon.h b/cpp/core/jni/JniCommon.h
@@ -170,18 +170,15 @@ static inline arrow::Compression::type getCompressionType(JNIEnv* env, jstring c
   if (codecJstr == NULL) {
     return arrow::Compression::UNCOMPRESSED;
   }
-  auto codecU = env->GetStringUTFChars(codecJstr, JNI_FALSE);
+  auto codec = env->GetStringUTFChars(codecJstr, JNI_FALSE);
 
-  std::string codecL;
-  std::transform(codecU, codecU + std::strlen(codecU), std::back_inserter(codecL), ::tolower);
+  // Convert codec string into lowercase.
+  std::string codecLower;
+  std::transform(codec, codec + std::strlen(codec), std::back_inserter(codecLower), ::tolower);
+  GLUTEN_ASSIGN_OR_THROW(auto compressionType, arrow::util::Codec::GetCompressionType(codecLower));
 
-  GLUTEN_ASSIGN_OR_THROW(auto compression_type, arrow::util::Codec::GetCompressionType(codecL));
-
-  if (compression_type == arrow::Compression::LZ4) {
-    compression_type = arrow::Compression::LZ4_FRAME;
-  }
-  env->ReleaseStringUTFChars(codecJstr, codecU);
-  return compression_type;
+  env->ReleaseStringUTFChars(codecJstr, codec);
+  return compressionType;
 }
 
 static inline gluten::CodecBackend getCodecBackend(JNIEnv* env, jstring codecJstr) {

diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
@@ -785,8 +785,8 @@ JNIEXPORT jlong JNICALL Java_io_glutenproject_vectorized_ShuffleWriterJniWrapper
     shuffleWriterOptions.buffer_size = bufferSize;
   }
 
+  shuffleWriterOptions.compression_type = getCompressionType(env, codecJstr);
   if (codecJstr != NULL) {
-    shuffleWriterOptions.compression_type = getCompressionType(env, codecJstr);
     shuffleWriterOptions.codec_backend = getCodecBackend(env, codecBackendJstr);
     shuffleWriterOptions.compression_mode = getCompressionMode(env, compressionModeJstr);
   }
@@ -991,8 +991,8 @@ JNIEXPORT jlong JNICALL Java_io_glutenproject_vectorized_ShuffleReaderJniWrapper
   ShuffleReaderOptions options = ShuffleReaderOptions::defaults();
   options.ipc_read_options.memory_pool = pool;
   options.ipc_read_options.use_threads = false;
+  options.compression_type = getCompressionType(env, compressionType);
   if (compressionType != nullptr) {
-    options.compression_type = getCompressionType(env, compressionType);
     options.codec_backend = getCodecBackend(env, compressionBackend);
   }
   std::shared_ptr<arrow::Schema> schema =