import quincy 17.2.0

[ceph.git] / ceph / src / arrow / java / vector / src / test / java / org / apache / arrow / vector / TestFixedSizeListVector.java

/*
 * 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.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;

import java.math.BigDecimal;
import java.util.Arrays;
import java.util.List;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.complex.FixedSizeListVector;
import org.apache.arrow.vector.complex.ListVector;
import org.apache.arrow.vector.complex.impl.UnionFixedSizeListReader;
import org.apache.arrow.vector.complex.impl.UnionFixedSizeListWriter;
import org.apache.arrow.vector.complex.impl.UnionListReader;
import org.apache.arrow.vector.complex.reader.FieldReader;
import org.apache.arrow.vector.types.Types.MinorType;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.apache.arrow.vector.util.TransferPair;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

public class TestFixedSizeListVector {

  private BufferAllocator allocator;

  @Before
  public void init() {
    allocator = new DirtyRootAllocator(Long.MAX_VALUE, (byte) 100);
  }

  @After
  public void terminate() throws Exception {
    allocator.close();
  }

  @Test
  public void testIntType() {
    try (FixedSizeListVector vector = FixedSizeListVector.empty("list", 2, allocator)) {
      IntVector nested = (IntVector) vector.addOrGetVector(FieldType.nullable(MinorType.INT.getType())).getVector();
      vector.allocateNew();

      for (int i = 0; i < 10; i++) {
        vector.setNotNull(i);
        nested.set(i * 2, i);
        nested.set(i * 2 + 1, i + 10);
      }
      vector.setValueCount(10);

      UnionFixedSizeListReader reader = vector.getReader();
      for (int i = 0; i < 10; i++) {
        reader.setPosition(i);
        Assert.assertTrue(reader.isSet());
        Assert.assertTrue(reader.next());
        assertEquals(i, reader.reader().readInteger().intValue());
        Assert.assertTrue(reader.next());
        assertEquals(i + 10, reader.reader().readInteger().intValue());
        Assert.assertFalse(reader.next());
        assertEquals(Arrays.asList(i, i + 10), reader.readObject());
      }
    }
  }

  @Test
  public void testFloatTypeNullable() {
    try (FixedSizeListVector vector = FixedSizeListVector.empty("list", 2, allocator)) {
      Float4Vector nested = (Float4Vector) vector.addOrGetVector(FieldType.nullable(MinorType.FLOAT4.getType()))
          .getVector();
      vector.allocateNew();

      for (int i = 0; i < 10; i++) {
        if (i % 2 == 0) {
          vector.setNotNull(i);
          nested.set(i * 2, i + 0.1f);
          nested.set(i * 2 + 1, i + 10.1f);
        }
      }
      vector.setValueCount(10);

      UnionFixedSizeListReader reader = vector.getReader();
      for (int i = 0; i < 10; i++) {
        reader.setPosition(i);
        if (i % 2 == 0) {
          Assert.assertTrue(reader.isSet());
          Assert.assertTrue(reader.next());
          assertEquals(i + 0.1f, reader.reader().readFloat(), 0.00001);
          Assert.assertTrue(reader.next());
          assertEquals(i + 10.1f, reader.reader().readFloat(), 0.00001);
          Assert.assertFalse(reader.next());
          assertEquals(Arrays.asList(i + 0.1f, i + 10.1f), reader.readObject());
        } else {
          Assert.assertFalse(reader.isSet());
          Assert.assertNull(reader.readObject());
        }
      }
    }
  }

  @Test
  public void testNestedInList() {
    try (ListVector vector = ListVector.empty("list", allocator)) {
      FixedSizeListVector tuples = (FixedSizeListVector) vector.addOrGetVector(
          FieldType.nullable(new ArrowType.FixedSizeList(2))).getVector();
      IntVector innerVector = (IntVector) tuples.addOrGetVector(FieldType.nullable(MinorType.INT.getType()))
          .getVector();
      vector.allocateNew();

      for (int i = 0; i < 10; i++) {
        if (i % 2 == 0) {
          int position = vector.startNewValue(i);
          for (int j = 0; j < i % 7; j++) {
            tuples.setNotNull(position + j);
            innerVector.set((position + j) * 2, j);
            innerVector.set((position + j) * 2 + 1, j + 1);
          }
          vector.endValue(i, i % 7);
        }
      }
      vector.setValueCount(10);

      UnionListReader reader = vector.getReader();
      for (int i = 0; i < 10; i++) {
        reader.setPosition(i);
        if (i % 2 == 0) {
          for (int j = 0; j < i % 7; j++) {
            Assert.assertTrue(reader.next());
            FieldReader innerListReader = reader.reader();
            for (int k = 0; k < 2; k++) {
              Assert.assertTrue(innerListReader.next());
              assertEquals(k + j, innerListReader.reader().readInteger().intValue());
            }
            Assert.assertFalse(innerListReader.next());
          }
          Assert.assertFalse(reader.next());
        } else {
          Assert.assertFalse(reader.isSet());
          Assert.assertNull(reader.readObject());
        }
      }
    }
  }

  @Test
  public void testTransferPair() {
    try (FixedSizeListVector from = new FixedSizeListVector(
        "from", allocator, new FieldType(true, new ArrowType.FixedSizeList(2), null), null);
         FixedSizeListVector to = new FixedSizeListVector(
             "to", allocator, new FieldType(true, new ArrowType.FixedSizeList(2), null), null)) {
      Float4Vector nested = (Float4Vector) from.addOrGetVector(FieldType.nullable(MinorType.FLOAT4.getType()))
          .getVector();
      from.allocateNew();

      for (int i = 0; i < 10; i++) {
        if (i % 2 == 0) {
          from.setNotNull(i);
          nested.set(i * 2, i + 0.1f);
          nested.set(i * 2 + 1, i + 10.1f);
        }
      }
      from.setValueCount(10);

      TransferPair pair = from.makeTransferPair(to);

      pair.copyValueSafe(0, 1);
      pair.copyValueSafe(2, 2);
      to.copyFromSafe(4, 3, from);

      to.setValueCount(10);

      UnionFixedSizeListReader reader = to.getReader();

      reader.setPosition(0);
      Assert.assertFalse(reader.isSet());
      Assert.assertNull(reader.readObject());

      reader.setPosition(1);
      Assert.assertTrue(reader.isSet());
      Assert.assertTrue(reader.next());
      assertEquals(0.1f, reader.reader().readFloat(), 0.00001);
      Assert.assertTrue(reader.next());
      assertEquals(10.1f, reader.reader().readFloat(), 0.00001);
      Assert.assertFalse(reader.next());
      assertEquals(Arrays.asList(0.1f, 10.1f), reader.readObject());

      reader.setPosition(2);
      Assert.assertTrue(reader.isSet());
      Assert.assertTrue(reader.next());
      assertEquals(2.1f, reader.reader().readFloat(), 0.00001);
      Assert.assertTrue(reader.next());
      assertEquals(12.1f, reader.reader().readFloat(), 0.00001);
      Assert.assertFalse(reader.next());
      assertEquals(Arrays.asList(2.1f, 12.1f), reader.readObject());

      reader.setPosition(3);
      Assert.assertTrue(reader.isSet());
      Assert.assertTrue(reader.next());
      assertEquals(4.1f, reader.reader().readFloat(), 0.00001);
      Assert.assertTrue(reader.next());
      assertEquals(14.1f, reader.reader().readFloat(), 0.00001);
      Assert.assertFalse(reader.next());
      assertEquals(Arrays.asList(4.1f, 14.1f), reader.readObject());

      for (int i = 4; i < 10; i++) {
        reader.setPosition(i);
        Assert.assertFalse(reader.isSet());
        Assert.assertNull(reader.readObject());
      }
    }
  }

  @Test
  public void testConsistentChildName() throws Exception {
    try (FixedSizeListVector listVector = FixedSizeListVector.empty("sourceVector", 2, allocator)) {
      String emptyListStr = listVector.getField().toString();
      Assert.assertTrue(emptyListStr.contains(ListVector.DATA_VECTOR_NAME));

      listVector.addOrGetVector(FieldType.nullable(MinorType.INT.getType()));
      String emptyVectorStr = listVector.getField().toString();
      Assert.assertTrue(emptyVectorStr.contains(ListVector.DATA_VECTOR_NAME));
    }
  }

  @Test
  public void testUnionFixedSizeListWriterWithNulls() throws Exception {
    /* Write to a decimal list vector
     * each list of size 3 and having its data values alternating between null and a non-null.
     * Read and verify
     */
    try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /*listSize=*/3, allocator)) {

      UnionFixedSizeListWriter writer = vector.getWriter();
      writer.allocate();

      final int valueCount = 100;

      for (int i = 0; i < valueCount; i++) {
        writer.startList();
        writer.decimal().writeDecimal(new BigDecimal(i));
        writer.writeNull();
        writer.decimal().writeDecimal(new BigDecimal(i * 3));
        writer.endList();
      }
      vector.setValueCount(valueCount);

      for (int i = 0; i < valueCount; i++) {
        List<BigDecimal> values = (List<BigDecimal>) vector.getObject(i);
        assertEquals(3, values.size());
        assertEquals(new BigDecimal(i), values.get(0));
        assertEquals(null, values.get(1));
        assertEquals(new BigDecimal(i * 3), values.get(2));
      }
    }
  }

  @Test
  public void testUnionFixedSizeListWriter() throws Exception {
    try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", 3, allocator)) {

      UnionFixedSizeListWriter writer1 = vector1.getWriter();
      writer1.allocate();

      int[] values1 = new int[] {1, 2, 3};
      int[] values2 = new int[] {4, 5, 6};
      int[] values3 = new int[] {7, 8, 9};

      //set some values
      writeListVector(vector1, writer1, values1);
      writeListVector(vector1, writer1, values2);
      writeListVector(vector1, writer1, values3);
      writer1.setValueCount(3);

      assertEquals(3, vector1.getValueCount());

      int[] realValue1 = convertListToIntArray(vector1.getObject(0));
      assertTrue(Arrays.equals(values1, realValue1));
      int[] realValue2 = convertListToIntArray(vector1.getObject(1));
      assertTrue(Arrays.equals(values2, realValue2));
      int[] realValue3 = convertListToIntArray(vector1.getObject(2));
      assertTrue(Arrays.equals(values3, realValue3));
    }
  }

  @Test
  public void testWriteDecimal() throws Exception {
    try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /*listSize=*/3, allocator)) {

      UnionFixedSizeListWriter writer = vector.getWriter();
      writer.allocate();

      final int valueCount = 100;

      for (int i = 0; i < valueCount; i++) {
        writer.startList();
        writer.decimal().writeDecimal(new BigDecimal(i));
        writer.decimal().writeDecimal(new BigDecimal(i * 2));
        writer.decimal().writeDecimal(new BigDecimal(i * 3));
        writer.endList();
      }
      vector.setValueCount(valueCount);

      for (int i = 0; i < valueCount; i++) {
        List<BigDecimal> values = (List<BigDecimal>) vector.getObject(i);
        assertEquals(3, values.size());
        assertEquals(new BigDecimal(i), values.get(0));
        assertEquals(new BigDecimal(i * 2), values.get(1));
        assertEquals(new BigDecimal(i * 3), values.get(2));
      }
    }
  }

  @Test
  public void testDecimalIndexCheck() throws Exception {
    try (final FixedSizeListVector vector = FixedSizeListVector.empty("vector", /*listSize=*/3, allocator)) {

      UnionFixedSizeListWriter writer = vector.getWriter();
      writer.allocate();

      IllegalStateException e = assertThrows(IllegalStateException.class, () -> {
        writer.startList();
        writer.decimal().writeDecimal(new BigDecimal(1));
        writer.decimal().writeDecimal(new BigDecimal(2));
        writer.decimal().writeDecimal(new BigDecimal(3));
        writer.decimal().writeDecimal(new BigDecimal(4));
        writer.endList();
      });
      assertEquals("values at index 0 is greater than listSize 3", e.getMessage());
    }
  }


  @Test(expected = IllegalStateException.class)
  public void testWriteIllegalData() throws Exception {
    try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", 3, allocator)) {

      UnionFixedSizeListWriter writer1 = vector1.getWriter();
      writer1.allocate();

      int[] values1 = new int[] {1, 2, 3};
      int[] values2 = new int[] {4, 5, 6, 7, 8};

      //set some values
      writeListVector(vector1, writer1, values1);
      writeListVector(vector1, writer1, values2);
      writer1.setValueCount(3);

      assertEquals(3, vector1.getValueCount());
      int[] realValue1 = convertListToIntArray(vector1.getObject(0));
      assertTrue(Arrays.equals(values1, realValue1));
      int[] realValue2 = convertListToIntArray(vector1.getObject(1));
      assertTrue(Arrays.equals(values2, realValue2));
    }
  }

  @Test
  public void testSplitAndTransfer() throws Exception {
    try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", 3, allocator)) {

      UnionFixedSizeListWriter writer1 = vector1.getWriter();
      writer1.allocate();

      int[] values1 = new int[] {1, 2, 3};
      int[] values2 = new int[] {4, 5, 6};
      int[] values3 = new int[] {7, 8, 9};

      //set some values
      writeListVector(vector1, writer1, values1);
      writeListVector(vector1, writer1, values2);
      writeListVector(vector1, writer1, values3);
      writer1.setValueCount(3);

      TransferPair transferPair = vector1.getTransferPair(allocator);
      transferPair.splitAndTransfer(0, 2);
      FixedSizeListVector targetVector = (FixedSizeListVector) transferPair.getTo();

      assertEquals(2, targetVector.getValueCount());
      int[] realValue1 = convertListToIntArray(targetVector.getObject(0));
      assertTrue(Arrays.equals(values1, realValue1));
      int[] realValue2 = convertListToIntArray(targetVector.getObject(1));
      assertTrue(Arrays.equals(values2, realValue2));

      targetVector.clear();
    }
  }

  @Test
  public void testZeroWidthVector() {
    try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", 0, allocator)) {

      UnionFixedSizeListWriter writer1 = vector1.getWriter();
      writer1.allocate();

      int[] values1 = new int[] {};
      int[] values2 = new int[] {};
      int[] values3 = null;
      int[] values4 = new int[] {};

      //set some values
      writeListVector(vector1, writer1, values1);
      writeListVector(vector1, writer1, values2);
      writeListVector(vector1, writer1, values3);
      writeListVector(vector1, writer1, values4);
      writer1.setValueCount(4);

      assertEquals(4, vector1.getValueCount());

      int[] realValue1 = convertListToIntArray(vector1.getObject(0));
      assertArrayEquals(values1, realValue1);
      int[] realValue2 = convertListToIntArray(vector1.getObject(1));
      assertArrayEquals(values2, realValue2);
      assertNull(vector1.getObject(2));
      int[] realValue4 = convertListToIntArray(vector1.getObject(3));
      assertArrayEquals(values4, realValue4);
    }
  }

  @Test
  public void testVectorWithNulls() {
    try (final FixedSizeListVector vector1 = FixedSizeListVector.empty("vector", 4, allocator)) {

      UnionFixedSizeListWriter writer1 = vector1.getWriter();
      writer1.allocate();

      List<Integer> values1 = Arrays.asList(null, 1, 2, 3);
      List<Integer> values2 = Arrays.asList(4, null, 5, 6);
      List<Integer> values3 = null;
      List<Integer> values4 = Arrays.asList(7, 8, null, 9);

      //set some values
      writeListVector(vector1, writer1, values1);
      writeListVector(vector1, writer1, values2);
      writeListVector(vector1, writer1, values3);
      writeListVector(vector1, writer1, values4);
      writer1.setValueCount(4);

      assertEquals(4, vector1.getValueCount());

      List realValue1 = vector1.getObject(0);
      assertEquals(values1, realValue1);
      List realValue2 = vector1.getObject(1);
      assertEquals(values2, realValue2);
      List realValue3 = vector1.getObject(2);
      assertEquals(values3, realValue3);
      List realValue4 = vector1.getObject(3);
      assertEquals(values4, realValue4);
    }
  }

  private int[] convertListToIntArray(List list) {
    int[] values = new int[list.size()];
    for (int i = 0; i < list.size(); i++) {
      values[i] = (int) list.get(i);
    }
    return values;
  }

  private void writeListVector(FixedSizeListVector vector, UnionFixedSizeListWriter writer, int[] values) {
    writer.startList();
    if (values != null) {
      for (int v : values) {
        writer.integer().writeInt(v);
      }
    } else {
      vector.setNull(writer.getPosition());
    }
    writer.endList();
  }

  private void writeListVector(FixedSizeListVector vector, UnionFixedSizeListWriter writer, List<Integer> values) {
    writer.startList();
    if (values != null) {
      for (Integer v : values) {
        if (v == null) {
          writer.writeNull();
        } else {
          writer.integer().writeInt(v);
        }
      }
    } else {
      vector.setNull(writer.getPosition());
    }
    writer.endList();
  }

}