[SPARK-48892][ML] Avoid per-row param read in Tokenizer

[zhengruifeng]

What changes were proposed in this pull request?

Inspired by #47258, I am checking other ML implementations, and find that we can also optimize Tokenizer in the same way

Why are the changes needed?

the function createTransformFunc is to build the udf for UnaryTransformer.transform:

spark/mllib/src/main/scala/org/apache/spark/ml/Transformer.scala

Line 118 in d679dab

val transformUDF = udf(this.createTransformFunc)

existing implementation read the params for each row.

Does this PR introduce any user-facing change?

No

How was this patch tested?

CI and manually tests:

create test dataset

spark.range(1000000).select(uuid().as("uuid")).write.mode("overwrite").parquet("/tmp/regex_tokenizer.parquet")

duration

val df = spark.read.parquet("/tmp/regex_tokenizer.parquet")
import org.apache.spark.ml.feature._
val tokenizer = new RegexTokenizer().setPattern("-").setInputCol("uuid")
Seq.range(0, 1000).foreach(i => tokenizer.transform(df).count()) // warm up
val tic = System.currentTimeMillis; Seq.range(0, 1000).foreach(i => tokenizer.transform(df).count()); System.currentTimeMillis - tic

result (before this PR)

scala> val tic = System.currentTimeMillis; Seq.range(0, 1000).foreach(i => tokenizer.transform(df).count()); System.currentTimeMillis - tic
val tic: Long = 1720613235068
val res5: Long = 50397

result (after this PR)

scala> val tic = System.currentTimeMillis; Seq.range(0, 1000).foreach(i => tokenizer.transform(df).count()); System.currentTimeMillis - tic
val tic: Long = 1720612871256
val res5: Long = 43748

Was this patch authored or co-authored using generative AI tooling?

No

[srowen]

Looks fine, it adds some complexity, but not much

[JoshRosen]

I'll wager that the expensive part was probably the configuration check itself plus the regex compilation, but not the branching (since those would predict well). Therefore I predict that you can get almost the whole speedup if you did something like

override protected def createTransformFunc: String => Seq[String] = {
  val re = $(pattern).r
  val _toLowercase = $(toLowercase)
  val _gaps = $(gaps)
  val minLength = $(minTokenLength)
  { originStr =>
      // scalastyle:off caselocale
      val str = if (_toLowerCase) originStr.toLowerCase() else originStr
      // scalastyle:on caselocale
      val tokens = if (_gaps) re.split(str).toImmutableArraySeq else re.findAllIn(str).toSeq
      tokens.filter(_.length >= minLength)
    }
}

Basically, I think it might be overkill or unnecessary to fully inline and expand the cross product like this. My suggested approach is easier to understand and probably nearly equivalent in performance.

[zhengruifeng]

I'll wager that the expensive part was probably the configuration check itself plus the regex compilation, but not the branching (since those would predict well). Therefore I predict that you can get almost the whole speedup if you did something like

override protected def createTransformFunc: String => Seq[String] = {
  val re = $(pattern).r
  val _toLowercase = $(toLowercase)
  val _gaps = $(gaps)
  val minLength = $(minTokenLength)
  { originStr =>
      // scalastyle:off caselocale
      val str = if (_toLowerCase) originStr.toLowerCase() else originStr
      // scalastyle:on caselocale
      val tokens = if (_gaps) re.split(str).toImmutableArraySeq else re.findAllIn(str).toSeq
      tokens.filter(_.length >= minLength)
    }
}

Basically, I think it might be overkill or unnecessary to fully inline and expand the cross product like this. My suggested approach is easier to understand and probably nearly equivalent in performance.

make sense, let me make the changes simple

[zhengruifeng]

merged to master