GH-948: Use buffer indexing for UUID vector

Author: jhrotko

vector/src/main/java/org/apache/arrow/vector/UuidVector.java

@@ -23,10 +23,13 @@
 import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.util.ArrowBufPointer;
+import org.apache.arrow.memory.util.ByteFunctionHelpers;
 import org.apache.arrow.memory.util.hash.ArrowBufHasher;
+import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.complex.impl.UuidReaderImpl;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.holders.ExtensionHolder;
 import org.apache.arrow.vector.holders.NullableUuidHolder;
 import org.apache.arrow.vector.holders.UuidHolder;
 import org.apache.arrow.vector.types.pojo.Field;
@@ -132,7 +135,8 @@ public int hashCode(int index) {
 
   @Override
   public int hashCode(int index, ArrowBufHasher hasher) {
-    return getUnderlyingVector().hashCode(index, hasher);
+    int start = this.getStartOffset(index);
+    return ByteFunctionHelpers.hash(hasher, this.getDataBuffer(), start, start + UUID_BYTE_WIDTH);
   }
 
   /**
@@ -146,17 +150,19 @@ public int isSet(int index) {
   }
 
   /**
-   * Gets the UUID value at the given index as an ArrowBuf.
+   * Reads the UUID value at the given index into a UuidHolder.
    *
-   * @param index the index to retrieve
-   * @return a buffer slice containing the 16-byte UUID
-   * @throws IllegalStateException if the value at the index is null and null checking is enabled
+   * @param index the index to read from
+   * @param holder the holder to populate with the UUID data
    */
-  public ArrowBuf get(int index) throws IllegalStateException {
+  public void get(int index, UuidHolder holder) {
+    Preconditions.checkArgument(index >= 0, "Cannot get negative index in UUID vector.");
     if (NullCheckingForGet.NULL_CHECKING_ENABLED && this.isSet(index) == 0) {
-      throw new IllegalStateException("Value at index is null");
+      holder.isSet = 0;
     } else {
-      return getBufferSlicePostNullCheck(index);
+      holder.isSet = 1;
+      holder.buffer = getDataBuffer();
+      holder.start = getStartOffset(index);
     }
   }
 
@@ -167,23 +173,24 @@ public ArrowBuf get(int index) throws IllegalStateException {
    * @param holder the holder to populate with the UUID data
    */
   public void get(int index, NullableUuidHolder holder) {
+    Preconditions.checkArgument(index >= 0, "Cannot get negative index in UUID vector.");
     if (NullCheckingForGet.NULL_CHECKING_ENABLED && this.isSet(index) == 0) {
       holder.isSet = 0;
     } else {
       holder.isSet = 1;
-      holder.buffer = getBufferSlicePostNullCheck(index);
+      holder.buffer = getDataBuffer();
+      holder.start = getStartOffset(index);
     }
   }
 
   /**
-   * Reads the UUID value at the given index into a UuidHolder.
+   * Calculates the byte offset for a given index in the data buffer.
    *
-   * @param index the index to read from
-   * @param holder the holder to populate with the UUID data
+   * @param index the index of the UUID value
+   * @return the byte offset in the data buffer
    */
-  public void get(int index, UuidHolder holder) {
-    holder.isSet = 1;
-    holder.buffer = getBufferSlicePostNullCheck(index);
+  public final int getStartOffset(int index) {
+    return index * UUID_BYTE_WIDTH;
   }
 
   /**
@@ -207,7 +214,7 @@ public void set(int index, UUID value) {
    * @param holder the holder containing the UUID data
    */
   public void set(int index, UuidHolder holder) {
-    this.set(index, holder.isSet, holder.buffer);
+    this.set(index, holder.buffer, holder.start);
   }
 
   /**
@@ -217,28 +224,11 @@ public void set(int index, UuidHolder holder) {
    * @param holder the holder containing the UUID data
    */
   public void set(int index, NullableUuidHolder holder) {
-    this.set(index, holder.isSet, holder.buffer);
-  }
-
-  /**
-   * Sets the UUID value at the given index with explicit null flag.
-   *
-   * @param index the index to set
-   * @param isSet 1 if the value is set, 0 if null
-   * @param buffer the buffer containing the 16-byte UUID data
-   */
-  public void set(int index, int isSet, ArrowBuf buffer) {
-    getUnderlyingVector().set(index, isSet, buffer);
-  }
-
-  /**
-   * Sets the UUID value at the given index from an ArrowBuf.
-   *
-   * @param index the index to set
-   * @param value the buffer containing the 16-byte UUID data
-   */
-  public void set(int index, ArrowBuf value) {
-    getUnderlyingVector().set(index, value);
+    if (holder.isSet == 0) {
+      getUnderlyingVector().setNull(index);
+    } else {
+      this.set(index, holder.buffer, holder.start);
+    }
   }
 
   /**
@@ -249,10 +239,12 @@ public void set(int index, ArrowBuf value) {
    * @param sourceOffset the offset in the source buffer where the UUID data starts
    */
   public void set(int index, ArrowBuf source, int sourceOffset) {
-    // Copy bytes from source buffer to target vector data buffer
-    ArrowBuf dataBuffer = getUnderlyingVector().getDataBuffer();
-    dataBuffer.setBytes((long) index * UUID_BYTE_WIDTH, source, sourceOffset, UUID_BYTE_WIDTH);
-    getUnderlyingVector().setIndexDefined(index);
+    Preconditions.checkNotNull(source, "Cannot set UUID vector, the source buffer is null.");
+
+    BitVectorHelper.setBit(getUnderlyingVector().getValidityBuffer(), index);
+    getUnderlyingVector()
+            .getDataBuffer()
+            .setBytes((long) index * UUID_BYTE_WIDTH, source, sourceOffset, UUID_BYTE_WIDTH);
   }
 
   /**
@@ -286,25 +278,34 @@ public void setSafe(int index, UUID value) {
    * @param holder the holder containing the UUID data, or null to set a null value
    */
   public void setSafe(int index, NullableUuidHolder holder) {
-    if (holder != null) {
-      getUnderlyingVector().setSafe(index, holder.isSet, holder.buffer);
-    } else {
+    if (holder == null || holder.isSet == 0) {
       getUnderlyingVector().setNull(index);
+    } else {
+      this.setSafe(index, holder.buffer, holder.start);
     }
   }
 
   /**
    * Sets the UUID value at the given index from a UuidHolder, expanding capacity if needed.
    *
    * @param index the index to set
-   * @param holder the holder containing the UUID data, or null to set a null value
+   * @param holder the holder containing the UUID data
    */
   public void setSafe(int index, UuidHolder holder) {
-    if (holder != null) {
-      getUnderlyingVector().setSafe(index, holder.isSet, holder.buffer);
-    } else {
-      getUnderlyingVector().setNull(index);
-    }
+    this.setSafe(index, holder.buffer, holder.start);
+  }
+
+  /**
+   * Sets the UUID value at the given index by copying from a source buffer, expanding capacity if
+   * needed.
+   *
+   * @param index the index to set
+   * @param buffer the source buffer to copy from
+   * @param start the offset in the source buffer where the UUID data starts
+   */
+  public void setSafe(int index, ArrowBuf buffer, int start) {
+    getUnderlyingVector().handleSafe(index);
+    this.set(index, buffer, start);
   }
 
   /**
@@ -400,15 +401,9 @@ public TransferPair getTransferPair(BufferAllocator allocator) {
     return getTransferPair(this.getField().getName(), allocator);
   }
 
-  private ArrowBuf getBufferSlicePostNullCheck(int index) {
-    return getUnderlyingVector()
-        .getDataBuffer()
-        .slice((long) index * UUID_BYTE_WIDTH, UUID_BYTE_WIDTH);
-  }
-
   @Override
   public int getTypeWidth() {
-    return getUnderlyingVector().getTypeWidth();
+    return UUID_BYTE_WIDTH;
   }
 
   /** {@link TransferPair} for {@link UuidVector}. */

vector/src/main/java/org/apache/arrow/vector/complex/impl/NullableUuidHolderReaderImpl.java

@@ -0,0 +1,124 @@
+/*
+ * 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.arrow.vector.complex.impl;
+
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.holders.NullableUuidHolder;
+import org.apache.arrow.vector.holders.UuidHolder;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.util.UuidUtility;
+
+/**
+ * Reader implementation for reading UUID values from a {@link NullableUuidHolder}.
+ *
+ * <p>This reader wraps a single UUID holder value and provides methods to read from it. Unlike
+ * {@link UuidReaderImpl} which reads from a vector, this reader operates on a holder instance.
+ *
+ * @see NullableUuidHolder
+ * @see UuidReaderImpl
+ */
+public class NullableUuidHolderReaderImpl extends AbstractFieldReader {
+  private final NullableUuidHolder holder;
+
+  /**
+   * Constructs a reader for the given UUID holder.
+   *
+   * @param holder the UUID holder to read from
+   */
+  public NullableUuidHolderReaderImpl(NullableUuidHolder holder) {
+    this.holder = holder;
+  }
+
+  @Override
+  public int size() {
+    throw new UnsupportedOperationException(
+        "size() is not supported on NullableUuidHolderReaderImpl. "
+            + "This reader wraps a single UUID holder value, not a collection. "
+            + "Use UuidReaderImpl for vector-based UUID reading.");
+  }
+
+  @Override
+  public boolean next() {
+    throw new UnsupportedOperationException(
+        "next() is not supported on NullableUuidHolderReaderImpl. "
+            + "This reader wraps a single UUID holder value, not an iterator. "
+            + "Use UuidReaderImpl for vector-based UUID reading.");
+  }
+
+  @Override
+  public void setPosition(int index) {
+    throw new UnsupportedOperationException(
+        "setPosition() is not supported on NullableUuidHolderReaderImpl. "
+            + "This reader wraps a single UUID holder value, not a vector. "
+            + "Use UuidReaderImpl for vector-based UUID reading.");
+  }
+
+  @Override
+  public Types.MinorType getMinorType() {
+    return Types.MinorType.EXTENSIONTYPE;
+  }
+
+  @Override
+  public boolean isSet() {
+    return holder.isSet == 1;
+  }
+
+  @Override
+  public void read(ExtensionHolder h) {
+    if (h instanceof NullableUuidHolder) {
+      NullableUuidHolder nullableHolder = (NullableUuidHolder) h;
+      nullableHolder.buffer = this.holder.buffer;
+      nullableHolder.isSet = this.holder.isSet;
+      nullableHolder.start = this.holder.start;
+    } else if (h instanceof UuidHolder) {
+      UuidHolder uuidHolder = (UuidHolder) h;
+      uuidHolder.buffer = this.holder.buffer;
+      uuidHolder.start = this.holder.start;
+    } else {
+      throw new IllegalArgumentException(
+          "Unsupported holder type: "
+              + h.getClass().getName()
+              + ". "
+              + "Only NullableUuidHolder and UuidHolder are supported for UUID values. "
+              + "Provided holder type cannot be used to read UUID data.");
+    }
+  }
+
+  @Override
+  public Object readObject() {
+    if (!isSet()) {
+      return null;
+    }
+    // Convert UUID bytes to Java UUID object
+    try {
+      return UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
+    } catch (Exception e) {
+      throw new RuntimeException(
+          String.format(
+              "Failed to read UUID from buffer. Invalid Arrow buffer state: "
+                  + "capacity=%d, readableBytes=%d, readerIndex=%d, writerIndex=%d, refCnt=%d. "
+                  + "The buffer must contain exactly 16 bytes of valid UUID data.",
+              holder.buffer.capacity(),
+              holder.buffer.readableBytes(),
+              holder.buffer.readerIndex(),
+              holder.buffer.writerIndex(),
+              holder.buffer.refCnt()),
+          e);
+    }
+  }
+}
+

vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterImpl.java

@@ -50,8 +50,15 @@ public void writeExtension(Object value) {
       vector.setSafe(getPosition(), (ArrowBuf) value);
     } else if (value instanceof java.util.UUID) {
       vector.setSafe(getPosition(), (java.util.UUID) value);
+    } else if (value instanceof ExtensionHolder) {
+      write((ExtensionHolder) value);
     } else {
-      throw new IllegalArgumentException("Unsupported value type for UUID: " + value.getClass());
+      throw new IllegalArgumentException(
+          "Unsupported value type for UUID: "
+              + value.getClass().getName()
+              + ". "
+              + "Supported types are: byte[] (16 bytes), ArrowBuf (16 bytes), or java.util.UUID. "
+              + "Convert your value to one of these types before writing.");
     }
     vector.setValueCount(getPosition() + 1);
   }

vector/src/main/java/org/apache/arrow/vector/holders/NullableUuidHolder.java

@@ -32,4 +32,7 @@
 public class NullableUuidHolder extends ExtensionHolder {
   /** Buffer containing 16-byte UUID data. */
   public ArrowBuf buffer;
+
+  /** Offset in the buffer where the UUID data starts. */
+  public int start = 0;
 }

vector/src/main/java/org/apache/arrow/vector/holders/UuidHolder.java

@@ -31,6 +31,9 @@ public class UuidHolder extends ExtensionHolder {
   /** Buffer containing 16-byte UUID data. */
   public ArrowBuf buffer;
 
+  /** Offset in the buffer where the UUID data starts. */
+  public int start = 0;
+
   /** Constructs a UuidHolder with isSet = 1. */
   public UuidHolder() {
     this.isSet = 1;

vector/src/test/java/org/apache/arrow/vector/TestListVector.java

@@ -1259,12 +1259,12 @@ public void testListVectorReaderForExtensionType() throws Exception {
       FieldReader uuidReader = reader.reader();
       UuidHolder holder = new UuidHolder();
       uuidReader.read(holder);
-      UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0);
+      UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
       assertEquals(u1, actualUuid);
       reader.next();
       uuidReader = reader.reader();
       uuidReader.read(holder);
-      actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0);
+      actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
       assertEquals(u2, actualUuid);
     }
   }
@@ -1300,12 +1300,12 @@ public void testCopyFromForExtensionType() throws Exception {
       FieldReader uuidReader = reader.reader();
       UuidHolder holder = new UuidHolder();
       uuidReader.read(holder);
-      UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0);
+      UUID actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
       assertEquals(u1, actualUuid);
       reader.next();
       uuidReader = reader.reader();
       uuidReader.read(holder);
-      actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, 0);
+      actualUuid = UuidUtility.uuidFromArrowBuf(holder.buffer, holder.start);
       assertEquals(u2, actualUuid);
     }
   }