Skip to content

fix: make EnsembleByKey transformSchema match output - #2575

Merged
Rana Singh (ranadeepsingh) merged 6 commits into
microsoft:masterfrom
fallintoplace:fix/ensemble-by-key-schema
Aug 12, 2026
Merged

fix: make EnsembleByKey transformSchema match output#2575
Rana Singh (ranadeepsingh) merged 6 commits into
microsoft:masterfrom
fallintoplace:fix/ensemble-by-key-schema

Conversation

@fallintoplace

@fallintoplace Minh Vu (fallintoplace) commented Jul 26, 2026

Copy link
Copy Markdown
Contributor

Problem

EnsembleByKey.transformSchema describes the opposite shape from transform:

  • with collapseGroup=true, transform returns grouping keys and aggregate columns, while transformSchema keeps every input column
  • with collapseGroup=false, transform joins aggregates back to the input, while transformSchema keeps only fields selected from the aggregation column set

This can break pipeline schema validation because the declared schema differs from the DataFrame produced at runtime. The old implementation also fails when transformSchema is called before transform without explicit colNames, and declares FloatType for scalar means that Spark produces as DoubleType.

Changes

  • initialize default output names consistently from both transform and transformSchema
  • derive grouping fields from keys
  • return grouping fields followed by aggregate fields for collapsed groups
  • match Spark join ordering for non-collapsed groups: keys, remaining non-overwritten input fields, then aggregates
  • declare scalar mean outputs as DoubleType for both float and double inputs
  • assert exact schema equality for double and float inputs in both collapse modes, with distinct key and aggregation columns and no explicit output name

Testing

  • core/testOnly com.microsoft.azure.synapse.ml.stages.EnsembleByKeySuite (38 tests passed)
  • core/testOnly com.microsoft.azure.synapse.ml.stages.EnsembleByKeyResolutionSuite (7 tests passed)
  • core/Test/scalastyle (0 findings)
  • core/Compile/scalastyle (0 findings)

Behavior change disclosure

Beyond the schema fix, this PR changes the join in collapseGroup=false mode from a null-unsafe equi-join to a null-safe one (<=>).

Rows whose grouping key is null were previously dropped from the output entirely. They are now retained. This is a data-loss fix rather than a regression, but it is an observable row-set change for anyone relying on the old behavior as an implicit null filter. Covered by an explicit test.

Copilot AI lite review requested due to automatic review settings July 26, 2026 02:05
@azure-pipelines

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

@github-actions

Copy link
Copy Markdown

Hey Minh Vu (@fallintoplace) 👋!
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 AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR updates EnsembleByKey’s transformSchema logic so it matches the actual output schema produced by transform, across both collapseGroup=true and collapseGroup=false modes, and adds a regression test to validate schema consistency.

Changes:

  • Rebuild transformSchema output fields using keys for grouping fields and consistent output ordering for both collapse modes.
  • In non-collapsing mode, model Spark join column ordering and drop/replace conflicting output column names in the predicted schema.
  • Add a test asserting transformSchema(schema) equals transform(df).schema for both collapse modes.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File Description
core/src/main/scala/com/microsoft/azure/synapse/ml/stages/EnsembleByKey.scala Adjusts transformSchema field selection and ordering to align with transform behavior.
core/src/test/scala/com/microsoft/azure/synapse/ml/stages/EnsembleByKeySuite.scala Adds schema-equality coverage for both collapseGroup modes.

Comment thread core/src/main/scala/com/microsoft/azure/synapse/ml/stages/EnsembleByKey.scala Outdated
@fallintoplace Minh Vu (fallintoplace) changed the title fix: align EnsembleByKey transform schema fix: make EnsembleByKey transformSchema match output Jul 26, 2026
@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

@codecov-commenter

Codecov Comments Bot (codecov-commenter) commented Jul 31, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 98.20789% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 86.99%. Comparing base (7088a6a) to head (9ba8644).

Files with missing lines Patch % Lines
...rosoft/azure/synapse/ml/stages/EnsembleByKey.scala 98.20% 5 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master    #2575      +/-   ##
==========================================
+ Coverage   86.81%   86.99%   +0.17%     
==========================================
  Files         336      336              
  Lines       18275    18518     +243     
  Branches     1673     1725      +52     
==========================================
+ Hits        15866    16110     +244     
+ Misses       2409     2408       -1     

☔ 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

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

@ranadeepsingh

Copy link
Copy Markdown
Collaborator

CI follow-up for f4bc2c9: Azure build 229036776 completed with 67 successful checks. The only failing job was Databricks E2E databricks-gpu, which aborted before running tests (Tests: succeeded 0, failed 0) because cluster startup terminated with INSTANCE_POOL_MAX_CAPACITY_REACHED / Instance pool is full. All CPU Databricks jobs, core/stages unit tests, Python/R tests, style, publish, and Fabric E2E passed. This is infrastructure-only, so I am not issuing another /azp run.

@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

Rana Singh (ranadeepsingh) added a commit to fallintoplace/SynapseML that referenced this pull request Aug 2, 2026
AB#2575

## Summary
Add Round 3 code review artifact for SynapseML PR microsoft#2575 focusing on edge cases and robustness.

## Prompting Intent
Perform Round 3 of the DIRECT sequential six-round review gauntlet for SynapseML PR microsoft#2575. Focus on edge cases, robustness, error handling, boundary conditions, concurrency/session behavior, and failure modes.

## Linked Sources
- Prompt: C:\Users\singhrana\.copilot\session-state\c2ea157b-c36a-4cbc-896d-e8caa3ea05bd\files\pr-2575\prompts\review-round-3.md
- Review artifact: C:\Users\singhrana\Documents\SynapseML-pr-2575\reviews\pr-2575\task-2575-attempt-1-review-3-gemini-3.6-flash.md

## Rationale
Recorded findings for edge cases, null join behavior, session case sensitivity, and vector aggregate nullability in the Round 3 review artifact as instructed.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Rana Singh (ranadeepsingh) added a commit to fallintoplace/SynapseML that referenced this pull request Aug 2, 2026
## Summary
Make EnsembleByKey schema resolution match Spark output for collapsed and
non-collapsed transforms, including scalar and vector means, nested and qualified
keys, duplicate attributes, null keys, overwrite behavior, output naming, metadata,
nullability, and unsupported inputs. Add generated-Python compatibility coverage
and preserve the complete six-round review audit trail.

## Prompting Intent
Take end-to-end ownership of GitHub PR microsoft#2575 in an isolated worktree. Independently
verify and harden EnsembleByKey schema behavior, run the required six-round DIRECT
sequential multi-model review, fix and rerun every finding until clean, validate
Scala and generated-language surfaces, avoid generated target edits and history
rewrites, and prepare the existing PR branch for normal push and Azure validation.

## Linked Sources
- Pull request: microsoft#2575
- Round 1 review: reviews/pr-2575/task-2575-attempt-1-review-1-claude-opus-5.md
- Round 2 review: reviews/pr-2575/task-2575-attempt-1-review-2-gpt-5.6-sol.md
- Round 3 review: reviews/pr-2575/task-2575-attempt-1-review-3-gemini-3.6-flash.md
- Round 4 review: reviews/pr-2575/task-2575-attempt-1-review-4-claude-opus-5.md
- Round 5 review: reviews/pr-2575/task-2575-attempt-1-review-5-gpt-5.6-sol.md
- Round 6 review: reviews/pr-2575/task-2575-attempt-1-review-6-gemini-3.6-flash.md

## Rationale
Keep business logic Scala-first and derive the Python surface through SynapseML's
internal-wrapper override pattern. Resolve configured fields against schema ordinals
and analyzed Spark attributes so transformSchema remains optimistic where qualifier
or ExprId identity is unavailable while transform rejects genuine runtime ambiguity.
Use positional normalization and public DataFrame APIs for Spark 3.5/4.1 compatibility,
and use null-safe joins to preserve non-collapsed rows with null grouping keys.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

Rana Singh (ranadeepsingh) added a commit to fallintoplace/SynapseML that referenced this pull request Aug 4, 2026
## Summary
Remove the six generated review audit files from the committed PR contents while
leaving all reviewed implementation, documentation, and regression tests unchanged.

## Prompting Intent
The engineer asked to prevent the generated review files from being merged. Preserve
the existing PR history, avoid force-pushing, and remove only those committed artifacts.

## Linked Sources
- Pull request: microsoft#2575
- User request: remove the review files from the committed files

## Rationale
Use a follow-up deletion commit rather than rewriting the existing commit so the
shared PR branch remains safe and auditable. The review process remains reflected in
the conversation and prior commit history, but the final merged tree contains only
the product changes and their tests.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

Rana Singh (ranadeepsingh) added a commit to fallintoplace/SynapseML that referenced this pull request Aug 4, 2026
## Summary
Merge upstream master at `617ad0f8fc8932d93fc6185cb2547e5ab2c2719a`
into the existing EnsembleByKey PR branch.

## Prompting Intent
The engineer asked to update the published PR branch with the latest master changes.
Preserve branch history and use a normal push rather than rebasing or force-pushing.

## Linked Sources
- Pull request: microsoft#2575
- Upstream repository: https://github.com/microsoft/SynapseML

## Rationale
Use a merge commit because the branch is already published and the engineer selected
the history-preserving option. This incorporates current upstream changes without
rewriting commits reviewed and validated on the shared PR branch.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

@azure-pipelines

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

Copilot AI review requested due to automatic review settings August 11, 2026 09:10
Rana Singh (ranadeepsingh) added a commit to fallintoplace/SynapseML that referenced this pull request Aug 11, 2026
AB#2575

## Summary
Add Round 3 code review artifact for SynapseML PR microsoft#2575 focusing on edge cases and robustness.

## Prompting Intent
Perform Round 3 of the DIRECT sequential six-round review gauntlet for SynapseML PR microsoft#2575. Focus on edge cases, robustness, error handling, boundary conditions, concurrency/session behavior, and failure modes.

## Linked Sources
- Prompt: C:\Users\singhrana\.copilot\session-state\c2ea157b-c36a-4cbc-896d-e8caa3ea05bd\files\pr-2575\prompts\review-round-3.md
- Review artifact: C:\Users\singhrana\Documents\SynapseML-pr-2575\reviews\pr-2575\task-2575-attempt-1-review-3-gemini-3.6-flash.md

## Rationale
Recorded findings for edge cases, null join behavior, session case sensitivity, and vector aggregate nullability in the Round 3 review artifact as instructed.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Rana Singh (ranadeepsingh) added a commit to fallintoplace/SynapseML that referenced this pull request Aug 11, 2026
## Summary
Make EnsembleByKey schema resolution match Spark output for collapsed and
non-collapsed transforms, including scalar and vector means, nested and qualified
keys, duplicate attributes, null keys, overwrite behavior, output naming, metadata,
nullability, and unsupported inputs. Add generated-Python compatibility coverage
and preserve the complete six-round review audit trail.

## Prompting Intent
Take end-to-end ownership of GitHub PR microsoft#2575 in an isolated worktree. Independently
verify and harden EnsembleByKey schema behavior, run the required six-round DIRECT
sequential multi-model review, fix and rerun every finding until clean, validate
Scala and generated-language surfaces, avoid generated target edits and history
rewrites, and prepare the existing PR branch for normal push and Azure validation.

## Linked Sources
- Pull request: microsoft#2575
- Round 1 review: reviews/pr-2575/task-2575-attempt-1-review-1-claude-opus-5.md
- Round 2 review: reviews/pr-2575/task-2575-attempt-1-review-2-gpt-5.6-sol.md
- Round 3 review: reviews/pr-2575/task-2575-attempt-1-review-3-gemini-3.6-flash.md
- Round 4 review: reviews/pr-2575/task-2575-attempt-1-review-4-claude-opus-5.md
- Round 5 review: reviews/pr-2575/task-2575-attempt-1-review-5-gpt-5.6-sol.md
- Round 6 review: reviews/pr-2575/task-2575-attempt-1-review-6-gemini-3.6-flash.md

## Rationale
Keep business logic Scala-first and derive the Python surface through SynapseML's
internal-wrapper override pattern. Resolve configured fields against schema ordinals
and analyzed Spark attributes so transformSchema remains optimistic where qualifier
or ExprId identity is unavailable while transform rejects genuine runtime ambiguity.
Use positional normalization and public DataFrame APIs for Spark 3.5/4.1 compatibility,
and use null-safe joins to preserve non-collapsed rows with null grouping keys.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Rana Singh (ranadeepsingh) added a commit to fallintoplace/SynapseML that referenced this pull request Aug 11, 2026
## Summary
Remove the six generated review audit files from the committed PR contents while
leaving all reviewed implementation, documentation, and regression tests unchanged.

## Prompting Intent
The engineer asked to prevent the generated review files from being merged. Preserve
the existing PR history, avoid force-pushing, and remove only those committed artifacts.

## Linked Sources
- Pull request: microsoft#2575
- User request: remove the review files from the committed files

## Rationale
Use a follow-up deletion commit rather than rewriting the existing commit so the
shared PR branch remains safe and auditable. The review process remains reflected in
the conversation and prior commit history, but the final merged tree contains only
the product changes and their tests.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review details

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

@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
AB#2575

## Summary
Add Round 3 code review artifact for SynapseML PR microsoft#2575 focusing on edge cases and robustness.

## Prompting Intent
Perform Round 3 of the DIRECT sequential six-round review gauntlet for SynapseML PR microsoft#2575. Focus on edge cases, robustness, error handling, boundary conditions, concurrency/session behavior, and failure modes.

## Linked Sources
- Prompt: C:\Users\singhrana\.copilot\session-state\c2ea157b-c36a-4cbc-896d-e8caa3ea05bd\files\pr-2575\prompts\review-round-3.md
- Review artifact: C:\Users\singhrana\Documents\SynapseML-pr-2575\reviews\pr-2575\task-2575-attempt-1-review-3-gemini-3.6-flash.md

## Rationale
Recorded findings for edge cases, null join behavior, session case sensitivity, and vector aggregate nullability in the Round 3 review artifact as instructed.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
## Summary
Make EnsembleByKey schema resolution match Spark output for collapsed and
non-collapsed transforms, including scalar and vector means, nested and qualified
keys, duplicate attributes, null keys, overwrite behavior, output naming, metadata,
nullability, and unsupported inputs. Add generated-Python compatibility coverage
and preserve the complete six-round review audit trail.

## Prompting Intent
Take end-to-end ownership of GitHub PR microsoft#2575 in an isolated worktree. Independently
verify and harden EnsembleByKey schema behavior, run the required six-round DIRECT
sequential multi-model review, fix and rerun every finding until clean, validate
Scala and generated-language surfaces, avoid generated target edits and history
rewrites, and prepare the existing PR branch for normal push and Azure validation.

## Linked Sources
- Pull request: microsoft#2575
- Round 1 review: reviews/pr-2575/task-2575-attempt-1-review-1-claude-opus-5.md
- Round 2 review: reviews/pr-2575/task-2575-attempt-1-review-2-gpt-5.6-sol.md
- Round 3 review: reviews/pr-2575/task-2575-attempt-1-review-3-gemini-3.6-flash.md
- Round 4 review: reviews/pr-2575/task-2575-attempt-1-review-4-claude-opus-5.md
- Round 5 review: reviews/pr-2575/task-2575-attempt-1-review-5-gpt-5.6-sol.md
- Round 6 review: reviews/pr-2575/task-2575-attempt-1-review-6-gemini-3.6-flash.md

## Rationale
Keep business logic Scala-first and derive the Python surface through SynapseML's
internal-wrapper override pattern. Resolve configured fields against schema ordinals
and analyzed Spark attributes so transformSchema remains optimistic where qualifier
or ExprId identity is unavailable while transform rejects genuine runtime ambiguity.
Use positional normalization and public DataFrame APIs for Spark 3.5/4.1 compatibility,
and use null-safe joins to preserve non-collapsed rows with null grouping keys.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
## Summary
Remove the six generated review audit files from the committed PR contents while
leaving all reviewed implementation, documentation, and regression tests unchanged.

## Prompting Intent
The engineer asked to prevent the generated review files from being merged. Preserve
the existing PR history, avoid force-pushing, and remove only those committed artifacts.

## Linked Sources
- Pull request: microsoft#2575
- User request: remove the review files from the committed files

## Rationale
Use a follow-up deletion commit rather than rewriting the existing commit so the
shared PR branch remains safe and auditable. The review process remains reflected in
the conversation and prior commit history, but the final merged tree contains only
the product changes and their tests.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot AI review requested due to automatic review settings August 11, 2026 17:29

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review details

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

@ranadeepsingh

Copy link
Copy Markdown
Collaborator

/azp run

@azure-pipelines

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

@ranadeepsingh
Rana Singh (ranadeepsingh) merged commit feea272 into microsoft:master Aug 12, 2026
78 checks passed
@ranadeepsingh

Copy link
Copy Markdown
Collaborator

Thanks for this contribution, Minh Vu (@fallintoplace), and for sticking with it through the review cycle — this one is now merged. 🎉

It was a genuinely good catch. EnsembleByKey.transformSchema was advertising a schema that did not match what transform actually produced, which is the kind of bug that stays invisible until something downstream depends on the declared schema — pipeline validation, or a PipelineModel being inspected before it ever sees data. Making the two agree is the right fix, and doing it in transformSchema rather than papering over it at the call site is what makes it hold up.

The added coverage is what made it easy to land with confidence: EnsembleByKeySuite and EnsembleByKeyResolutionSuite came out to 46 Scala tests between them, plus the new Python test, and all of it ran green in CI before merge.

Appreciate you taking the time to contribute to SynapseML. Hope to see more from you.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants