diff --git a/core/src/main/scala/org/apache/spark/util/random/SamplingUtils.scala b/core/src/main/scala/org/apache/spark/util/random/SamplingUtils.scala
index f17ac7fb6a8e3..0a8c59bbaae96 100644
--- a/core/src/main/scala/org/apache/spark/util/random/SamplingUtils.scala
+++ b/core/src/main/scala/org/apache/spark/util/random/SamplingUtils.scala
@@ -56,12 +56,13 @@ private[spark] object SamplingUtils {
       val rand = new XORShiftRandom(seed)
       while (input.hasNext) {
         val item = input.next()
-        val replacementIndex = l < Int.MaxValue match {
-          case true => rand.nextInt(l.toInt)
-          case false => rand.nextInt()
+        val replacementIndex = if (l < Int.MaxValue) {
+          rand.nextInt(l.toInt)
+        } else {
+          rand.nextLong()
         }
         if (replacementIndex < k) {
-          reservoir(replacementIndex) = item
+          reservoir(replacementIndex.toInt) = item
         }
         l += 1
       }