From 11fb00a71e713fe10df633909f361fc44f73ddc8 Mon Sep 17 00:00:00 2001 From: Xusen Yin Date: Thu, 7 May 2015 22:38:39 +0800 Subject: [PATCH] change it into an Estimator --- .../main/scala/org/apache/spark/ml/feature/Bucketizer.scala | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala index b3d8b17cecdfd..ef41ac996195f 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala @@ -18,11 +18,11 @@ package org.apache.spark.ml.feature import org.apache.spark.annotation.AlphaComponent -import org.apache.spark.ml.Transformer import org.apache.spark.ml.attribute.NominalAttribute import org.apache.spark.ml.param._ import org.apache.spark.ml.param.shared.{HasInputCol, HasOutputCol} import org.apache.spark.ml.util.SchemaUtils +import org.apache.spark.ml.{Estimator, Model} import org.apache.spark.sql._ import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.{DoubleType, StructType} @@ -32,7 +32,8 @@ import org.apache.spark.sql.types.{DoubleType, StructType} * `Bucketizer` maps a column of continuous features to a column of feature buckets. */ @AlphaComponent -final class Bucketizer extends Transformer with HasInputCol with HasOutputCol { +final class Bucketizer(override val parent: Estimator[Bucketizer] = null) + extends Model[Bucketizer] with HasInputCol with HasOutputCol { /** * The given buckets should match 1) its size is larger than zero; 2) it is ordered in a non-DESC