[core] Improve vector capacity growth for huge batches

paimon-common/src/main/java/org/apache/paimon/data/columnar/writable/AbstractWritableVector.java

@@ -31,6 +31,11 @@ public abstract class AbstractWritableVector implements WritableColumnVector, Se
 
     private static final long serialVersionUID = 1L;
 
+    static final int MAX_ROUNDED_ARRAY_LENGTH = Integer.MAX_VALUE - 15;
+    static final int DEFAULT_HUGE_VECTOR_THRESHOLD = 1 << 20;
+    static final double DEFAULT_HUGE_VECTOR_RESERVE_RATIO = 1.2;
+    static final int DEFAULT_HUGE_VECTOR_SHRINK_RATIO = 4;
+
     // If the whole column vector has no nulls, this is true, otherwise false.
     protected boolean noNulls = true;
 
@@ -41,13 +46,16 @@ public abstract class AbstractWritableVector implements WritableColumnVector, Se
 
     protected int capacity;
 
+    private final int initialCapacity;
+
     /**
      * The Dictionary for this column. If it's not null, will be used to decode the value in get().
      */
     protected Dictionary dictionary;
 
     public AbstractWritableVector(int capacity) {
         this.capacity = capacity;
+        this.initialCapacity = capacity;
     }
 
     /** Update the dictionary. */
@@ -91,8 +99,10 @@ public int getCapacity() {
 
     @Override
     public void reset() {
-        // To reduce copy, Ww don't result the capacity to initial capacity here. Which means the
-        // capacity will be the same as expand.
+        // Keep normal expanded capacity for reuse, but release buffers inflated by one-off spikes.
+        if (shouldShrinkCapacity()) {
+            capacity = initialCapacity;
+        }
         noNulls = true;
         isAllNull = false;
         elementsAppended = 0;
@@ -103,7 +113,7 @@ public void reserve(int requiredCapacity) {
         if (requiredCapacity < 0) {
             throw new IllegalArgumentException("Invalid capacity: " + requiredCapacity);
         } else if (requiredCapacity > capacity) {
-            int newCapacity = (int) Math.min(Integer.MAX_VALUE, requiredCapacity * 2L);
+            int newCapacity = calculateNewCapacity(requiredCapacity);
             if (requiredCapacity <= newCapacity) {
                 try {
                     reserveInternal(newCapacity);
@@ -113,11 +123,24 @@ public void reserve(int requiredCapacity) {
                 }
             } else {
                 throw new UnsupportedOperationException(
-                        "Cannot allocate :" + newCapacity + " elements");
+                        "Cannot allocate " + requiredCapacity + " elements");
             }
             capacity = newCapacity;
         }
     }
 
+    private boolean shouldShrinkCapacity() {
+        return capacity > DEFAULT_HUGE_VECTOR_THRESHOLD
+                && capacity > (long) elementsAppended * DEFAULT_HUGE_VECTOR_SHRINK_RATIO;
+    }
+
+    static int calculateNewCapacity(int requiredCapacity) {
+        long newCapacity =
+                requiredCapacity < DEFAULT_HUGE_VECTOR_THRESHOLD
+                        ? requiredCapacity * 2L
+                        : (long) (requiredCapacity * DEFAULT_HUGE_VECTOR_RESERVE_RATIO);
+        return (int) Math.min(MAX_ROUNDED_ARRAY_LENGTH, newCapacity);
+    }
+
     protected abstract void reserveInternal(int newCapacity);
 }

paimon-common/src/test/java/org/apache/paimon/data/columnar/heap/HeapVectorCapacityTest.java

@@ -0,0 +1,59 @@
+/*
+ * 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.columnar.heap;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for heap vector capacity reuse. */
+class HeapVectorCapacityTest {
+
+    @Test
+    void testHugeVectorCapacityIsReleasedOnReset() {
+        int threshold = 1 << 20;
+        HeapIntVector vector = new HeapIntVector(4);
+
+        vector.reserve(threshold);
+        assertThat(vector.getCapacity()).isGreaterThan(threshold);
+        assertThat(vector.vector.length).isEqualTo(vector.getCapacity());
+
+        vector.addElementsAppended(1);
+        vector.reset();
+
+        assertThat(vector.getCapacity()).isEqualTo(4);
+        assertThat(vector.vector).hasSize(4);
+    }
+
+    @Test
+    void testHugeVectorCapacityIsRetainedWhenUsageIsNotTooSmall() {
+        int threshold = 1 << 20;
+        HeapIntVector vector = new HeapIntVector(4);
+
+        vector.reserve(threshold);
+        int expandedCapacity = vector.getCapacity();
+        // Use ceiling division (shrink ratio = 4)
+        // so capacity is not strictly greater than usage * 4.
+        vector.addElementsAppended((expandedCapacity + 3) / 4);
+        vector.reset();
+
+        assertThat(vector.getCapacity()).isEqualTo(expandedCapacity);
+        assertThat(vector.vector).hasSize(expandedCapacity);
+    }
+}

paimon-common/src/test/java/org/apache/paimon/data/columnar/writable/AbstractWritableVectorTest.java

@@ -0,0 +1,142 @@
+/*
+ * 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.columnar.writable;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link AbstractWritableVector} capacity growth. */
+class AbstractWritableVectorTest {
+
+    @Test
+    void testSmallVectorDoublesRequiredCapacity() {
+        TestVector vector = new TestVector(4);
+
+        vector.reserve(5);
+
+        assertThat(vector.getCapacity()).isEqualTo(10);
+        assertThat(vector.reservedCapacity).isEqualTo(10);
+    }
+
+    @Test
+    void testHugeVectorUsesConservativeGrowth() {
+        int threshold = AbstractWritableVector.DEFAULT_HUGE_VECTOR_THRESHOLD;
+        TestVector vector = new TestVector(threshold - 1);
+
+        vector.reserve(threshold);
+
+        assertThat(vector.getCapacity())
+                .isEqualTo(
+                        (int)
+                                (threshold
+                                        * AbstractWritableVector
+                                                .DEFAULT_HUGE_VECTOR_RESERVE_RATIO));
+    }
+
+    @Test
+    void testHugeVectorCapacityResetsToInitialCapacityWhenUsageIsSmall() {
+        int threshold = AbstractWritableVector.DEFAULT_HUGE_VECTOR_THRESHOLD;
+        TestVector vector = new TestVector(4);
+
+        vector.reserve(threshold);
+        vector.addElementsAppended(1);
+        vector.reset();
+
+        assertThat(vector.getCapacity()).isEqualTo(4);
+        assertThat(vector.getElementsAppended()).isZero();
+    }
+
+    @Test
+    void testHugeVectorCapacityIsRetainedWhenUsageIsNotTooSmall() {
+        int threshold = AbstractWritableVector.DEFAULT_HUGE_VECTOR_THRESHOLD;
+        TestVector vector = new TestVector(4);
+
+        vector.reserve(threshold);
+        int expandedCapacity = vector.getCapacity();
+        // Use ceiling division so capacity is not strictly greater than usage * 4.
+        int retainedUsage =
+                (expandedCapacity + AbstractWritableVector.DEFAULT_HUGE_VECTOR_SHRINK_RATIO - 1)
+                        / AbstractWritableVector.DEFAULT_HUGE_VECTOR_SHRINK_RATIO;
+        vector.addElementsAppended(retainedUsage);
+        vector.reset();
+
+        assertThat(vector.getCapacity()).isEqualTo(expandedCapacity);
+        assertThat(vector.getElementsAppended()).isZero();
+    }
+
+    @Test
+    void testNormalExpandedCapacityIsRetainedAfterReset() {
+        TestVector vector = new TestVector(4);
+
+        vector.reserve(5);
+        vector.reset();
+
+        assertThat(vector.getCapacity()).isEqualTo(10);
+    }
+
+    @Test
+    void testReserveRejectsCapacityBeyondMaxRoundedArrayLength() {
+        TestVector vector = new TestVector(4);
+
+        assertThatThrownBy(() -> vector.reserve(Integer.MAX_VALUE))
+                .isInstanceOf(UnsupportedOperationException.class)
+                .hasMessage("Cannot allocate " + Integer.MAX_VALUE + " elements");
+    }
+
+    private static class TestVector extends AbstractWritableVector {
+
+        private int reservedCapacity;
+
+        private TestVector(int capacity) {
+            super(capacity);
+            this.reservedCapacity = capacity;
+        }
+
+        @Override
+        public void setNullAt(int rowId) {}
+
+        @Override
+        public void setNulls(int rowId, int count) {}
+
+        @Override
+        public void fillWithNulls() {}
+
+        @Override
+        public WritableIntVector reserveDictionaryIds(int capacity) {
+            return null;
+        }
+
+        @Override
+        public WritableIntVector getDictionaryIds() {
+            return null;
+        }
+
+        @Override
+        public boolean isNullAt(int i) {
+            return false;
+        }
+
+        @Override
+        protected void reserveInternal(int newCapacity) {
+            this.reservedCapacity = newCapacity;
+        }
+    }
+}