getRowsColumnRange always has cells grouped by row

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/SnapshotTransaction.java

@@ -20,6 +20,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.LinkedHashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
@@ -476,19 +477,19 @@ private Iterator<Map.Entry<Cell, byte[]>> filterDeletedValues(
     private Iterator<Map.Entry<Cell, Value>> getRowColumnRangePostFiltered(
             TableReference tableRef, RowColumnRangeIterator iterator, int batchHint) {
         return Iterators.concat(Iterators.transform(Iterators.partition(iterator, batchHint), batch -> {
-            ImmutableMap.Builder<Cell, Value> rawBuilder = ImmutableMap.builder();
-            batch.forEach(rawBuilder::put);
-            Map<Cell, Value> raw = rawBuilder.build();
+            // N.B. This batch could be spread across multiple rows, and those rows might extend into other
+            // batches. We are given cells for a row grouped together, so easiest way to ensure they stay together
+            // is to preserve the original order.
+            Map<Cell, Value> raw = new LinkedHashMap<>();
+            batch.forEach(entry -> raw.put(entry.getKey(), entry.getValue()));
             validatePreCommitRequirementsOnReadIfNecessary(tableRef, getStartTimestamp());
             if (raw.isEmpty()) {
                 return Collections.emptyIterator();
             }
-            SortedMap<Cell, Value> postFiltered = ImmutableSortedMap.copyOf(
-                    getWithPostFilteringSync(
-                            tableRef,
-                            raw,
-                            x -> x));
-            return postFiltered.entrySet().iterator();
+            return getWithPostFilteringSync(
+                    tableRef,
+                    raw,
+                    x -> x).iterator();
         }));
     }
 

atlasdb-tests-shared/src/test/java/com/palantir/atlasdb/transaction/impl/SnapshotTransactionTest.java

@@ -41,6 +41,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -87,6 +88,7 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimaps;
+import com.google.common.primitives.UnsignedBytes;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
@@ -1184,6 +1186,7 @@ public void testRowsColumnRangesSingleIteratorVersion() {
         runTestForGetRowsColumnRangeSingleIteratorVersion(10, 10, 5);
         runTestForGetRowsColumnRangeSingleIteratorVersion(10, 10, 10);
         runTestForGetRowsColumnRangeSingleIteratorVersion(100, 100, 99);
+        runTestForGetRowsColumnRangeSingleIteratorVersion(100, 11, 0);
     }
 
     private void runTestForGetRowsColumnRangeSingleIteratorVersion(
@@ -1213,19 +1216,35 @@ private void runTestForGetRowsColumnRangeSingleIteratorVersion(
         });
         keyValueService.addGarbageCollectionSentinelValues(TABLE, expectedDeletedCells);
 
+        List<byte[]> shuffledRows = expectedRows;
+        Collections.shuffle(shuffledRows);
         List<Cell> cells = serializableTxManager.runTaskReadOnly(tx ->
                 ImmutableList.copyOf(Iterators.transform(
                         tx.getRowsColumnRange(
                                 TABLE,
-                                expectedRows,
+                                shuffledRows,
                                 new ColumnRangeSelection(null, null),
                                 10),
                         Map.Entry::getKey)));
+        expectedCells.sort(Comparator
+                .comparing(
+                        (Cell cell) -> indexOfArrayEquality(shuffledRows, cell.getRowName()),
+                        Comparator.naturalOrder())
+                .thenComparing(Cell::getColumnName, UnsignedBytes.lexicographicalComparator()));
         Assertions.assertThat(cells).isEqualTo(expectedCells);
 
         keyValueService.truncateTable(TABLE);
     }
 
+    private int indexOfArrayEquality(List<byte[]> list, byte[] value) {
+        for (int i = 0; i < list.size(); i++) {
+            if (Arrays.equals(list.get(i), value)) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
     @Test
     public void commitThrowsIfRolledBackAtCommitTime_expiredLocks() {
         final Cell cell = Cell.create(PtBytes.toBytes("row1"), PtBytes.toBytes("column1"));