fix: use 64-bit bulk sparse row offsets

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/dataset/DatasetAggregator.scala

@@ -6,7 +6,7 @@ package com.microsoft.azure.synapse.ml.lightgbm.dataset
 import com.microsoft.azure.synapse.ml.lightgbm.dataset.DatasetUtils.getRowAsDoubleArray
 import com.microsoft.azure.synapse.ml.lightgbm.swig._
 import com.microsoft.azure.synapse.ml.lightgbm.{ColumnParams, LightGBMUtils, NativeOmpCallSite, PartitionTaskContext}
-import com.microsoft.ml.lightgbm.{SWIGTYPE_p_int, lightgbmlib, lightgbmlibConstants}
+import com.microsoft.ml.lightgbm.{SWIGTYPE_p_long_long, lightgbmlib, lightgbmlibConstants}
 import org.apache.spark.internal.Logging
 import org.apache.spark.ml.linalg.SQLDataTypes.VectorType
 import org.apache.spark.ml.linalg.{DenseVector, SparseVector}
@@ -42,6 +42,24 @@ private[lightgbm] object ChunkedArrayUtils {
     }
   }
 
+  /** Widens per-row counts before they become cumulative CSR offsets. */
+  def copyIntChunkedArrayToLong(chunkedArray: IntChunkedArray,
+                                mainArray: LongSwigArray,
+                                startIndex: Long): Unit = {
+    val chunkSize = chunkedArray.getChunkSize
+    val fullChunks = chunkedArray.getChunksCount - 1
+    for (chunk <- 0L until fullChunks) {
+      for (index <- 0L until chunkSize) {
+        mainArray.setItem(startIndex + chunk * chunkSize + index,
+          chunkedArray.getItem(chunk, index, -1).toLong)
+      }
+    }
+    for (index <- 0L until chunkedArray.getLastChunkAddCount) {
+      mainArray.setItem(startIndex + fullChunks * chunkSize + index,
+        chunkedArray.getItem(fullChunks, index, -1).toLong)
+    }
+  }
+
   def insertTransposedChunkedArray[T: Numeric](chunkedArray: ChunkedArray[T],
                                                numCols: Int,
                                                mainArray: BaseSwigArray[T],
@@ -461,7 +479,7 @@ private[lightgbm] abstract class BaseSparseAggregatedColumns(chunkSize: Int)
   extends BaseAggregatedColumns(chunkSize) {
   protected var indexes: IntSwigArray = _
   protected var values: DoubleSwigArray = _
-  protected var indexPointers: IntSwigArray = _
+  protected var indexPointers: LongSwigArray = _
 
   /**
     * Aggregated variables for knowing how large full array should be allocated to
@@ -489,7 +507,7 @@ private[lightgbm] abstract class BaseSparseAggregatedColumns(chunkSize: Int)
     val indptrCount = this.indptrCount.get()
     indexes = new IntSwigArray(indexesCount)
     values = new DoubleSwigArray(indexesCount)
-    indexPointers = new IntSwigArray(indptrCount)
+    indexPointers = new LongSwigArray(indptrCount)
     indexPointers.setItem(0, 0)
     updateConcurrentMapOffsets(pIdToIndexesCountOffset)
     updateConcurrentMapOffsets(pIdToIndptrCountOffset, 1L)
@@ -499,7 +517,7 @@ private[lightgbm] abstract class BaseSparseAggregatedColumns(chunkSize: Int)
 
   def getValues: DoubleSwigArray = values
 
-  def getIndexPointers: IntSwigArray = indexPointers
+  def getIndexPointers: LongSwigArray = indexPointers
 
   override def cleanup(): Unit = {
     Option(labels).foreach(_.delete())
@@ -510,26 +528,26 @@ private[lightgbm] abstract class BaseSparseAggregatedColumns(chunkSize: Int)
     Option(indexPointers).foreach(_.delete())
   }
 
-  private def indexPointerArrayIncrement(indptrArray: SWIGTYPE_p_int): Unit = {
+  private[lightgbm] def indexPointerArrayIncrement(indptrArray: SWIGTYPE_p_long_long): Unit = {
     // Update indptr array indexes in sparse matrix
     (1L until indptrCount.get()).foreach { index =>
-      val indptrPrevValue = lightgbmlib.intArray_getitem(indptrArray, index - 1)
-      val indptrCurrValue = lightgbmlib.intArray_getitem(indptrArray, index)
-      lightgbmlib.intArray_setitem(indptrArray, index, indptrPrevValue + indptrCurrValue)
+      val indptrPrevValue = lightgbmlib.longArray_getitem(indptrArray, index - 1)
+      val indptrCurrValue = lightgbmlib.longArray_getitem(indptrArray, index)
+      lightgbmlib.longArray_setitem(indptrArray, index, indptrPrevValue + indptrCurrValue)
     }
   }
 
   def generateDataset(ctx: PartitionTaskContext, referenceDataset: Option[LightGBMDataset]): LightGBMDataset = {
     indexPointerArrayIncrement(getIndexPointers.array)
     val pointer = lightgbmlib.voidpp_handle()
     val numRows = indptrCount.get() - 1
-    logInfo(s"LightGBM task ${ctx.taskId} part ${ctx.partitionId} generating dense dataset" +
+    logInfo(s"LightGBM task ${ctx.taskId} part ${ctx.partitionId} generating sparse dataset" +
       s" with $numRows rows and $numCols columns")
     // Generate the dataset for features
     LightGBMUtils.registerNativeOmpThreads(NativeOmpCallSite.SparseDataset, ctx.trainingCtx.datasetParams)
     LightGBMUtils.validate(lightgbmlib.LGBM_DatasetCreateFromCSR(
-      lightgbmlib.int_to_voidp_ptr(indexPointers.array),
-      lightgbmlibConstants.C_API_DTYPE_INT32,
+      lightgbmlib.int64_t_to_voidp_ptr(indexPointers.array),
+      lightgbmlibConstants.C_API_DTYPE_INT64,
       indexes.array,
       lightgbmlib.double_to_voidp_ptr(values.array),
       lightgbmlibConstants.C_API_DTYPE_FLOAT64,
@@ -558,7 +576,7 @@ private[lightgbm] final class SparseAggregatedColumns(chunkSize: Int)
     val sparseChunkedColumns = chunkedCols.asInstanceOf[SparseChunkedColumns]
     sparseChunkedColumns.indexes.coalesceTo(indexes)
     sparseChunkedColumns.values.coalesceTo(values)
-    sparseChunkedColumns.indexPointers.coalesceTo(indexPointers)
+    ChunkedArrayUtils.copyIntChunkedArrayToLong(sparseChunkedColumns.indexPointers, indexPointers, 0L)
   }
 }
 
@@ -585,7 +603,7 @@ private[lightgbm] final class SparseSyncAggregatedColumns(chunkSize: Int)
     val sparseChunkedCols = chunkedCols.asInstanceOf[SparseChunkedColumns]
     ChunkedArrayUtils.copyChunkedArray(sparseChunkedCols.indexes, indexes, featureIndexes.head)
     ChunkedArrayUtils.copyChunkedArray(sparseChunkedCols.values, values, featureIndexes.head)
-    ChunkedArrayUtils.copyChunkedArray(sparseChunkedCols.indexPointers, indexPointers, featureIndexes(1))
+    ChunkedArrayUtils.copyIntChunkedArrayToLong(sparseChunkedCols.indexPointers, indexPointers, featureIndexes(1))
   }
 
 }

lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala

@@ -0,0 +1,134 @@
+// Copyright (C) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in project root for information.
+
+package com.microsoft.azure.synapse.ml.lightgbm.split1
+
+import com.microsoft.azure.synapse.ml.lightgbm._
+import com.microsoft.azure.synapse.ml.lightgbm.dataset._
+import com.microsoft.azure.synapse.ml.lightgbm.swig.{IntChunkedArray, LongSwigArray}
+import org.apache.spark.ml.linalg.{SQLDataTypes, Vectors}
+import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
+import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
+
+// scalastyle:off magic.number
+class BulkSparseOffsetsSuite extends LightGBMTestUtils {
+  private class PrefixColumns extends BaseSparseAggregatedColumns(2) {
+    def increment(array: LongSwigArray, count: Long): Unit = {
+      indptrCount.set(count)
+      indexPointerArrayIncrement(array.array)
+    }
+  }
+
+  private val max = Int.MaxValue.toLong
+  private val cases = Seq(
+    Seq(0L, 2L, 0L, 1L), Seq(0L, max - 1L, 1L), Seq(0L, max, 1L),
+    Seq(0L, max, max), Seq(0L, 0L), Seq(0L), Seq.empty[Long])
+
+  cases.zipWithIndex.foreach { case (counts, index) =>
+    test(s"bulk sparse offsets use a 64-bit prefix: case $index") {
+      LightGBMUtils.initializeNativeLibrary()
+      val array = new LongSwigArray(math.max(counts.length, 1))
+      try {
+        counts.zipWithIndex.foreach { case (count, i) => array.setItem(i, count) }
+        new PrefixColumns().increment(array, counts.length)
+        val expected = counts.scanLeft(0L)(_ + _).tail
+        assert(counts.indices.map(i => array.getItem(i)) == expected)
+      } finally {
+        array.delete()
+      }
+    }
+  }
+
+  test("bulk sparse row counts widen across chunks at a nonzero destination offset") {
+    LightGBMUtils.initializeNativeLibrary()
+    val source = new IntChunkedArray(2)
+    val target = new LongSwigArray(7)
+    try {
+      (0 until 7).foreach(i => target.setItem(i, -7L))
+      Seq(0, Int.MaxValue, 1, 0, 2).foreach(source.add)
+      ChunkedArrayUtils.copyIntChunkedArrayToLong(source, target, 1L)
+      assert((0 until 7).map(i => target.getItem(i)) == Seq(-7L, 0L, max, 1L, 0L, 2L, -7L))
+    } finally {
+      source.delete()
+      target.delete()
+    }
+  }
+
+  test("empty sparse row counts leave the destination unchanged") {
+    LightGBMUtils.initializeNativeLibrary()
+    val source = new IntChunkedArray(2)
+    val target = new LongSwigArray(1)
+    try {
+      target.setItem(0, 9L)
+      ChunkedArrayUtils.copyIntChunkedArrayToLong(source, target, 0L)
+      assert(target.getItem(0) == 9L)
+    } finally {
+      source.delete()
+      target.delete()
+    }
+  }
+
+  Seq(false, true).foreach { singleDataset =>
+    test(s"bulk sparse aggregation widens row counts with singleDataset=$singleDataset") {
+      import spark.implicits._
+      val offsets = spark.range(0L, 1L, 1L, 1).mapPartitions { _ =>
+        LightGBMUtils.initializeNativeLibrary()
+        val schema = StructType(Seq(StructField("label", DoubleType),
+          StructField("features", SQLDataTypes.VectorType)))
+        val rows = Seq(Vectors.sparse(4, Array(0, 3), Array(1.0, 2.0)),
+          Vectors.sparse(4, Array.emptyIntArray, Array.emptyDoubleArray),
+          Vectors.sparse(4, Array(2), Array(3.0))).map { features =>
+          new GenericRowWithSchema(Array[Any](0.0, features), schema)
+        }
+        val chunks = new SparseChunkedColumns(new PeekingIterator(rows.iterator),
+          ColumnParams("label", "features", None, None, None), schema, 2, singleDataset)
+        val aggregate: BaseSparseAggregatedColumns = if (singleDataset) {
+          new SparseSyncAggregatedColumns(2)
+        } else {
+          new SparseAggregatedColumns(2)
+        }
+        try {
+          aggregate.incrementCount(chunks, org.apache.spark.TaskContext.getPartitionId())
+          aggregate.addRows(chunks)
+          val indptr = aggregate.getIndexPointers
+          aggregate.indexPointerArrayIncrement(indptr.array)
+          Iterator((0 until 4).map(i => indptr.getItem(i)).toArray)
+        } finally {
+          chunks.release()
+          aggregate.cleanup()
+        }
+      }.collect()
+      assert(offsets.length == 1)
+      assert(offsets.head.toSeq == Seq(0L, 2L, 2L, 3L))
+    }
+
+    test(s"bulk sparse native training and persistence match streaming with singleDataset=$singleDataset") {
+      import spark.implicits._
+      val data = (0 until 64).map { index =>
+        (index, (index % 2).toDouble,
+          Vectors.sparse(4, Array(0, 2), Array((index % 7 + 1).toDouble, (index % 2 + 1).toDouble)))
+      }.toDF("id", "label", "features").repartition(1).cache()
+      try {
+        def estimator(mode: String): LightGBMClassifier = new LightGBMClassifier()
+          .setDataTransferMode(mode).setMatrixType("sparse").setUseSingleDatasetMode(singleDataset)
+          .setNumTasks(1).setNumThreads(1).setMaxStreamingOMPThreads(1)
+          .setMicroBatchSize(16).setNumIterations(2).setNumLeaves(4)
+          .setDefaultListenPort(getAndIncrementPort())
+        val bulk = estimator("bulk").fit(data)
+        val streaming = estimator("streaming").fit(data)
+        val path = tmpDir.resolve(s"bulk-sparse-$singleDataset").toString
+        bulk.saveNativeModel(path, overwrite = false)
+        val restored = LightGBMClassificationModel.loadNativeModelFromFile(path)
+        def probabilities(model: LightGBMClassificationModel): Seq[Double] = {
+          model.transform(data).orderBy("id").select("probability").collect()
+            .map(_.getAs[org.apache.spark.ml.linalg.Vector](0)(1)).toSeq
+        }
+        val predictions = probabilities(bulk)
+        assert(predictions.zip(probabilities(restored)).forall { case (a, b) => math.abs(a - b) < 1e-9 })
+        assert(predictions.zip(probabilities(streaming)).forall { case (a, b) => math.abs(a - b) < 1e-9 })
+      } finally {
+        data.unpersist()
+      }
+    }
+  }
+}