From 0bbdf1dbfacbd88067f1bc95070d459cb206c429 Mon Sep 17 00:00:00 2001 From: Lenin Mookiah Date: Sun, 25 Oct 2026 07:51:04 -0700 Subject: [PATCH 1/2] fix: use 64-bit bulk sparse row offsets Fixes #2786 --- .../lightgbm/dataset/DatasetAggregator.scala | 44 ++++-- .../split1/BulkSparseOffsetsSuite.scala | 133 ++++++++++++++++++ 2 files changed, 164 insertions(+), 13 deletions(-) create mode 100644 lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala diff --git a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/dataset/DatasetAggregator.scala b/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/dataset/DatasetAggregator.scala index 53dbc20f3b..acd9b7bc80 100644 --- a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/dataset/DatasetAggregator.scala +++ b/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,12 +528,12 @@ 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) } } @@ -523,13 +541,13 @@ private[lightgbm] abstract class BaseSparseAggregatedColumns(chunkSize: Int) 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)) } } diff --git a/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala b/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala new file mode 100644 index 0000000000..8c998f2221 --- /dev/null +++ b/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala @@ -0,0 +1,133 @@ +// 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") { + val offsets = sc.parallelize(Seq(0), 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() + } + } + } +} From 6b5dc4a1e8eab1b5d30d01e8a8946e7a697fdf9b Mon Sep 17 00:00:00 2001 From: Lenin Mookiah Date: Sun, 25 Oct 2026 08:27:55 -0700 Subject: [PATCH 2/2] test: use Dataset input for bulk sparse aggregation regression --- .../synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala b/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala index 8c998f2221..ac57552408 100644 --- a/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala +++ b/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BulkSparseOffsetsSuite.scala @@ -70,7 +70,8 @@ class BulkSparseOffsetsSuite extends LightGBMTestUtils { Seq(false, true).foreach { singleDataset => test(s"bulk sparse aggregation widens row counts with singleDataset=$singleDataset") { - val offsets = sc.parallelize(Seq(0), 1).mapPartitions { _ => + 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)))