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.search; 19 20 import static org.junit.Assert.assertEquals; 21 22 import java.util.Arrays; 23 import java.util.Collection; 24 25 import org.apache.arrow.algorithm.sort.DefaultVectorComparators; 26 import org.apache.arrow.algorithm.sort.VectorValueComparator; 27 import org.apache.arrow.memory.BufferAllocator; 28 import org.apache.arrow.memory.RootAllocator; 29 import org.apache.arrow.vector.IntVector; 30 import org.junit.After; 31 import org.junit.Before; 32 import org.junit.Test; 33 import org.junit.runner.RunWith; 34 import org.junit.runners.Parameterized; 35 36 /** 37 * Test cases for {@link VectorRangeSearcher}. 38 */ 39 @RunWith(Parameterized.class) 40 public class TestVectorRangeSearcher { 41 42 private BufferAllocator allocator; 43 44 private int repeat; 45 TestVectorRangeSearcher(int repeat)46 public TestVectorRangeSearcher(int repeat) { 47 this.repeat = repeat; 48 } 49 50 @Before prepare()51 public void prepare() { 52 allocator = new RootAllocator(1024 * 1024); 53 } 54 55 @After shutdown()56 public void shutdown() { 57 allocator.close(); 58 } 59 60 @Test testGetLowerBounds()61 public void testGetLowerBounds() { 62 final int maxValue = 100; 63 try (IntVector intVector = new IntVector("int vec", allocator)) { 64 // allocate vector 65 intVector.allocateNew(maxValue * repeat); 66 intVector.setValueCount(maxValue * repeat); 67 68 // prepare data in sorted order 69 // each value is repeated some times 70 for (int i = 0; i < maxValue; i++) { 71 for (int j = 0; j < repeat; j++) { 72 if (i == 0) { 73 intVector.setNull(i * repeat + j); 74 } else { 75 intVector.set(i * repeat + j, i); 76 } 77 } 78 } 79 80 // do search 81 VectorValueComparator<IntVector> comparator = DefaultVectorComparators.createDefaultComparator(intVector); 82 for (int i = 0; i < maxValue; i++) { 83 int result = VectorRangeSearcher.getFirstMatch(intVector, comparator, intVector, i * repeat); 84 assertEquals(i * repeat, result); 85 } 86 } 87 } 88 89 @Test testGetLowerBoundsNegative()90 public void testGetLowerBoundsNegative() { 91 final int maxValue = 100; 92 try (IntVector intVector = new IntVector("int vec", allocator); 93 IntVector negVector = new IntVector("neg vec", allocator)) { 94 // allocate vector 95 intVector.allocateNew(maxValue * repeat); 96 intVector.setValueCount(maxValue * repeat); 97 98 negVector.allocateNew(maxValue); 99 negVector.setValueCount(maxValue); 100 101 // prepare data in sorted order 102 // each value is repeated some times 103 for (int i = 0; i < maxValue; i++) { 104 for (int j = 0; j < repeat; j++) { 105 if (i == 0) { 106 intVector.setNull(i * repeat + j); 107 } else { 108 intVector.set(i * repeat + j, i); 109 } 110 } 111 negVector.set(i, maxValue + i); 112 } 113 114 // do search 115 VectorValueComparator<IntVector> comparator = DefaultVectorComparators.createDefaultComparator(intVector); 116 for (int i = 0; i < maxValue; i++) { 117 int result = VectorRangeSearcher.getFirstMatch(intVector, comparator, negVector, i); 118 assertEquals(-1, result); 119 } 120 } 121 } 122 123 @Test testGetUpperBounds()124 public void testGetUpperBounds() { 125 final int maxValue = 100; 126 try (IntVector intVector = new IntVector("int vec", allocator)) { 127 // allocate vector 128 intVector.allocateNew(maxValue * repeat); 129 intVector.setValueCount(maxValue * repeat); 130 131 // prepare data in sorted order 132 // each value is repeated some times 133 for (int i = 0; i < maxValue; i++) { 134 for (int j = 0; j < repeat; j++) { 135 if (i == 0) { 136 intVector.setNull(i * repeat + j); 137 } else { 138 intVector.set(i * repeat + j, i); 139 } 140 } 141 } 142 143 // do search 144 VectorValueComparator<IntVector> comparator = DefaultVectorComparators.createDefaultComparator(intVector); 145 for (int i = 0; i < maxValue; i++) { 146 int result = VectorRangeSearcher.getLastMatch(intVector, comparator, intVector, i * repeat); 147 assertEquals((i + 1) * repeat - 1, result); 148 } 149 } 150 } 151 152 @Test testGetUpperBoundsNegative()153 public void testGetUpperBoundsNegative() { 154 final int maxValue = 100; 155 try (IntVector intVector = new IntVector("int vec", allocator); 156 IntVector negVector = new IntVector("neg vec", allocator)) { 157 // allocate vector 158 intVector.allocateNew(maxValue * repeat); 159 intVector.setValueCount(maxValue * repeat); 160 161 negVector.allocateNew(maxValue); 162 negVector.setValueCount(maxValue); 163 164 // prepare data in sorted order 165 // each value is repeated some times 166 for (int i = 0; i < maxValue; i++) { 167 for (int j = 0; j < repeat; j++) { 168 if (i == 0) { 169 intVector.setNull(i * repeat + j); 170 } else { 171 intVector.set(i * repeat + j, i); 172 } 173 } 174 negVector.set(i, maxValue + i); 175 } 176 177 // do search 178 VectorValueComparator<IntVector> comparator = DefaultVectorComparators.createDefaultComparator(intVector); 179 for (int i = 0; i < maxValue; i++) { 180 int result = VectorRangeSearcher.getLastMatch(intVector, comparator, negVector, i); 181 assertEquals(-1, result); 182 } 183 } 184 } 185 186 @Parameterized.Parameters(name = "repeat = {0}") getRepeat()187 public static Collection<Object[]> getRepeat() { 188 return Arrays.asList( 189 new Object[]{1}, 190 new Object[]{2}, 191 new Object[]{5}, 192 new Object[]{10} 193 ); 194 } 195 } 196