From d9f2a5012d53ed72288733e04144295153d91437 Mon Sep 17 00:00:00 2001
From: Shay Nativ <shay.nativ@gmail.com>
Date: Wed, 29 Mar 2017 14:32:23 +0300
Subject: [PATCH] [ML] Node: Added getter for impurityStats

---
 mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala
index 07e98a142b10e..2b8264dcaa507 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala
@@ -43,6 +43,11 @@ sealed abstract class Node extends Serializable {
    */
   private[ml] def impurityStats: ImpurityCalculator
 
+  /** Getter method for the impurity stats */
+  def getImpurityStats(): Array[Double] = {
+    impurityStats.stats
+  }
+
   /** Recursive prediction helper method */
   private[ml] def predictImpl(features: Vector): LeafNode