Skip to content

[SPARK-58281][ML][CONNECT] Avoid parent overcounting in PipelineModel size estimates#57451

Open
zhengruifeng wants to merge 6 commits into
apache:masterfrom
zhengruifeng:pipeline-model-size-estimate-dev3
Open

[SPARK-58281][ML][CONNECT] Avoid parent overcounting in PipelineModel size estimates#57451
zhengruifeng wants to merge 6 commits into
apache:masterfrom
zhengruifeng:pipeline-model-size-estimate-dev3

Conversation

@zhengruifeng

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

Adds a PipelineModel.estimatedSize override that sums model-stage estimates individually and directly estimates non-model transformer stages. This prevents a reflection walk of the complete pipeline object graph.

Adds a regression test that fits a pipeline containing StringIndexer and asserts its size estimate remains below 16 KiB.

Why are the changes needed?

SPARK-57521 clears the top-level model parent before the default size walk. A copied PipelineModel still contains copied model stages whose estimator parents are retained, so a reflection walk can still reach shared Spark-session state. Estimating nested model stages independently applies their parentless-copy behavior at every model boundary.

Does this PR introduce any user-facing change?

Yes. Pipeline model cache-size estimates no longer include shared state reachable through nested stage parents, preventing unnecessary cache overcounting.

How was this patch tested?

  • Added a fitted StringIndexer pipeline size-estimation regression test.
  • 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

@zhengruifeng zhengruifeng changed the title [WIP][ML] Avoid parent overcounting in PipelineModel size estimates [SPARK-58281][ML][CONNECT] Avoid parent overcounting in PipelineModel size estimates Jul 23, 2026
@zhengruifeng
zhengruifeng marked this pull request as ready for review July 23, 2026 02:58
"copy should create an instance with the same parent")
}

test("PipelineModel estimated size") {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new regression test covers only a single-stage pipeline whose one stage is a Model (StringIndexer). The case _ => 0L branch; the novel, behavior-changing part of this PR (the old default walk counted non-model stage bytes; the new code skips them) is never exercised. The assertion is also upper-bound-only (estimatedSize < 16 KiB), so a zero-returning implementation would pass it. Consider adding a pipeline that mixes a Model stage with a non-model Transformer stage, plus a lower-bound assertion and a hasParent-preserved check, mirroring the SPARK-57521 tests in ModelSuite (ModelSuite.scala:36-47).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

make sense

@uros-b uros-b left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @zhengruifeng, I left just one comment - otherwise LGTM

@zhengruifeng
zhengruifeng force-pushed the pipeline-model-size-estimate-dev3 branch from f01d197 to 88e5126 Compare July 24, 2026 01:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants