pass BinaryRow to resetIfTooLarge for ownership check

Author: zhongheng.gy

paimon-common/src/main/java/org/apache/paimon/data/RowHelper.java

@@ -39,7 +39,7 @@ public class RowHelper implements Serializable {
     /**
      * Threshold in bytes for releasing the internal reuse buffer. When big records are written, the
      * BinaryRowWriter's internal segment can grow very large via grow(). The {@link
-     * #resetIfTooLarge()} method checks this threshold and releases the bloated
+     * #resetIfTooLarge(BinaryRow)} method checks this threshold and releases the bloated
      * reuseRow/reuseWriter to avoid holding onto oversized buffers indefinitely.
      */
     private static final int REUSE_RELEASE_THRESHOLD = 4 * 1024 * 1024; // 4MB
@@ -90,15 +90,22 @@ public void copyInto(InternalRow row) {
     }
 
     /**
-     * Release the internal reuse buffer if the segment exceeds the threshold AND the last written
-     * record is small. This hysteresis avoids thrashing when records are consistently large, while
-     * still reclaiming memory when the workload transitions back to small records.
+     * Release the internal reuse buffer if the given row is the reuse row produced by this helper,
+     * the backing segment exceeds the threshold, and the current record is small. The identity
+     * check ({@code currentRow == reuseRow}) ensures we only act when the caller actually used this
+     * helper's buffer — if the input was already a {@link BinaryRow}, {@code toBinaryRow()} returns
+     * it directly and the helper state is stale, so we must skip cleanup.
+     *
+     * <p>The hysteresis ({@code currentRow.getSizeInBytes() < threshold}) avoids thrashing when
+     * records are consistently large, while still reclaiming memory when the workload transitions
+     * back to small records.
      */
-    public void resetIfTooLarge() {
-        if (reuseWriter != null
+    public void resetIfTooLarge(BinaryRow currentRow) {
+        if (currentRow == reuseRow
+                && reuseWriter != null
                 && reuseWriter.getSegments() != null
                 && reuseWriter.getSegments().size() > REUSE_RELEASE_THRESHOLD
-                && reuseRow.getSizeInBytes() < REUSE_RELEASE_THRESHOLD) {
+                && currentRow.getSizeInBytes() < REUSE_RELEASE_THRESHOLD) {
             reuseRow = null;
             reuseWriter = null;
         }

paimon-common/src/main/java/org/apache/paimon/data/serializer/InternalRowSerializer.java

@@ -59,15 +59,11 @@ public InternalRowSerializer duplicate() {
 
     @Override
     public void serialize(InternalRow row, DataOutputView target) throws IOException {
+        BinaryRow binaryRow = toBinaryRow(row);
         try {
-            binarySerializer.serialize(toBinaryRow(row), target);
+            binarySerializer.serialize(binaryRow, target);
         } finally {
-            // Must use finally here: toBinaryRow() may inflate RowHelper's internal buffer
-            // for large records (e.g. 100MB+). The serialization can exit via EOFException
-            // thrown by SimpleCollectingOutputView.nextSegment() when the sort buffer is
-            // full, which is caught by BinaryInMemorySortBuffer.write() as a normal signal.
-            // Without finally, the bloated buffer would never be released on that path.
-            rowHelper.resetIfTooLarge();
+            rowHelper.resetIfTooLarge(binaryRow);
         }
     }
 
@@ -141,12 +137,11 @@ public InternalRow createReuseInstance() {
     @Override
     public int serializeToPages(InternalRow row, AbstractPagedOutputView target)
             throws IOException {
+        BinaryRow binaryRow = toBinaryRow(row);
         try {
-            return binarySerializer.serializeToPages(toBinaryRow(row), target);
+            return binarySerializer.serializeToPages(binaryRow, target);
         } finally {
-            // Same as serialize(): must use finally because EOFException may bypass normal
-            // return when the sort buffer is full.
-            rowHelper.resetIfTooLarge();
+            rowHelper.resetIfTooLarge(binaryRow);
         }
     }
 

paimon-common/src/test/java/org/apache/paimon/data/RowHelperTest.java

@@ -18,6 +18,7 @@
 
 package org.apache.paimon.data;
 
+import org.apache.paimon.memory.MemorySegment;
 import org.apache.paimon.types.DataTypes;
 import org.apache.paimon.types.RowKind;
 
@@ -27,69 +28,89 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 
-/** Tests for {@link RowHelper}, focusing on the resetIfTooLarge() behavior. */
+/** Tests for {@link RowHelper}, focusing on the resetIfTooLarge(BinaryRow) behavior. */
 class RowHelperTest {
 
     @Test
-    void testResetIfTooLargeReleasesAfterTransitionToSmallRecord() {
+    void testReleasesWhenTransitionFromLargeToSmallRecord() {
         RowHelper helper = new RowHelper(Arrays.asList(DataTypes.STRING(), DataTypes.BYTES()));
 
         // Write a large record (> 4MB) to inflate the internal buffer
-        byte[] largePayload = new byte[5 * 1024 * 1024]; // 5MB
-        Arrays.fill(largePayload, (byte) 'x');
+        byte[] largePayload = new byte[5 * 1024 * 1024];
         GenericRow largeRow = GenericRow.of(BinaryString.fromString("key"), largePayload);
         largeRow.setRowKind(RowKind.INSERT);
         helper.copyInto(largeRow);
+        BinaryRow reuseAfterLarge = helper.reuseRow();
+        assertThat(reuseAfterLarge).isNotNull();
 
+        // Hysteresis: should NOT release when current record is large
+        helper.resetIfTooLarge(reuseAfterLarge);
         assertThat(helper.reuseRow()).isNotNull();
 
-        // Hysteresis: resetIfTooLarge() should NOT release when last record is large
-        helper.resetIfTooLarge();
-        assertThat(helper.reuseRow()).isNotNull();
-
-        // Now write a small record — buffer is still oversized from the large record
+        // Write a small record — buffer is still oversized from the large record
         GenericRow smallRow = GenericRow.of(BinaryString.fromString("s"), new byte[10]);
         smallRow.setRowKind(RowKind.INSERT);
         helper.copyInto(smallRow);
+        BinaryRow reuseAfterSmall = helper.reuseRow();
 
-        // resetIfTooLarge() should release now: buffer > 4MB but last record < 4MB
-        helper.resetIfTooLarge();
+        // Should release now: buffer > 4MB but current record < 4MB
+        helper.resetIfTooLarge(reuseAfterSmall);
         assertThat(helper.reuseRow()).isNull();
     }
 
     @Test
-    void testResetIfTooLargeKeepsSmallBuffer() {
+    void testKeepsSmallBuffer() {
         RowHelper helper = new RowHelper(Arrays.asList(DataTypes.STRING(), DataTypes.INT()));
 
-        // Write a small record (< 4MB)
         GenericRow smallRow = GenericRow.of(BinaryString.fromString("hello"), 42);
         smallRow.setRowKind(RowKind.INSERT);
         helper.copyInto(smallRow);
+        BinaryRow reuse = helper.reuseRow();
+        assertThat(reuse).isNotNull();
 
+        // Small buffer should NOT be released
+        helper.resetIfTooLarge(reuse);
         assertThat(helper.reuseRow()).isNotNull();
+    }
 
-        // resetIfTooLarge() should NOT release the small buffer
-        helper.resetIfTooLarge();
+    @Test
+    void testSkipsWhenCurrentRowIsNotReuseRow() {
+        RowHelper helper = new RowHelper(Arrays.asList(DataTypes.STRING(), DataTypes.BYTES()));
+
+        // Write a large record to inflate the buffer
+        byte[] largePayload = new byte[5 * 1024 * 1024];
+        GenericRow largeRow = GenericRow.of(BinaryString.fromString("key"), largePayload);
+        largeRow.setRowKind(RowKind.INSERT);
+        helper.copyInto(largeRow);
+        assertThat(helper.reuseRow()).isNotNull();
+
+        // Simulate the BinaryRow input path: toBinaryRow() returns the input directly,
+        // not the helper's reuseRow. Pass a different BinaryRow instance.
+        BinaryRow externalRow = new BinaryRow(2);
+        externalRow.pointTo(MemorySegment.wrap(new byte[32]), 0, 32);
+
+        // Should NOT release because externalRow != reuseRow
+        helper.resetIfTooLarge(externalRow);
         assertThat(helper.reuseRow()).isNotNull();
     }
 
     @Test
-    void testResetIfTooLargeBeforeCopyInto() {
+    void testSafeToCallWithNullReuseRow() {
         RowHelper helper = new RowHelper(Arrays.asList(DataTypes.STRING()));
-
-        // reuseRow is null before any copyInto
         assertThat(helper.reuseRow()).isNull();
 
-        // resetIfTooLarge() should be safe to call when reuseRow is null
-        helper.resetIfTooLarge();
+        // Should be safe — reuseRow is null, no NPE
+        BinaryRow someRow = new BinaryRow(1);
+        someRow.pointTo(MemorySegment.wrap(new byte[16]), 0, 16);
+        helper.resetIfTooLarge(someRow);
         assertThat(helper.reuseRow()).isNull();
     }
 
     @Test
-    void testReuseIsRecreatedAfterRelease() {
+    void testReuseRecreatedAfterRelease() {
         RowHelper helper = new RowHelper(Arrays.asList(DataTypes.STRING(), DataTypes.BYTES()));
 
-        // Write a large record to inflate the buffer, then a small record to trigger release
+        // Inflate buffer, then transition to small
         byte[] largePayload = new byte[5 * 1024 * 1024];
         GenericRow largeRow = GenericRow.of(BinaryString.fromString("key"), largePayload);
         largeRow.setRowKind(RowKind.INSERT);
@@ -99,16 +120,15 @@ void testReuseIsRecreatedAfterRelease() {
         smallRow.setRowKind(RowKind.INSERT);
         helper.copyInto(smallRow);
 
-        // Buffer is oversized + last record is small → release
-        helper.resetIfTooLarge();
+        helper.resetIfTooLarge(helper.reuseRow());
         assertThat(helper.reuseRow()).isNull();
 
         // Write another small record — reuseRow should be recreated
         helper.copyInto(smallRow);
         assertThat(helper.reuseRow()).isNotNull();
 
-        // Small buffer should survive resetIfTooLarge()
-        helper.resetIfTooLarge();
+        // Small buffer should survive
+        helper.resetIfTooLarge(helper.reuseRow());
         assertThat(helper.reuseRow()).isNotNull();
     }
 }