[common] Fix RowHelper internal buffer never shrinking for large records

Author: zhongheng.gy

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

@@ -36,6 +36,14 @@ public class RowHelper implements Serializable {
 
     private static final long serialVersionUID = 1L;
 
+    /**
+     * 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
+     * reuseRow/reuseWriter to avoid holding onto oversized buffers indefinitely.
+     */
+    private static final int REUSE_RELEASE_THRESHOLD = 4 * 1024 * 1024; // 4MB
+
     private final FieldGetter[] fieldGetters;
     private final ValueSetter[] valueSetters;
     private final boolean[] writeNulls;
@@ -81,6 +89,21 @@ public void copyInto(InternalRow row) {
         reuseWriter.complete();
     }
 
+    /**
+     * 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.
+     */
+    public void resetIfTooLarge() {
+        if (reuseWriter != null
+                && reuseWriter.getSegments() != null
+                && reuseWriter.getSegments().size() > REUSE_RELEASE_THRESHOLD
+                && reuseRow.getSizeInBytes() < REUSE_RELEASE_THRESHOLD) {
+            reuseRow = null;
+            reuseWriter = null;
+        }
+    }
+
     public BinaryRow reuseRow() {
         return reuseRow;
     }

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

@@ -59,7 +59,16 @@ public InternalRowSerializer duplicate() {
 
     @Override
     public void serialize(InternalRow row, DataOutputView target) throws IOException {
-        binarySerializer.serialize(toBinaryRow(row), target);
+        try {
+            binarySerializer.serialize(toBinaryRow(row), 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();
+        }
     }
 
     @Override
@@ -132,7 +141,13 @@ public InternalRow createReuseInstance() {
     @Override
     public int serializeToPages(InternalRow row, AbstractPagedOutputView target)
             throws IOException {
-        return binarySerializer.serializeToPages(toBinaryRow(row), target);
+        try {
+            return binarySerializer.serializeToPages(toBinaryRow(row), target);
+        } finally {
+            // Same as serialize(): must use finally because EOFException may bypass normal
+            // return when the sort buffer is full.
+            rowHelper.resetIfTooLarge();
+        }
     }
 
     @Override

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

@@ -0,0 +1,114 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.paimon.data;
+
+import org.apache.paimon.types.DataTypes;
+import org.apache.paimon.types.RowKind;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link RowHelper}, focusing on the resetIfTooLarge() behavior. */
+class RowHelperTest {
+
+    @Test
+    void testResetIfTooLargeReleasesAfterTransitionToSmallRecord() {
+        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');
+        GenericRow largeRow = GenericRow.of(BinaryString.fromString("key"), largePayload);
+        largeRow.setRowKind(RowKind.INSERT);
+        helper.copyInto(largeRow);
+
+        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
+        GenericRow smallRow = GenericRow.of(BinaryString.fromString("s"), new byte[10]);
+        smallRow.setRowKind(RowKind.INSERT);
+        helper.copyInto(smallRow);
+
+        // resetIfTooLarge() should release now: buffer > 4MB but last record < 4MB
+        helper.resetIfTooLarge();
+        assertThat(helper.reuseRow()).isNull();
+    }
+
+    @Test
+    void testResetIfTooLargeKeepsSmallBuffer() {
+        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);
+
+        assertThat(helper.reuseRow()).isNotNull();
+
+        // resetIfTooLarge() should NOT release the small buffer
+        helper.resetIfTooLarge();
+        assertThat(helper.reuseRow()).isNotNull();
+    }
+
+    @Test
+    void testResetIfTooLargeBeforeCopyInto() {
+        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();
+        assertThat(helper.reuseRow()).isNull();
+    }
+
+    @Test
+    void testReuseIsRecreatedAfterRelease() {
+        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
+        byte[] largePayload = new byte[5 * 1024 * 1024];
+        GenericRow largeRow = GenericRow.of(BinaryString.fromString("key"), largePayload);
+        largeRow.setRowKind(RowKind.INSERT);
+        helper.copyInto(largeRow);
+
+        GenericRow smallRow = GenericRow.of(BinaryString.fromString("small"), new byte[10]);
+        smallRow.setRowKind(RowKind.INSERT);
+        helper.copyInto(smallRow);
+
+        // Buffer is oversized + last record is small → release
+        helper.resetIfTooLarge();
+        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();
+        assertThat(helper.reuseRow()).isNotNull();
+    }
+}