[SPARK-16933] [ML] Fix AFTAggregator in AFTSurvivalRegression serializes unnecessary data.

[yanboliang]

What changes were proposed in this pull request?

Similar to LeastSquaresAggregator in #14109, AFTAggregator used for AFTSurvivalRegression ends up serializing the parameters and featuresStd, which is not necessary and can cause performance issues for high dimensional data. This patch removes this serialization. This PR is highly inspired by #14109.

How was this patch tested?

I tested this locally and verified the serialization reduction.

Before patch

After patch

[SparkQA]

Test build #63314 has finished for PR 14519 at commit d152a3a.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[yanboliang]

cc @sethah @dbtsai

[srowen]

Let's put this into #14109

[dbtsai]

In line 506,

  private val gradientSumArray = Array.ofDim[Double](parameters.length)

the code will evaluate the lazy parameters in the driver.

[dbtsai]

BTW, after thinking a bit, some of the lazy is not needed. lazy is for avoiding doing computation in the driver; however
@transient private val parameters = bcParameters.value should work without lazy. Also, sigma or intercept may not need lazy. Thanks.

[yanboliang]

@dbtsai I addressed the parameters.length issue. But I can not remove lazy from @transient private lazy val parameters = bcParameters.value and intercept/sigma. Otherwise, it complains NullPointerException. If I removed both @transient and lazy, it works well, but this does not coincide with our requirements. It's a little weird and I'm still work on to figure out the root cause, can you give me some suggestion? Thanks.

[dbtsai]

You are right. In scala, when we use @transient private val, that lazy evaluation will be only evaluated once even after serialization/deserialization cycle. As a result, after the AFTAggregator is broadcasted into executors, the variable will be be evaluated again, and will be default to null.

[SparkQA]

Test build #63362 has finished for PR 14519 at commit 287f153.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[sethah]

@yanboliang Can you do a quick test to make sure the shuffle write size is the expected size? For example, in logistic regression only the gradient should be serialized which is an array of numFeatures doubles. The expected shuffle write size is then roughly numFeatures * 8 bytes for each task. It would be nice to check before/after.

[dbtsai]

LGTM. Will be nice to see the compassion of shuffle write size, and then will be ready to merge. Thanks.

[yanboliang]

I tested this locally and verified the serialization reduction. I posted the shuffle size comparison diagram in the PR description. I will merge this into master. Thanks for your review! @dbtsai @sethah