@@ -31,6 +31,7 @@ import org.roaringbitmap.{ArrayContainer, BitmapContainer, RoaringArray, Roaring
31
31
import org .apache .spark ._
32
32
import org .apache .spark .api .python .PythonBroadcast
33
33
import org .apache .spark .broadcast .HttpBroadcast
34
+ import org .apache .spark .network .util .ByteUnit
34
35
import org .apache .spark .network .nio .{GetBlock , GotBlock , PutBlock }
35
36
import org .apache .spark .scheduler .{CompressedMapStatus , HighlyCompressedMapStatus }
36
37
import org .apache .spark .storage ._
@@ -51,18 +52,18 @@ class KryoSerializer(conf: SparkConf)
51
52
52
53
private val bufferSizeKb = conf.getSizeAsKb(" spark.kryoserializer.buffer" , " 64k" )
53
54
54
- if (bufferSizeKb >= 2048 * 1024 ) {
55
+ if (bufferSizeKb >= ByteUnit . GiB .toKiB( 2 ) ) {
55
56
throw new IllegalArgumentException (" spark.kryoserializer.buffer must be less than " +
56
- s " 2048 mb, got: + ${bufferSizeKb/ 1024 } mb. " )
57
+ s " 2048 mb, got: + ${ByteUnit . KiB .toMiB( bufferSizeKb) } mb. " )
57
58
}
58
- private val bufferSize = (bufferSizeKb * 1024 ).toInt
59
+ private val bufferSize = ByteUnit . KiB .toBytes (bufferSizeKb).toInt
59
60
60
61
val maxBufferSizeMb = conf.getSizeAsMb(" spark.kryoserializer.buffer.max" , " 64m" ).toInt
61
- if (maxBufferSizeMb >= 2048 ) {
62
+ if (maxBufferSizeMb >= ByteUnit . GiB .toMiB( 2 ) ) {
62
63
throw new IllegalArgumentException (" spark.kryoserializer.buffer.max must be less than " +
63
64
s " 2048 mb, got: + $maxBufferSizeMb mb. " )
64
65
}
65
- private val maxBufferSize = maxBufferSizeMb * 1024 * 1024
66
+ private val maxBufferSize = ByteUnit . MiB .toBytes( maxBufferSizeMb).toInt
66
67
67
68
private val referenceTracking = conf.getBoolean(" spark.kryo.referenceTracking" , true )
68
69
private val registrationRequired = conf.getBoolean(" spark.kryo.registrationRequired" , false )
0 commit comments