Getting the leaf probabilities of a tree model in spark

Question

I\'m trying to refactor a trained spark tree-based model (RandomForest or GBT classifiers) in such a way it can be exported in environments without spark. The toDebugS

Answer 1

I'm trying to refactor a trained spark tree-based model (RandomForest or GBT classifiers) in such a way it can be exported in environments without spark. The

Given growing number of tools designed for real-time serving of Spark (and other) models, that's probably reinventing the wheel.

However if you want to access model internals from plain Python it is best to load its serialized form.

Let's say you have:

from pyspark.ml.classification import RandomForestClassificationModel

rf_model: RandomForestClassificationModel
path: str  # Absolute path

And you save the model:

rf_model.write().save(path)

You can load it back using Parquet reader that supports mixes of struct and list types. Model writer writes both node data:

node_data = spark.read.parquet("{}/data".format(path))

node_data.printSchema()

root
 |-- treeID: integer (nullable = true)
 |-- nodeData: struct (nullable = true)
 |    |-- id: integer (nullable = true)
 |    |-- prediction: double (nullable = true)
 |    |-- impurity: double (nullable = true)
 |    |-- impurityStats: array (nullable = true)
 |    |    |-- element: double (containsNull = true)
 |    |-- rawCount: long (nullable = true)
 |    |-- gain: double (nullable = true)
 |    |-- leftChild: integer (nullable = true)
 |    |-- rightChild: integer (nullable = true)
 |    |-- split: struct (nullable = true)
 |    |    |-- featureIndex: integer (nullable = true)
 |    |    |-- leftCategoriesOrThreshold: array (nullable = true)
 |    |    |    |-- element: double (containsNull = true)
 |    |    |-- numCategories: integer (nullable = true)

and tree metadata:

tree_meta = spark.read.parquet("{}/treesMetadata".format(path))

tree_meta.printSchema()                            
root
 |-- treeID: integer (nullable = true)
 |-- metadata: string (nullable = true)
 |-- weights: double (nullable = true)

where the former one provides all the information you need, as the prediction process is basically an aggregation of impurtityStats *.

You could also access this data directly using underlying Java objects

from  collections import namedtuple
import numpy as np

LeafNode = namedtuple("LeafNode", ("prediction", "impurity"))
InternalNode = namedtuple(
    "InternalNode", ("left", "right", "prediction", "impurity", "split"))
CategoricalSplit = namedtuple("CategoricalSplit", ("feature_index", "categories"))
ContinuousSplit = namedtuple("ContinuousSplit", ("feature_index", "threshold"))

def jtree_to_python(jtree):
    def jsplit_to_python(jsplit):
        if jsplit.getClass().toString().endswith(".ContinuousSplit"):
            return ContinuousSplit(jsplit.featureIndex(), jsplit.threshold())
        else:
            jcat = jsplit.toOld().categories()
            return CategoricalSplit(
                jsplit.featureIndex(),
                [jcat.apply(i) for i in range(jcat.length())])

    def jnode_to_python(jnode):
        prediction = jnode.prediction()        
        stats = np.array(list(jnode.impurityStats().stats()))

        if jnode.numDescendants() != 0:  # InternalNode
            left = jnode_to_python(jnode.leftChild())
            right = jnode_to_python(jnode.rightChild())
            split = jsplit_to_python(jnode.split())

            return InternalNode(left, right, prediction, stats, split)            

        else:
            return LeafNode(prediction, stats) 

    return jnode_to_python(jtree.rootNode())

which can be applied to RandomForestModel like this:

nodes = [jtree_to_python(t) for t in rf_model._java_obj.trees()]

Furthermore such structure can be easily used to make predictions, for both individual trees (warning: Python 3.7+ ahead. For legacy usage please refer to functools documentation):

from functools import singledispatch

@singledispatch
def should_go_left(split, vector): pass

@should_go_left.register
def _(split: CategoricalSplit, vector):
    return vector[split.feature_index] in split.categories

@should_go_left.register
def _(split: ContinuousSplit, vector):
    return vector[split.feature_index] <= split.threshold

@singledispatch
def predict(node, vector): pass

@predict.register
def _(node: LeafNode, vector):
    return node.prediction, node.impurity

@predict.register
def _(node: InternalNode, vector):
    return predict(
        node.left if should_go_left(node.split, vector) else node.right,
        vector
    )

and forests:

from typing import Iterable, Union

def predict_probability(nodes: Iterable[Union[InternalNode, LeafNode]], vector):
    total = np.array([
        v / v.sum() for _, v in  (
            predict(node, vector) for node in nodes
        )
    ]).sum(axis=0)
    return total / total.sum()

That however depends on the internal API (and weakness of Scala package-scoped access modifiers) and might break in the future.

---

* DataFrame as loaded from data path can be easily transformed to a structure compatible with predict and predict_probability functions defined above.

from pyspark.sql.dataframe import DataFrame 
from itertools import groupby
from operator import itemgetter


def model_data_to_tree(tree_data: DataFrame):
    def dict_to_tree(node_id, nodes):
        node = nodes[node_id]
        prediction = node.prediction
        impurity = np.array(node.impurityStats)

        if node.leftChild == -1 and node.rightChild == -1:
            return LeafNode(prediction, impurity)
        else:
            left = dict_to_tree(node.leftChild, nodes)
            right = dict_to_tree(node.rightChild, nodes)
            feature_index = node.split.featureIndex
            left_value = node.split.leftCategoriesOrThreshold

            split = (
                CategoricalSplit(feature_index, left_value)
                if node.split.numCategories != -1
                else ContinuousSplit(feature_index, left_value[0])
            )

            return InternalNode(left, right, prediction, impurity, split)

    tree_id = itemgetter("treeID")
    rows = tree_data.collect()
    return ([
        dict_to_tree(0, {node.nodeData.id: node.nodeData for node in nodes})
        for tree, nodes in groupby(sorted(rows, key=tree_id), key=tree_id)
    ] if "treeID" in tree_data.columns
    else [dict_to_tree(0, {node.id: node for node in rows})])