From eb6f205698ffa22be404a87e31be0f72fedf3332 Mon Sep 17 00:00:00 2001
From: HMS17 <hms243@cornell.edu>
Date: Tue, 16 Dec 2025 08:41:04 -0500
Subject: [PATCH] [BI-2755] - improve default sub-unit sort

---
 .../brapi/v2/services/BrAPITrialService.java         | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/main/java/org/breedinginsight/brapi/v2/services/BrAPITrialService.java b/src/main/java/org/breedinginsight/brapi/v2/services/BrAPITrialService.java
index 9ca961920..9f5a93410 100644
--- a/src/main/java/org/breedinginsight/brapi/v2/services/BrAPITrialService.java
+++ b/src/main/java/org/breedinginsight/brapi/v2/services/BrAPITrialService.java
@@ -935,8 +935,16 @@ private boolean isSubEntityDataset(List<BrAPIObservationUnit> ous){
 
     private void sortDefaultForObservationUnit(List<BrAPIObservationUnit> ous) {
         Comparator<BrAPIObservationUnit> studyNameComparator = Comparator.comparing(BrAPIObservationUnit::getStudyName, new IntOrderComparator());
-        Comparator<BrAPIObservationUnit> ouNameComparator = Comparator.comparing(BrAPIObservationUnit::getObservationUnitName, new IntOrderComparator());
-        ous.sort( (studyNameComparator).thenComparing(ouNameComparator));
+
+        if (isSubEntityDataset(ous)) {
+            Comparator<BrAPIObservationUnit> subUnitComparator = Comparator.comparing(BrAPIObservationUnit::getObservationUnitName, new IntOrderComparator());
+            Comparator<BrAPIObservationUnit> ouNameComparator = Comparator.comparing(row -> (row.getAdditionalInfo().get(BrAPIAdditionalInfoFields.EXP_UNIT_ID).toString()), new IntOrderComparator());
+            ous.sort((studyNameComparator).thenComparing(ouNameComparator).thenComparing(subUnitComparator));
+        }
+        else {
+            Comparator<BrAPIObservationUnit> ouNameComparator = Comparator.comparing(BrAPIObservationUnit::getObservationUnitName, new IntOrderComparator());
+            ous.sort((studyNameComparator).thenComparing(ouNameComparator));
+        }
     }
 
     private void sortDefaultForExportRows(@NotNull List<Map<String, Object>> exportRows) {