decision tree learning algorithm for regression

1# Decision tree learning algorithm for regression
2
3from pyspark.ml.linalg import Vectors
4df = spark.createDataFrame([
5  (1.0, Vectors.dense(1.0)),
6  (0.0, Vectors.sparse(1, [], []))], ["label", "features"])
7dt = DecisionTreeRegressor(maxDepth=2, varianceCol="variance")
8model = dt.fit(df)
9model.depth
10# 1
11model.numNodes
12# 3
13model.featureImportances
14# SparseVector(1, {0: 1.0}
15model.numFeatures
16# 1
17test0 = spark.createDataFrame([(Vectors.dense(-1.0),)], ["features"])
18model.transform(test0).head().prediction
19# 0.0
20test1 = spark.createDataFrame([(Vectors.sparse(1, [0], [1.0]),)], ["features"])
21model.transform(test1).head().prediction
22# 1.0
23dtr_path = temp_path + "/dtr"
24dt.save(dtr_path)
25dt2 = DecisionTreeRegressor.load(dtr_path)
26dt2.getMaxDepth()
27# 2
28model_path = temp_path + "/dtr_model"
29model.save(model_path)
30model2 = DecisionTreeRegressionModel.load(model_path)
31model.numNodes == model2.numNodes
32# True
33model.depth == model2.depth
34# True
35model.transform(test1).head().variance
36# 0.0