diff --git a/paimon-common/src/main/java/org/apache/paimon/data/columnar/writable/AbstractWritableVector.java b/paimon-common/src/main/java/org/apache/paimon/data/columnar/writable/AbstractWritableVector.java index fd4981b403d0..59e5ce91d1fb 100644 --- a/paimon-common/src/main/java/org/apache/paimon/data/columnar/writable/AbstractWritableVector.java +++ b/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,6 +46,8 @@ 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(). */ @@ -48,6 +55,7 @@ public abstract class AbstractWritableVector implements WritableColumnVector, Se 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); } diff --git a/paimon-common/src/test/java/org/apache/paimon/data/columnar/heap/HeapVectorCapacityTest.java b/paimon-common/src/test/java/org/apache/paimon/data/columnar/heap/HeapVectorCapacityTest.java new file mode 100644 index 000000000000..907601281b16 --- /dev/null +++ b/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); + } +} diff --git a/paimon-common/src/test/java/org/apache/paimon/data/columnar/writable/AbstractWritableVectorTest.java b/paimon-common/src/test/java/org/apache/paimon/data/columnar/writable/AbstractWritableVectorTest.java new file mode 100644 index 000000000000..9e63d4933f84 --- /dev/null +++ b/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; + } + } +}