diff options
Diffstat (limited to 'python/pyspark/ml/clustering.py')
| -rw-r--r-- | python/pyspark/ml/clustering.py | 14 |
1 files changed, 9 insertions, 5 deletions
diff --git a/python/pyspark/ml/clustering.py b/python/pyspark/ml/clustering.py index 2db5b82c44..e22d5c8ea4 100644 --- a/python/pyspark/ml/clustering.py +++ b/python/pyspark/ml/clustering.py @@ -87,12 +87,14 @@ class KMeans(JavaEstimator, HasFeaturesCol, HasPredictionCol, HasMaxIter, HasTol .. versionadded:: 1.5.0 """ - k = Param(Params._dummy(), "k", "number of clusters to create") + k = Param(Params._dummy(), "k", "number of clusters to create", + typeConverter=TypeConverters.toInt) initMode = Param(Params._dummy(), "initMode", "the initialization algorithm. This can be either \"random\" to " + "choose random points as initial cluster centers, or \"k-means||\" " + - "to use a parallel variant of k-means++") - initSteps = Param(Params._dummy(), "initSteps", "steps for k-means initialization mode") + "to use a parallel variant of k-means++", TypeConverters.toString) + initSteps = Param(Params._dummy(), "initSteps", "steps for k-means initialization mode", + typeConverter=TypeConverters.toInt) @keyword_only def __init__(self, featuresCol="features", predictionCol="prediction", k=2, @@ -227,10 +229,12 @@ class BisectingKMeans(JavaEstimator, HasFeaturesCol, HasPredictionCol, HasMaxIte .. versionadded:: 2.0.0 """ - k = Param(Params._dummy(), "k", "number of clusters to create") + k = Param(Params._dummy(), "k", "number of clusters to create", + typeConverter=TypeConverters.toInt) minDivisibleClusterSize = Param(Params._dummy(), "minDivisibleClusterSize", "the minimum number of points (if >= 1.0) " + - "or the minimum proportion") + "or the minimum proportion", + typeConverter=TypeConverters.toFloat) @keyword_only def __init__(self, featuresCol="features", predictionCol="prediction", maxIter=20, |