perf: cache offsetBufferAddress in CometPlainVector for variable-width vectors (apache#4364)

Author: 0lai0

spark/src/main/java/org/apache/comet/vector/CometPlainVector.java

@@ -33,6 +33,7 @@
 /** A column vector whose elements are plainly decoded. */
 public class CometPlainVector extends CometDecodedVector {
   private final long valueBufferAddress;
+  private final long offsetBufferAddress;
   private final boolean isBaseFixedWidthVector;
 
   private byte booleanByteCache;
@@ -58,6 +59,12 @@ public CometPlainVector(ValueVector vector, boolean isUuid, boolean isReused) {
     }
 
     isBaseFixedWidthVector = valueVector instanceof BaseFixedWidthVector;
+    if (vector instanceof BaseVariableWidthVector) {
+      this.offsetBufferAddress =
+          ((BaseVariableWidthVector) vector).getOffsetBuffer().memoryAddress();
+    } else {
+      this.offsetBufferAddress = -1;
+    }
     this.isReused = isReused;
   }
 
@@ -123,13 +130,11 @@ public double getDouble(int rowId) {
   @Override
   public UTF8String getUTF8String(int rowId) {
     if (isNullAt(rowId)) return null;
-    if (!isBaseFixedWidthVector) {
-      BaseVariableWidthVector varWidthVector = (BaseVariableWidthVector) valueVector;
-      long offsetBufferAddress = varWidthVector.getOffsetBuffer().memoryAddress();
+    if (offsetBufferAddress != -1) {
       int offset = Platform.getInt(null, offsetBufferAddress + rowId * 4L);
       int length = Platform.getInt(null, offsetBufferAddress + (rowId + 1L) * 4L) - offset;
       return UTF8String.fromAddress(null, valueBufferAddress + offset, length);
-    } else {
+    } else if (isBaseFixedWidthVector) {
       BaseFixedWidthVector fixedWidthVector = (BaseFixedWidthVector) valueVector;
       int length = fixedWidthVector.getTypeWidth();
       int offset = rowId * length;
@@ -142,6 +147,8 @@ public UTF8String getUTF8String(int rowId) {
       } else {
         return UTF8String.fromString(convertToUuid(result).toString());
       }
+    } else {
+      throw new IllegalStateException("Unsupported UTF8 vector type: " + valueVector.getName());
     }
   }
 
@@ -150,17 +157,15 @@ public byte[] getBinary(int rowId) {
     if (isNullAt(rowId)) return null;
     int offset;
     int length;
-    if (valueVector instanceof BaseVariableWidthVector) {
-      BaseVariableWidthVector varWidthVector = (BaseVariableWidthVector) valueVector;
-      long offsetBufferAddress = varWidthVector.getOffsetBuffer().memoryAddress();
+    if (offsetBufferAddress != -1) {
       offset = Platform.getInt(null, offsetBufferAddress + rowId * 4L);
       length = Platform.getInt(null, offsetBufferAddress + (rowId + 1L) * 4L) - offset;
     } else if (valueVector instanceof BaseFixedWidthVector) {
       BaseFixedWidthVector fixedWidthVector = (BaseFixedWidthVector) valueVector;
       length = fixedWidthVector.getTypeWidth();
       offset = rowId * length;
     } else {
-      throw new RuntimeException("Unsupported binary vector type: " + valueVector.getName());
+      throw new IllegalStateException("Unsupported binary vector type: " + valueVector.getName());
     }
     byte[] result = new byte[length];
     Platform.copyMemory(

spark/src/test/java/org/apache/comet/vector/TestCometPlainVector.java

@@ -0,0 +1,77 @@
+/*
+ * 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.comet.vector;
+
+import java.nio.charset.StandardCharsets;
+
+import org.junit.Test;
+
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.VarBinaryVector;
+import org.apache.arrow.vector.VarCharVector;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+public class TestCometPlainVector {
+
+  @Test
+  public void testGetUTF8StringWithVariableWidthVector() {
+    try (RootAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
+      VarCharVector vector = new VarCharVector("strings", allocator);
+      vector.allocateNew();
+      vector.setSafe(0, bytes("alpha"));
+      vector.setSafe(1, bytes(""));
+      vector.setSafe(2, bytes("spark"));
+      vector.setValueCount(4); // row 3 is null (validity bit not set)
+
+      try (CometPlainVector cv = new CometPlainVector(vector, false)) {
+        assertEquals("alpha", cv.getUTF8String(0).toString());
+        assertEquals("", cv.getUTF8String(1).toString());
+        assertEquals("spark", cv.getUTF8String(2).toString());
+        assertNull(cv.getUTF8String(3));
+      }
+    }
+  }
+
+  @Test
+  public void testGetBinaryWithVariableWidthVector() {
+    try (RootAllocator allocator = new RootAllocator(Integer.MAX_VALUE)) {
+      VarBinaryVector vector = new VarBinaryVector("bytes", allocator);
+      vector.allocateNew();
+      vector.setSafe(0, new byte[] {1, 2, 3}, 0, 3);
+      vector.setSafe(1, new byte[0], 0, 0);
+      vector.setSafe(2, new byte[] {4, 5}, 0, 2);
+      vector.setValueCount(4); // row 3 is null (validity bit not set)
+
+      try (CometPlainVector cv = new CometPlainVector(vector, false)) {
+        assertArrayEquals(new byte[] {1, 2, 3}, cv.getBinary(0));
+        assertArrayEquals(new byte[0], cv.getBinary(1));
+        assertArrayEquals(new byte[] {4, 5}, cv.getBinary(2));
+        assertNull(cv.getBinary(3));
+      }
+    }
+  }
+
+  private static byte[] bytes(String s) {
+    return s.getBytes(StandardCharsets.UTF_8);
+  }
+}