add Estimator and Transformer

mengxr · mengxr · commit a3015cf6d53b · 2014-12-31T15:39:38.000-08:00
diff --git a/python/pyspark/ml/__init__.py b/python/pyspark/ml/__init__.py
@@ -3,7 +3,7 @@
 from pyspark import SparkContext
 from pyspark.ml.param import Param
 
-__all__ = ["Pipeline"]
+__all__ = ["Pipeline", "Transformer", "Estimator"]
 
 # An implementation of PEP3102 for Python 2.
 _keyword_only_secret = 70861589
@@ -60,3 +60,15 @@ def transform(self, dataset):
         for t in self.transformers:
             dataset = t.transform(dataset)
         return dataset
+
+
+class Estimator(object):
+
+    def fit(self, dataset, params={}):
+        raise NotImplementedError()
+
+
+class Transformer(object):
+
+    def transform(self, dataset, paramMap={}):
+        raise NotImplementedError()
diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py
@@ -1,9 +1,9 @@
 from pyspark.sql import SchemaRDD
-from pyspark.ml import _jvm
+from pyspark.ml import Estimator, Transformer, _jvm
 from pyspark.ml.param import Param
 
 
-class LogisticRegression(object):
+class LogisticRegression(Estimator):
     """
     Logistic regression.
     """
@@ -45,7 +45,7 @@ def fit(self, dataset, params=None):
         return LogisticRegressionModel(java_model)
 
 
-class LogisticRegressionModel(object):
+class LogisticRegressionModel(Transformer):
     """
     Model fitted by LogisticRegression.
     """
