fixup

microsoft · Jun 2, 2021 · 15fa45d · 15fa45d
1 parent 90e5f0a
commit 15fa45d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/src/main/scala/com/microsoft/ml/spark/lightgbm/dataset/DatasetAggregator.scala b/src/main/scala/com/microsoft/ml/spark/lightgbm/dataset/DatasetAggregator.scala
@@ -45,7 +45,7 @@ class SparseDatasetAggregator(columnParams: ColumnParams, chunkSize: Int,
   @volatile var threadRowStartIndex = 0L
   @volatile var threadInitScoreStartIndex = 0L
   @volatile var threadIndexesStartIndex = 0L
-  @volatile var threadIndptrStartIndex = 0L
+  @volatile var threadIndptrStartIndex = 1L
 
   def setNumCols(numCols: Int): Unit = {
     this.numCols = numCols
@@ -57,6 +57,10 @@ class SparseDatasetAggregator(columnParams: ColumnParams, chunkSize: Int,
                      indptrCount: Long): Unit = {
     if (synchronized) {
       this.synchronized {
+        // Add extra 0 for start of indptr in parallel case
+        if (this.indptrCount == 0) {
+          this.indptrCount += 1
+        }
         innerIncrementCount(rowCount, initScoreCount, indexesCount, indptrCount)
       }
     } else {
@@ -86,6 +90,7 @@ class SparseDatasetAggregator(columnParams: ColumnParams, chunkSize: Int,
     indexesArray = Some(lightgbmlib.new_intArray(this.indexesCount))
     valuesArray = Some(lightgbmlib.new_doubleArray(this.indexesCount))
     indptrArray = Some(lightgbmlib.new_intArray(this.indptrCount))
+    lightgbmlib.intArray_setitem(indptrArray.get, 0, 0)
     groupColumnValuesArray = new Array[Row](this.rowCount.toInt)
   }
 

diff --git a/src/main/scala/com/microsoft/ml/spark/lightgbm/dataset/DatasetUtils.scala b/src/main/scala/com/microsoft/ml/spark/lightgbm/dataset/DatasetUtils.scala
@@ -245,7 +245,9 @@ object DatasetUtils {
     var indptrChunkedArray = new int32ChunkedArray(chunkSize)
     val groupColumnValues: ListBuffer[Row] = new ListBuffer[Row]()
     var rowCount = 0
-    indptrChunkedArray.add(0)
+    if (!useSingleDataset) {
+      indptrChunkedArray.add(0)
+    }
     while (rowsIter.hasNext) {
       rowCount += 1
       val row = rowsIter.next()
@@ -276,6 +278,21 @@ object DatasetUtils {
       indexesChunkedArray, valuesChunkedArray, indptrChunkedArray, groupColumnValues)
   }
 
+  def clearSparseArrays(labelsChunkedArray: floatChunkedArray,
+                        weightChunkedArrayOpt: Option[floatChunkedArray],
+                        initScoreChunkedArrayOpt: Option[doubleChunkedArray],
+                        indexesChunkedArray: int32ChunkedArray,
+                        valuesChunkedArray: doubleChunkedArray,
+                        indptrChunkedArray: int32ChunkedArray): Unit = {
+    // Clear memory
+    labelsChunkedArray.delete()
+    weightChunkedArrayOpt.foreach(_.delete())
+    initScoreChunkedArrayOpt.foreach(_.delete())
+    indexesChunkedArray.delete()
+    valuesChunkedArray.delete()
+    indptrChunkedArray.delete()
+  }
+
   def aggregateDenseStreamedData(rowsIter: Iterator[Row], columnParams: ColumnParams,
                                  referenceDataset: Option[LightGBMDataset], schema: StructType,
                                  log: Logger, trainParams: TrainParams): Option[LightGBMDataset] = {