This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 2621d6d24f23 [SPARK-58633][ML] Use integer label indexes in
StringIndexerModel
2621d6d24f23 is described below
commit 2621d6d24f23664d0ee94bcf814f247a1aa055ae
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Mon Aug 10 12:04:02 2026 +0800
[SPARK-58633][ML] Use integer label indexes in StringIndexerModel
### What changes were proposed in this pull request?
This PR changes `StringIndexerModel` label lookup maps from
`OpenHashMap[String, Double]` to `OpenHashMap[String, Int]`. The indexer
UDF returns integer
indexes, and the output column is cast to `DoubleType` at the end of the
transformation.
### Why are the changes needed?
String index values are inherently integral. Storing them as `Int` instead
of `Double` reduces
the memory footprint of the lookup maps captured by `StringIndexerModel`
transformations while
preserving the public double-valued output.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The existing `StringIndexerSuite` covers normal, unseen, skipped, kept,
null, numeric, and
multi-column inputs as well as the `DoubleType` output contract. It passed
all 28 tests:
`build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 'mllib/testOnly
org.apache.spark.ml.feature.StringIndexerSuite'`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57838 from zhengruifeng/SPARK-58633.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../main/scala/org/apache/spark/ml/feature/StringIndexer.scala | 10 +++++-----
1 file changed, 5 insertions(+), 5 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 bfb78adfd42d..c363d3ca2f32 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
@@ -344,7 +344,7 @@ class StringIndexerModel (
private def filterInvalidData(
dataset: Dataset[_],
inputColNames: Seq[String],
- labelsToIndexArray: Array[OpenHashMap[String, Double]]): Dataset[_] = {
+ labelsToIndexArray: Array[OpenHashMap[String, Int]]): Dataset[_] = {
val conditions: Seq[Column] = inputColNames.indices.map { i =>
val inputColName = inputColNames(i)
val labelToIndex = labelsToIndexArray(i)
@@ -365,9 +365,9 @@ class StringIndexerModel (
private def getIndexer(
labels: Seq[String],
- labelToIndex: OpenHashMap[String, Double],
+ labelToIndex: OpenHashMap[String, Int],
keepInvalid: Boolean) = {
- val unknownIndex = labels.length.toDouble
+ val unknownIndex = labels.length
if (keepInvalid) {
udf { label: String =>
if (label == null) {
@@ -397,7 +397,7 @@ class StringIndexerModel (
val (inputColNames, outputColNames) = getInOutCols()
val labelsToIndexArray = labelsArray.map { labels =>
- val map = new OpenHashMap[String, Double](labels.length)
+ val map = new OpenHashMap[String, Int](labels.length)
labels.zipWithIndex.foreach { case (label, idx) =>
map.update(label, idx)
}
@@ -429,7 +429,7 @@ class StringIndexerModel (
val indexer = getIndexer(labels.toImmutableArraySeq, labelToIndex,
keepInvalid)
- outputColumns(i) = indexer(dataset(inputColName).cast(StringType))
+ outputColumns(i) =
indexer(dataset(inputColName).cast(StringType)).cast(DoubleType)
.as(outputColName, metadata)
} catch {
case _: AnalysisException =>
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
