[Issue #18] Fix IndexKey timestamp via deferred binding after routing-key init

src/main/java/io/pixelsdb/pixels/sink/event/RowChangeEvent.java

@@ -54,12 +54,21 @@ public class RowChangeEvent
     private TableMetadata tableMetadata = null;
     private Map<String, SinkProto.ColumnValue> beforeValueMap;
     private Map<String, SinkProto.ColumnValue> afterValueMap;
+    /**
+     * Primary-key bytes used for bucket routing / PK-change detection.
+     * Independent of TransService timestamp.
+     */
+    @Getter
+    private ByteString beforeRoutingKey;
+    @Getter
+    private ByteString afterRoutingKey;
     @Getter
     private IndexProto.IndexKey beforeKey;
     @Getter
     private IndexProto.IndexKey afterKey;
 
-    private boolean indexKeyInited = false;
+    private boolean routingKeyInited = false;
+    private boolean indexKeyBound = false;
 
     @Getter
     private long tableId;
@@ -76,7 +85,7 @@ public RowChangeEvent(SinkProto.RowRecord rowRecord) throws SinkException
         this.tableMetadata = tableMetadataRegistry.getMetadata(
                 rowRecord.getSource().getDb(), rowRecord.getSource().getTable());
         init();
-        initIndexKey();
+        initRoutingKey();
     }
 
     public RowChangeEvent(SinkProto.RowRecord rowRecord, TypeDescription schema) throws SinkException
@@ -135,9 +144,12 @@ private void initColumnValueMap(SinkProto.RowValue rowValue, Map<String, SinkPro
                 .forEach(i -> map.put(schema.getFieldNames().get(i), rowValue.getValuesList().get(i)));
     }
 
-    public void initIndexKey() throws SinkException
+    /**
+     * Builds primary-key bytes for routing. Safe before TransService timestamp is known.
+     */
+    public void initRoutingKey() throws SinkException
     {
-        if (indexKeyInited)
+        if (routingKeyInited)
         {
             return;
         }
@@ -149,76 +161,113 @@ public void initIndexKey() throws SinkException
 
         if (!this.tableMetadata.hasPrimaryIndex())
         {
+            routingKeyInited = true;
             return;
         }
         if (hasBeforeData())
         {
-            this.beforeKey = generateIndexKey(tableMetadata, beforeValueMap);
+            this.beforeRoutingKey = generateRoutingKey(tableMetadata, beforeValueMap);
         }
 
         if (hasAfterData())
         {
-            this.afterKey = generateIndexKey(tableMetadata, afterValueMap);
+            this.afterRoutingKey = generateRoutingKey(tableMetadata, afterValueMap);
         }
 
-        indexKeyInited = true;
+        routingKeyInited = true;
     }
 
-    public void updateIndexKey() throws SinkException
+    /**
+     * Builds Retina {@link IndexProto.IndexKey} from routing bytes and {@link #timeStamp}.
+     * Call after the TransService timestamp is assigned.
+     */
+    public void bindIndexKey() throws SinkException
     {
+        initRoutingKey();
+
+        if (this.tableMetadata == null || !this.tableMetadata.hasPrimaryIndex())
+        {
+            return;
+        }
+
+        SinglePointIndex index = tableMetadata.getIndex();
+        long tableIdForKey = tableMetadata.getTable().getId();
         if (hasBeforeData())
         {
-            this.beforeKey = generateIndexKey(tableMetadata, beforeValueMap);
+            this.beforeKey = toIndexKey(beforeRoutingKey, index.getId(), tableIdForKey);
         }
 
         if (hasAfterData())
         {
-            this.afterKey = generateIndexKey(tableMetadata, afterValueMap);
+            this.afterKey = toIndexKey(afterRoutingKey, index.getId(), tableIdForKey);
         }
+
+        indexKeyBound = true;
+    }
+
+    /**
+     * @deprecated Use {@link #initRoutingKey()} for bucketing and {@link #bindIndexKey()} after
+     * timestamp assignment. Kept for callers that only need routing.
+     */
+    @Deprecated
+    public void initIndexKey() throws SinkException
+    {
+        initRoutingKey();
+    }
+
+    /**
+     * @deprecated Use {@link #bindIndexKey()}.
+     */
+    @Deprecated
+    public void updateIndexKey() throws SinkException
+    {
+        bindIndexKey();
+    }
+
+    public boolean isIndexKeyBound()
+    {
+        return indexKeyBound;
     }
 
     public int getBeforeBucketFromIndex()
     {
-        assert indexKeyInited;
+        assert routingKeyInited;
         if (hasBeforeData())
         {
-            return getBucketFromIndexKey(beforeKey);
+            return getBucketIdFromByteBuffer(beforeRoutingKey);
         }
         throw new IllegalCallerException("Event dosen't have before data");
     }
 
     public boolean isPkChanged() throws SinkException
     {
-        if (!indexKeyInited)
+        if (!routingKeyInited)
         {
-            initIndexKey();
+            initRoutingKey();
         }
 
         if (getOp() != SinkProto.OperationType.UPDATE)
         {
             return false;
         }
 
-        ByteString beforeKey = getBeforeKey().getKey();
-        ByteString afterKey = getAfterKey().getKey();
-
-        return !beforeKey.equals(afterKey);
+        return !beforeRoutingKey.equals(afterRoutingKey);
     }
 
     public int getAfterBucketFromIndex()
     {
-        assert indexKeyInited;
+        assert routingKeyInited;
         if (hasAfterData())
         {
-            return getBucketFromIndexKey(afterKey);
+            return getBucketIdFromByteBuffer(afterRoutingKey);
         }
         throw new IllegalCallerException("Event dosen't have after data");
     }
 
-    private IndexProto.IndexKey generateIndexKey(TableMetadata tableMetadata, Map<String, SinkProto.ColumnValue> rowValue)
+    private ByteString generateRoutingKey(
+            TableMetadata tableMetadata, Map<String, SinkProto.ColumnValue> rowValue)
     {
         List<String> keyColumnNames = tableMetadata.getKeyColumnNames();
-        SinglePointIndex index = tableMetadata.getIndex();
         int len = keyColumnNames.size();
         List<ByteString> keyColumnValues = new ArrayList<>(len);
         int keySize = 0;
@@ -234,12 +283,16 @@ private IndexProto.IndexKey generateIndexKey(TableMetadata tableMetadata, Map<St
         {
             byteBuffer.put(value.toByteArray());
         }
+        return ByteString.copyFrom(byteBuffer.rewind());
+    }
 
+    private IndexProto.IndexKey toIndexKey(ByteString routingKey, long indexId, long tableIdForKey)
+    {
         return IndexProto.IndexKey.newBuilder()
                 .setTimestamp(timeStamp)
-                .setKey(ByteString.copyFrom(byteBuffer.rewind()))
-                .setIndexId(index.getId())
-                .setTableId(tableMetadata.getTable().getId())
+                .setKey(routingKey)
+                .setIndexId(indexId)
+                .setTableId(tableIdForKey)
                 .build();
     }
 

src/main/java/io/pixelsdb/pixels/sink/event/RowChangeEventFactory.java

@@ -95,7 +95,7 @@ public static RowChangeEvent create(
         }
 
         RowChangeEvent event = new RowChangeEvent(rowRecord, schema, metadata);
-        event.initIndexKey();
+        event.initRoutingKey();
         return event;
     }
 }

src/main/java/io/pixelsdb/pixels/sink/writer/AbstractBucketedWriter.java

@@ -28,7 +28,7 @@ public void writeRowChangeEvent(RowChangeEvent event, C context) throws SinkExce
             return;
         }
 
-        event.initIndexKey();
+        event.initRoutingKey();
 
         switch (event.getOp())
         {
@@ -92,7 +92,8 @@ private RowChangeEvent buildDeleteEvent(RowChangeEvent event) throws SinkExcepti
 
         RowChangeEvent deleteEvent =
                 new RowChangeEvent(builder.build(), event.getSchema());
-        deleteEvent.initIndexKey();
+        deleteEvent.setTimeStamp(event.getTimeStamp());
+        deleteEvent.initRoutingKey();
         return deleteEvent;
     }
 
@@ -105,7 +106,8 @@ private RowChangeEvent buildInsertEvent(RowChangeEvent event) throws SinkExcepti
 
         RowChangeEvent insertEvent =
                 new RowChangeEvent(builder.build(), event.getSchema());
-        insertEvent.initIndexKey();
+        insertEvent.setTimeStamp(event.getTimeStamp());
+        insertEvent.initRoutingKey();
         return insertEvent;
     }
 }

src/main/java/io/pixelsdb/pixels/sink/writer/NoneWriter.java

@@ -93,10 +93,11 @@ public boolean writeRow(RowChangeEvent rowChangeEvent)
         metricsFacade.recordRowChange(rowChangeEvent.getTable(), rowChangeEvent.getOp());
         try
         {
-            rowChangeEvent.initIndexKey();
-            if (rowChangeEvent.getAfterKey() != null)
+            rowChangeEvent.initRoutingKey();
+            if (rowChangeEvent.getAfterRoutingKey() != null)
             {
-                metricsFacade.recordPrimaryKeyUpdateDistribution(rowChangeEvent.getTable(), rowChangeEvent.getAfterKey().getKey());
+                metricsFacade.recordPrimaryKeyUpdateDistribution(
+                        rowChangeEvent.getTable(), rowChangeEvent.getAfterRoutingKey());
             }
 
 

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

@@ -47,6 +47,9 @@ public static RetinaProto.TableUpdateData buildTableUpdateData(
                 .setTableName(tableName);
         for (RowChangeEvent event : events)
         {
+            // Bind IndexKey once with the same timestamp written into TableUpdateData.
+            event.setTimeStamp(timestamp);
+            event.bindIndexKey();
             addRowChange(event, builder);
         }
         return builder.build();

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

@@ -236,6 +236,7 @@ protected void writeRowChangeEvent(SinkContext ctx, RowChangeEvent event) throws
     {
         if (ctx != null)
         {
+            // Routing uses key bytes only; IndexKey is bound later in RetinaPayloadBuilder.
             event.setTimeStamp(ctx.getTimestamp());
         }
         retinaBucketDispatcher.writeRowChangeEvent(event, ctx);

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

@@ -51,12 +51,7 @@ public void flush(List<RowChangeEvent> batch)
         try
         {
             List<RetinaProto.TableUpdateData> tableUpdateData = new LinkedList<>();
-            for (RowChangeEvent event : batch)
-            {
-                event.setTimeStamp(pixelsTransContext.getTimestamp());
-                event.updateIndexKey();
-            }
-
+            // Timestamp + IndexKey binding happen inside RetinaPayloadBuilder.
             RetinaProto.TableUpdateData update =
                     buildTableUpdateDataFromBatch(pixelsTransContext, batch);
             if (update != null)
@@ -97,9 +92,6 @@ public void flush(List<RowChangeEvent> batch)
                         }
                     }
             );
-        } catch (SinkException e)
-        {
-            throw new RuntimeException(e);
         } finally
         {
             writeLock.unlock();

src/test/java/io/pixelsdb/pixels/sink/cdc/CdcOrderContractTest.java

@@ -157,7 +157,7 @@ private RowChangeEvent event(
                 metadataFixture.metadata().getTypeDescription(),
                 metadataFixture.metadata());
         event.setTimeStamp(7001L);
-        event.initIndexKey();
+        event.initRoutingKey();
         return event;
     }
 
@@ -202,9 +202,9 @@ List<Long> primaryKeys()
         {
             return emitted.stream()
                     .map(event -> event.hasAfterData()
-                            ? event.getAfterKey()
-                            : event.getBeforeKey())
-                    .map(key -> readLong(key.getKey()))
+                            ? event.getAfterRoutingKey()
+                            : event.getBeforeRoutingKey())
+                    .map(CdcOrderContractTest::readLong)
                     .toList();
         }
     }

src/test/java/io/pixelsdb/pixels/sink/cdc/RecordingRetinaWriter.java

@@ -82,7 +82,6 @@ public boolean writeRow(RowChangeEvent row)
             int ordinal = recordedRequests.size() + 1;
             long timestamp = TIMESTAMP_BASE + ordinal;
             row.setTimeStamp(timestamp);
-            row.updateIndexKey();
 
             RetinaProto.TableUpdateData tableUpdate =
                     RetinaPayloadBuilder.buildTableUpdateData(

src/test/java/io/pixelsdb/pixels/sink/event/RowBatchTest.java

@@ -31,7 +31,7 @@ void integerRowBatchTest()
         TypeDescription schema = TypeDescription.createSchemaFromStrings(
                 java.util.List.of("id"), java.util.List.of("int"));
         VectorizedRowBatch rowBatch = schema.createRowBatch(
-                3, TypeDescription.Mode.CREATE_INT_VECTOR_FOR_INT);
+                3, TypeDescription.VectorLayout.NONE);
         IntColumnVector vector = (IntColumnVector) rowBatch.cols[0];
         vector.add(10);
         vector.add(20);
@@ -52,7 +52,7 @@ void varcharRowBatchTest()
     {
         TypeDescription schema = TypeDescription.createSchemaFromStrings(
                 java.util.List.of("name"), java.util.List.of("varchar(100)"));
-        VectorizedRowBatch rowBatch = schema.createRowBatch(3, TypeDescription.Mode.CREATE_INT_VECTOR_FOR_INT);
+        VectorizedRowBatch rowBatch = schema.createRowBatch(3, TypeDescription.VectorLayout.NONE);
         BinaryColumnVector v = (BinaryColumnVector) rowBatch.cols[0];
         v.add("rr");
         v.add("zz");

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

@@ -141,6 +141,15 @@ private static RowChangeEvent event(
             SinkProto.RowValue before,
             SinkProto.RowValue after,
             TableMetadata metadata) throws Exception
+    {
+        return eventWithRoutingOnly(operation, before, after, metadata);
+    }
+
+    private static RowChangeEvent eventWithRoutingOnly(
+            SinkProto.OperationType operation,
+            SinkProto.RowValue before,
+            SinkProto.RowValue after,
+            TableMetadata metadata) throws Exception
     {
         SinkProto.RowRecord.Builder record = SinkProto.RowRecord.newBuilder()
                 .setOp(operation)
@@ -158,8 +167,7 @@ private static RowChangeEvent event(
 
         RowChangeEvent event = new RowChangeEvent(
                 record.build(), metadata.getTypeDescription(), metadata);
-        event.setTimeStamp(TIMESTAMP);
-        event.initIndexKey();
+        event.initRoutingKey();
         return event;
     }
 
@@ -202,4 +210,24 @@ private static ByteString bytes(String value)
     {
         return ByteString.copyFromUtf8(value);
     }
+
+    @Test
+    void shouldBindIndexKeyWithPayloadTimestamp() throws Exception
+    {
+        TableMetadata metadata = tableMetadata();
+        RowChangeEvent insert = eventWithRoutingOnly(
+                SinkProto.OperationType.INSERT, null, row("1", "insert"), metadata);
+
+        assertEquals(bytes("1"), insert.getAfterRoutingKey());
+        assertEquals(null, insert.getAfterKey());
+
+        RetinaProto.TableUpdateData tableUpdate =
+                RetinaPayloadBuilder.buildTableUpdateData(
+                        TABLE_NAME, TIMESTAMP, List.of(insert));
+
+        assertEquals(TIMESTAMP, tableUpdate.getTimestamp());
+        assertEquals(TIMESTAMP, tableUpdate.getInsertData(0).getIndexKeys(0).getTimestamp());
+        assertEquals(bytes("1"), tableUpdate.getInsertData(0).getIndexKeys(0).getKey());
+        assertEquals(TIMESTAMP, insert.getAfterKey().getTimestamp());
+    }
 }