[Issue #21] Coalesce transaction-local primary-key changes in Retina batches

docs/retina-transaction-batches.md

@@ -0,0 +1,25 @@
+# Retina transaction batches
+
+The Retina payload builder folds repeated primary keys within one source transaction and one RPC batch. It keeps the earliest before image and final after image:
+
+| Changes | Net write |
+| --- | --- |
+| U → U | U with final values |
+| I → U | I with final values |
+| I → D | None |
+| U → D | D |
+| D → I | U with final values |
+| I → U → D | None |
+| U → D → I | U with final values |
+
+Different transactions and RPC batches remain separate. Primary-key changes must be split into an old-key delete and new-key insert by the bucket dispatcher before batching. Snapshots and events without a source transaction ID are not folded. Each operation group has unique keys, so repeated keys without a safe transaction boundary execute in separate, ordered groups.
+
+Transaction counters use the original event count after a successful RPC acknowledgment. An entirely canceled batch sends an empty RPC to retain the acknowledgment boundary; failed RPCs do not advance counters.
+
+Coalescing uses only batch-local state and introduces no locks. Existing per-writer flush serialization and Retina stream processing retain ordering between RPCs; different buckets can still execute concurrently. Native visibility concurrency is handled separately in [Pixels PR #1412](https://github.com/pixelsdb/pixels/pull/1412).
+
+Run focused regressions with:
+
+```sh
+mvn -Dtest=RetinaPayloadBuilderTest,TableCrossTxCoalescingTest test
+```

src/main/java/io/pixelsdb/pixels/sink/writer/retina/RetinaPayloadBuilder.java

@@ -19,6 +19,9 @@
 import io.pixelsdb.pixels.sink.event.RowChangeEvent;
 import io.pixelsdb.pixels.sink.exception.SinkException;
 
+import com.google.protobuf.ByteString;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -32,12 +35,13 @@ private RetinaPayloadBuilder()
     }
 
     /**
-     * Builds table updates while preserving the source operation order.
+     * Folds repeated keys within each source transaction, then groups net operations.
      *
      * <p>Retina processes the operation fields inside one {@code TableUpdateData}
      * in its own fixed order. Therefore, only adjacent events with the same
-     * operation may share a payload. The order of the returned payloads matches
-     * the order of {@code events}.</p>
+     * operation and source transaction may share a payload. Within a transaction,
+     * independent keys keep their first-occurrence order, with each key reduced
+     * to its earliest before image and final after image.</p>
      */
     public static List<RetinaProto.TableUpdateData> buildTableUpdateDataList(
             String tableName,
@@ -49,16 +53,28 @@ public static List<RetinaProto.TableUpdateData> buildTableUpdateDataList(
             throw new IllegalArgumentException("events is empty");
         }
 
+        events = TransactionBatchCoalescer.coalesce(events);
         List<RetinaProto.TableUpdateData> updates = new ArrayList<>();
         int groupStart = 0;
-        for (int i = 1; i <= events.size(); i++)
+        Set<ByteString> groupKeys = new HashSet<>();
+        for (int i = 0; i < events.size(); i++)
         {
-            if (i == events.size() || events.get(i).getOp() != events.get(groupStart).getOp())
+            RowChangeEvent event = events.get(i);
+            event.initRoutingKey();
+            ByteString key = event.isDelete() ? event.getBeforeRoutingKey() : event.getAfterRoutingKey();
+            if (i > groupStart && (event.getOp() != events.get(groupStart).getOp()
+                    || !event.getTransaction().getId().equals(events.get(groupStart).getTransaction().getId())
+                    || groupKeys.contains(key)))
             {
-                updates.add(buildTableUpdateData(
-                        tableName, timestamp, events.subList(groupStart, i)));
+                updates.add(buildTableUpdateData(tableName, timestamp, events.subList(groupStart, i)));
                 groupStart = i;
+                groupKeys.clear();
             }
+            groupKeys.add(key);
+        }
+        if (groupStart < events.size())
+        {
+            updates.add(buildTableUpdateData(tableName, timestamp, events.subList(groupStart, events.size())));
         }
         return updates;
     }

src/main/java/io/pixelsdb/pixels/sink/writer/retina/TableCrossTxWriter.java

@@ -144,12 +144,14 @@ private void addTransactionBatch(
             List<String> fullTableNames,
             List<Integer> tableUpdateCount)
     {
-        List<RetinaProto.TableUpdateData> updates =
-                buildTableUpdateDataFromBatch(txId, transactionBatch);
-        if (updates.isEmpty())
+        if (SinkContextManager.getInstance().getSinkContext(txId) == null)
         {
             return;
         }
+        List<RetinaProto.TableUpdateData> updates =
+                buildTableUpdateDataFromBatch(txId, transactionBatch);
+        // A canceled I->D still consumes source events. Keep its counters, and
+        // acknowledge them with the RPC (an empty RPC is valid when all rows cancel).
         tableUpdateData.addAll(updates);
         txIds.add(txId);
         fullTableNames.add(transactionBatch.get(0).getFullTableName());

src/main/java/io/pixelsdb/pixels/sink/writer/retina/TransactionBatchCoalescer.java

@@ -0,0 +1,133 @@
+/*
+ * Copyright 2026 PixelsDB.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.pixelsdb.pixels.sink.writer.retina;
+
+import com.google.protobuf.ByteString;
+import io.pixelsdb.pixels.sink.SinkProto;
+import io.pixelsdb.pixels.sink.event.RowChangeEvent;
+import io.pixelsdb.pixels.sink.exception.SinkException;
+
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+/** Transaction-local net effects for one detached RPC batch. No shared state or locks. */
+final class TransactionBatchCoalescer
+{
+    private TransactionBatchCoalescer() { }
+
+    static List<RowChangeEvent> coalesce(List<RowChangeEvent> events) throws SinkException
+    {
+        List<RowChangeEvent> result = new ArrayList<>();
+        Map<ByteString, Change> changes = new LinkedHashMap<>();
+        RowChangeEvent scope = null;
+        for (RowChangeEvent event : events)
+        {
+            // Snapshot / unidentified transactions have no safe source transaction boundary.
+            if (event.isSnapshot() || event.getTransaction().getId().isBlank())
+            {
+                append(changes, result);
+                scope = null;
+                result.add(event);
+                continue;
+            }
+            if (scope != null && !sameScope(scope, event))
+            {
+                append(changes, result);
+            }
+            scope = event;
+            event.initRoutingKey();
+            if (event.isUpdate() && event.isPkChanged())
+            {
+                throw new SinkException("Primary-key changes must be split before Retina batching");
+            }
+            ByteString key = event.isDelete() ? event.getBeforeRoutingKey() : event.getAfterRoutingKey();
+            if (key == null || event.getOp() == SinkProto.OperationType.UNRECOGNIZED)
+            {
+                throw new SinkException("Missing primary key or unsupported Retina operation");
+            }
+            Change change = changes.get(key);
+            if (change == null)
+            {
+                changes.put(key, new Change(event));
+            }
+            else
+            {
+                change.accept(event);
+            }
+        }
+        append(changes, result);
+        return result;
+    }
+
+    private static boolean sameScope(RowChangeEvent a, RowChangeEvent b)
+    {
+        return a.getTransaction().getId().equals(b.getTransaction().getId())
+                && a.getRowRecord().getSource().equals(b.getRowRecord().getSource())
+                && a.getTableId() == b.getTableId()
+                && a.getTableMetadata().getPrimaryIndexKeyId() == b.getTableMetadata().getPrimaryIndexKeyId();
+    }
+
+    private static void append(Map<ByteString, Change> changes, List<RowChangeEvent> result)
+            throws SinkException
+    {
+        for (Change change : changes.values())
+        {
+            RowChangeEvent event = change.finish();
+            if (event != null) result.add(event);
+        }
+        changes.clear();
+    }
+
+    private static final class Change
+    {
+        private final RowChangeEvent first;
+        private RowChangeEvent last;
+        private boolean exists;
+
+        private Change(RowChangeEvent event)
+        {
+            this.first = event;
+            this.last = event;
+            this.exists = event.hasAfterData();
+        }
+
+        private void accept(RowChangeEvent event) throws SinkException
+        {
+            if (event.isInsert() == exists)
+            {
+                throw new SinkException("Invalid operation sequence for one primary key in a transaction");
+            }
+            last = event;
+            exists = event.hasAfterData();
+        }
+
+        private RowChangeEvent finish() throws SinkException
+        {
+            if (first == last) return first;
+            boolean existedBefore = first.hasBeforeData();
+            if (!existedBefore && !exists) return null;
+            SinkProto.RowRecord.Builder record = last.getRowRecord().toBuilder();
+            if (existedBefore) record.setBefore(first.getBefore());
+            else record.clearBefore();
+            if (!exists) record.clearAfter();
+            record.setOp(!exists ? SinkProto.OperationType.DELETE
+                    : existedBefore ? SinkProto.OperationType.UPDATE : SinkProto.OperationType.INSERT);
+            return new RowChangeEvent(record.build(), first.getSchema(), first.getTableMetadata());
+        }
+    }
+}

src/test/java/io/pixelsdb/pixels/sink/writer/retina/RetinaPayloadBuilderTest.java

@@ -275,4 +275,133 @@ void shouldBindIndexKeyWithPayloadTimestamp() throws Exception
         assertEquals(bytes("1"), tableUpdate.getInsertData(0).getIndexKeys(0).getKey());
         assertEquals(TIMESTAMP, insert.getAfterKey().getTimestamp());
     }
+
+    @Test
+    void shouldFoldPostgresTransactionSequences() throws Exception
+    {
+        String[] sequences = {"UU", "IU", "ID", "UD", "DI", "IUD", "UDI", "IDI", "DID"};
+        SinkProto.OperationType[] expected = {SinkProto.OperationType.UPDATE, SinkProto.OperationType.INSERT,
+                null, SinkProto.OperationType.DELETE, SinkProto.OperationType.UPDATE, null,
+                SinkProto.OperationType.UPDATE, SinkProto.OperationType.INSERT, SinkProto.OperationType.DELETE};
+        for (int n = 0; n < sequences.length; n++)
+        {
+            List<RowChangeEvent> input = transactionSequence(sequences[n], "tx");
+            List<RowChangeEvent> net = TransactionBatchCoalescer.coalesce(input);
+            if (expected[n] == null) assertEquals(List.of(), net, sequences[n]);
+            else
+            {
+                assertEquals(1, net.size(), sequences[n]);
+                assertEquals(expected[n], net.get(0).getOp(), sequences[n]);
+                if (net.get(0).hasBeforeData()) assertEquals(input.get(0).getBefore(), net.get(0).getBefore());
+                if (net.get(0).hasAfterData()) assertEquals(input.get(input.size() - 1).getAfter(), net.get(0).getAfter());
+            }
+            List<RetinaProto.TableUpdateData> payload = RetinaPayloadBuilder.buildTableUpdateDataList(TABLE_NAME, TIMESTAMP, input);
+            assertEquals(expected[n] == null ? 0 : 1, payload.size(), sequences[n]);
+        }
+    }
+
+    @Test
+    void shouldPreserveTransactionsAndRpcBoundaries() throws Exception
+    {
+        RowChangeEvent a = transactionSequence("U", "a").get(0);
+        RowChangeEvent b = transactionSequence("U", "b").get(0);
+        assertEquals(2, RetinaPayloadBuilder.buildTableUpdateDataList(TABLE_NAME, TIMESTAMP, List.of(a, b)).size());
+        List<RowChangeEvent> canceled = transactionSequence("ID", "c");
+        assertEquals(1, RetinaPayloadBuilder.buildTableUpdateDataList(TABLE_NAME, TIMESTAMP, canceled.subList(0, 1)).get(0).getInsertDataCount());
+        assertEquals(1, RetinaPayloadBuilder.buildTableUpdateDataList(TABLE_NAME, TIMESTAMP, canceled.subList(1, 2)).get(0).getDeleteDataCount());
+    }
+
+    @Test
+    void shouldFoldInterleavedKeysAndLeaveSourceRecordsUnchanged() throws Exception
+    {
+        List<RowChangeEvent> changes = transactionSequence("UU", "tx");
+        RowChangeEvent other = withTransaction(event(SinkProto.OperationType.UPDATE, row("2", "old"), row("2", "new"), tableMetadata()), "tx", 2);
+        List<RetinaProto.TableUpdateData> updates = RetinaPayloadBuilder.buildTableUpdateDataList(TABLE_NAME, TIMESTAMP, List.of(changes.get(0), other, changes.get(1)));
+        assertEquals(1, updates.size());
+        assertEquals(2, updates.get(0).getUpdateDataCount());
+        assertEquals(bytes("2"), updates.get(0).getUpdateData(1).getIndexKeys(0).getKey());
+        assertEquals(bytes("2"), updates.get(0).getUpdateData(0).getColValues(1));
+        assertEquals(bytes("1"), changes.get(0).getAfterData().get(1));
+    }
+
+    @Test
+    void shouldRejectUnsplitPrimaryKeyChangesAndInvalidSequences() throws Exception
+    {
+        RowChangeEvent changedKey = withTransaction(event(SinkProto.OperationType.UPDATE, row("1", "old"), row("2", "new"), tableMetadata()), "tx", 1);
+        assertThrows(io.pixelsdb.pixels.sink.exception.SinkException.class,
+                () -> TransactionBatchCoalescer.coalesce(List.of(changedKey)));
+        assertThrows(io.pixelsdb.pixels.sink.exception.SinkException.class,
+                () -> TransactionBatchCoalescer.coalesce(transactionSequence("II", "tx")));
+    }
+
+    static List<RowChangeEvent> transactionSequence(String sequence, String tx) throws Exception
+    {
+        List<RowChangeEvent> result = new java.util.ArrayList<>();
+        TableMetadata metadata = tableMetadata();
+        for (int i = 0; i < sequence.length(); i++)
+        {
+            char op = sequence.charAt(i);
+            result.add(withTransaction(event(op == 'I' ? SinkProto.OperationType.INSERT : op == 'U' ? SinkProto.OperationType.UPDATE : SinkProto.OperationType.DELETE,
+                    op == 'I' ? null : row("1", Integer.toString(i)),
+                    op == 'D' ? null : row("1", Integer.toString(i + 1)), metadata), tx, i + 1));
+        }
+        return result;
+    }
+
+    @Test
+    void canceledTransactionMustNotCombineSeparateRunsOfTheSameKey() throws Exception
+    {
+        List<RowChangeEvent> events = new java.util.ArrayList<>();
+        events.addAll(transactionSequence("U", "a"));
+        events.addAll(transactionSequence("ID", "b"));
+        events.addAll(transactionSequence("U", "a"));
+        assertEquals(2, RetinaPayloadBuilder.buildTableUpdateDataList(TABLE_NAME, TIMESTAMP, events).size());
+    }
+
+    @Test
+    void shouldCoalesceSplitPrimaryKeyChangeAndKeepFinalNullValue() throws Exception
+    {
+        TableMetadata metadata = tableMetadata();
+        RowChangeEvent deleteOld = withTransaction(event(SinkProto.OperationType.DELETE, row("1", "old"), null, metadata), "tx", 1);
+        RowChangeEvent insertNew = withTransaction(event(SinkProto.OperationType.INSERT, null, row("2", "new"), metadata), "tx", 1);
+        SinkProto.RowValue finalRow = SinkProto.RowValue.newBuilder()
+                .addValues(SinkProto.ColumnValue.newBuilder().setValue(bytes("2")))
+                .addValues(SinkProto.ColumnValue.newBuilder().setIsNull(true)).build();
+        RowChangeEvent updateNew = withTransaction(event(SinkProto.OperationType.UPDATE, row("2", "new"), finalRow, metadata), "tx", 2);
+        List<RetinaProto.TableUpdateData> result = RetinaPayloadBuilder.buildTableUpdateDataList(TABLE_NAME, TIMESTAMP, List.of(deleteOld, insertNew, updateNew));
+        assertEquals(2, result.size());
+        assertEquals(bytes("1"), result.get(0).getDeleteData(0).getIndexKeys(0).getKey());
+        assertEquals(bytes("2"), result.get(1).getInsertData(0).getIndexKeys(0).getKey());
+        assertEquals(List.of(false, true), result.get(1).getInsertData(0).getIsNullList());
+    }
+
+    @Test
+    void concurrentBatchesMustKeepIndependentNetEffects() throws Exception
+    {
+        java.util.concurrent.ExecutorService executor = java.util.concurrent.Executors.newFixedThreadPool(8);
+        try
+        {
+            List<java.util.concurrent.Future<?>> results = new java.util.ArrayList<>();
+            for (int i = 0; i < 100; i++)
+            {
+                final int id = i;
+                results.add(executor.submit(() -> {
+                    List<RetinaProto.TableUpdateData> result = RetinaPayloadBuilder.buildTableUpdateDataList(
+                            TABLE_NAME, TIMESTAMP + id, transactionSequence("UU", "tx" + id));
+                    assertEquals(1, result.size());
+                    assertEquals(1, result.get(0).getUpdateDataCount());
+                    assertEquals(TIMESTAMP + id, result.get(0).getUpdateData(0).getIndexKeys(0).getTimestamp());
+                    return null;
+                }));
+            }
+            for (java.util.concurrent.Future<?> result : results) result.get();
+        }
+        finally { executor.shutdownNow(); }
+    }
+
+    private static RowChangeEvent withTransaction(RowChangeEvent event, String tx, long order) throws Exception
+    {
+        return new RowChangeEvent(event.getRowRecord().toBuilder().setTransaction(
+                SinkProto.TransactionInfo.newBuilder().setId(tx).setTotalOrder(order)).build(), event.getSchema(), event.getTableMetadata());
+    }
 }