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