content: /*
   * 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;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import org.apache.arrow.memory.ArrowBuf;
  import org.apache.arrow.memory.BufferAllocator;
  import org.apache.arrow.vector.complex.MapVector;
  import org.apache.arrow.vector.complex.UnionVector;
  import org.apache.arrow.vector.complex.VectorWithOrdinal;
  import org.apache.arrow.vector.complex.impl.UnionWriter;
  import org.apache.arrow.vector.holders.NullableBitHolder;
  import org.apache.arrow.vector.holders.NullableFloat4Holder;
  import org.apache.arrow.vector.holders.NullableIntHolder;
  import org.apache.arrow.vector.holders.NullableUInt4Holder;
  import org.apache.arrow.vector.types.Types.MinorType;
  import org.apache.arrow.vector.types.UnionMode;
  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.util.TransferPair;
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  
  public class TestUnionVector {
    private static final String EMPTY_SCHEMA_PATH = "";
  
    private BufferAllocator allocator;
  
    @BeforeEach
    public void init() {
      allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100);
    }
  
    @AfterEach
    public void terminate() throws Exception {
      allocator.close();
    }
  
    @Test
    public void testUnionVector() throws Exception {
  
      final NullableUInt4Holder uInt4Holder = new NullableUInt4Holder();
      uInt4Holder.value = 100;
      uInt4Holder.isSet = 1;
  
      try (UnionVector unionVector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
        unionVector.allocateNew();
  
        // write some data
        unionVector.setType(0, MinorType.UINT4);
        unionVector.setSafe(0, uInt4Holder);
        unionVector.setType(2, MinorType.UINT4);
        unionVector.setSafe(2, uInt4Holder);
        unionVector.setValueCount(4);
  
        // check that what we wrote is correct
        assertEquals(4, unionVector.getValueCount());
  
        assertEquals(false, unionVector.isNull(0));
        assertEquals(100, unionVector.getObject(0));
  
        assertNull(unionVector.getObject(1));
  
        assertEquals(false, unionVector.isNull(2));
        assertEquals(100, unionVector.getObject(2));
  
        assertNull(unionVector.getObject(3));
      }
    }
  
    @Test
    public void testUnionVectorMapValue() throws Exception {
      try (UnionVector unionVector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
        unionVector.allocateNew();
  
        UnionWriter writer = (UnionWriter) unionVector.getWriter();
  
        // populate map vector with the following two records
        // [
        //    null,
        //    [[1: 2], [3: 4], [5: null]]
        // ]
  
        writer.setPosition(0);
        writer.writeNull();
  
        writer.setPosition(1);
        writer.startMap();
  
        writer.startEntry();
        writer.key().integer().writeInt(1);
        writer.value().integer().writeInt(2);
        writer.endEntry();
  
        writer.startEntry();
        writer.key().integer().writeInt(3);
        writer.value().integer().writeInt(4);
        writer.endEntry();
  
        writer.startEntry();
        writer.key().integer().writeInt(5);
        writer.endEntry();
  
        writer.endMap();
  
        unionVector.setValueCount(2);
  
        // check that what we wrote is correct
        assertEquals(2, unionVector.getValueCount());
  
        // first entry
        assertNull(unionVector.getObject(0));
  
        // second entry
        List<Map<String, Integer>> resultList = (List<Map<String, Integer>>) unionVector.getObject(1);
        assertEquals(3, resultList.size());
  
        Map<String, Integer> resultMap = resultList.get(0);
        assertEquals(1, (int) resultMap.get(MapVector.KEY_NAME));
        assertEquals(2, (int) resultMap.get(MapVector.VALUE_NAME));
  
        resultMap = resultList.get(1);
        assertEquals(3, (int) resultMap.get(MapVector.KEY_NAME));
        assertEquals(4, (int) resultMap.get(MapVector.VALUE_NAME));
  
        resultMap = resultList.get(2);
        assertEquals(5, (int) resultMap.get(MapVector.KEY_NAME));
        assertNull(resultMap.get(MapVector.VALUE_NAME));
      }
    }
  
    @Test
    public void testTransfer() throws Exception {
      try (UnionVector srcVector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
        srcVector.allocateNew();
  
        // write some data
        srcVector.setType(0, MinorType.INT);
        srcVector.setSafe(0, newIntHolder(5));
        srcVector.setType(1, MinorType.BIT);
        srcVector.setSafe(1, newBitHolder(false));
        srcVector.setType(3, MinorType.INT);
        srcVector.setSafe(3, newIntHolder(10));
        srcVector.setType(5, MinorType.BIT);
        srcVector.setSafe(5, newBitHolder(false));
        srcVector.setValueCount(6);
  
        try (UnionVector destVector =
            new UnionVector(
                EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
          TransferPair pair = srcVector.makeTransferPair(destVector);
  
          // Creating the transfer should transfer the type of the field at least.
          assertEquals(srcVector.getField(), destVector.getField());
  
          // transfer
          pair.transfer();
  
          assertEquals(srcVector.getField(), destVector.getField());
  
          // now check the values are transferred
          assertEquals(6, destVector.getValueCount());
  
          assertFalse(destVector.isNull(0));
          assertEquals(5, destVector.getObject(0));
  
          assertFalse(destVector.isNull(1));
          assertEquals(false, destVector.getObject(1));
  
          assertNull(destVector.getObject(2));
  
          assertFalse(destVector.isNull(3));
          assertEquals(10, destVector.getObject(3));
  
          assertNull(destVector.getObject(4));
  
          assertFalse(destVector.isNull(5));
          assertEquals(false, destVector.getObject(5));
        }
      }
    }
  
    @Test
    public void testSplitAndTransfer() throws Exception {
      try (UnionVector sourceVector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
  
        sourceVector.allocateNew();
  
        /* populate the UnionVector */
        sourceVector.setType(0, MinorType.INT);
        sourceVector.setSafe(0, newIntHolder(5));
        sourceVector.setType(1, MinorType.INT);
        sourceVector.setSafe(1, newIntHolder(10));
        sourceVector.setType(2, MinorType.INT);
        sourceVector.setSafe(2, newIntHolder(15));
        sourceVector.setType(3, MinorType.INT);
        sourceVector.setSafe(3, newIntHolder(20));
        sourceVector.setType(4, MinorType.INT);
        sourceVector.setSafe(4, newIntHolder(25));
        sourceVector.setType(5, MinorType.INT);
        sourceVector.setSafe(5, newIntHolder(30));
        sourceVector.setType(6, MinorType.INT);
        sourceVector.setSafe(6, newIntHolder(35));
        sourceVector.setType(7, MinorType.INT);
        sourceVector.setSafe(7, newIntHolder(40));
        sourceVector.setType(8, MinorType.INT);
        sourceVector.setSafe(8, newIntHolder(45));
        sourceVector.setType(9, MinorType.INT);
        sourceVector.setSafe(9, newIntHolder(50));
        sourceVector.setValueCount(10);
  
        /* check the vector output */
        assertEquals(10, sourceVector.getValueCount());
        assertEquals(false, sourceVector.isNull(0));
        assertEquals(5, sourceVector.getObject(0));
        assertEquals(false, sourceVector.isNull(1));
        assertEquals(10, sourceVector.getObject(1));
        assertEquals(false, sourceVector.isNull(2));
        assertEquals(15, sourceVector.getObject(2));
        assertEquals(false, sourceVector.isNull(3));
        assertEquals(20, sourceVector.getObject(3));
        assertEquals(false, sourceVector.isNull(4));
        assertEquals(25, sourceVector.getObject(4));
        assertEquals(false, sourceVector.isNull(5));
        assertEquals(30, sourceVector.getObject(5));
        assertEquals(false, sourceVector.isNull(6));
        assertEquals(35, sourceVector.getObject(6));
        assertEquals(false, sourceVector.isNull(7));
        assertEquals(40, sourceVector.getObject(7));
        assertEquals(false, sourceVector.isNull(8));
        assertEquals(45, sourceVector.getObject(8));
        assertEquals(false, sourceVector.isNull(9));
        assertEquals(50, sourceVector.getObject(9));
  
        try (UnionVector toVector =
            new UnionVector(
                EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
  
          final TransferPair transferPair = sourceVector.makeTransferPair(toVector);
  
          final int[][] transferLengths = {{0, 3}, {3, 1}, {4, 2}, {6, 1}, {7, 1}, {8, 2}};
  
          for (final int[] transferLength : transferLengths) {
            final int start = transferLength[0];
            final int length = transferLength[1];
  
            transferPair.splitAndTransfer(start, length);
  
            /* check the toVector output after doing the splitAndTransfer */
            for (int i = 0; i < length; i++) {
              assertEquals(
                  sourceVector.getObject(start + i),
                  toVector.getObject(i),
                  "Different data at indexes: " + (start + i) + "and " + i);
            }
          }
        }
      }
    }
  
    @Test
    public void testSplitAndTransferWithMixedVectors() throws Exception {
      try (UnionVector sourceVector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
  
        sourceVector.allocateNew();
  
        /* populate the UnionVector */
        sourceVector.setType(0, MinorType.INT);
        sourceVector.setSafe(0, newIntHolder(5));
  
        sourceVector.setType(1, MinorType.FLOAT4);
        sourceVector.setSafe(1, newFloat4Holder(5.5f));
  
        sourceVector.setType(2, MinorType.INT);
        sourceVector.setSafe(2, newIntHolder(10));
  
        sourceVector.setType(3, MinorType.FLOAT4);
        sourceVector.setSafe(3, newFloat4Holder(10.5f));
  
        sourceVector.setType(4, MinorType.INT);
        sourceVector.setSafe(4, newIntHolder(15));
  
        sourceVector.setType(5, MinorType.FLOAT4);
        sourceVector.setSafe(5, newFloat4Holder(15.5f));
  
        sourceVector.setType(6, MinorType.INT);
        sourceVector.setSafe(6, newIntHolder(20));
  
        sourceVector.setType(7, MinorType.FLOAT4);
        sourceVector.setSafe(7, newFloat4Holder(20.5f));
  
        sourceVector.setType(8, MinorType.INT);
        sourceVector.setSafe(8, newIntHolder(30));
  
        sourceVector.setType(9, MinorType.FLOAT4);
        sourceVector.setSafe(9, newFloat4Holder(30.5f));
        sourceVector.setValueCount(10);
  
        /* check the vector output */
        assertEquals(10, sourceVector.getValueCount());
        assertEquals(false, sourceVector.isNull(0));
        assertEquals(5, sourceVector.getObject(0));
        assertEquals(false, sourceVector.isNull(1));
        assertEquals(5.5f, sourceVector.getObject(1));
        assertEquals(false, sourceVector.isNull(2));
        assertEquals(10, sourceVector.getObject(2));
        assertEquals(false, sourceVector.isNull(3));
        assertEquals(10.5f, sourceVector.getObject(3));
        assertEquals(false, sourceVector.isNull(4));
        assertEquals(15, sourceVector.getObject(4));
        assertEquals(false, sourceVector.isNull(5));
        assertEquals(15.5f, sourceVector.getObject(5));
        assertEquals(false, sourceVector.isNull(6));
        assertEquals(20, sourceVector.getObject(6));
        assertEquals(false, sourceVector.isNull(7));
        assertEquals(20.5f, sourceVector.getObject(7));
        assertEquals(false, sourceVector.isNull(8));
        assertEquals(30, sourceVector.getObject(8));
        assertEquals(false, sourceVector.isNull(9));
        assertEquals(30.5f, sourceVector.getObject(9));
  
        try (UnionVector toVector =
            new UnionVector(
                EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
  
          final TransferPair transferPair = sourceVector.makeTransferPair(toVector);
  
          final int[][] transferLengths = {{0, 2}, {2, 1}, {3, 2}, {5, 3}, {8, 2}};
  
          for (final int[] transferLength : transferLengths) {
            final int start = transferLength[0];
            final int length = transferLength[1];
  
            transferPair.splitAndTransfer(start, length);
  
            /* check the toVector output after doing the splitAndTransfer */
            for (int i = 0; i < length; i++) {
              assertEquals(
                  sourceVector.getObject(start + i),
                  toVector.getObject(i),
                  "Different values at index: " + i);
            }
          }
        }
      }
    }
  
    @Test
    public void testGetFieldTypeInfo() throws Exception {
      Map<String, String> metadata = new HashMap<>();
      metadata.put("key1", "value1");
  
      int[] typeIds = new int[2];
      typeIds[0] = MinorType.INT.ordinal();
      typeIds[1] = MinorType.VARCHAR.ordinal();
  
      List<Field> children = new ArrayList<>();
      children.add(new Field("int", FieldType.nullable(MinorType.INT.getType()), null));
      children.add(new Field("varchar", FieldType.nullable(MinorType.VARCHAR.getType()), null));
  
      final FieldType fieldType =
          new FieldType(
              false,
              new ArrowType.Union(UnionMode.Sparse, typeIds),
              /* dictionary= */ null,
              metadata);
      final Field field = new Field("union", fieldType, children);
  
      MinorType minorType = MinorType.UNION;
      UnionVector vector = (UnionVector) minorType.getNewVector(field, allocator, null);
      vector.initializeChildrenFromFields(children);
  
      assertTrue(vector.getField().equals(field));
  
      // Union has 2 child vectors
      assertEquals(2, vector.size());
  
      // Check child field 0
      VectorWithOrdinal intChild = vector.getChildVectorWithOrdinal("int");
      assertEquals(0, intChild.ordinal);
      assertEquals(intChild.vector.getField(), children.get(0));
  
      // Check child field 1
      VectorWithOrdinal varcharChild = vector.getChildVectorWithOrdinal("varchar");
      assertEquals(1, varcharChild.ordinal);
      assertEquals(varcharChild.vector.getField(), children.get(1));
    }
  
    @Test
    public void testGetBufferAddress() throws Exception {
      try (UnionVector vector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
        boolean error = false;
  
        vector.allocateNew();
  
        /* populate the UnionVector */
        vector.setType(0, MinorType.INT);
        vector.setSafe(0, newIntHolder(5));
  
        vector.setType(1, MinorType.FLOAT4);
        vector.setSafe(1, newFloat4Holder(5.5f));
  
        vector.setType(2, MinorType.INT);
        vector.setSafe(2, newIntHolder(10));
  
        vector.setType(3, MinorType.FLOAT4);
        vector.setSafe(3, newFloat4Holder(10.5f));
  
        vector.setValueCount(10);
  
        /* check the vector output */
        assertEquals(10, vector.getValueCount());
        assertEquals(false, vector.isNull(0));
        assertEquals(5, vector.getObject(0));
        assertEquals(false, vector.isNull(1));
        assertEquals(5.5f, vector.getObject(1));
        assertEquals(false, vector.isNull(2));
        assertEquals(10, vector.getObject(2));
        assertEquals(false, vector.isNull(3));
        assertEquals(10.5f, vector.getObject(3));
  
        List<ArrowBuf> buffers = vector.getFieldBuffers();
  
        try {
          vector.getOffsetBufferAddress();
        } catch (UnsupportedOperationException ue) {
          error = true;
        } finally {
          assertTrue(error);
          error = false;
        }
  
        try {
          vector.getDataBufferAddress();
        } catch (UnsupportedOperationException ue) {
          error = true;
        } finally {
          assertTrue(error);
        }
  
        assertEquals(1, buffers.size());
      }
    }
  
    @Test
    public void testSetGetNull() {
      try (UnionVector srcVector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
        srcVector.allocateNew();
  
        final NullableIntHolder holder = new NullableIntHolder();
        holder.isSet = 1;
        holder.value = 5;
  
        // write some data
        srcVector.setType(0, MinorType.INT);
        srcVector.setSafe(0, holder);
  
        assertFalse(srcVector.isNull(0));
  
        holder.isSet = 0;
        srcVector.setSafe(0, holder);
  
        assertNull(srcVector.getObject(0));
      }
    }
  
    @Test
    public void testCreateNewVectorWithoutTypeExceptionThrown() {
      try (UnionVector vector =
          new UnionVector(
              EMPTY_SCHEMA_PATH, allocator, /* field type */ null, /* call-back */ null)) {
        IllegalArgumentException e1 =
            assertThrows(IllegalArgumentException.class, () -> vector.getTimeStampMilliTZVector());
        assertEquals(
            "No TimeStampMilliTZ present. Provide ArrowType argument to create a new vector",
            e1.getMessage());
  
        IllegalArgumentException e2 =
            assertThrows(IllegalArgumentException.class, () -> vector.getDurationVector());
        assertEquals(
            "No Duration present. Provide ArrowType argument to create a new vector",
            e2.getMessage());
  
        IllegalArgumentException e3 =
            assertThrows(IllegalArgumentException.class, () -> vector.getFixedSizeBinaryVector());
        assertEquals(
            "No FixedSizeBinary present. Provide ArrowType argument to create a new vector",
            e3.getMessage());
  
        IllegalArgumentException e4 =
            assertThrows(IllegalArgumentException.class, () -> vector.getDecimalVector());
        assertEquals(
            "No Decimal present. Provide ArrowType argument to create a new vector", e4.getMessage());
      }
    }
  
    private static NullableIntHolder newIntHolder(int value) {
      final NullableIntHolder holder = new NullableIntHolder();
      holder.isSet = 1;
      holder.value = value;
      return holder;
    }
  
    private static NullableBitHolder newBitHolder(boolean value) {
      final NullableBitHolder holder = new NullableBitHolder();
      holder.isSet = 1;
      holder.value = value ? 1 : 0;
      return holder;
    }
  
    private static NullableFloat4Holder newFloat4Holder(float value) {
      final NullableFloat4Holder holder = new NullableFloat4Holder();
      holder.isSet = 1;
      holder.value = value;
      return holder;
    }
  }
  