fix(lightgbm): de-duplicate feature names so training no longer fails on repeated columns

[Rana Singh (ranadeepsingh)]

Why

Fixes #2242.

LightGBM refuses to build a Dataset whose feature names repeat. Spark can hand it repeated names in two ways — through the AttributeGroup metadata of the features column, or through a user-supplied slotNames — and either one killed training:

Dataset set feature names call failed in LightGBM with error:
Feature (Column_) appears more than one time.

The names are perfectly usable; only their uniqueness is a LightGBM requirement. So the fix is to make them unique before LightGBM ever sees them, rather than failing the job.

What changed

slotNames param ──┐
                  ├─► getSlotNamesWithMetadata ─► ensureUniqueFeatureNames ─► LightGBM
AttributeGroup ───┘                                (rename dupes, keep order)

Change	Detail
De-duplicate feature names	Both metadata-derived and explicit names now pass through ensureUniqueFeatureNames, which suffixes repeats while preserving order (LightGBM feature names are positional). A warning lists what was renamed.
Name the streaming reference Dataset	createReferenceDatasetFromSample now takes the feature names and applies them, so the reference Dataset matches the per-partition Datasets. Split into createDatasetFromSamples / setFeatureNamesIfProvided / serializeAndCleanup.
Validate slot names unconditionally	validateSlotNames previously ran only when AttributeGroup metadata was present, so explicit slotNames were never checked.
Validate before the first native call	Naming the reference Dataset introduces an LGBM_DatasetSetFeatureNames call earlier in trainOneDataBatch than the old validation point, so validation moved up to where featuresSchema resolves.
Fix native handle leaks	createDatasetFromSamples frees its voidpp handle in a finally; the reference Dataset is freed in a finally spanning naming and serialization.

Correctness notes

Three of these were found by testing, not by reading — each is covered by a test confirmed to fail against the simpler implementation.

Subtlety	Why the obvious approach is wrong
Generated names can collide with later originals	Remembering only emitted names turns ["Column_", "Column_", "Column__1"] into a collision on Column__1 — the exact error this PR fixes. All original names are reserved up front.
LightGBM normalizes spaces to underscores before checking duplicates	"a b" and "a_b" are two strings in Scala but one feature natively. Uniqueness is now decided on the normalized form while the original names are what get emitted.
Validation must precede the native call	LightGBM's own rejection is Do not support special JSON characters in feature name. — it never says which column. Validating late let that opaque error win and broke …bad column names fails early. It now fails in 0.3s with an actionable IllegalArgumentException.

Two smaller corrections: the invalid-name message listed \ as disallowed, but neither the regex nor LightGBM's CheckAllowedJSON rejects it (which accepts exactly " , : [ ] { }) — the message now matches reality, and the regex is deliberately unchanged. And when the name count disagrees with numCols, names are skipped with a warning rather than risking an out-of-bounds native read.

No bulk-mode retry. An earlier revision caught the failure and retried with dataTransferMode=bulk. Removed: BulkPartitionTask calls setFeatureNames with the same duplicated names, so the retry fails identically — and it mutated the estimator's param mid-fit, visible to concurrent callers. Deduplicating up front removes the need for a fallback.

Behavior changes

• Repeated feature names are renamed instead of fatal.
• validateSlotNames now also checks explicitly supplied slotNames. A job passing slotNames containing " , : [ ] { } with no attribute metadata previously reached LightGBM unchecked; it now fails fast with the existing, actionable error. This is strictly a better diagnostic — those names were already rejected by LightGBM, just opaquely and later.

How was this patch tested?

Added to VerifyLightGBMCommon:

Test	Covers
Verify duplicate feature names are handled correctly	Repeats via AttributeGroup metadata
Verify explicit slotNames parameter is used	Explicit slotNames still applied
Verify duplicate explicit slotNames are made unique	Repeats via user-supplied slotNames
Verify a generated slot name cannot collide with a later original name	Collision case — verified to fail against the naive implementation
Verify names differing only by space vs underscore are made unique	Native normalization — verified to fail before the fix

Locally on Java 11 with the full CI datasets: VerifyLightGBMCommon 9/9; the stream/bulk regressor, ranker and classifier suites plus NetworkManagerSuite and TrainUtilsSuite 52/52 (including …bad column names fails early); lightgbm/scalastyle and lightgbm/Test/scalastyle 0 errors.

Checklist

• Tests added
• No dependency changes
• Bug fix, not a new feature — no website samples needed

[Rana Singh (ranadeepsingh)]

/azp run

[github-actions]

Hey Rana Singh (@ranadeepsingh) 👋!
Thank you so much for contributing to our repository 🙌.
Someone from SynapseML Team will be reviewing this pull request soon.

We use semantic commit messages to streamline the release process.
Before your pull request can be merged, you should make sure your first commit and PR title start with a semantic prefix.
This helps us to create release messages and credit you for your hard work!

Examples of commit messages with semantic prefixes:

• fix: Fix LightGBM crashes with empty partitions
• feat: Make HTTP on Spark back-offs configurable
• docs: Update Spark Serving usage
• build: Add codecov support
• perf: improve LightGBM memory usage
• refactor: make python code generation rely on classes
• style: Remove nulls from CNTKModel
• test: Add test coverage for CNTKModel

To test your commit locally, please follow our guild on building from source.
Check out the developer guide for additional guidance on testing your change.

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 87.03%. Comparing base (c570407) to head (d9a915e).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2508      +/-   ##
==========================================
+ Coverage   86.76%   87.03%   +0.27%     
==========================================
  Files         338      338              
  Lines       18785    18817      +32     
  Branches     1804     1803       -1     
==========================================
+ Hits        16299    16378      +79     
+ Misses       2486     2439      -47     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Copilot]

Pull request overview

This PR fixes LightGBM training failures caused by duplicate feature names (including collisions after LightGBM’s space-to-underscore normalization) by ensuring slot/feature names are unique and validated before the first native call that consumes them. It also propagates feature names onto the streaming reference Dataset and refactors reference-dataset creation for clearer lifecycle management.

Changes:

• De-duplicate feature/slot names (from both slotNames and AttributeGroup metadata) using LightGBM-style normalization rules.
• Validate slot names earlier in fit to fail fast with actionable errors before native feature-name calls.
• Name the streaming reference Dataset consistently and add regression tests covering duplicate/collision scenarios.

Show a summary per file

File	Description
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/LightGBMBase.scala	Ensure unique feature names (including normalized collisions) and validate names before native calls.
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/dataset/ReferenceDatasetUtils.scala	Refactor reference Dataset creation and (now) apply feature names to reference datasets.
lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/VerifyLightGBMCommon.scala	Add tests for duplicate feature names from metadata and explicit slotNames, including collision edge cases.

Review details

Tip

Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Suppressed comments (1)

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

• serializeAndCleanup frees the Dataset handle. If the caller also adds defensive cleanup (needed for failures before serialization), this risks double-free. Prefer making Dataset lifetime ownership explicit: serialize here, but free the Dataset in the caller's finally so all failure paths free exactly once.

    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")
    LightGBMUtils.validate(lightgbmlib.LGBM_DatasetFree(datasetHandle), "Free Dataset")
    toByteArray(bufferHandlePtr, bufferLen)

• Files reviewed: 3/3 changed files
• Comments generated: 2
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Copilot]

Review details

Suppressed comments (2)

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

• getSlotNamesWithMetadata (and therefore ensureUniqueFeatureNames) is invoked multiple times during a single fit in streaming mode (e.g., once inside calculateRowStatistics and again when building the TrainingContext in executeTraining). Since ensureUniqueFeatureNames logs a warning when it renames duplicates, the same warning can be emitted more than once per batch, which is noisy and can look like multiple independent problems. Consider computing the (unique) feature name array once per batch and threading it through to all consumers.

    // Get feature names to set on the reference dataset (ensures unique names for Spark 3.5+)
    val featureNames = getSlotNamesWithMetadata(featuresSchema)

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

• validateSlotNames currently validates the de-duplicated names returned by getSlotNamesWithMetadata. If the input contains invalid characters and also has duplicates, ensureUniqueFeatureNames can append a suffix (e.g., bad,name_1), and the thrown IllegalArgumentException will list names the user never supplied, making the error harder to act on. Validate the raw names first (either explicit slotNames or metadata-derived names) and only then de-duplicate for LightGBM.

This issue also appears on line 585 of the same file.

  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) {

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Copilot]

Review details

Suppressed comments (1)

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

• getSlotNamesWithMetadata assumes every Attribute entry is non-null. Spark attribute metadata can contain null slots (this file already handles case (null, _) in getCategoricalIndexes), and a null here would NPE during name extraction and break training.

Guard null attributes and fall back to the default index-based name when an entry is null.

          val colNames = attributes.indices.map(_.toString).toArray
          attributes.foreach(attr =>
            attr.index.foreach(index => colNames(index) = attr.name.getOrElse(index.toString)))
          // Ensure unique feature names to avoid LightGBM error:

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Rana Singh (ranadeepsingh)]

Good catch — you were right, and it was worse than a theoretical risk. Fixed by centralizing the guard in LightGBMDataset.setFeatureNames as you suggested.

I reverted the guard and ran the new tests to confirm the exposure is real rather than assume it:

\
java.lang.Exception: Dataset set feature names call failed in LightGBM with error:
basic_string: construction from null is not valid
\\

LGBM_DatasetSetFeatureNames reads numCols entries, runs off the end of the Java String[], gets null, and std::string construction from null throws — killing the executor task and failing the whole fit. It reproduces in both streaming and bulk transfer modes, exactly the paths you identified. slotNames is user-supplied and never length-checked upstream (validateSlotNames only screens for special characters), so any user passing a mismatched length hits this.

Changes:

• Length guard moved into LightGBMDataset.setFeatureNames, covering BulkPartitionTask, StreamingPartitionTask, and ReferenceDatasetUtils.createReferenceDatasetFromSample.
• ReferenceDatasetUtils.setFeatureNamesIfProvided keeps its own guard since it operates on a raw SWIGTYPE_p_void handle and cannot delegate.
• Two regression tests, one per transfer mode.

Validation: VerifyLightGBMCommon 11/11 pass with the guard, 2 fail without it. lightgbm/scalastyle and lightgbm/Test/scalastyle clean.

/azp run

[Copilot]

Review details

Suppressed comments (2)

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

• The code skips naming when featureNamesArray.length != numCols, but the preceding comment only explains the short-array risk. Since the condition also skips longer arrays, the comment should be updated to reflect the full mismatch behavior to avoid confusion.

        // 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.

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

• The guard skips naming when names.length != numCols, but the comment only justifies the short-array case. This is misleading for future maintainers (the current condition also skips when the array is longer than numCols). Update the comment to describe the full mismatch behavior and rationale.

      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.

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).