ceph/src/arrow/java/algorithm/src/test/java/org/apache/arrow/algorithm/sort/TestCompositeVectorComparator.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.algorithm.sort;
  19
  20 import static org.junit.jupiter.api.Assertions.assertTrue;
  21
  22 import java.util.Arrays;
  23
  24 import org.apache.arrow.memory.BufferAllocator;
  25 import org.apache.arrow.memory.RootAllocator;
  26 import org.apache.arrow.vector.IntVector;
  27 import org.apache.arrow.vector.ValueVector;
  28 import org.apache.arrow.vector.VarCharVector;
  29 import org.apache.arrow.vector.VectorSchemaRoot;
  30 import org.junit.After;
  31 import org.junit.Before;
  32 import org.junit.Test;
  33
  34 /**
  35  * Test cases for {@link CompositeVectorComparator}.
  36  */
  37 public class TestCompositeVectorComparator {
  38
  39   private BufferAllocator allocator;
  40
  41   @Before
  42   public void prepare() {
  43     allocator = new RootAllocator(1024 * 1024);
  44   }
  45
  46   @After
  47   public void shutdown() {
  48     allocator.close();
  49   }
  50
  51   @Test
  52   public void testCompareVectorSchemaRoot() {
  53     final int vectorLength = 10;
  54     IntVector intVec1 = new IntVector("int1", allocator);
  55     VarCharVector strVec1 = new VarCharVector("str1", allocator);
  56
  57     IntVector intVec2 = new IntVector("int2", allocator);
  58     VarCharVector strVec2 = new VarCharVector("str2", allocator);
  59
  60     try (VectorSchemaRoot batch1 = new VectorSchemaRoot(Arrays.asList(intVec1, strVec1));
  61          VectorSchemaRoot batch2 = new VectorSchemaRoot(Arrays.asList(intVec2, strVec2))) {
  62
  63       intVec1.allocateNew(vectorLength);
  64       strVec1.allocateNew(vectorLength * 10, vectorLength);
  65       intVec2.allocateNew(vectorLength);
  66       strVec2.allocateNew(vectorLength * 10, vectorLength);
  67
  68       for (int i = 0; i < vectorLength; i++) {
  69         intVec1.set(i, i);
  70         strVec1.set(i, new String("a" + i).getBytes());
  71         intVec2.set(i, i);
  72         strVec2.set(i, new String("a5").getBytes());
  73       }
  74
  75       VectorValueComparator<IntVector> innerComparator1 =
  76               DefaultVectorComparators.createDefaultComparator(intVec1);
  77       innerComparator1.attachVectors(intVec1, intVec2);
  78       VectorValueComparator<VarCharVector> innerComparator2 =
  79               DefaultVectorComparators.createDefaultComparator(strVec1);
  80       innerComparator2.attachVectors(strVec1, strVec2);
  81
  82       VectorValueComparator<ValueVector> comparator = new CompositeVectorComparator(
  83           new VectorValueComparator[]{innerComparator1, innerComparator2}
  84       );
  85
  86       // verify results
  87
  88       // both elements are equal, the result is equal
  89       assertTrue(comparator.compare(5, 5) == 0);
  90
  91       // the first element being equal, the second is smaller, and the result is smaller
  92       assertTrue(comparator.compare(1, 1) < 0);
  93       assertTrue(comparator.compare(2, 2) < 0);
  94       assertTrue(comparator.compare(3, 3) < 0);
  95
  96       // the first element being equal, the second is larger, and the result is larger
  97       assertTrue(comparator.compare(7, 7) > 0);
  98       assertTrue(comparator.compare(8, 8) > 0);
  99       assertTrue(comparator.compare(9, 9) > 0);
 100
 101       // the first element is smaller, the result is always smaller
 102       assertTrue(comparator.compare(1, 2) < 0);
 103       assertTrue(comparator.compare(3, 7) < 0);
 104       assertTrue(comparator.compare(4, 9) < 0);
 105
 106       // the first element is larger, the result is always larger
 107       assertTrue(comparator.compare(2, 0) > 0);
 108       assertTrue(comparator.compare(8, 7) > 0);
 109       assertTrue(comparator.compare(4, 1) > 0);
 110     }
 111   }
 112 }