From f91a0341e4e3fcb37d5cab2a6902209b7b43d718 Mon Sep 17 00:00:00 2001
From: tsreaper <tsreaper96@gmail.com>
Date: Fri, 2 Feb 2024 20:46:22 +0800
Subject: [PATCH] [orc] Optimize ensureSize for orc ColumnVector (#2845)

---
 .../format/orc/writer/RowDataVectorizer.java  | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/paimon-format/src/main/java/org/apache/paimon/format/orc/writer/RowDataVectorizer.java b/paimon-format/src/main/java/org/apache/paimon/format/orc/writer/RowDataVectorizer.java
index 5349935ca5cc..b05b5e208995 100644
--- a/paimon-format/src/main/java/org/apache/paimon/format/orc/writer/RowDataVectorizer.java
+++ b/paimon-format/src/main/java/org/apache/paimon/format/orc/writer/RowDataVectorizer.java
@@ -194,8 +194,10 @@ private static void setColumn(
         listColumnVector.lengths[rowId] = arrayData.size();
         listColumnVector.offsets[rowId] = listColumnVector.childCount;
         listColumnVector.childCount += listColumnVector.lengths[rowId];
-        listColumnVector.child.ensureSize(
-                listColumnVector.childCount, listColumnVector.offsets[rowId] != 0);
+        ensureSize(
+                listColumnVector.child,
+                listColumnVector.childCount,
+                listColumnVector.offsets[rowId] != 0);
 
         InternalRow convertedRowData = convert(arrayData, arrayType.getElementType());
         for (int i = 0; i < arrayData.size(); i++) {
@@ -221,10 +223,14 @@ private static void setColumn(
         mapColumnVector.lengths[rowId] = mapData.size();
         mapColumnVector.offsets[rowId] = mapColumnVector.childCount;
         mapColumnVector.childCount += mapColumnVector.lengths[rowId];
-        mapColumnVector.keys.ensureSize(
-                mapColumnVector.childCount, mapColumnVector.offsets[rowId] != 0);
-        mapColumnVector.values.ensureSize(
-                mapColumnVector.childCount, mapColumnVector.offsets[rowId] != 0);
+        ensureSize(
+                mapColumnVector.keys,
+                mapColumnVector.childCount,
+                mapColumnVector.offsets[rowId] != 0);
+        ensureSize(
+                mapColumnVector.values,
+                mapColumnVector.childCount,
+                mapColumnVector.offsets[rowId] != 0);
 
         InternalRow convertedKeyRowData = convert(keyArray, mapType.getKeyType());
         InternalRow convertedValueRowData = convert(valueArray, mapType.getValueType());
@@ -258,6 +264,13 @@ private static void setColumn(
         }
     }
 
+    private static void ensureSize(ColumnVector cv, int size, boolean preserveData) {
+        int currentLength = cv.isNull.length;
+        if (currentLength < size) {
+            cv.ensureSize(Math.max(currentLength * 2, size), preserveData);
+        }
+    }
+
     /**
      * Converting ArrayData to RowData for calling {@link RowDataVectorizer#setColumn(int,
      * ColumnVector, DataType, InternalRow, int)} recursively with array.