[WIP][ML][SPARK-51379] Move treeAggregate's final aggregation from driver to executor

[zhengruifeng]

What changes were proposed in this pull request?

Move treeAggregate's final aggregation from driver to executor.

ee20fbb introduced an optimization that:

Move final iteration of aggregation of RDD.treeAggregate to an executor with one partition and fetch that result to the driver

This PR tries to apply this optimization so that less memory is required for ML's treeAggregate.

Why are the changes needed?

to save driver memory

Does this PR introduce any user-facing change?

no

How was this patch tested?

ci and manually check

preparing data:

df = spark.read.format("libsvm").load("data/mllib/sample_libsvm_data.txt")
for i in range(10):
    df = df.union(df)

df.count()
df.repartition(1024).write.mode("overwrite").parquet("/tmp/test_data")

training a lr

from pyspark.ml.classification import *
df = spark.read.parquet("/tmp/test_data")
lr = LogisticRegression()
model = lr.fit(df)

before:

after:

In each iteration, the data sent to driver is reduced from 136.1 KiB to 21.3 KiB

Was this patch authored or co-authored using generative AI tooling?

no