From ebed932a3846288051576696a68752279de62db7 Mon Sep 17 00:00:00 2001
From: Holden Karau <holden@pigscanfly.ca>
Date: Wed, 29 Jul 2015 08:37:55 -0700
Subject: [PATCH] Add Experimental tag and some scaladocs. Also don't require
 that the inputCol has the metadata on it, instead have the labelsCol
 specified when creating the inverse.

---
 .../spark/ml/feature/StringIndexer.scala      | 20 ++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
index 56bf6cec5a604..4992d5db1f57a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -154,18 +154,32 @@ class StringIndexerModel private[ml] (
     copyValues(copied, extra)
   }
 
+  /**
+   * Return a model to perform the inverse transformation.
+   * Note: By default we keep the original columns during this transformation, so the inverse
+   * should only be used on new columns such as predicted labels.
+   */
   def invert(inputCol: String, outputCol: String): StringIndexerInverseTransformer = {
-    new StringIndexerInverseTransformer()
+    val labelsCol: String = $(this.outputCol)
+    new StringIndexerInverseTransformer(labelsCol)
       .setInputCol(inputCol)
       .setOutputCol(outputCol)
   }
 }
 
+/**
+ * :: Experimental ::
+ * Transform a provided column back to the original input types using the metadata on the
+ * labelsCol. Note: By default we keep the original columns during this transformation,
+ * so the inverse should only be used on new columns such as predicted labels.
+ */
+@Experimental
 class StringIndexerInverseTransformer private[ml] (
-  override val uid: String) extends Transformer
+  override val uid: String,
+  val labelsCol: String) extends Transformer
     with HasInputCol with HasOutputCol {
 
-  def this() = this(Identifiable.randomUID("strIdxInv"))
+  def this(labelsCol: String) = this(Identifiable.randomUID("strIdxInv"), labelsCol)
 
   /** @group setParam */
   def setInputCol(value: String): this.type = set(inputCol, value)