[SPARK-58279][ML][CONNECT] Estimate size of TargetEncoder, VectorIndexer, CountVectorizer, and MinHashLSH models#57448
Conversation
zhengruifeng
changed the title
zhengruifeng
changed the title
| // For ml connect only | ||
| private[ml] def this() = this("", -1, Map.empty) | ||
|
|
||
| private[spark] override def estimatedSize: Long = { |
There was a problem hiding this comment.
Note: VectorIndexerModel.estimatedSize counts only categoryMaps and omits partialFeatureAttributes, an eager private val (not lazy) built in the primary constructor as an Array[Attribute] of numFeatures allocated attribute objects. Because the old default estimated a parentless copy (new VectorIndexerModel(uid, numFeatures, categoryMaps), which re-materializes partialFeatureAttributes), that array was previously counted and is now dropped from the estimate. numFeatures is not bounded by categoryMaps.size, so a high-dimensional model (e.g. numFeatures=1000, few categorical) under-counts by ~1000 attribute objects. The merged sibling StringIndexerModel (SPARK-58251) sets the precedent by explicitly counting its derived eager field: SizeEstimator.estimate((labelsArray, labelsToIndexArray)). Fix is one line: estimateMatadataSize + SizeEstimator.estimate((categoryMaps, partialFeatureAttributes)).
There was a problem hiding this comment.
Good point. partialFeatureAttributes is eagerly constructed derived state, so the current override can undercount it compared with the previous parentless-copy estimate. I will created a sub-task https://issues.apache.org/jira/browse/SPARK-58293 to defer its initialization and work on that follow-up soon. The follow-up will ensure size estimation accounts for the state when it is materialized.
uros-b
left a comment
uros-b
left a comment
There was a problem hiding this comment.
Thank you @zhengruifeng, I left just one comment - otherwise LGTM
…xer, CountVectorizer, and MinHashLSH models ### What changes were proposed in this pull request? Adds explicit `estimatedSize` implementations for `TargetEncoderModel`, `VectorIndexerModel`, `CountVectorizerModel`, and `MinHashLSHModel`. Each override accounts for model metadata and the model-owned lookup state: encoder statistics, category maps, vocabulary, or random hash coefficients. Adds fitted-model regression tests in the owning algorithm suites with a 4 KiB upper bound. ### Why are the changes needed? The default reflection-based estimate can traverse substantially more state than the fitted model owns. Direct accounting keeps model-size estimation small, predictable, and focused on the learned state. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? - `JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt 'mllib/Test/compile'`\n\n### Was this patch authored or co-authored using generative AI tooling?\n\nGenerated-by: Codex GPT-5 Closes #57448 from zhengruifeng/feature-lookup-size-estimates-dev3. Authored-by: Ruifeng Zheng <ruifengz@apache.org> Signed-off-by: Ruifeng Zheng <ruifengz@apache.org> (cherry picked from commit b3a4544) Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>
3 participants
What changes were proposed in this pull request?
Adds explicit
estimatedSizeimplementations forTargetEncoderModel,VectorIndexerModel,CountVectorizerModel, andMinHashLSHModel. Each override accounts for model metadata and the model-owned lookup state: encoder statistics, category maps, vocabulary, or random hash coefficients.Adds fitted-model regression tests in the owning algorithm suites with a 4 KiB upper bound.
Why are the changes needed?
The default reflection-based estimate can traverse substantially more state than the fitted model owns. Direct accounting keeps model-size estimation small, predictable, and focused on the learned state.
Does this PR introduce any user-facing change?
No.
How was this patch tested?
JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt 'mllib/Test/compile'\n\n### Was this patch authored or co-authored using generative AI tooling?\n\nGenerated-by: Codex GPT-5