/
githubmirror
/
spark
Обзор
Документация
Войти
/
githubmirror
/
spark
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
python/pyspark/ml/tests/test_clustering.py
519 строк
18 KB
Tian Gao
[SPARK-55020][PYTHON] Disable gc when executing gRPC command
28 янв 2026, 04:01
28 янв 2026, 04:01
e17d273
Код
Авторство
О чём код?
# # 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 tempfile import numpy as np from pyspark.ml.linalg import Vectors, SparseVector from pyspark.ml.clustering import ( KMeans, KMeansModel, KMeansSummary, BisectingKMeans, BisectingKMeansModel, BisectingKMeansSummary, GaussianMixture, GaussianMixtureModel, GaussianMixtureSummary, LDA, LDAModel, LocalLDAModel, DistributedLDAModel, PowerIterationClustering, ) from pyspark.testing.sqlutils import ReusedSQLTestCase class ClusteringTestsMixin: def test_kmeans(self): df = ( self.spark.createDataFrame( [ (1, 1.0, Vectors.dense([-0.1, -0.05])), (2, 2.0, Vectors.dense([-0.01, -0.1])), (3, 3.0, Vectors.dense([0.9, 0.8])), (4, 1.0, Vectors.dense([0.75, 0.935])), (5, 1.0, Vectors.dense([-0.83, -0.68])), (6, 1.0, Vectors.dense([-0.91, -0.76])), ], ["index", "weight", "features"], ) .coalesce(1) .sortWithinPartitions("index") .select("weight", "features") ) km = KMeans( k=2, maxIter=2, weightCol="weight", ) self.assertEqual(km.getK(), 2) self.assertEqual(km.getMaxIter(), 2) self.assertEqual(km.getWeightCol(), "weight") model = km.fit(df) self.assertEqual(km.uid, model.uid) centers = model.clusterCenters() self.assertEqual(len(centers), 2) self.assertEqual(model.numFeatures, 2) self.assertTrue(np.allclose(centers[0], [-0.372, -0.338], atol=1e-3), centers[0]) self.assertTrue(np.allclose(centers[1], [0.8625, 0.83375], atol=1e-3), centers[1]) output = model.transform(df) expected_cols = ["weight", "features", "prediction"] self.assertEqual(output.columns, expected_cols) self.assertEqual(output.count(), 6) self.assertTrue(np.allclose(model.predict(Vectors.dense(0.0, 5.0)), 1, atol=1e-4)) def check_summary(): # Model summary self.assertTrue(model.hasSummary) summary = model.summary self.assertTrue(isinstance(summary, KMeansSummary)) self.assertEqual(summary.k, 2) self.assertEqual(summary.numIter, 2) self.assertEqual(summary.clusterSizes, [4, 2]) self.assertTrue(np.allclose(summary.trainingCost, 1.35710375, atol=1e-4)) self.assertEqual(summary.featuresCol, "features") self.assertEqual(summary.predictionCol, "prediction") self.assertEqual(summary.cluster.columns, ["prediction"]) self.assertEqual(summary.cluster.count(), 6) self.assertEqual(summary.predictions.columns, expected_cols) self.assertEqual(summary.predictions.count(), 6) # check summary before model offloading occurs check_summary() # save & load with tempfile.TemporaryDirectory(prefix="kmeans_model") as d: km.write().overwrite().save(d) km2 = KMeans.load(d) self.assertEqual(str(km), str(km2)) model.write().overwrite().save(d) model2 = KMeansModel.load(d) self.assertEqual(str(model), str(model2)) self.assertFalse(model2.hasSummary) with self.assertRaisesRegex(Exception, "No training summary available"): model2.summary def test_bisecting_kmeans(self): df = ( self.spark.createDataFrame( [ (1, 1.0, Vectors.dense([-0.1, -0.05])), (2, 2.0, Vectors.dense([-0.01, -0.1])), (3, 3.0, Vectors.dense([0.9, 0.8])), (4, 1.0, Vectors.dense([0.75, 0.935])), (5, 1.0, Vectors.dense([-0.83, -0.68])), (6, 1.0, Vectors.dense([-0.91, -0.76])), ], ["index", "weight", "features"], ) .coalesce(1) .sortWithinPartitions("index") .select("weight", "features") ) bkm = BisectingKMeans( k=2, maxIter=2, minDivisibleClusterSize=1.0, weightCol="weight", ) self.assertEqual(bkm.getK(), 2) self.assertEqual(bkm.getMaxIter(), 2) self.assertEqual(bkm.getMinDivisibleClusterSize(), 1.0) self.assertEqual(bkm.getWeightCol(), "weight") model = bkm.fit(df) self.assertEqual(bkm.uid, model.uid) centers = model.clusterCenters() self.assertEqual(len(centers), 2) self.assertEqual(model.numFeatures, 2) self.assertTrue(np.allclose(centers[0], [-0.372, -0.338], atol=1e-3), centers[0]) self.assertTrue(np.allclose(centers[1], [0.8625, 0.83375], atol=1e-3), centers[1]) output = model.transform(df) expected_cols = ["weight", "features", "prediction"] self.assertEqual(output.columns, expected_cols) self.assertEqual(output.count(), 6) self.assertTrue(np.allclose(model.predict(Vectors.dense(0.0, 5.0)), 1, atol=1e-4)) # BisectingKMeans-specific method: computeCost self.assertTrue(np.allclose(model.computeCost(df), 1.164325125, atol=1e-4)) # Model summary self.assertTrue(model.hasSummary) summary = model.summary self.assertTrue(isinstance(summary, BisectingKMeansSummary)) self.assertEqual(summary.k, 2) self.assertEqual(summary.numIter, 2) self.assertEqual(summary.clusterSizes, [4, 2]) self.assertTrue(np.allclose(summary.trainingCost, 1.3571037499999998, atol=1e-4)) self.assertEqual(summary.featuresCol, "features") self.assertEqual(summary.predictionCol, "prediction") self.assertEqual(summary.cluster.columns, ["prediction"]) self.assertEqual(summary.cluster.count(), 6) self.assertEqual(summary.predictions.columns, expected_cols) self.assertEqual(summary.predictions.count(), 6) # save & load with tempfile.TemporaryDirectory(prefix="bisecting_kmeans") as d: bkm.write().overwrite().save(d) bkm2 = BisectingKMeans.load(d) self.assertEqual(str(bkm), str(bkm2)) model.write().overwrite().save(d) model2 = BisectingKMeansModel.load(d) self.assertEqual(str(model), str(model2)) def test_gaussian_mixture(self): spark = self.spark df = ( spark.createDataFrame( [ (1, 1.0, Vectors.dense([-0.1, -0.05])), (2, 2.0, Vectors.dense([-0.01, -0.1])), (3, 3.0, Vectors.dense([0.9, 0.8])), (4, 1.0, Vectors.dense([0.75, 0.935])), (5, 1.0, Vectors.dense([-0.83, -0.68])), (6, 1.0, Vectors.dense([-0.91, -0.76])), ], ["index", "weight", "features"], ) .coalesce(1) .sortWithinPartitions("index") .select("weight", "features") ) gmm = GaussianMixture( k=2, maxIter=2, weightCol="weight", seed=1, ) self.assertEqual(gmm.getK(), 2) self.assertEqual(gmm.getMaxIter(), 2) self.assertEqual(gmm.getWeightCol(), "weight") self.assertEqual(gmm.getSeed(), 1) model = gmm.fit(df) self.assertEqual(gmm.uid, model.uid) self.assertEqual(model.numFeatures, 2) self.assertEqual(len(model.weights), 2) self.assertTrue( np.allclose(model.weights, [0.541014115744985, 0.4589858842550149], atol=1e-4), model.weights, ) # check the gaussians gaussians = model.gaussians self.assertEqual(len(gaussians), 2) self.assertTrue( np.allclose( gaussians[0].mean.toArray(), [0.28586084899633746, 0.28513455726904297], atol=1e-4, ), gaussians[0].mean, ) self.assertTrue( np.allclose( gaussians[0].cov.toArray(), [[0.41732752, 0.38378601], [0.38378601, 0.36454957]], atol=1e-4, ), gaussians[0].cov, ) self.assertTrue( np.allclose( gaussians[1].mean.toArray(), [0.04795771063097124, 0.06212817950777127], atol=1e-4, ), gaussians[1].mean, ) self.assertTrue( np.allclose( gaussians[1].cov.toArray(), [[0.50359595, 0.44696663], [0.44696663, 0.40424231]], atol=1e-4, ), gaussians[1].cov, ) self.assertEqual(model.gaussiansDF.columns, ["mean", "cov"]) self.assertEqual(model.gaussiansDF.count(), 2) vec = Vectors.dense(0.0, 5.0) pred = model.predict(vec) self.assertTrue(np.allclose(pred, 0, atol=1e-4), pred) pred = model.predictProbability(vec) self.assertTrue(np.allclose(pred.toArray(), [0.5, 0.5], atol=1e-4), pred) output = model.transform(df) expected_cols = ["weight", "features", "probability", "prediction"] self.assertEqual(output.columns, expected_cols) self.assertEqual(output.count(), 6) def check_summary(): # Model summary self.assertTrue(model.hasSummary) summary = model.summary self.assertTrue(isinstance(summary, GaussianMixtureSummary)) self.assertEqual(summary.k, 2) self.assertEqual(summary.numIter, 2) self.assertEqual(len(summary.clusterSizes), 2) self.assertEqual(summary.clusterSizes, [3, 3]) ll = summary.logLikelihood self.assertTrue(ll < 0, ll) self.assertTrue(np.allclose(ll, -1.311264553744033, atol=1e-4), ll) self.assertEqual(summary.featuresCol, "features") self.assertEqual(summary.predictionCol, "prediction") self.assertEqual(summary.probabilityCol, "probability") self.assertEqual(summary.cluster.columns, ["prediction"]) self.assertEqual(summary.cluster.count(), 6) self.assertEqual(summary.predictions.columns, expected_cols) self.assertEqual(summary.predictions.count(), 6) self.assertEqual(summary.probability.columns, ["probability"]) self.assertEqual(summary.predictions.count(), 6) # save & load with tempfile.TemporaryDirectory(prefix="gaussian_mixture") as d: gmm.write().overwrite().save(d) gmm2 = GaussianMixture.load(d) self.assertEqual(str(gmm), str(gmm2)) model.write().overwrite().save(d) model2 = GaussianMixtureModel.load(d) self.assertEqual(str(model), str(model2)) def test_local_lda(self): spark = self.spark df = ( spark.createDataFrame( [ [1, Vectors.dense([0.0, 1.0])], [2, SparseVector(2, {0: 1.0})], ], ["id", "features"], ) .coalesce(1) .sortWithinPartitions("id") ) lda = LDA(k=2, optimizer="online", seed=1) lda.setMaxIter(1) self.assertEqual(lda.getK(), 2) self.assertEqual(lda.getOptimizer(), "online") self.assertEqual(lda.getMaxIter(), 1) self.assertEqual(lda.getSeed(), 1) model = lda.fit(df) self.assertEqual(lda.uid, model.uid) self.assertIsInstance(model, LDAModel) self.assertIsInstance(model, LocalLDAModel) self.assertNotIsInstance(model, DistributedLDAModel) self.assertFalse(model.isDistributed()) dc = model.estimatedDocConcentration() self.assertTrue(np.allclose(dc.toArray(), [0.5, 0.5], atol=1e-4), dc) topics = model.topicsMatrix() self.assertTrue( np.allclose( topics.toArray(), [[1.20296728, 1.15740442], [0.99357675, 1.02993164]], atol=1e-4 ), topics, ) ll = model.logLikelihood(df) self.assertTrue(np.allclose(ll, -3.2125122434040088, atol=1e-4), ll) lp = model.logPerplexity(df) self.assertTrue(np.allclose(lp, 1.6062561217020044, atol=1e-4), lp) dt = model.describeTopics() self.assertEqual(dt.columns, ["topic", "termIndices", "termWeights"]) self.assertEqual(dt.count(), 2) # LocalLDAModel specific methods self.assertEqual(model.vocabSize(), 2) output = model.transform(df) expected_cols = ["id", "features", "topicDistribution"] self.assertEqual(output.columns, expected_cols) self.assertEqual(output.count(), 2) # save & load with tempfile.TemporaryDirectory(prefix="local_lda") as d: lda.write().overwrite().save(d) lda2 = LDA.load(d) self.assertEqual(str(lda), str(lda2)) model.write().overwrite().save(d) model2 = LocalLDAModel.load(d) self.assertEqual(str(model), str(model2)) def test_distributed_lda(self): spark = self.spark df = ( spark.createDataFrame( [ [1, Vectors.dense([0.0, 1.0])], [2, SparseVector(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) self.assertEqual(lda.uid, model.uid) self.assertIsInstance(model, LDAModel) self.assertNotIsInstance(model, LocalLDAModel) self.assertIsInstance(model, DistributedLDAModel) self.assertTrue(model.isDistributed()) self.assertEqual(model.vocabSize(), 2) dc = model.estimatedDocConcentration() self.assertTrue(np.allclose(dc.toArray(), [26.0, 26.0], atol=1e-4), dc) topics = model.topicsMatrix() self.assertTrue( np.allclose( topics.toArray(), [[0.39149926, 0.60850074], [0.60991237, 0.39008763]], atol=1e-4 ), topics, ) ll = model.logLikelihood(df) self.assertTrue(np.allclose(ll, -3.719138517085772, atol=1e-4), ll) lp = model.logPerplexity(df) self.assertTrue(np.allclose(lp, 1.859569258542886, atol=1e-4), lp) dt = model.describeTopics() self.assertEqual(dt.columns, ["topic", "termIndices", "termWeights"]) self.assertEqual(dt.count(), 2) # DistributedLDAModel specific methods ll = model.trainingLogLikelihood() self.assertTrue(np.allclose(ll, -1.3847360462201639, atol=1e-4), ll) lp = model.logPrior() self.assertTrue(np.allclose(lp, -69.59963186898915, atol=1e-4), lp) model.getCheckpointFiles() output = model.transform(df) expected_cols = ["id", "features", "topicDistribution"] self.assertEqual(output.columns, expected_cols) self.assertEqual(output.count(), 2) # Test toLocal() localModel = model.toLocal() self.assertIsInstance(localModel, LDAModel) self.assertIsInstance(localModel, LocalLDAModel) self.assertNotIsInstance(localModel, DistributedLDAModel) self.assertFalse(localModel.isDistributed()) output = localModel.transform(df) expected_cols = ["id", "features", "topicDistribution"] self.assertEqual(output.columns, expected_cols) self.assertEqual(output.count(), 2) # save & load with tempfile.TemporaryDirectory(prefix="distributed_lda") as d: lda.write().overwrite().save(d) lda2 = LDA.load(d) self.assertEqual(str(lda), str(lda2)) model.write().overwrite().save(d) model2 = DistributedLDAModel.load(d) self.assertEqual(str(model), str(model2)) def test_power_iteration_clustering(self): spark = self.spark data = [ (1, 0, 0.5), (2, 0, 0.5), (2, 1, 0.7), (3, 0, 0.5), (3, 1, 0.7), (3, 2, 0.9), (4, 0, 0.5), (4, 1, 0.7), (4, 2, 0.9), (4, 3, 1.1), (5, 0, 0.5), (5, 1, 0.7), (5, 2, 0.9), (5, 3, 1.1), (5, 4, 1.3), ] df = spark.createDataFrame(data, ["src", "dst", "weight"]).repartition(1) pic = PowerIterationClustering(k=2, weightCol="weight") pic.setMaxIter(40) self.assertEqual(pic.getK(), 2) self.assertEqual(pic.getMaxIter(), 40) self.assertEqual(pic.getWeightCol(), "weight") assignments = pic.assignClusters(df) self.assertEqual(assignments.columns, ["id", "cluster"]) self.assertEqual(assignments.count(), 6) # save & load with tempfile.TemporaryDirectory(prefix="power_iteration_clustering") as d: pic.write().overwrite().save(d) pic2 = PowerIterationClustering.load(d) self.assertEqual(str(pic), str(pic2)) self.assertEqual(pic.uid, pic2.uid) self.assertEqual(pic.getK(), pic2.getK()) self.assertEqual(pic.getMaxIter(), pic2.getMaxIter()) self.assertEqual(pic.getWeightCol(), pic2.getWeightCol()) class ClusteringTests(ClusteringTestsMixin, ReusedSQLTestCase): pass if __name__ == "__main__": from pyspark.testing import main main()