Make UUID implementation available for downstream projects

vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java → vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidReaderImpl.java

@@ -16,8 +16,8 @@
  */
 package org.apache.arrow.vector.complex.impl;
 
-import org.apache.arrow.vector.UuidVector;
-import org.apache.arrow.vector.holder.UuidHolder;
+import org.apache.arrow.vector.extension.UuidHolder;
+import org.apache.arrow.vector.extension.UuidVector;
 import org.apache.arrow.vector.holders.ExtensionHolder;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.Field;

vector/src/test/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java → vector/src/main/java/org/apache/arrow/vector/complex/impl/UuidWriterFactory.java

@@ -17,12 +17,13 @@
 package org.apache.arrow.vector.complex.impl;
 
 import org.apache.arrow.vector.ExtensionTypeVector;
-import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.complex.writer.FieldWriter;
+import org.apache.arrow.vector.extension.UuidVector;
 
 public class UuidWriterFactory implements ExtensionTypeWriterFactory {
 
   @Override
-  public AbstractFieldWriter getWriterImpl(ExtensionTypeVector extensionTypeVector) {
+  public FieldWriter getWriterImpl(ExtensionTypeVector extensionTypeVector) {
     if (extensionTypeVector instanceof UuidVector) {
       return new UuidWriterImpl((UuidVector) extensionTypeVector);
     }

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

@@ -16,10 +16,11 @@
  */
 package org.apache.arrow.vector.complex.impl;
 
-import java.nio.ByteBuffer;
 import java.util.UUID;
-import org.apache.arrow.vector.UuidVector;
-import org.apache.arrow.vector.holder.UuidHolder;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.extension.NullableUuidHolder;
+import org.apache.arrow.vector.extension.UuidHolder;
+import org.apache.arrow.vector.extension.UuidVector;
 import org.apache.arrow.vector.holders.ExtensionHolder;
 
 public class UuidWriterImpl extends AbstractExtensionTypeWriter<UuidVector> {
@@ -30,18 +31,25 @@ public UuidWriterImpl(UuidVector vector) {
 
   @Override
   public void writeExtension(Object value) {
-    UUID uuid = (UUID) value;
-    ByteBuffer bb = ByteBuffer.allocate(16);
-    bb.putLong(uuid.getMostSignificantBits());
-    bb.putLong(uuid.getLeastSignificantBits());
-    vector.setSafe(getPosition(), bb.array());
+    if (value instanceof UUID) {
+      vector.setSafe(getPosition(), (UUID) value);
+    } else if (value instanceof byte[]) {
+      vector.setSafe(getPosition(), (byte[]) value);
+    } else if (value instanceof ArrowBuf) {
+      vector.setSafe(getPosition(), (ArrowBuf) value);
+    } else {
+      throw new IllegalArgumentException("Unsupported value type for UUID: " + value.getClass());
+    }
     vector.setValueCount(getPosition() + 1);
   }
 
   @Override
   public void write(ExtensionHolder holder) {
-    UuidHolder uuidHolder = (UuidHolder) holder;
-    vector.setSafe(getPosition(), uuidHolder.value);
+    if (holder instanceof UuidHolder) {
+      vector.setSafe(getPosition(), ((UuidHolder) holder).buffer);
+    } else if (holder instanceof NullableUuidHolder) {
+      vector.setSafe(getPosition(), ((NullableUuidHolder) holder).buffer);
+    }
     vector.setValueCount(getPosition() + 1);
   }
 }

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

@@ -0,0 +1,49 @@
+/*
+ * 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.extension;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+
+/** Nullable holder for UUID extension type values. */
+public class NullableUuidHolder extends ExtensionHolder {
+  /** Buffer containing the UUID bytes (16 bytes) when isSet = 1, undefined when isSet = 0. */
+  public ArrowBuf buffer;
+
+  /** Default constructor initializes the holder as null (isSet = 0). */
+  public NullableUuidHolder() {
+    this.isSet = 0;
+  }
+
+  /**
+   * Reason for not supporting the operation is that ValueHolders are potential scalar replacements
+   * and hence we don't want any methods to be invoked on them.
+   */
+  @Override
+  public int hashCode() {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Reason for not supporting the operation is that ValueHolders are potential scalar replacements
+   * and hence we don't want any methods to be invoked on them.
+   */
+  @Override
+  public String toString() {
+    throw new UnsupportedOperationException();
+  }
+}

vector/src/test/java/org/apache/arrow/vector/holder/UuidHolder.java → vector/src/main/java/org/apache/arrow/vector/extension/UuidHolder.java

@@ -14,10 +14,21 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.arrow.vector.holder;
+package org.apache.arrow.vector.extension;
 
+import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.vector.holders.ExtensionHolder;
 
+/** Holder for UUID extension type values. */
 public class UuidHolder extends ExtensionHolder {
-  public byte[] value;
+
+  /** Width of the UUID in bytes. */
+  public static int WIDTH = 16;
+
+  /** Buffer containing the UUID bytes (16 bytes). */
+  public ArrowBuf buffer;
+
+  public UuidHolder() {
+    this.isSet = 1;
+  }
 }

vector/src/test/java/org/apache/arrow/vector/types/pojo/UuidType.java → vector/src/main/java/org/apache/arrow/vector/extension/UuidType.java

@@ -14,15 +14,28 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.arrow.vector.types.pojo;
+package org.apache.arrow.vector.extension;
 
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.FixedSizeBinaryVector;
-import org.apache.arrow.vector.UuidVector;
+import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.ArrowType.ExtensionType;
+import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
+import org.apache.arrow.vector.types.pojo.FieldType;
 
 public class UuidType extends ExtensionType {
+  private static final AtomicBoolean registered = new AtomicBoolean(false);
+  public static final UuidType INSTANCE = new UuidType();
+
+  /** Register the extension type so it can be used globally. */
+  public static void ensureRegistered() {
+    if (!registered.getAndSet(true)) {
+      // The values don't matter, we just need an instance
+      ExtensionTypeRegistry.register(INSTANCE);
+    }
+  }
 
   @Override
   public ArrowType storageType() {
@@ -45,14 +58,19 @@ public ArrowType deserialize(ArrowType storageType, String serializedData) {
       throw new UnsupportedOperationException(
           "Cannot construct UuidType from underlying type " + storageType);
     }
-    return new UuidType();
+    return INSTANCE;
   }
 
   @Override
   public String serialize() {
     return "";
   }
 
+  @Override
+  public boolean isComplex() {
+    return false;
+  }
+
   @Override
   public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
     return new UuidVector(name, allocator, new FixedSizeBinaryVector(name, allocator, 16));

vector/src/test/java/org/apache/arrow/vector/UuidVector.java → vector/src/main/java/org/apache/arrow/vector/extension/UuidVector.java

@@ -14,18 +14,21 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.arrow.vector;
+package org.apache.arrow.vector.extension;
 
 import java.nio.ByteBuffer;
 import java.util.UUID;
+import org.apache.arrow.memory.ArrowBuf;
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.memory.util.hash.ArrowBufHasher;
+import org.apache.arrow.vector.ExtensionTypeVector;
+import org.apache.arrow.vector.FixedSizeBinaryVector;
+import org.apache.arrow.vector.ValueIterableVector;
+import org.apache.arrow.vector.ValueVector;
 import org.apache.arrow.vector.complex.impl.UuidReaderImpl;
 import org.apache.arrow.vector.complex.reader.FieldReader;
-import org.apache.arrow.vector.holder.UuidHolder;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.types.pojo.UuidType;
 import org.apache.arrow.vector.util.TransferPair;
 
 public class UuidVector extends ExtensionTypeVector<FixedSizeBinaryVector>
@@ -39,10 +42,19 @@ public UuidVector(
   }
 
   public UuidVector(String name, BufferAllocator allocator) {
-    super(name, allocator, new FixedSizeBinaryVector(name, allocator, 16));
+    super(name, allocator, new FixedSizeBinaryVector(name, allocator, UuidHolder.WIDTH));
     this.field = new Field(name, FieldType.nullable(new UuidType()), null);
   }
 
+  /** Constructor with field and allocator. */
+  public UuidVector(Field field, BufferAllocator allocator) {
+    super(
+        field.getName(),
+        allocator,
+        new FixedSizeBinaryVector(field.getName(), allocator, UuidHolder.WIDTH));
+    this.field = field;
+  }
+
   @Override
   public UUID getObject(int index) {
     final ByteBuffer bb = ByteBuffer.wrap(getUnderlyingVector().getObject(index));
@@ -59,6 +71,16 @@ public int hashCode(int index, ArrowBufHasher hasher) {
     return getUnderlyingVector().hashCode(index, hasher);
   }
 
+  public int isSet(int index) {
+    return getUnderlyingVector().isSet(index);
+  }
+
+  /**
+   * Set the value at the index of the vector to the given value.
+   *
+   * @param index position in the vector
+   * @param uuid given value
+   */
   public void set(int index, UUID uuid) {
     ByteBuffer bb = ByteBuffer.allocate(16);
     bb.putLong(uuid.getMostSignificantBits());
@@ -87,21 +109,78 @@ protected FieldReader getReaderImpl() {
     return new UuidReaderImpl(this);
   }
 
+  public void setSafe(int index, ArrowBuf value) {
+    getUnderlyingVector().setSafe(index, value);
+  }
+
   public void setSafe(int index, byte[] value) {
-    getUnderlyingVector().setIndexDefined(index);
     getUnderlyingVector().setSafe(index, value);
   }
 
+  /** Set value at index using UUID. */
+  public void setSafe(int index, UUID uuid) {
+    ByteBuffer bb = ByteBuffer.allocate(16);
+    bb.putLong(uuid.getMostSignificantBits());
+    bb.putLong(uuid.getLeastSignificantBits());
+    getUnderlyingVector().setSafe(index, bb.array());
+  }
+
+  /** Set value at index using NullableUuidHolder. */
+  public void setSafe(int index, NullableUuidHolder holder) {
+    if (holder != null) {
+      getUnderlyingVector().setSafe(index, holder.isSet, holder.buffer);
+    } else {
+      getUnderlyingVector().setNull(index);
+    }
+  }
+
+  /** Set value at index using UuidHolder. */
+  public void setSafe(int index, UuidHolder holder) {
+    if (holder != null) {
+      getUnderlyingVector().setSafe(index, holder.isSet, holder.buffer);
+    } else {
+      getUnderlyingVector().setNull(index);
+    }
+  }
+
+  /** Get value at index into UuidHolder. */
   public void get(int index, UuidHolder holder) {
-    holder.value = getUnderlyingVector().get(index);
+    holder.buffer =
+        getUnderlyingVector()
+            .getDataBuffer()
+            .slice((long) index * UuidHolder.WIDTH, UuidHolder.WIDTH);
     holder.isSet = 1;
   }
 
+  /**
+   * Get the slice at the given index, if present, and assign it to the buffer of the provided
+   * NullableUuidHolder.
+   *
+   * @param index position in the vector
+   * @param holder the holder to store the value in
+   */
+  public void get(int index, NullableUuidHolder holder) {
+    if (isNull(index)) {
+      holder.isSet = 0;
+    } else {
+      holder.buffer =
+          getUnderlyingVector()
+              .getDataBuffer()
+              .slice((long) index * UuidHolder.WIDTH, UuidHolder.WIDTH);
+      holder.isSet = 1;
+    }
+  }
+
   public class TransferImpl implements TransferPair {
     UuidVector to;
     ValueVector targetUnderlyingVector;
     TransferPair tp;
 
+    /**
+     * Constructor for TransferImpl.
+     *
+     * @param to the target vector
+     */
     public TransferImpl(UuidVector to) {
       this.to = to;
       targetUnderlyingVector = this.to.getUnderlyingVector();

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

@@ -39,7 +39,8 @@
 import org.apache.arrow.vector.complex.impl.UuidWriterFactory;
 import org.apache.arrow.vector.complex.reader.FieldReader;
 import org.apache.arrow.vector.complex.writer.BaseWriter.ExtensionWriter;
-import org.apache.arrow.vector.holder.UuidHolder;
+import org.apache.arrow.vector.extension.UuidHolder;
+import org.apache.arrow.vector.extension.UuidType;
 import org.apache.arrow.vector.holders.DurationHolder;
 import org.apache.arrow.vector.holders.FixedSizeBinaryHolder;
 import org.apache.arrow.vector.holders.TimeStampMilliTZHolder;
@@ -48,7 +49,6 @@
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.types.pojo.UuidType;
 import org.apache.arrow.vector.util.TransferPair;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -1259,13 +1259,13 @@ public void testListVectorReaderForExtensionType() throws Exception {
       FieldReader uuidReader = reader.reader();
       UuidHolder holder = new UuidHolder();
       uuidReader.read(holder);
-      ByteBuffer bb = ByteBuffer.wrap(holder.value);
+      ByteBuffer bb = holder.buffer.nioBuffer();
       UUID actualUuid = new UUID(bb.getLong(), bb.getLong());
       assertEquals(u1, actualUuid);
       reader.next();
       uuidReader = reader.reader();
       uuidReader.read(holder);
-      bb = ByteBuffer.wrap(holder.value);
+      bb = holder.buffer.nioBuffer();
       actualUuid = new UUID(bb.getLong(), bb.getLong());
       assertEquals(u2, actualUuid);
     }

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

@@ -35,14 +35,15 @@
 import org.apache.arrow.vector.complex.impl.NullableStructWriter;
 import org.apache.arrow.vector.complex.writer.Float8Writer;
 import org.apache.arrow.vector.complex.writer.IntWriter;
+import org.apache.arrow.vector.extension.UuidType;
+import org.apache.arrow.vector.extension.UuidVector;
 import org.apache.arrow.vector.holders.ComplexHolder;
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType;
 import org.apache.arrow.vector.types.pojo.ArrowType.Struct;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
-import org.apache.arrow.vector.types.pojo.UuidType;
 import org.apache.arrow.vector.util.TransferPair;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;