GH-946: Add Variant extension type support

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

@@ -0,0 +1,68 @@
+/*
+ * 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.NullableVariantHolder;
+import org.apache.arrow.vector.types.Types;
+
+public class NullableVariantHolderReaderImpl extends AbstractFieldReader {
+  private final NullableVariantHolder holder;
+
+  public NullableVariantHolderReaderImpl(NullableVariantHolder holder) {
+    this.holder = holder;
+  }
+
+  @Override
+  public int size() {
+    throw new UnsupportedOperationException("You can't call size on a Holder value reader.");
+  }
+
+  @Override
+  public boolean next() {
+    throw new UnsupportedOperationException("You can't call next on a single value reader.");
+  }
+
+  @Override
+  public void setPosition(int index) {
+    throw new UnsupportedOperationException("You can't call setPosition on a single value reader.");
+  }
+
+  @Override
+  public Types.MinorType getMinorType() {
+    return Types.MinorType.EXTENSIONTYPE;
+  }
+
+  @Override
+  public boolean isSet() {
+    return holder.isSet == 1;
+  }
+
+  /**
+   * Reads the variant holder data into the provided holder.
+   *
+   * @param h the holder to read into
+   */
+  public void read(NullableVariantHolder h) {
+    h.metadataStart = this.holder.metadataStart;
+    h.metadataEnd = this.holder.metadataEnd;
+    h.metadataBuffer = this.holder.metadataBuffer;
+    h.valueStart = this.holder.valueStart;
+    h.valueEnd = this.holder.valueEnd;
+    h.valueBuffer = this.holder.valueBuffer;
+    h.isSet = this.isSet() ? 1 : 0;
+  }
+}

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

@@ -0,0 +1,72 @@
+/*
+ * 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.extension.VariantVector;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.holders.NullableVariantHolder;
+import org.apache.arrow.vector.holders.VariantHolder;
+import org.apache.arrow.vector.types.Types;
+import org.apache.arrow.vector.types.pojo.Field;
+
+public class VariantReaderImpl extends AbstractFieldReader {
+  private final VariantVector vector;
+
+  public VariantReaderImpl(VariantVector vector) {
+    this.vector = vector;
+  }
+
+  @Override
+  public Types.MinorType getMinorType() {
+    return this.vector.getMinorType();
+  }
+
+  @Override
+  public Field getField() {
+    return this.vector.getField();
+  }
+
+  @Override
+  public boolean isSet() {
+    return !this.vector.isNull(this.idx());
+  }
+
+  @Override
+  public void read(ExtensionHolder holder) {
+    if (holder instanceof VariantHolder) {
+      vector.get(idx(), (VariantHolder) holder);
+    } else if (holder instanceof NullableVariantHolder) {
+      vector.get(idx(), (NullableVariantHolder) holder);
+    } else {
+      throw new IllegalArgumentException(
+          "Unsupported holder type for VariantReader: " + holder.getClass());
+    }
+  }
+
+  public void read(VariantHolder h) {
+    this.vector.get(this.idx(), h);
+  }
+
+  public void read(NullableVariantHolder h) {
+    this.vector.get(this.idx(), h);
+  }
+
+  @Override
+  public Object readObject() {
+    return this.vector.getObject(this.idx());
+  }
+}

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

@@ -0,0 +1,45 @@
+/*
+ * 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.ExtensionTypeVector;
+import org.apache.arrow.vector.extension.VariantVector;
+
+/**
+ * Factory for creating {@link VariantWriterImpl} instances.
+ *
+ * <p>This factory is used to create writers for Variant extension type vectors.
+ *
+ * @see VariantWriterImpl
+ * @see org.apache.arrow.vector.extension.VariantType
+ */
+public class VariantWriterFactory implements ExtensionTypeWriterFactory {
+
+  /**
+   * Creates a writer implementation for the given extension type vector.
+   *
+   * @param extensionTypeVector the vector to create a writer for
+   * @return a {@link VariantWriterImpl} if the vector is a {@link VariantVector}, null otherwise
+   */
+  @Override
+  public AbstractFieldWriter getWriterImpl(ExtensionTypeVector extensionTypeVector) {
+    if (extensionTypeVector instanceof VariantVector) {
+      return new VariantWriterImpl((VariantVector) extensionTypeVector);
+    }
+    return null;
+  }
+}

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

@@ -0,0 +1,87 @@
+/*
+ * 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.extension.VariantVector;
+import org.apache.arrow.vector.holders.ExtensionHolder;
+import org.apache.arrow.vector.holders.NullableVariantHolder;
+import org.apache.arrow.vector.holders.VariantHolder;
+
+/**
+ * Writer implementation for VARIANT extension type vectors.
+ *
+ * <p>This writer handles writing variant data to a {@link VariantVector}. It accepts both {@link
+ * VariantHolder} and {@link NullableVariantHolder} objects containing metadata and value buffers
+ * and writes them to the appropriate position in the vector.
+ */
+public class VariantWriterImpl extends AbstractExtensionTypeWriter<VariantVector> {
+
+  private static final String UNSUPPORTED_TYPE_TEMPLATE = "Unsupported type for Variant: %s";
+
+  /**
+   * Constructs a new VariantWriterImpl for the given vector.
+   *
+   * @param vector the variant vector to write to
+   */
+  public VariantWriterImpl(VariantVector vector) {
+    super(vector);
+  }
+
+  /**
+   * Writes an extension type value to the vector.
+   *
+   * <p>This method validates that the object is an {@link ExtensionHolder} and delegates to {@link
+   * #write(ExtensionHolder)}.
+   *
+   * @param object the object to write, must be an {@link ExtensionHolder}
+   * @throws IllegalArgumentException if the object is not an {@link ExtensionHolder}
+   */
+  @Override
+  public void writeExtension(Object object) {
+    if (object instanceof ExtensionHolder) {
+      write((ExtensionHolder) object);
+    } else {
+      throw new IllegalArgumentException(
+          String.format(UNSUPPORTED_TYPE_TEMPLATE, object.getClass().getName()));
+    }
+  }
+
+  /**
+   * Writes a variant holder to the vector at the current position.
+   *
+   * <p>The holder can be either a {@link VariantHolder} (non-nullable, always set) or a {@link
+   * NullableVariantHolder} (nullable, may be null). The data is written using {@link
+   * VariantVector#setSafe(int, NullableVariantHolder)} which handles buffer allocation and copying.
+   *
+   * @param extensionHolder the variant holder to write, must be a {@link VariantHolder} or {@link
+   *     NullableVariantHolder}
+   * @throws IllegalArgumentException if the holder is neither a {@link VariantHolder} nor a {@link
+   *     NullableVariantHolder}
+   */
+  @Override
+  public void write(ExtensionHolder extensionHolder) {
+    if (extensionHolder instanceof VariantHolder) {
+      vector.setSafe(getPosition(), (VariantHolder) extensionHolder);
+    } else if (extensionHolder instanceof NullableVariantHolder) {
+      vector.setSafe(getPosition(), (NullableVariantHolder) extensionHolder);
+    } else {
+      throw new IllegalArgumentException(
+          String.format(UNSUPPORTED_TYPE_TEMPLATE, extensionHolder.getClass().getName()));
+    }
+    vector.setValueCount(getPosition() + 1);
+  }
+}

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

@@ -0,0 +1,80 @@
+/*
+ * 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.BufferAllocator;
+import org.apache.arrow.vector.FieldVector;
+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 final class VariantType extends ExtensionType {
+
+  public static final VariantType INSTANCE = new VariantType();
+
+  public static final String EXTENSION_NAME = "parquet.variant";
+
+  static {
+    ExtensionTypeRegistry.register(INSTANCE);
+  }
+
+  private VariantType() {}
+
+  @Override
+  public ArrowType storageType() {
+    return ArrowType.Struct.INSTANCE;
+  }
+
+  @Override
+  public String extensionName() {
+    return EXTENSION_NAME;
+  }
+
+  @Override
+  public boolean extensionEquals(ExtensionType other) {
+    return other instanceof VariantType;
+  }
+
+  @Override
+  public String serialize() {
+    return "";
+  }
+
+  @Override
+  public ArrowType deserialize(ArrowType storageType, String serializedData) {
+    if (!storageType.equals(this.storageType())) {
+      throw new UnsupportedOperationException(
+          "Cannot construct VariantType from underlying type " + storageType);
+    }
+    return INSTANCE;
+  }
+
+  @Override
+  public FieldVector getNewVector(String name, FieldType fieldType, BufferAllocator allocator) {
+    return new VariantVector(name, allocator);
+  }
+
+  @Override
+  public boolean isComplex() {
+    // The type itself is not complex meaning we need separate functions to convert/extract
+    // different types.
+    // Meanwhile, the containing vector is complex in terms of containing multiple values (metadata
+    // and value)
+    return false;
+  }
+}