Made Assignment picklable

Author: Sergei Lebedev

python/pyspark/mllib/clustering.py

@@ -618,8 +618,7 @@ def assignments(self):
         """
         Returns the cluster assignments of this model.
         """
-        return self.call("getAssignments").map(
-            lambda x: (PowerIterationClustering.Assignment(*x)))
+        return self.call("getAssignments").map(lambda x: Assignment(*x))
 
     @classmethod
     @since('1.5.0')
@@ -633,6 +632,12 @@ def load(cls, sc, path):
         return PowerIterationClusteringModel(wrapper)
 
 
+#: Represents an (id, cluster) tuple.
+#:
+#: .. versionadded:: 1.5.0
+Assignment = namedtuple("Assignment", ["id", "cluster"])
+
+
 class PowerIterationClustering(object):
     """
     Power Iteration Clustering (PIC), a scalable graph clustering algorithm
@@ -671,12 +676,8 @@ def train(cls, rdd, k, maxIterations=100, initMode="random"):
                               rdd.map(_convert_to_vector), int(k), int(maxIterations), initMode)
         return PowerIterationClusteringModel(model)
 
-    class Assignment(namedtuple("Assignment", ["id", "cluster"])):
-        """
-        Represents an (id, cluster) tuple.
-
-        .. versionadded:: 1.5.0
-        """
+    # Backward-compatible alias.
+    Assignment = Assignment
 
 
 class StreamingKMeansModel(KMeansModel):