This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 9f027a981f20 [SPARK-58546][ML] Avoid redundant StringIndexer label
lookups
9f027a981f20 is described below
commit 9f027a981f20283bdc493ee3ed550fdcafceeb3b
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Tue Aug 4 20:24:53 2026 +0800
[SPARK-58546][ML] Avoid redundant StringIndexer label lookups
### What changes were proposed in this pull request?
This PR passes `keepInvalid` into `StringIndexerModel.getIndexer`, computes
it once during transformation, and uses one `OpenHashMap.get` lookup for
non-null labels.
### Why are the changes needed?
The previous UDF performed `contains` followed by map access for known
labels, causing two lookups per value. This removes the redundant lookup while
preserving invalid-label behavior.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Not run (not requested). Verified with `git diff --check`, ASCII scan, and
changed-line length scan.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57745 from zhengruifeng/single-lookup-string-indexer.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 1f2b47d48843e955f71ff006598b327182576b93)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../apache/spark/ml/feature/StringIndexer.scala | 46 +++++++++++-----------
1 file changed, 24 insertions(+), 22 deletions(-)
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
index c0f7e312dcbb..bfb78adfd42d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -363,28 +363,32 @@ class StringIndexerModel (
.where(conditions.reduce(_ and _))
}
- private def getIndexer(labels: Seq[String], labelToIndex:
OpenHashMap[String, Double]) = {
- val keepInvalid = (getHandleInvalid == StringIndexer.KEEP_INVALID)
-
- udf { label: String =>
- if (label == null) {
- if (keepInvalid) {
- labels.length
+ private def getIndexer(
+ labels: Seq[String],
+ labelToIndex: OpenHashMap[String, Double],
+ keepInvalid: Boolean) = {
+ val unknownIndex = labels.length.toDouble
+ if (keepInvalid) {
+ udf { label: String =>
+ if (label == null) {
+ unknownIndex
} else {
+ labelToIndex.get(label).getOrElse(unknownIndex)
+ }
+ }.asNondeterministic()
+ } else {
+ udf { label: String =>
+ if (label == null) {
throw new SparkException("StringIndexer encountered NULL value. To
handle or skip " +
"NULLS, try setting StringIndexer.handleInvalid.")
- }
- } else {
- if (labelToIndex.contains(label)) {
- labelToIndex(label)
- } else if (keepInvalid) {
- labels.length
} else {
- throw new SparkException(s"Unseen label: $label. To handle unseen
labels, " +
- s"set Param handleInvalid to ${StringIndexer.KEEP_INVALID}.")
+ labelToIndex.get(label).getOrElse {
+ throw new SparkException(s"Unseen label: $label. To handle unseen
labels, " +
+ s"set Param handleInvalid to ${StringIndexer.KEEP_INVALID}.")
+ }
}
- }
- }.asNondeterministic()
+ }.asNondeterministic()
+ }
}
@Since("2.0.0")
@@ -400,6 +404,7 @@ class StringIndexerModel (
map
}
val outputColumns = new Array[Column](outputColNames.length)
+ val keepInvalid = getHandleInvalid == StringIndexer.KEEP_INVALID
// Skips invalid rows if `handleInvalid` is set to
`StringIndexer.SKIP_INVALID`.
val filteredDataset = if (getHandleInvalid == StringIndexer.SKIP_INVALID) {
@@ -416,16 +421,13 @@ class StringIndexerModel (
try {
dataset.col(inputColName)
- val filteredLabels = getHandleInvalid match {
- case StringIndexer.KEEP_INVALID => labels :+ "__unknown"
- case _ => labels
- }
+ val filteredLabels = if (keepInvalid) labels :+ "__unknown" else labels
val metadata = NominalAttribute.defaultAttr
.withName(outputColName)
.withValues(filteredLabels)
.toMetadata()
- val indexer = getIndexer(labels.toImmutableArraySeq, labelToIndex)
+ val indexer = getIndexer(labels.toImmutableArraySeq, labelToIndex,
keepInvalid)
outputColumns(i) = indexer(dataset(inputColName).cast(StringType))
.as(outputColName, metadata)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
