python/pyspark/ml/tests/connect/test_connect_model_offloading.py - master - githubmirror/spark

python/pyspark/ml/tests/connect/test_connect_model_offloading.py
312 строк10 KB
Tian Gao
[SPARK-54906][PYTHON][TESTS] Unify test entry for all pyspark tests
07 янв 2026, 04:21
07 янв 2026, 04:21c27ede3
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import json

import numpy as np

from pyspark.sql import functions as sf
from pyspark.ml.linalg import Vectors
from pyspark.ml.classification import (
    LinearSVC,
    LinearSVCSummary,
    LinearSVCTrainingSummary,
)
from pyspark.ml.regression import (
    LinearRegression,
    LinearRegressionSummary,
    LinearRegressionTrainingSummary,
)
from pyspark.ml.clustering import (
    LDA,
    LDAModel,
    LocalLDAModel,
    DistributedLDAModel,
)
from pyspark.ml.fpm import (
    FPGrowth,
    FPGrowthModel,
)
from pyspark.testing.connectutils import ReusedConnectTestCase


class ModelOffloadingTests(ReusedConnectTestCase):
    def test_linear_svc_offloading(self):
        # force clean up the ml cache
        self.spark.client._cleanup_ml_cache()

        df = (
            self.spark.createDataFrame(
                [
                    (1.0, 1.0, Vectors.dense(0.0, 5.0)),
                    (0.0, 2.0, Vectors.dense(1.0, 2.0)),
                    (1.0, 3.0, Vectors.dense(2.0, 1.0)),
                    (0.0, 4.0, Vectors.dense(3.0, 3.0)),
                ],
                ["label", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("weight")
        )
        vec = Vectors.dense(0.0, 5.0)

        svc = LinearSVC(maxIter=1, regParam=1.0)
        self.assertEqual(svc.getMaxIter(), 1)
        self.assertEqual(svc.getRegParam(), 1.0)

        model = svc.fit(df)

        def check_model(m):
            self.assertEqual(svc.uid, m.uid)
            self.assertEqual(m.numClasses, 2)
            self.assertEqual(m.predict(vec), 1.0)

            self.assertTrue(m.hasSummary)
            summary = m.summary()

            self.assertIsInstance(summary, LinearSVCSummary)
            self.assertIsInstance(summary, LinearSVCTrainingSummary)
            self.assertEqual(summary.labels, [0.0, 1.0])

        # model is cached!
        # '{"id":"xxx","class":"org.apache.spark.ml.classification.LinearSVCModel","size":xxx}'
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 1, cached)
        self.assertEqual(
            json.loads(cached[0])["class"],
            "org.apache.spark.ml.classification.LinearSVCModel",
            cached,
        )

        check_model(model)

        # model is offloaded!
        self.spark.client._delete_ml_cache([model._java_obj._ref_id], evict_only=True)
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 0, cached)

        check_model(model)

    def test_linear_regression_offloading(self):
        # force clean up the ml cache
        self.spark.client._cleanup_ml_cache()

        df = (
            self.spark.createDataFrame(
                [
                    (1.0, 1.0, Vectors.dense(0.0, 5.0)),
                    (0.0, 2.0, Vectors.dense(1.0, 2.0)),
                    (1.5, 3.0, Vectors.dense(2.0, 1.0)),
                    (0.7, 4.0, Vectors.dense(1.5, 3.0)),
                ],
                ["label", "weight", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("weight")
        )
        vec = Vectors.dense(0.0, 5.0)

        lr = LinearRegression(
            regParam=0.0,
            maxIter=2,
            solver="normal",
            weightCol="weight",
        )
        self.assertEqual(lr.getRegParam(), 0)
        self.assertEqual(lr.getMaxIter(), 2)

        model = lr.fit(df)

        def check_model(m):
            self.assertEqual(lr.uid, m.uid)
            self.assertEqual(m.numFeatures, 2)
            self.assertTrue(np.allclose(m.predict(vec), 0.21249999999999963, atol=1e-4))

            summary = m.summary
            self.assertTrue(isinstance(summary, LinearRegressionSummary))
            self.assertTrue(isinstance(summary, LinearRegressionTrainingSummary))
            self.assertEqual(summary.predictions.count(), 4)

        # model is cached!
        # '{"id":"xxx","class":"org.apache.spark.ml.regression.LinearRegressionModel","size":xxx}'
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 1, cached)
        self.assertEqual(
            json.loads(cached[0])["class"],
            "org.apache.spark.ml.regression.LinearRegressionModel",
            cached,
        )

        check_model(model)

        # model is offloaded!
        self.spark.client._delete_ml_cache([model._java_obj._ref_id], evict_only=True)
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 0, cached)

        check_model(model)

    def test_lda_offloading(self):
        # force clean up the ml cache
        self.spark.client._cleanup_ml_cache()

        df = (
            self.spark.createDataFrame(
                [
                    [1, Vectors.dense([0.0, 1.0])],
                    [2, Vectors.sparse(2, {0: 1.0})],
                ],
                ["id", "features"],
            )
            .coalesce(1)
            .sortWithinPartitions("id")
        )

        lda = LDA(k=2, optimizer="em", seed=1)
        lda.setMaxIter(1)

        self.assertEqual(lda.getK(), 2)
        self.assertEqual(lda.getOptimizer(), "em")
        self.assertEqual(lda.getMaxIter(), 1)
        self.assertEqual(lda.getSeed(), 1)

        model = lda.fit(df)

        def check_model(m):
            self.assertEqual(lda.uid, m.uid)
            self.assertIsInstance(m, LDAModel)
            self.assertNotIsInstance(m, LocalLDAModel)
            self.assertIsInstance(m, DistributedLDAModel)
            self.assertTrue(m.isDistributed())
            self.assertEqual(m.vocabSize(), 2)

            output = m.transform(df)
            expected_cols = ["id", "features", "topicDistribution"]
            self.assertEqual(output.columns, expected_cols)
            self.assertEqual(output.count(), 2)

        # model is cached!
        # '{"id":"xxx","class":"org.apache.spark.ml.clustering.DistributedLDAModel","size":xxx}'
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 1, cached)
        self.assertEqual(
            json.loads(cached[0])["class"],
            "org.apache.spark.ml.clustering.DistributedLDAModel",
            cached,
        )

        check_model(model)

        # both model and local_model are is cached!
        local_model = model.toLocal()
        # '{"id":"xxx","class":"org.apache.spark.ml.clustering.LocalLDAModel","size":xxx}'
        # '{"id":"xxx","class":"org.apache.spark.ml.clustering.DistributedLDAModel","size":xxx}'
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 2, cached)
        self.assertEqual(
            sorted([json.loads(c)["class"] for c in cached]),
            [
                "org.apache.spark.ml.clustering.DistributedLDAModel",
                "org.apache.spark.ml.clustering.LocalLDAModel",
            ],
            cached,
        )

        def check_local_model(m):
            self.assertIsInstance(m, LDAModel)
            self.assertIsInstance(m, LocalLDAModel)
            self.assertNotIsInstance(m, DistributedLDAModel)
            self.assertFalse(m.isDistributed())
            self.assertEqual(m.vocabSize(), 2)

            output = m.transform(df)
            expected_cols = ["id", "features", "topicDistribution"]
            self.assertEqual(output.columns, expected_cols)
            self.assertEqual(output.count(), 2)

        check_local_model(local_model)

        # both model and local_model are offloaded!
        self.spark.client._delete_ml_cache([model._java_obj._ref_id], evict_only=True)
        self.spark.client._delete_ml_cache([local_model._java_obj._ref_id], evict_only=True)
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 0, cached)

        check_model(model)
        check_local_model(local_model)

    def test_fp_growth_offloading(self):
        # force clean up the ml cache
        self.spark.client._cleanup_ml_cache()

        df = self.spark.createDataFrame(
            [
                ["r z h k p"],
                ["z y x w v u t s"],
                ["s x o n r"],
                ["x z y m t s q e"],
                ["z"],
                ["x z y r q t p"],
            ],
            ["items"],
        ).select(sf.split("items", " ").alias("items"))

        fp = FPGrowth(minSupport=0.2, minConfidence=0.7)
        fp.setNumPartitions(1)

        model = fp.fit(df)

        def check_model(m):
            self.assertIsInstance(m, FPGrowthModel)
            self.assertEqual(fp.uid, m.uid)
            self.assertEqual(m.freqItemsets.columns, ["items", "freq"])
            self.assertEqual(m.freqItemsets.count(), 54)

            self.assertEqual(
                m.associationRules.columns,
                ["antecedent", "consequent", "confidence", "lift", "support"],
            )
            self.assertEqual(m.associationRules.count(), 89)

            output = m.transform(df)
            self.assertEqual(output.columns, ["items", "prediction"])
            self.assertEqual(output.count(), 6)

        # model is cached!
        # '{"id":"xxx","class":"org.apache.spark.ml.fpm.FPGrowthModel","size":xxx}'
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 1, cached)
        self.assertEqual(
            json.loads(cached[0])["class"],
            "org.apache.spark.ml.fpm.FPGrowthModel",
            cached,
        )

        check_model(model)

        # model is offloaded!
        self.spark.client._delete_ml_cache([model._java_obj._ref_id], evict_only=True)
        cached = self.spark.client._get_ml_cache_info()
        self.assertEqual(len(cached), 0, cached)

        check_model(model)


if __name__ == "__main__":
    from pyspark.testing import main

    main()