Skip to content

Commit

Permalink
chore: Use proper constructor of IndexShuffleBlockResolver (apache#610)
Browse files Browse the repository at this point in the history
  • Loading branch information
viirya authored and kazuyukitanimura committed Jul 1, 2024
1 parent 097f8bd commit ae443d8
Showing 1 changed file with 18 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

package org.apache.spark.sql.comet.execution.shuffle

import java.util.Collections
import java.util.concurrent.ConcurrentHashMap

import scala.collection.JavaConverters._
Expand Down Expand Up @@ -61,7 +62,23 @@ class CometShuffleManager(conf: SparkConf) extends ShuffleManager with Logging {

private lazy val shuffleExecutorComponents = loadShuffleExecutorComponents(conf)

override val shuffleBlockResolver = new IndexShuffleBlockResolver(conf)
override val shuffleBlockResolver: IndexShuffleBlockResolver = {
// The patch versions of Spark 3.4 have different constructor signatures:
// See https://github.com/apache/spark/commit/5180694705be3508bd21dd9b863a59b8cb8ba193
// We look for proper constructor by reflection.
classOf[IndexShuffleBlockResolver].getDeclaredConstructors
.filter(c => List(2, 3).contains(c.getParameterCount()))
.map { c =>
c.getParameterCount match {
case 2 =>
c.newInstance(conf, null).asInstanceOf[IndexShuffleBlockResolver]
case 3 =>
c.newInstance(conf, null, Collections.emptyMap())
.asInstanceOf[IndexShuffleBlockResolver]
}
}
.head
}

/**
* (override) Obtains a [[ShuffleHandle]] to pass to tasks.
Expand Down

0 comments on commit ae443d8

Please sign in to comment.