[TreePartitioner] Improve TreePartitioner - Change return type from ITypedPartition to specific TreePartition - Change return type of #getTreeNode to TreePartitionNode to meet the method name