/*
 * 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.algorithm.sort;

import static org.junit.jupiter.api.Assertions.assertEquals;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.complex.StructVector;
import org.apache.arrow.vector.testing.ValueVectorDataPopulator;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.FieldType;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

/** Test cases for {@link GeneralOutOfPlaceVectorSorter}. */
public class TestGeneralOutOfPlaceVectorSorter {

  private BufferAllocator allocator;

  @BeforeEach
  public void prepare() {
    allocator = new RootAllocator(1024 * 1024);
  }

  @AfterEach
  public void shutdown() {
    allocator.close();
  }

  VectorValueComparator<StructVector> getComparator(StructVector structVector) {
    IntVector child0 = structVector.getChild("column0", IntVector.class);
    VectorValueComparator<IntVector> childComp0 =
        DefaultVectorComparators.createDefaultComparator(child0);
    childComp0.attachVector(child0);

    IntVector child1 = structVector.getChild("column1", IntVector.class);
    VectorValueComparator<IntVector> childComp1 =
        DefaultVectorComparators.createDefaultComparator(child1);
    childComp1.attachVector(child1);

    VectorValueComparator<StructVector> comp =
        new VectorValueComparator<StructVector>() {

          @Override
          public int compareNotNull(int index1, int index2) {
            // compare values by lexicographic order
            int result0 = childComp0.compare(index1, index2);
            if (result0 != 0) {
              return result0;
            }
            return childComp1.compare(index1, index2);
          }

          @Override
          public VectorValueComparator createNew() {
            return this;
          }
        };

    return comp;
  }

  @Test
  public void testSortStructVector() {
    final int vectorLength = 7;
    try (StructVector srcVector = StructVector.empty("src struct", allocator);
        StructVector dstVector = StructVector.empty("dst struct", allocator)) {

      IntVector srcChild0 =
          srcVector.addOrGet(
              "column0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class);
      IntVector srcChild1 =
          srcVector.addOrGet(
              "column1", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class);

      IntVector dstChild0 =
          dstVector.addOrGet(
              "column0", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class);
      IntVector dstChild1 =
          dstVector.addOrGet(
              "column1", FieldType.nullable(new ArrowType.Int(32, true)), IntVector.class);

      // src struct vector values:
      // [
      //   (2, 1)
      //   (3, 4)
      //   (5, 4)
      //   (null, 3)
      //   (7, null)
      //   (null, null)
      //   (6, 6)
      // ]

      ValueVectorDataPopulator.setVector(srcChild0, 2, 3, 5, null, 7, null, 6);
      ValueVectorDataPopulator.setVector(srcChild1, 1, 4, 4, 3, null, null, 6);
      srcVector.setIndexDefined(0);
      srcVector.setIndexDefined(1);
      srcVector.setIndexDefined(2);
      srcVector.setIndexDefined(3);
      srcVector.setIndexDefined(4);
      srcVector.setIndexDefined(6);
      srcVector.setValueCount(vectorLength);

      dstChild0.allocateNew(vectorLength);
      dstChild1.allocateNew(vectorLength);
      dstVector.setValueCount(vectorLength);

      // construct the comparator
      VectorValueComparator<StructVector> comp = getComparator(srcVector);

      // sort the vector
      GeneralOutOfPlaceVectorSorter<StructVector> sorter = new GeneralOutOfPlaceVectorSorter<>();
      sorter.sortOutOfPlace(srcVector, dstVector, comp);

      // validate results
      assertEquals(vectorLength, dstVector.getValueCount());
      assertEquals(
          "["
              + "null, "
              + "{\"column1\":3}, "
              + "{\"column0\":2,\"column1\":1}, "
              + "{\"column0\":3,\"column1\":4}, "
              + "{\"column0\":5,\"column1\":4}, "
              + "{\"column0\":6,\"column1\":6}, "
              + "{\"column0\":7}"
              + "]",
          dstVector.toString());
    }
  }
}
