Spark MLlib (ML Pipeline) の決定木モデルの木構造をダンプするコード

package biz.k11i.spark.misc

import org.apache.spark.ml.classification.DecisionTreeClassificationModel
import org.apache.spark.ml.tree.{CategoricalSplit, ContinuousSplit, InternalNode, LeafNode}

/**
  * spark.ml の DecisionTreeClassificationModel の木構造を標準出力に書き出す。
  */
object DecisionTreePrinter {
  def printTree(model: DecisionTreeClassificationModel): Unit = {
    model.rootNode match {
      case node: InternalNode => printNodes(node, 0)
      case leaf: LeafNode => printLeaf(leaf, 0)
    }
  }

  def printNodes(node: InternalNode, numIndents: Int): Unit = {
    val indents = "  " * numIndents

    node.split match {
      case cat: CategoricalSplit => println(s"${indents}category, featureIndex=${cat.featureIndex}, left=${cat.leftCategories.mkString(",")}, right=${cat.rightCategories.mkString(",")}")
      case con: ContinuousSplit => println(s"${indents}continuous, featureIndex=${con.featureIndex}, threshold=${con.threshold}")
    }

    Seq(node.leftChild, node.rightChild).foreach {
      case internalNode: InternalNode => printNodes(internalNode, numIndents + 1)
      case leafNode: LeafNode => printLeaf(leafNode, numIndents + 1)
    }
  }

  def printLeaf(node: LeafNode, indent: Int): Unit = {
    val indents = "  " * indent
    println(s"${indents}prediction=${node.prediction}")
  }
}