fix: generate TabularSHAP samples lazily

[Lenin Mookiah (leninworld)]

Related Issues/PRs

Related to #2206. The reported eager-buffering failure is reproduced with synthetic notebook-shaped data; the reporter's original notebook run has not been independently reproduced.

What changes are proposed in this pull request?

Generate TabularSHAP sample rows incrementally instead of materializing a sample array for each instance/background pair. This avoids the large buffers created when Spark fuses the broadcast join and array expansion.

• Contributor's change: use Dataset.mapPartitions with a row encoder. Preserve the sampler, sample budget, weights, coalition vectors, column order, target-class handling and public API.
• Maintainer follow-up: preserve nulls recursively inside background arrays, maps and structs; add a real Spark sampling-path regression and public nested-null regressions; explain remaining scoring cost in the notebook. The submitted row encoder otherwise rejected valid nested background nulls.
• Preserve the contributor's commits. Merge the current target rather than rewriting the fork branch.

The lazy iterator test remains as a separate consumption check. The added Spark-path test executes createSamples with two broadcast background rows and checks encoding before the first result, without heap-size or timing assertions.

How is this patch tested?

• I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.

Maintainer validation: Windows amd64, JDK 11, Spark 3.5.0, Scala 2.12.17.

Command/check	Result
sbt core/compile core/Test/compile	Passed
sbt core/scalastyle core/Test/scalastyle	Passed, zero errors/warnings
sbt "core/testOnly com.microsoft.azure.synapse.ml.explainers.split1.TabularSHAPExplainerSuite com.microsoft.azure.synapse.ml.explainers.split1.VectorSHAPExplainerSuite"	13 passed, none skipped
sbt core/codegen	Passed; generated Python wrapper import/public-setter smoke passed with PySpark 3.5.0
black --check --workers 2 --extend-exclude 'docs/' . (22.3.0)	Passed, 209 files unchanged; explicit asyncio loop initialization used for the host Python
Edited notebook JSON parsing	Passed

Fail-before/pass-after: the nested-null public-transform test fails on the submitted head at the row encoder and passes after the follow-up. Replaying the original eager createSamples implementation makes the Spark-path regression fail: it encodes 20,000 feature values before returning one sample. The lazy implementation passes. Scalar-null/non-null controls and genuine incompatible-type rejection were also checked.

Memory and runtime: a local public-Spark harness used 13 mixed string/integer features, 5 explained rows, 100 background rows, 5,000 samples and 8 local workers. With a 1 GiB heap, the pre-PR sampling implementation terminates with OutOfMemoryError: Java heap space; the lazy path completes all 5 rows in 21.7 seconds, with finite explanations and the expected baseline/sum. This is a replay of the old sampling method with unchanged shared code, not a full old-release installation.

With a 4 GiB heap, the same workload's two warm runs were 10.44/10.16 seconds before and 11.38/10.33 seconds after (about 5% slower on their median); cold runs were 17.72/20.91 seconds. A smaller 20-background-row workload showed about 17% warm overhead. This is a measured memory/runtime trade-off, not a universal speedup. Maintainer changes add schema work once per plan, not per sample.

The contributor separately reported Docker/Linux fail-before/pass-after runs on the original implementation and released 1.0.4/1.1.3, plus deterministic SHAP parity checks:

CI for 9e66546: Azure build 240826424, triggered with /azp run, completed against merge 856d25b (parents: current master d51118a and this head). Published results contain 3,746 passed, 22 skipped, zero failed. Both added regressions ran and passed in TabularSHAP's 8/8 test run; VectorSHAP passed 5/5, and Databricks CPU/GPU notebook tests passed 43/43 and 4/4.

Azure reports partiallySucceeded because three non-gating dependency-cache tasks encountered SSL errors in UnitTests misc; that job's 502 tests and test publishing passed. All 22 skipped test cases also skipped in the successful prior master build 240603209. The website coverage publisher reported no coverage file; website tests passed. Fabric E2E and Internal compatibility were intentionally excluded by existing fork-PR conditions.

GitHub compile/style checks, website build/tests, dependency review and both configured CodeQL analyses passed. Website deployment is intentionally skipped for PRs. The aggregate CodeQL check is neutral because a separate master API-upload configuration is absent on this PR. Current-head automated review has zero open or suppressed findings; its earlier CI-pending note is now superseded by these results.

Spark Connect and port-runtime validation were not performed locally. VectorSHAP's existing sampling implementation is outside this focused change. Contributor sign-off and required human review remain outstanding.

Does this PR change any dependencies?

• No. You can skip this section.
• Yes. Make sure the dependencies are resolved correctly, and list changes here.

Does this PR add a new feature? If so, have you added samples on website?

• No. You can skip this section.
• Yes. Make sure you have added samples following below steps.

1. Find the corresponding markdown file for your new feature in website/docs/documentation folder.
   Make sure you choose the correct class estimators/transformers and namespace.
2. Follow the pattern in markdown file and add another section for your new API, including pyspark, scala (and .NET potentially) samples.
3. Make sure the DocTable points to correct API link.
4. Navigate to website folder, and run npm start to make sure the website renders correctly.
5. Don't forget to add <!--pytest-codeblocks:cont--> before each python code blocks to enable auto-tests for python samples.
6. Make sure the WebsiteSamplesTests job pass in the pipeline.

[azure-pipelines]

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

[github-actions]

Hey Lenin Mookiah (@leninworld) 👋!
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.

[Copilot]

🟡 Changes recommended

The new RDD-based path breaks Spark Connect compatibility and violates repository requirements.

1 open finding

• Replace unsupported Dataset.rdd with mapPartitions for Spark Connect · New

What changed in this PR

Makes TabularSHAP sample generation lazy to avoid executor heap exhaustion.

Changes:

• Replaces array-based sampling with partition-level lazy iteration.
• Preserves nullable background values.
• Adds regressions for laziness and null handling.

File	Description
TabularSHAP.scala	Implements lazy sample generation.
TabularSHAPExplainerSuite.scala	Tests lazy generation and nullable backgrounds.

🧠 Review effort: Balanced

---

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

[Copilot]

🟡 Changes recommended

The laziness test covers only a helper and does not protect the production transformer path from rematerialization.

1 open finding

• Add path-level regression test preventing sample iterator buffering · New

1 resolved since last review

• Replace unsupported Dataset.rdd with mapPartitions for Spark Connect

🧠 Review effort: Balanced

[Copilot]

🔵 Needs a closer look

The nullable-background regression can pass vacuously when no results are returned.

1 open finding

• Add path-level regression test preventing sample iterator buffering

Previously missed (1)

In code that hasn't changed since last review

Test does not assert expected row count

core/​src/​test/​scala/​com/​microsoft/​azure/​synapse/​ml/​explainers/​split1/​TabularSHAPExplainerSuite.scala:150

This regression passes vacuously if the nullable-background path returns no explanations, because foreach performs no assertions on an empty result. Assert the expected row count before validating the values so the test proves all three input rows survive the new expansion path.

🧠 Review effort: Balanced

[Copilot]

🟢 Approval recommended

No blocking issues remain, and reported targeted tests, parity checks, and memory-constrained validation support the focused change.

1 open finding

• Add path-level regression test preventing sample iterator buffering

🧠 Review effort: Balanced

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

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

[Copilot]

🔵 Needs a closer look

Current-head CI remains pending for the changed Spark sampling and row-encoding path despite reported successful local validation.

0 open findings

1 resolved since last review

• Add path-level regression test preventing sample iterator buffering

🧠 Review effort: Balanced

[Rana Singh (ranadeepsingh)]

Thanks Lenin Mookiah (@leninworld) for fixing the eager sample buffering! Your lazy sampling path avoids an out-of-memory failure in a synthetic workload shaped like #2206. I have kept the issue linkage conditional because I have not independently reproduced the reporter's original notebook.

I preserved your commits and lazy sampling design, merged the current target without rewriting your branch, and added 9e66546:

• Make the sample encoder's nested struct/array/map types nullable so valid background nulls remain supported.
• Add a real Spark sampling-path regression and a public-transform nested-null regression.
• Add a short notebook note about streaming's memory benefit and remaining model-scoring cost.

Both regressions have fail-before/pass-after evidence. The old sampling method replay ran out of heap at 1 GiB; the lazy path completed. At 4 GiB, two warm runs showed about 5% median overhead at the larger workload and 17% at the smaller one. The description records this as a measured memory/runtime trade-off, not a universal speedup.

Azure build 240826424, triggered with /azp run, has finished for this head's verified merge with master: 3,746 passed, 22 pre-existing skips, zero failed. Its partiallySucceeded status is from non-gating cache SSL warnings; the affected job's tests passed. Published TabularSHAP results confirm both new regressions passed:

• TabularSHAP createSamples streams encoded rows before expanding the sample budget
• TabularKernelSHAP preserves nulls inside background features

All 13 TabularSHAP/VectorSHAP tests passed. Required GitHub checks pass, and the current-head automated review has no open or suppressed findings. Its CI-pending note is now superseded by the completed build.

Please review these follow-ups and sign off if they fit your intent. If they do not fit your vision for the PR, feel free to revert my commit, or let me know and I will revert it. Contributor sign-off and required human review are still outstanding; I have not merged the PR.

[Lenin Mookiah (leninworld)]

Thanks Lenin Mookiah (Lenin Mookiah (@leninworld)) for fixing the eager sample buffering! Your lazy sampling path avoids an out-of-memory failure in a synthetic workload shaped like #2206. I have kept the issue linkage conditional because I have not independently reproduced the reporter's original notebook.

I preserved your commits and lazy sampling design, merged the current target without rewriting your branch, and added 9e66546:

• Make the sample encoder's nested struct/array/map types nullable so valid background nulls remain supported.
• Add a real Spark sampling-path regression and a public-transform nested-null regression.
• Add a short notebook note about streaming's memory benefit and remaining model-scoring cost.

Both regressions have fail-before/pass-after evidence. The old sampling method replay ran out of heap at 1 GiB; the lazy path completed. At 4 GiB, two warm runs showed about 5% median overhead at the larger workload and 17% at the smaller one. The description records this as a measured memory/runtime trade-off, not a universal speedup.

Azure build 240826424, triggered with /azp run, has finished for this head's verified merge with master: 3,746 passed, 22 pre-existing skips, zero failed. Its partiallySucceeded status is from non-gating cache SSL warnings; the affected job's tests passed. Published TabularSHAP results confirm both new regressions passed:

• TabularSHAP createSamples streams encoded rows before expanding the sample budget
• TabularKernelSHAP preserves nulls inside background features

All 13 TabularSHAP/VectorSHAP tests passed. Required GitHub checks pass, and the current-head automated review has no open or suppressed findings. Its CI-pending note is now superseded by the completed build.

Please review these follow-ups and sign off if they fit your intent. If they do not fit your vision for the PR, feel free to revert my commit, or let me know and I will revert it. Contributor sign-off and required human review are still outstanding; I have not merged the PR.

Thanks again, Rana Singh (@ranadeepsingh), for the follow-up. I independently rechecked 9e66546 in Linux amd64 Docker with Java 11.0.32, Spark 3.5.0, and Scala 2.12.18, and the additions fit the intent of the PR:

• TabularSHAPExplainerSuite passed 8/8 and VectorSHAPExplainerSuite passed 5/5. Core compilation, test compilation, packaging, and both Scalastyle checks also passed.
• Without recursive nullability, the nested-null regression fails with The 0th field 'value' of input row cannot be null. With the original array/explode sampling path, it instead produces an incorrect base value of approximately 0 rather than 10.
• Replaying the original eager createSamples implementation makes the new Spark-path regression fail after encoding 20,000 feature values before returning its first sample.
• With a 1 GiB driver heap and 16 cores, the notebook-shaped synthetic run—5 explained rows × 100 background rows × 5,000 samples—completed all five rows in 77.9 seconds without an out-of-memory error. Fully enumerated parity cases matched the pre-PR implementation within 3e-21.
  The recursive nullability correction, complementary Spark-path coverage, and documentation updates all align with the intended bounded-memory change. Thank you for strengthening the implementation and its regression coverage.