Reduce duplicates

Author: boneanxs

shims/common/pom.xml

@@ -39,6 +39,13 @@
       <artifactId>spark-hive_${scala.binary.version}</artifactId>
       <scope>provided</scope>
     </dependency>
+
+    <!--test-->
+    <dependency>
+      <groupId>org.scalatest</groupId>
+      <artifactId>scalatest_${scala.binary.version}</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>

shims/common/src/main/scala/org/apache/gluten/config/GlutenConfig.scala

@@ -17,8 +17,8 @@
 package org.apache.gluten.config
 
 import org.apache.spark.internal.Logging
-import org.apache.spark.network.util.{ByteUnit, JavaUtils}
-import org.apache.spark.sql.internal.{SQLConf, SQLConfProvider}
+import org.apache.spark.network.util.ByteUnit
+import org.apache.spark.sql.internal.{GlutenConfigUtil, SQLConf, SQLConfProvider}
 
 import com.google.common.collect.ImmutableList
 import org.apache.hadoop.security.UserGroupInformation
@@ -447,7 +447,7 @@ object GlutenConfig {
    */
   def getNativeSessionConf(
       backendName: String,
-      conf: scala.collection.Map[String, String]): util.Map[String, String] = {
+      conf: Map[String, String]): util.Map[String, String] = {
     val nativeConfMap = new util.HashMap[String, String]()
     val keys = Set(
       DEBUG_ENABLED.key,
@@ -505,31 +505,17 @@ object GlutenConfig {
       (SPARK_SHUFFLE_SPILL_COMPRESS, SPARK_SHUFFLE_SPILL_COMPRESS_DEFAULT.toString)
     )
     keyWithDefault.forEach(e => nativeConfMap.put(e._1, conf.getOrElse(e._1, e._2)))
-
-    conf
-      .get(SPARK_UNSAFE_SORTER_SPILL_READER_BUFFER_SIZE)
-      .foreach(
-        v =>
-          nativeConfMap
-            .put(
-              SPARK_UNSAFE_SORTER_SPILL_READER_BUFFER_SIZE,
-              JavaUtils.byteStringAs(v, ByteUnit.BYTE).toString))
-    conf
-      .get(SPARK_SHUFFLE_SPILL_DISK_WRITE_BUFFER_SIZE)
-      .foreach(
-        v =>
-          nativeConfMap
-            .put(
-              SPARK_SHUFFLE_SPILL_DISK_WRITE_BUFFER_SIZE,
-              JavaUtils.byteStringAs(v, ByteUnit.BYTE).toString))
-    conf
-      .get(SPARK_SHUFFLE_FILE_BUFFER)
-      .foreach(
-        v =>
-          nativeConfMap
-            .put(
-              SPARK_SHUFFLE_FILE_BUFFER,
-              (JavaUtils.byteStringAs(v, ByteUnit.KiB) * 1024).toString))
+    GlutenConfigUtil.mapByteConfValue(
+      conf,
+      SPARK_UNSAFE_SORTER_SPILL_READER_BUFFER_SIZE,
+      ByteUnit.BYTE)(
+      v => nativeConfMap.put(SPARK_UNSAFE_SORTER_SPILL_READER_BUFFER_SIZE, v.toString))
+    GlutenConfigUtil.mapByteConfValue(
+      conf,
+      SPARK_SHUFFLE_SPILL_DISK_WRITE_BUFFER_SIZE,
+      ByteUnit.BYTE)(v => nativeConfMap.put(SPARK_SHUFFLE_SPILL_DISK_WRITE_BUFFER_SIZE, v.toString))
+    GlutenConfigUtil.mapByteConfValue(conf, SPARK_SHUFFLE_FILE_BUFFER, ByteUnit.KiB)(
+      v => nativeConfMap.put(SPARK_SHUFFLE_FILE_BUFFER, (v * 1024).toString))
 
     conf
       .get(LEGACY_TIME_PARSER_POLICY.key)

shims/common/src/main/scala/org/apache/spark/sql/internal/GlutenConfigUtil.scala

@@ -18,6 +18,8 @@ package org.apache.spark.sql.internal
 
 import org.apache.gluten.config._
 
+import org.apache.spark.network.util.{ByteUnit, JavaUtils}
+
 object GlutenConfigUtil {
   private def getConfString(configProvider: ConfigProvider, key: String, value: String): String = {
     Option(ConfigEntry.findEntry(key))
@@ -42,4 +44,9 @@ object GlutenConfigUtil {
         }
     }.toMap
   }
+
+  def mapByteConfValue(conf: Map[String, String], key: String, unit: ByteUnit)(
+      f: Long => Unit): Unit = {
+    conf.get(key).foreach(v => f(JavaUtils.byteStringAs(v, unit)))
+  }
 }

shims/common/src/test/scala/org/apache/spark/sql/internal/GlutenConfigUtilSuite.scala

@@ -0,0 +1,35 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.internal
+
+import org.apache.spark.network.util.ByteUnit
+
+import org.scalatest.funsuite.AnyFunSuite
+
+class GlutenConfigUtilSuite extends AnyFunSuite {
+
+  test("mapByteConfValue should return correct value") {
+    val conf = Map(
+      "spark.unsafe.sorter.spill.reader.buffer.size" -> "2m"
+    )
+
+    GlutenConfigUtil.mapByteConfValue(
+      conf,
+      "spark.unsafe.sorter.spill.reader.buffer.size",
+      ByteUnit.BYTE)(v => assert(2097152L.equals(v)))
+  }
+}