remove ensureRegistered method from UuidType

Author: jhrotko

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

@@ -16,7 +16,6 @@
  */
 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;
@@ -64,21 +63,10 @@ public class UuidType extends ExtensionType {
   /** Storage type for UUID: FixedSizeBinary(16). */
   public static final ArrowType STORAGE_TYPE = new ArrowType.FixedSizeBinary(UUID_BYTE_WIDTH);
 
-  /** Registers the UuidType in the extension type registry. */
-  private static final AtomicBoolean registered = new AtomicBoolean(false);
-
   static {
     ExtensionTypeRegistry.register(INSTANCE);
   }
 
-  /** 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(UuidType.INSTANCE);
-    }
-  }
-
   @Override
   public ArrowType storageType() {
     return STORAGE_TYPE;

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

@@ -20,6 +20,7 @@
 import org.apache.arrow.memory.BufferAllocator;
 import org.apache.arrow.vector.types.Types.MinorType;
 import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.ExtensionTypeRegistry;
 import org.apache.arrow.vector.types.pojo.FieldType;
 
 public class TestUtils {
@@ -62,4 +63,14 @@ public static String generateRandomString(int length) {
     }
     return sb.toString();
   }
+
+  /*
+   * Ensure the extension type is registered, as there might other tests trying to unregister the
+   * type. ex.: TestExtensionType#readUnderlyingType
+   */
+  public static void ensureRegistered(ArrowType.ExtensionType type) {
+    if (ExtensionTypeRegistry.lookup(type.extensionName()) == null) {
+      ExtensionTypeRegistry.register(type);
+    }
+  }
 }

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

@@ -16,6 +16,7 @@
  */
 package org.apache.arrow.vector;
 
+import static org.apache.arrow.vector.TestUtils.ensureRegistered;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -164,9 +165,9 @@ void testVectorOperations() {
 
   @Test
   void testIpcRoundTrip() {
-    UuidType.ensureRegistered();
+    UuidType type = UuidType.INSTANCE;
+    ensureRegistered(type);
 
-    UuidType type = new UuidType();
     Schema schema = new Schema(Collections.singletonList(Field.nullable("uuid", type)));
     byte[] serialized = schema.serializeAsMessage();
     Schema deserialized = Schema.deserializeMessage(ByteBuffer.wrap(serialized));
@@ -175,9 +176,9 @@ void testIpcRoundTrip() {
 
   @Test
   void testVectorIpcRoundTrip() throws IOException {
-    UuidType.ensureRegistered();
+    UuidType type = UuidType.INSTANCE;
+    ensureRegistered(type);
 
-    UuidType type = new UuidType();
     UUID uuid1 = UUID.randomUUID();
     UUID uuid2 = UUID.randomUUID();
 

vector/src/test/java/org/apache/arrow/vector/types/pojo/TestExtensionType.java

@@ -16,6 +16,7 @@
  */
 package org.apache.arrow.vector.types.pojo;
 
+import static org.apache.arrow.vector.TestUtils.ensureRegistered;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -60,9 +61,9 @@ public class TestExtensionType {
   /** Test that a custom UUID type can be round-tripped through a temporary file. */
   @Test
   public void roundtripUuid() throws IOException {
-    ExtensionTypeRegistry.register(new UuidType());
+    ensureRegistered(UuidType.INSTANCE);
     final Schema schema =
-        new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
+        new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE)));
     try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
         final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
       UUID u1 = UUID.randomUUID();
@@ -90,7 +91,7 @@ public void roundtripUuid() throws IOException {
         assertEquals(root.getSchema(), readerRoot.getSchema());
 
         final Field field = readerRoot.getSchema().getFields().get(0);
-        final UuidType expectedType = new UuidType();
+        final UuidType expectedType = UuidType.INSTANCE;
         assertEquals(
             field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
             expectedType.extensionName());
@@ -114,9 +115,9 @@ public void roundtripUuid() throws IOException {
   /** Test that a custom UUID type can be read as its underlying type. */
   @Test
   public void readUnderlyingType() throws IOException {
-    ExtensionTypeRegistry.register(new UuidType());
+    ensureRegistered(UuidType.INSTANCE);
     final Schema schema =
-        new Schema(Collections.singletonList(Field.nullable("a", new UuidType())));
+        new Schema(Collections.singletonList(Field.nullable("a", UuidType.INSTANCE)));
     try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
         final VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator)) {
       UUID u1 = UUID.randomUUID();
@@ -136,7 +137,7 @@ public void readUnderlyingType() throws IOException {
         writer.end();
       }
 
-      ExtensionTypeRegistry.unregister(new UuidType());
+      ExtensionTypeRegistry.unregister(UuidType.INSTANCE);
 
       try (final SeekableByteChannel channel =
               Files.newByteChannel(Paths.get(file.getAbsolutePath()));
@@ -154,7 +155,7 @@ public void readUnderlyingType() throws IOException {
                 .getByteWidth());
 
         final Field field = readerRoot.getSchema().getFields().get(0);
-        final UuidType expectedType = new UuidType();
+        final UuidType expectedType = UuidType.INSTANCE;
         assertEquals(
             field.getMetadata().get(ExtensionType.EXTENSION_METADATA_KEY_NAME),
             expectedType.extensionName());
@@ -255,7 +256,7 @@ public void roundtripLocation() throws IOException {
 
   @Test
   public void testVectorCompare() {
-    UuidType uuidType = new UuidType();
+    UuidType uuidType = UuidType.INSTANCE;
     ExtensionTypeRegistry.register(uuidType);
     try (final BufferAllocator allocator = new RootAllocator(Integer.MAX_VALUE);
         UuidVector a1 =