fix: preserve partition topology when counting Spark rows

[ranadeepsingh]

Summary

Fix per-partition row counting so it preserves the exact partition topology used by downstream LightGBM training tasks.

Fixes #2278.

This replaces only the proven intent behind #2282; it does not reuse that PR's corrupted diff or modify the old PR.

Root cause and primary failure

ClusterUtil.getNumRowsPerPartition counted a separate df.select(lit(0)) query. With adaptive execution enabled, Spark can coalesce that cheap projected query differently from the real training DataFrame.

A current-master AQE fixture produced:

• real DataFrame: 20 partitions
• projected counting query: 4 partitions
• real task IDs indexing the shortened array: primary ArrayIndexOutOfBoundsException, including Index 18 out of bounds for length 4

That is the same partition-ID-shaped failure reported in #2278. Once a worker dies on this index access, the remaining LightGBM workers report secondary connection failures.

Change

Count on the original DataFrame RDD rather than constructing a separately optimized literal projection. This intentionally trades projection pruning for the required topology correctness.

The regression test:

1. creates a 20-partition shuffle whose full rows remain above AQE's coalescing target;
2. proves the old literal projection coalesces to fewer partitions; and
3. verifies returned counts exactly match every real DataFrame partition.

Separately investigated hypotheses

• Feature width: malformed short dense rows can independently fail in SampledData.pushRow, but no evidence ties inconsistent feature widths to [BUG]java.lang.ArrayIndexOutOfBoundsException on multi-node cluster run #2278, so no unrelated validation change is included.
• Barrier/network startup: the existing barrier TrainValidationSplit test passes; connection errors follow the primary task failure rather than precede it.
• Native pointer lifecycle: the 300-iteration custom-objective path remains functional. The cleanup proposed in fix: java.lang.ArrayIndexOutOfBoundsException in multi-node cluster run #2282 cannot explain a partition-index exception, and no stable native-memory regression was established, so it is not included.

Validation

• core/testOnly com.microsoft.azure.synapse.ml.core.utils.VerifyClusterUtil
• core/compile, core/Test/compile
• core/scalastyle, core/Test/scalastyle
• lightgbm/Test/compile
• lightgbm/scalastyle, lightgbm/Test/scalastyle
• LightGBM streaming barrier TrainValidationSplit test
• LightGBM streaming custom-loss test (300 iterations)

Review

• No public method signatures changed.
• No serialization, authentication, network, file, or secret-handling behavior changed.
• The test uses an isolated Spark session and adds no external service dependency.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[github-actions]

Hey @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.

[ranadeepsingh]

/azp run

[azure-pipelines]

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

[Copilot]

Pull request overview

Fixes an AQE-induced mismatch where getNumRowsPerPartition could observe a different partition topology than the DataFrame used by downstream training tasks (e.g., LightGBM), leading to partition-indexing errors.

Changes:

• Update ClusterUtil.getNumRowsPerPartition to count rows using the original DataFrame’s RDD rather than a separately-optimized projection.
• Add a Scala regression test that constructs an AQE coalescing fixture and asserts per-partition counts match the real DataFrame topology.

Show a summary per file

File	Description
core/src/main/scala/com/microsoft/azure/synapse/ml/core/utils/ClusterUtil.scala	Switches row counting to use the DataFrame’s own RDD to preserve partition topology under AQE.
core/src/test/scala/com/microsoft/azure/synapse/ml/core/utils/VerifyClusterUtil.scala	Adds an AQE regression test validating the partition-topology-preserving row counts.

Review details

Suppressed comments (1)

core/src/test/scala/com/microsoft/azure/synapse/ml/core/utils/VerifyClusterUtil.scala:52

• To match the try { added at the start of the test, add a finally block that closes the session. Without a finally, a failed assertion would skip cleanup.

    assert(projected.length < expected.length,
      s"Fixture must expose adaptive coalescing: ${projected.length} projected vs ${expected.length} actual")
    assert(actual.sameElements(expected),
      s"Expected partition counts ${expected.mkString(",")}, got ${actual.mkString(",")}")
  }

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

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.78%. Comparing base (7d9fabc) to head (1ce1ec4).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2593      +/-   ##
==========================================
- Coverage   84.79%   84.78%   -0.02%     
==========================================
  Files         334      334              
  Lines       17806    17806              
  Branches     1623     1623              
==========================================
- Hits        15099    15096       -3     
- Misses       2707     2710       +3     

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

[ranadeepsingh]

/azp run

[azure-pipelines]

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