Fix IllegalStateException in UnifiedQueryPlanner.preserveCollation on composite-collation plans

[noCharger]

Description

UnifiedQueryPlanner.preserveCollation fails to plan any query whose top RelNode carries more than one collation in its trait set. It surfaces as a 500 Failed to plan query.

Root cause. The method reads the collation via:

RelCollation collation = logical.getTraitSet().getCollation();

RelTraitSet.getCollation() delegates to getTrait(RelCollationTraitDef.INSTANCE), which assumes a single collation. When a plan derives several collations, Calcite stores them as a RelCompositeTrait and getTrait refuses to return it, throwing:

java.lang.IllegalStateException: Trait index 1 has multiple values in this trait set; use getTraits(RelTraitDef) instead of getTrait(RelTraitDef)

Such a composite arises whenever a non-Sort top node is ordered by more than one equivalent key — for example ... | sort age | eval x = age: x mirrors age, so the top LogicalProject is sorted on both age and x and derives two collations. A plan with a single derived collation reads back fine, which is why this was not caught earlier. (A single multi-column literal-row plan such as makeresults format=csv data=... derives one multi-field collation, not a composite, so it does not trigger this.)

Fix. Read the collation through the composite-safe getTraits(RelCollationTraitDef.INSTANCE) and materialize a LogicalSort only when there is exactly one genuine (non-empty) sort collation. A composite (several incidental derived collations) is not a user sort, so the plan is left unwrapped instead of crashing. Behavior for the single-collation case is unchanged.

Related Issues

N/A

Check List

• New functionality includes testing.
• New functionality has been documented.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit cb87a7e)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ No major issues detected

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to cb87a7e

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Add null check for collation element Add a null check for collations.get(0) before calling getFieldCollations() to prevent potential NullPointerException. Even though the list size is checked, the element itself could theoretically be null. api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java [161-166] List<RelCollation> collations = logical.getTraitSet().getTraits(RelCollationTraitDef.INSTANCE); if (collations != null && collations.size() == 1 + && collations.get(0) != null && !collations.get(0).getFieldCollations().isEmpty()) { return LogicalSort.create(logical, collations.get(0), null, null); } Suggestion importance[1-10]: 3 __ Why: While adding a null check for collations.get(0) is technically defensive programming, it's unlikely that getTraits() would return a list containing null elements. The suggestion addresses a theoretical edge case but doesn't fix a critical issue or significantly improve code quality.	Low

---

Previous suggestions

Suggestions up to commit f860d87

Category	Suggestion	Impact
Possible issue	Add null check for collation object Add a null check before calling getFieldCollations() to prevent potential NullPointerException. While collations.get(0) is guarded by the size check, the returned RelCollation object itself could theoretically be null depending on the trait implementation. api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java [161-167] List<RelCollation> collations = logical.getTraitSet().getTraits(RelCollationTraitDef.INSTANCE); if (collations != null && collations.size() == 1 + && collations.get(0) != null && !collations.get(0).getFieldCollations().isEmpty()) { return LogicalSort.create(logical, collations.get(0), null, null); } Suggestion importance[1-10]: 5 __ Why: While adding a null check for collations.get(0) is defensive programming, the existing size check (collations.size() == 1) already ensures an element exists at index 0. The suggestion addresses a theoretical edge case but is unlikely to occur in practice with standard Calcite trait implementations.	Low

Suggestions up to commit 3f39114

Category	Suggestion	Impact
Possible issue	Add null-check for list element Add a null-check for collations.get(0) before calling getFieldCollations() to prevent potential NullPointerException. Even though the size check ensures the list has one element, the element itself could theoretically be null. api/src/main/java/org/opensearch/sql/api/UnifiedQueryPlanner.java [166-171] List<RelCollation> collations = logical.getTraitSet().getTraits(RelCollationTraitDef.INSTANCE); if (collations != null && collations.size() == 1 + && collations.get(0) != null && !collations.get(0).getFieldCollations().isEmpty()) { return LogicalSort.create(logical, collations.get(0), null, null); } Suggestion importance[1-10]: 3 __ Why: While adding a null-check for collations.get(0) is technically defensive programming, it's unlikely that the Calcite API would return a list containing a null RelCollation element. The suggestion addresses a theoretical edge case but has minimal practical impact given the API contract.	Low

[github-actions]

Persistent review updated to latest commit f860d87

[github-actions]

Persistent review updated to latest commit cb87a7e