Github user jkbradley commented on a diff in the pull request:
https://github.com/apache/spark/pull/14872#discussion_r76853720
--- Diff:
mllib/src/main/scala/org/apache/spark/ml/tree/impl/PartitionInfo.scala ---
@@ -0,0 +1,422 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.ml.tree.impl
+
+import org.apache.spark.util.collection.BitSet
+
+/**
+ * Intermediate data stored during learning.
+ * TODO(smurching): Rename; maybe TrainingInfo?
+ *
+ * Node indexing for nodeOffsets, activeNodes:
+ * Nodes are indexed left-to-right along the periphery of the tree, with
0-based indices.
+ * The periphery is the set of leaf nodes (active and inactive).
+ *
+ * @param columns Array of columns.
+ * Each column is sorted first by nodes (left-to-right
along the tree periphery);
+ * all columns share this first level of sorting.
+ * Within each node's group, each column is sorted based
on feature value;
+ * this second level of sorting differs across columns.
+ * @param nodeOffsets Offsets into the columns indicating the first level
of sorting (by node).
+ * The rows corresponding to node i are in the range
+ * [nodeOffsets(i), nodeOffsets(i+1)).
+ * @param activeNodes Nodes which are active (still being split).
+ * Inactive nodes are known to be leaves in the final
tree.
+ */
+private[impl] case class PartitionInfo(
+ columns: Array[FeatureVector],
+ nodeOffsets: Array[Int],
+ activeNodes: BitSet,
+ fullImpurityAggs: Array[ImpurityAggregatorSingle]) extends
Serializable {
+
+ // pre-allocated temporary buffers that we use to sort
+ // instances in left and right children during update
+ val tempVals: Array[Int] = new Array[Int](columns(0).values.length)
+ val tempIndices: Array[Int] = new Array[Int](columns(0).values.length)
+
+ /** For debugging */
+ override def toString: String = {
+ "PartitionInfo(" +
+ " columns: {\n" +
+ columns.mkString(",\n") +
+ " },\n" +
+ s" nodeOffsets: ${nodeOffsets.mkString(", ")},\n" +
+ s" activeNodes: ${activeNodes.iterator.mkString(", ")},\n" +
+ ")\n"
+ }
+
+ /**
+ * Update columns and nodeOffsets for the next level of the tree.
+ *
+ * Update columns:
+ * For each column,
+ * For each (previously) active node,
+ * Sort corresponding range of instances based on bit vector.
+ * Update nodeOffsets, activeNodes:
+ * Split offsets for nodes which split (which can be identified using
the bit vector).
+ *
+ * @param instanceBitVector Bit vector encoding splits for the next
level of the tree.
+ * These must follow a 2-level ordering, where the
first level is by node
+ * and the second level is by row index.
+ * bitVector(i) = false iff instance i goes to the
left child.
+ * For instances at inactive (leaf) nodes, the value
can be arbitrary.
+ * @return Updated partition info
+ */
+ def update(
+ instanceBitVector: BitSet,
+ newNumNodeOffsets: Int,
+ labels: Array[Byte],
+ metadata: DecisionTreeMetadata): PartitionInfo = {
+ // Create a 2-level representation of the new nodeOffsets (to be
flattened).
+ // These 2 levels correspond to original nodes and their children (if
split).
+ val newNodeOffsets = nodeOffsets.map(Array(_))
+ val newFullImpurityAggs = fullImpurityAggs.map(Array(_))
+
+ val newColumns = columns.zipWithIndex.map { case (col, index) =>
+ index match {
+ case 0 => first(col, instanceBitVector, metadata,
+ labels, newNodeOffsets, newFullImpurityAggs)
+ case _ => rest(col, instanceBitVector, newNodeOffsets)
+ }
+ col
+ }
+
+ // Identify the new activeNodes based on the 2-level representation of
the new nodeOffsets.
+ val newActiveNodes = new BitSet(newNumNodeOffsets - 1)
+ var newNodeOffsetsIdx = 0
+ var i = 0
+ while (i < newNodeOffsets.length) {
+ val offsets = newNodeOffsets(i)
+ if (offsets.length == 2) {
+ newActiveNodes.set(newNodeOffsetsIdx)
+ newActiveNodes.set(newNodeOffsetsIdx + 1)
+ newNodeOffsetsIdx += 2
+ } else {
+ newNodeOffsetsIdx += 1
+ }
+ i += 1
+ }
+ PartitionInfo(newColumns, newNodeOffsets.flatten, newActiveNodes,
newFullImpurityAggs.flatten)
+ }
+
+
+ /**
+ * Sort the very first column in the [[PartitionInfo.columns]]. While
+ * we sort the column, we also update [[PartitionInfo.nodeOffsets]]
+ * (by modifying @param newNodeOffsets) and
[[PartitionInfo.fullImpurityAggs]]
+ * (by modifying @param newFullImpurityAggs).
+ *
+ * @param col The very first column in [[PartitionInfo.columns]]
+ * @param metadata Used to create new [[ImpurityAggregatorSingle]] for a
new child
+ * node in the tree
+ * @param labels Labels are read as we sort column to populate stats
for each
+ * new ImpurityAggregatorSingle
+ */
+ private def first(
+ col: FeatureVector,
+ instanceBitVector: BitSet,
+ metadata: DecisionTreeMetadata,
+ labels: Array[Byte],
+ newNodeOffsets: Array[Array[Int]],
+ newFullImpurityAggs: Array[Array[ImpurityAggregatorSingle]]) = {
+ activeNodes.iterator.foreach { nodeIdx =>
--- End diff --
activeNodes will be replaced by an Array of nodes. Some will be leafs (no
split), and some will not. Could still use instanceBitVector to avoid
recomputing split decisions for each row.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
