Ensure GroupKey is preserved and commutative in join/combine

- Store original keys in KeyPair and use them for output emission in join/combine
- Always prefer GroupKey over plain keys when both are present for the same match
- Fixes timing-dependent bugs when joining or combining channels with mixed GroupKey/plain keys
- Updates buffer logic to store KeyPair objects, ensuring metadata like group size is retained
- Updates remainder and mismatch handling to work with new buffer structure
- Adds and passes comprehensive tests for GroupKey preservation and commutativity

Fixes #4104

Signed-off-by: Rob Syme <[email protected]>

Author: robsyme

modules/nextflow/src/main/groovy/nextflow/extension/CombineOp.groovy

@@ -48,6 +48,8 @@ class CombineOp {
 
     private Map<Object,List> rightValues = [:]
 
+    private Map<Object,Object> originalKeyMap = [:]
+
     private static final int LEFT = 0
 
     private static final int RIGHT = 1
@@ -93,7 +95,7 @@ class CombineOp {
         opts.onNext = {
             if( pivot ) {
                 def pair = makeKey(pivot, it)
-                emit(target, index, pair.keys, pair.values)
+                emit(target, index, pair)
             }
             else {
                 emit(target, index, NONE, it)
@@ -120,15 +122,44 @@ class CombineOp {
     }
 
     @PackageScope
-    synchronized void emit( DataflowWriteChannel target, int index, List p, v ) {
+    synchronized void emit( DataflowWriteChannel target, int index, KeyPair pair ) {
+        emit(target, index, pair.originalKeys, pair.keys, pair.values)
+    }
+
+    @PackageScope
+    synchronized void emit( DataflowWriteChannel target, int index, List originalKeys, List keys, v ) {
+        def p = keys  // Use normalized keys for matching
+
+        // Store/update the mapping from normalized key to original key
+        // Prefer GroupKey over plain keys
+        def existingOriginal = originalKeyMap.get(p)
+        if (existingOriginal == null) {
+            originalKeyMap[p] = originalKeys
+        } else {
+            // Check if any of the new original keys is a GroupKey
+            for (int i = 0; i < originalKeys.size(); i++) {
+                def newKey = originalKeys[i]
+                def oldKey = existingOriginal instanceof List ? existingOriginal[i] : existingOriginal
+                if (newKey instanceof GroupKey && !(oldKey instanceof GroupKey)) {
+                    originalKeyMap[p] = originalKeys
+                    break
+                }
+            }
+        }
+
+        // Use the best available original key (preferring GroupKey)
+        def bestOriginalKeys = originalKeyMap[p]
+        
+        // Ensure bestOriginalKeys is a List for the tuple method
+        def bestKeysList = bestOriginalKeys instanceof List ? bestOriginalKeys : [bestOriginalKeys]
 
         if( leftValues[p] == null ) leftValues[p] = []
         if( rightValues[p] == null ) rightValues[p] = []
 
         if( index == LEFT ) {
             log.trace "combine >> left >> by=$p; val=$v; right-values: ${rightValues[p]}"
             for ( Object x : rightValues[p] ) {
-                target.bind( tuple(p, v, x) )
+                target.bind( tuple(bestKeysList, v, x) )  // Use best original keys in output
             }
             leftValues[p].add(v)
             return
@@ -137,7 +168,7 @@ class CombineOp {
         if( index == RIGHT ) {
             log.trace "combine >> right >> by=$p; val=$v; right-values: ${leftValues[p]}"
             for ( Object x : leftValues[p] ) {
-                target.bind( tuple(p, x, v) )
+                target.bind( tuple(bestKeysList, x, v) )  // Use best original keys in output
             }
             rightValues[p].add(v)
             return
@@ -146,6 +177,12 @@ class CombineOp {
         throw new IllegalArgumentException("Not a valid spread operator index: $index")
     }
 
+    @PackageScope
+    synchronized void emit( DataflowWriteChannel target, int index, List p, v ) {
+        // Legacy method for when pivot is NONE
+        emit(target, index, p, p, v)
+    }
+
     DataflowWriteChannel apply() {
 
         target = CH.create()

modules/nextflow/src/main/groovy/nextflow/extension/DataflowHelper.groovy

@@ -361,13 +361,10 @@ class DataflowHelper {
             if( pivot != [0] )
                 throw new IllegalArgumentException("Not a valid `by` index: $pivot")
             result.addKey(entry)
-            result.values = []
             return result
         }
 
         def list = (List)entry
-        result.keys = new ArrayList(pivot.size())
-        result.values = new ArrayList(list.size())
 
         for( int i=0; i<list.size(); i++ ) {
             if( i in pivot )

modules/nextflow/src/main/groovy/nextflow/extension/JoinOp.groovy

@@ -58,6 +58,8 @@ class JoinOp {
 
     private Set uniqueKeys = new LinkedHashSet()
 
+    private Map<Object,Object> originalKeyMap = new HashMap()
+
     JoinOp( DataflowReadChannel source, DataflowReadChannel target, Map params = null ) {
         CheckHelper.checkParams('join', params, JOIN_PARAMS)
         this.source = source
@@ -173,6 +175,24 @@ class JoinOp {
         // get the index key for this object
         final item0 = DataflowHelper.makeKey(pivot, data)
 
+        // Store the mapping from normalized key to original key
+        // Prefer GroupKey over plain keys
+        def existingOriginal = originalKeyMap.get(item0.keys)
+        if (existingOriginal == null) {
+            originalKeyMap[item0.keys] = item0.originalKeys
+        } else {
+            // Check if any of the new original keys is a GroupKey
+            // If so, replace the existing mapping
+            for (int i = 0; i < item0.originalKeys.size(); i++) {
+                def newKey = item0.originalKeys[i]
+                def oldKey = existingOriginal instanceof List ? existingOriginal[i] : existingOriginal
+                if (newKey instanceof GroupKey && !(oldKey instanceof GroupKey)) {
+                    originalKeyMap[item0.keys] = item0.originalKeys
+                    break
+                }
+            }
+        }
+
         // check for unique keys
         checkForDuplicate(item0.keys, item0.values, index, false)
 
@@ -190,8 +210,8 @@ class JoinOp {
         def entries = channels[index]
 
         // add the received item to the list
-        // when it is used in the gather op add always as the first item
-        entries << item0.values
+        // Store the full KeyPair to preserve original keys
+        entries << item0
         setSingleton(index, item0.values.size()==0)
 
         // now check if it has received an element matching for each channel
@@ -200,15 +220,39 @@ class JoinOp {
         }
 
         def result = []
+        
+        // Find the best key (prefer GroupKey) from all channels
+        def bestOriginalKeys = null
+        for (Map.Entry<Integer,List> entry : channels.entrySet()) {
+            def channelItems = entry.getValue()
+            if (channelItems && channelItems.size() > 0) {
+                def keyPair = channelItems[0] as KeyPair
+                if (bestOriginalKeys == null) {
+                    bestOriginalKeys = keyPair.originalKeys
+                } else {
+                    // Check if this channel has a GroupKey version
+                    for (int i = 0; i < keyPair.originalKeys.size(); i++) {
+                        def candidateKey = keyPair.originalKeys[i]
+                        def currentKey = bestOriginalKeys instanceof List ? bestOriginalKeys[i] : bestOriginalKeys
+                        if (candidateKey instanceof GroupKey && !(currentKey instanceof GroupKey)) {
+                            bestOriginalKeys = keyPair.originalKeys
+                            break
+                        }
+                    }
+                }
+            }
+        }
+        
         // add the key
-        addToList(result, item0.keys)
+        addToList(result, bestOriginalKeys ?: item0.originalKeys)
 
         final itr = channels.iterator()
         while( itr.hasNext() ) {
             def entry = (Map.Entry<Integer,List>)itr.next()
 
             def list = entry.getValue()
-            addToList(result, list[0])
+            def keyPair = list[0] as KeyPair
+            addToList(result, keyPair.values)
 
             list.remove(0)
             if( list.size() == 0 ) {
@@ -221,6 +265,17 @@ class JoinOp {
         return result
     }
 
+    // Helper method to retrieve original data from buffer
+    private def getOriginalDataFromBuffer(Map<Object,Map<Integer,List>> buffer, Object key, int channelIndex) {
+        def channels = buffer.get(key)
+        if (channels == null) return null
+        def items = channels.get(channelIndex)
+        if (items == null || items.isEmpty()) return null
+        // Need to reconstruct the original data from the values and the key
+        // This is a simplified version - in reality we'd need to track the full original items
+        return null  // For now, we'll use a different approach
+    }
+
     private final void checkRemainder(Map<Object,Map<Integer,List>> buffers, int count, DataflowWriteChannel target ) {
        log.trace "Operator `join` remainder buffer: ${-> buffers}"
 
@@ -231,17 +286,43 @@ class JoinOp {
 
                 boolean fill=false
                 def result = new ArrayList(count+1)
-                addToList(result, key)
+                
+                // Find the best original key from available channels
+                def bestOriginalKey = null
+                for( int i=0; i<count; i++ ) {
+                    List items = entry[i]
+                    if( items && items.size() > 0 ) {
+                        def keyPair = items[0] as KeyPair
+                        if (bestOriginalKey == null) {
+                            bestOriginalKey = keyPair.originalKeys
+                        } else {
+                            // Check if this channel has a GroupKey version
+                            for (int j = 0; j < keyPair.originalKeys.size(); j++) {
+                                def candidateKey = keyPair.originalKeys[j]
+                                def currentKey = bestOriginalKey instanceof List ? bestOriginalKey[j] : bestOriginalKey
+                                if (candidateKey instanceof GroupKey && !(currentKey instanceof GroupKey)) {
+                                    bestOriginalKey = keyPair.originalKeys
+                                    break
+                                }
+                            }
+                        }
+                    }
+                }
+                
+                // Use the best available original key, or fall back to the map key
+                def originalKey = bestOriginalKey ?: originalKeyMap.get(key) ?: key
+                addToList(result, originalKey)
 
                 for( int i=0; i<count; i++ ) {
-                    List values = entry[i]
-                    if( values ) {
+                    List items = entry[i]
+                    if( items ) {
+                        def keyPair = items[0] as KeyPair
                         // track unique keys
-                        checkForDuplicate(key, values[0], i,false)
+                        checkForDuplicate(key, keyPair.values, i, false)
 
-                        addToList(result, values[0])
+                        addToList(result, keyPair.values)
                         fill |= true
-                        values.remove(0)
+                        items.remove(0)
                     }
                     else if( !singleton(i) ) {
                         addToList(result, null)
@@ -277,7 +358,17 @@ class JoinOp {
 
             result[key] = []
             for( Map.Entry entry : remainder0 ) {
-                result[key].add(csv0(entry.value,','))
+                def items = entry.value
+                def values = []
+                // Extract values from KeyPair objects
+                items.each { item ->
+                    if (item instanceof KeyPair) {
+                        values.add(item.values)
+                    } else {
+                        values.add(item)
+                    }
+                }
+                result[key].add(csv0(values,','))
             }
         }
 

modules/nextflow/src/main/groovy/nextflow/extension/KeyPair.groovy

@@ -31,12 +31,17 @@ import nextflow.extension.GroupKey
 @EqualsAndHashCode
 class KeyPair {
     List keys
+    List originalKeys
     List values
 
+    KeyPair() {
+        this.keys = []
+        this.originalKeys = []
+        this.values = []
+    }
+
     void addKey(el) {
-        if (keys == null) {
-            keys = []
-        }
+        originalKeys.add(el)
         keys.add(safeStr(el))
     }