[orc] Optimize ensureSize for orc ColumnVector (apache#2845)

zhu3pang · Feb 2, 2024 · f91a034 · f91a034
1 parent 0f117d4
commit f91a034
Showing 1 changed file with 19 additions and 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
@@ -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.