fix: LightGBM Ensure non-duplicate column names

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/LightGBMBase.scala

@@ -199,7 +199,7 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
 
   private def getSlotNamesWithMetadata(featuresSchema: StructField): Option[Array[String]] = {
     if (getSlotNames.nonEmpty) {
-      Some(getSlotNames)
+      Some(ensureUniqueFeatureNames(getSlotNames))
     } else {
       AttributeGroup.fromStructField(featuresSchema).attributes.flatMap(attributes =>
         if (attributes.isEmpty) {
@@ -208,28 +208,82 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
           val colNames = attributes.indices.map(_.toString).toArray
           attributes.foreach(attr =>
             attr.index.foreach(index => colNames(index) = attr.name.getOrElse(index.toString)))
-          Some(colNames)
+          // Ensure unique feature names to avoid LightGBM error:
+          // "Feature (Column_) appears more than one time"
+          // This can occur in Spark 3.5+ due to changes in AttributeGroup metadata handling
+          Some(ensureUniqueFeatureNames(colNames))
         }
       )
     }
   }
 
-  private def validateSlotNames(featuresSchema: StructField): Unit = {
-    val metadata = AttributeGroup.fromStructField(featuresSchema)
-    if (metadata.attributes.isDefined) {
-      val slotNamesOpt = getSlotNamesWithMetadata(featuresSchema)
-      val pattern = new Regex("[\",:\\[\\]{}]")
-      slotNamesOpt.foreach(slotNames => {
-        val badSlotNames = slotNames.flatMap(slotName =>
-          if (pattern.findFirstIn(slotName).isEmpty) None else Option(slotName))
-        if (!badSlotNames.isEmpty) {
-          throw new IllegalArgumentException(
-            s"Invalid slot names detected in features column: ${badSlotNames.mkString(",")}" +
-            " \n Special characters \" , : \\ [ ] { } will cause unexpected behavior in LGBM unless changed." +
-            " This error can be fixed by renaming the problematic columns prior to vector assembly.")
-        }
-      })
+  /**
+    * Maps a feature name to the form LightGBM compares internally. LightGBM replaces every space
+    * with an underscore before checking for duplicates, so "a b" and "a_b" are the same feature
+    * as far as the native library is concerned even though they differ in Scala.
+    */
+  private def normalizeFeatureName(name: String): String = name.replace(' ', '_')
+
+  /**
+    * Ensures all feature names are unique by appending a numeric suffix to repeated names.
+    * LightGBM rejects a Dataset whose feature names repeat, failing the native
+    * LGBM_DatasetSetFeatureNames call with "Feature (X) appears more than one time", and Spark
+    * can surface repeated names through AttributeGroup metadata on the features column.
+    *
+    * Uniqueness is decided on the normalized form (see normalizeFeatureName), because that is
+    * what LightGBM compares. The original names are what get emitted, so this only affects which
+    * names are considered to collide.
+    *
+    * Every original name is reserved up front, so a generated name can never collide with an
+    * original that appears later in the array. Generated names are reserved as they are handed
+    * out, so they cannot collide with each other either. Order is preserved, which matters
+    * because feature names are positional in LightGBM.
+    *
+    * @param names The array of feature names that may contain duplicates.
+    * @return An array of unique feature names, in the original order.
+    */
+  private def ensureUniqueFeatureNames(names: Array[String]): Array[String] = {
+    val reserved = scala.collection.mutable.HashSet[String](names.map(normalizeFeatureName): _*)
+    val emitted = scala.collection.mutable.HashSet[String]()
+    val renamed = scala.collection.mutable.ArrayBuffer[String]()
+
+    val uniqueNames = names.map { name =>
+      if (emitted.add(normalizeFeatureName(name))) {
+        name
+      } else {
+        // Terminates because only finitely many names are reserved.
+        val uniqueName = Iterator.from(1)
+          .map(suffix => s"${name}_$suffix")
+          .find(candidate => !reserved.contains(normalizeFeatureName(candidate)))
+          .get
+        reserved.add(normalizeFeatureName(uniqueName))
+        emitted.add(normalizeFeatureName(uniqueName))
+        renamed += name
+        uniqueName
+      }
+    }
+
+    if (renamed.nonEmpty) {
+      log.warn(s"Duplicate feature names detected and renamed: ${renamed.distinct.mkString(", ")}. " +
+        "Set the 'slotNames' parameter explicitly to control feature naming.")
     }
+
+    uniqueNames
+  }
+
+  private def validateSlotNames(featuresSchema: StructField): Unit = {
+    val slotNamesOpt = getSlotNamesWithMetadata(featuresSchema)
+    val pattern = new Regex("[\",:\\[\\]{}]")
+    slotNamesOpt.foreach(slotNames => {
+      val badSlotNames = slotNames.flatMap(slotName =>
+        if (pattern.findFirstIn(slotName).isEmpty) None else Option(slotName))
+      if (!badSlotNames.isEmpty) {
+        throw new IllegalArgumentException(
+          s"Invalid slot names detected in features column: ${badSlotNames.mkString(",")}" +
+          " \n Special characters \" , : [ ] { } will cause unexpected behavior in LGBM unless changed." +
+          " This error can be fixed by renaming the problematic columns prior to vector assembly.")
+      }
+    })
   }
 
   /**
@@ -414,6 +468,10 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
     val (numCols, numInitScoreClasses) = calculateColumnStatistics(preprocessedDF, measures)
 
     val featuresSchema = dataset.schema(getFeaturesCol)
+    // Validate before any native LightGBM call that consumes feature names (e.g. reference
+    // Dataset creation below), so an invalid name surfaces as an actionable error naming the
+    // offending columns rather than an opaque native failure.
+    validateSlotNames(featuresSchema)
     val generalTrainParams: BaseTrainParams = getTrainParams(numTasks, featuresSchema, numTasksPerExecutor)
     val trainParams = addCustomTrainParams(generalTrainParams, dataset)
     log.info(s"LightGBM batch $batchIndex of $batchCount, parameters: ${trainParams.toString()}")
@@ -422,7 +480,7 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
     val (serializedReferenceDataset: Option[Array[Byte]], partitionCounts: Option[Array[Long]]) =
       if (isStreamingMode) {
         val (referenceDataset, partitionCounts) =
-          calculateRowStatistics(trainingData, trainParams, numCols, measures)
+          calculateRowStatistics(trainingData, trainParams, numCols, featuresSchema, measures)
 
         // Save the reference Dataset so it's available to client and other batches
         if (getReferenceDataset.isEmpty) {
@@ -432,7 +490,6 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
         (Some(referenceDataset), Some(partitionCounts))
       } else (None, None)
 
-    validateSlotNames(featuresSchema)
     executeTraining(preprocessedDF,
                     validationData,
                     serializedReferenceDataset,
@@ -503,12 +560,14 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
     * @param dataframe The dataset to train on.
     * @param trainingParams The training parameters.
     * @param numCols The number of feature columns.
+    * @param featuresSchema The schema of the features column.
     * @param measures Instrumentation measures.
     * @return The serialized Dataset reference and an array of partition counts.
     */
   private def calculateRowStatistics(dataframe: DataFrame,
                                      trainingParams: BaseTrainParams,
                                      numCols: Int,
+                                     featuresSchema: StructField,
                                      measures: InstrumentationMeasures): (Array[Byte], Array[Long]) = {
     measures.markRowStatisticsStart()
 
@@ -523,6 +582,9 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
       trainingParams.generalParams.categoricalFeatures,
       trainingParams.executionParams.numThreads)
 
+    // Get feature names to set on the reference dataset (ensures unique names for Spark 3.5+)
+    val featureNames = getSlotNamesWithMetadata(featuresSchema)
+
     // Either get a reference dataset (as bytes) from params, or calculate it
     val precalculatedDataset = getReferenceDataset
     val serializedReference = if (precalculatedDataset.nonEmpty) {
@@ -541,6 +603,7 @@ trait LightGBMBase[TrainedModel <: Model[TrainedModel] with LightGBMModelParams]
         totalNumRows,
         numCols,
         collectedSampleData,
+        featureNames,
         measures,
         log)
     }

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

@@ -7,6 +7,7 @@ import com.microsoft.azure.synapse.ml.lightgbm.LightGBMUtils
 import com.microsoft.azure.synapse.ml.lightgbm.dataset.DatasetUtils.countCardinality
 import com.microsoft.lightgbm.SwigPtrWrapper
 import com.microsoft.ml.lightgbm._
+import org.slf4j.{Logger, LoggerFactory}
 
 import scala.reflect.ClassTag
 
@@ -179,8 +180,17 @@ class LightGBMDataset(val datasetPtr: SWIGTYPE_p_void) extends AutoCloseable {
     // Add in slot names if they exist
     featureNamesOpt.foreach { featureNamesArray =>
       if (featureNamesArray.nonEmpty) {
-        LightGBMUtils.validate(lightgbmlib.LGBM_DatasetSetFeatureNames(datasetPtr, featureNamesArray, numCols),
-          "Dataset set feature names")
+        // LGBM_DatasetSetFeatureNames reads numCols entries from the array, so a shorter array
+        // is an out-of-bounds native read. slotNames is user-supplied and unvalidated, so guard
+        // every dataset-naming path here rather than at individual call sites. LightGBM falls
+        // back to its own generated names when naming is skipped.
+        if (featureNamesArray.length != numCols) {
+          LightGBMDataset.Log.warn(s"Skipping feature names: got ${featureNamesArray.length} names " +
+            s"for $numCols feature columns.")
+        } else {
+          LightGBMUtils.validate(lightgbmlib.LGBM_DatasetSetFeatureNames(datasetPtr, featureNamesArray, numCols),
+            "Dataset set feature names")
+        }
       }
     }
     this
@@ -191,3 +201,7 @@ class LightGBMDataset(val datasetPtr: SWIGTYPE_p_void) extends AutoCloseable {
     LightGBMUtils.validate(lightgbmlib.LGBM_DatasetFree(datasetPtr), "Free Dataset")
   }
 }
+
+object LightGBMDataset {
+  private val Log: Logger = LoggerFactory.getLogger(classOf[LightGBMDataset])
+}

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

@@ -16,25 +16,42 @@ object ReferenceDatasetUtils {
                                        numRows: Long,
                                        numCols: Int,
                                        sampledRowData: Array[Row],
+                                       featureNames: Option[Array[String]],
                                        measures: InstrumentationMeasures,
                                        log: Logger): Array[Byte] = {
     log.info(s"Creating reference training dataset with ${sampledRowData.length} samples and config: $datasetParams")
 
-    // Pre-create allocated native pointers so it's easy to clean them up
-    val datasetVoidPtr = lightgbmlib.voidpp_handle()
     val lenPtr = lightgbmlib.new_intp()
     val bufferHandlePtr = lightgbmlib.voidpp_handle()
-
     val sampledData = SampledData(sampledRowData.length, numCols)
+
     try {
-      // create properly formatted sampled data
       measures.markSamplingStart()
       sampledRowData.zipWithIndex.foreach({case (row, index) => sampledData.pushRow(row, index, featuresCol)})
       measures.markSamplingStop()
 
-      // Create dataset from samples
-      // 1. Generate the dataset for features
-      val datasetVoidPtr = lightgbmlib.voidpp_handle()
+      val datasetHandle = createDatasetFromSamples(sampledData, numCols, numRows, datasetParams)
+      try {
+        setFeatureNamesIfProvided(datasetHandle, featureNames, numCols, log)
+        serializeReference(datasetHandle, bufferHandlePtr, lenPtr, log)
+      } finally {
+        // Free unconditionally so a failure while naming or serializing cannot leak the native
+        // Dataset. Deliberately not validated: throwing here would mask the original failure.
+        lightgbmlib.LGBM_DatasetFree(datasetHandle)
+      }
+    } finally {
+      sampledData.delete()
+      lightgbmlib.delete_voidpp(bufferHandlePtr)
+      lightgbmlib.delete_intp(lenPtr)
+    }
+  }
+
+  private def createDatasetFromSamples(sampledData: SampledData,
+                                       numCols: Int,
+                                       numRows: Long,
+                                       datasetParams: String): SWIGTYPE_p_void = {
+    val datasetVoidPtr = lightgbmlib.voidpp_handle()
+    try {
       LightGBMUtils.validate(lightgbmlib.LGBM_DatasetCreateFromSampledColumn(
         sampledData.getSampleData,
         sampledData.getSampleIndices,
@@ -45,31 +62,43 @@ object ReferenceDatasetUtils {
         numRows,
         datasetParams,
         datasetVoidPtr), "Dataset create from samples")
-
-
-      // 2. Serialize the raw dataset to a native buffer
-      val datasetHandle: SWIGTYPE_p_void = lightgbmlib.voidpp_value(datasetVoidPtr)
-      LightGBMUtils.validate(lightgbmlib.LGBM_DatasetSerializeReferenceToBinary(
-        datasetHandle,
-        bufferHandlePtr,
-        lenPtr), "Serialize ref")
-      val bufferLen: Int = lightgbmlib.intp_value(lenPtr)
-      log.info(s"Created serialized reference dataset of length $bufferLen")
-
-      // The dataset is now serialized to a buffer, so we don't need original
-      LightGBMUtils.validate(lightgbmlib.LGBM_DatasetFree(datasetHandle), "Free Dataset")
-
-      // This will also free the buffer
-      toByteArray(bufferHandlePtr, bufferLen)
-    }
-    finally {
-      sampledData.delete()
+      lightgbmlib.voidpp_value(datasetVoidPtr)
+    } finally {
+      // Frees the void** container only. The Dataset it points at is freed by LGBM_DatasetFree.
       lightgbmlib.delete_voidpp(datasetVoidPtr)
-      lightgbmlib.delete_voidpp(bufferHandlePtr)
-      lightgbmlib.delete_intp(lenPtr)
     }
   }
 
+  private def setFeatureNamesIfProvided(datasetHandle: SWIGTYPE_p_void,
+                                        featureNames: Option[Array[String]],
+                                        numCols: Int,
+                                        log: Logger): Unit = {
+    featureNames.foreach { names =>
+      if (names.length != numCols) {
+        // LGBM_DatasetSetFeatureNames reads numCols entries from the array, so a shorter array
+        // would be an out-of-bounds native read. Skip naming rather than risk it; LightGBM then
+        // falls back to its own generated names, which is the behavior prior to this feature.
+        log.warn(s"Skipping feature names on reference dataset: got ${names.length} names " +
+          s"for $numCols feature columns.")
+      } else if (names.nonEmpty) {
+        log.info(s"Setting ${names.length} feature names on reference dataset")
+        LightGBMUtils.validate(lightgbmlib.LGBM_DatasetSetFeatureNames(datasetHandle, names, numCols),
+          "Dataset set feature names")
+      }
+    }
+  }
+
+  private def serializeReference(datasetHandle: SWIGTYPE_p_void,
+                                 bufferHandlePtr: SWIGTYPE_p_p_void,
+                                 lenPtr: SWIGTYPE_p_int,
+                                 log: Logger): Array[Byte] = {
+    LightGBMUtils.validate(lightgbmlib.LGBM_DatasetSerializeReferenceToBinary(
+      datasetHandle, bufferHandlePtr, lenPtr), "Serialize ref")
+    val bufferLen: Int = lightgbmlib.intp_value(lenPtr)
+    log.info(s"Created serialized reference dataset of length $bufferLen")
+    toByteArray(bufferHandlePtr, bufferLen)
+  }
+
   def getInitializedReferenceDataset(ctx: PartitionTaskContext): LightGBMDataset = {
     // The definition is broadcast from Spark, so retrieve it
     val serializedDataset: Array[Byte] = ctx.trainingCtx.serializedReferenceDataset.get