Skip to content

Commit d95a81b

Browse files
ivanceaivancea
committed
Fixed tests
1 parent c266504 commit d95a81b

File tree

4 files changed

+152
-32
lines changed

4 files changed

+152
-32
lines changed

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/GroupingAggregatorFunctionTestCase.java‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import org.elasticsearch.compute.data.DoubleBlock;
2020
import org.elasticsearch.compute.data.ElementType;
2121
import org.elasticsearch.compute.data.FloatBlock;
22+
import org.elasticsearch.compute.data.IntArrayBlock;
23+
import org.elasticsearch.compute.data.IntBigArrayBlock;
2224
import org.elasticsearch.compute.data.IntBlock;
2325
import org.elasticsearch.compute.data.IntVector;
2426
import org.elasticsearch.compute.data.LongBlock;
@@ -648,8 +650,7 @@ public AddInput prepareProcessPage(SeenGroupIds ignoredSeenGroupIds, Page page)
648650
return seen;
649651
}, page);
650652

651-
@Override
652-
public void add(int positionOffset, IntBlock groupIds) {
653+
private void addBlock(int positionOffset, IntBlock groupIds) {
653654
for (int offset = 0; offset < groupIds.getPositionCount(); offset += emitChunkSize) {
654655
try (IntBlock.Builder builder = blockFactory().newIntBlockBuilder(emitChunkSize)) {
655656
int endP = Math.min(groupIds.getPositionCount(), offset + emitChunkSize);
@@ -682,6 +683,16 @@ public void add(int positionOffset, IntBlock groupIds) {
682683
}
683684
}
684685

686+
@Override
687+
public void add(int positionOffset, IntArrayBlock groupIds) {
688+
addBlock(positionOffset, groupIds);
689+
}
690+
691+
@Override
692+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
693+
addBlock(positionOffset, groupIds);
694+
}
695+
685696
@Override
686697
public void add(int positionOffset, IntVector groupIds) {
687698
int[] chunk = new int[emitChunkSize];

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/AddPageTests.java‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
import org.elasticsearch.compute.aggregation.GroupingAggregatorFunction;
1212
import org.elasticsearch.compute.data.BlockFactory;
1313
import org.elasticsearch.compute.data.BlockFactoryTests;
14+
import org.elasticsearch.compute.data.IntArrayBlock;
15+
import org.elasticsearch.compute.data.IntBigArrayBlock;
1416
import org.elasticsearch.compute.data.IntBlock;
1517
import org.elasticsearch.compute.data.IntVector;
1618
import org.elasticsearch.test.ESTestCase;
@@ -156,8 +158,7 @@ Added added(int positionOffset, int... ords) {
156158
private class TestAddInput implements GroupingAggregatorFunction.AddInput {
157159
private final List<Added> added = new ArrayList<>();
158160

159-
@Override
160-
public void add(int positionOffset, IntBlock groupIds) {
161+
private void addBlock(int positionOffset, IntBlock groupIds) {
161162
List<List<Integer>> result = new ArrayList<>(groupIds.getPositionCount());
162163
for (int p = 0; p < groupIds.getPositionCount(); p++) {
163164
int valueCount = groupIds.getValueCount(p);
@@ -172,9 +173,19 @@ public void add(int positionOffset, IntBlock groupIds) {
172173
added.add(new Added(positionOffset, result));
173174
}
174175

176+
@Override
177+
public void add(int positionOffset, IntArrayBlock groupIds) {
178+
addBlock(positionOffset, groupIds);
179+
}
180+
181+
@Override
182+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
183+
addBlock(positionOffset, groupIds);
184+
}
185+
175186
@Override
176187
public void add(int positionOffset, IntVector groupIds) {
177-
add(positionOffset, groupIds.asBlock());
188+
addBlock(positionOffset, groupIds.asBlock());
178189
}
179190

180191
@Override
@@ -187,7 +198,12 @@ private class CountingAddInput implements GroupingAggregatorFunction.AddInput {
187198
private int count;
188199

189200
@Override
190-
public void add(int positionOffset, IntBlock groupIds) {
201+
public void add(int positionOffset, IntArrayBlock groupIds) {
202+
count++;
203+
}
204+
205+
@Override
206+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
191207
count++;
192208
}
193209

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/BlockHashTests.java‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
import org.elasticsearch.compute.data.BytesRefVector;
2020
import org.elasticsearch.compute.data.DoubleBlock;
2121
import org.elasticsearch.compute.data.ElementType;
22+
import org.elasticsearch.compute.data.IntArrayBlock;
23+
import org.elasticsearch.compute.data.IntBigArrayBlock;
2224
import org.elasticsearch.compute.data.IntBlock;
2325
import org.elasticsearch.compute.data.IntVector;
2426
import org.elasticsearch.compute.data.LongBlock;
@@ -1276,7 +1278,13 @@ public void close() {
12761278
) {
12771279
hash1.add(page, new GroupingAggregatorFunction.AddInput() {
12781280
@Override
1279-
public void add(int positionOffset, IntBlock groupIds) {
1281+
public void add(int positionOffset, IntArrayBlock groupIds) {
1282+
groupIds.incRef();
1283+
output1.add(new Output(positionOffset, groupIds, null));
1284+
}
1285+
1286+
@Override
1287+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
12801288
groupIds.incRef();
12811289
output1.add(new Output(positionOffset, groupIds, null));
12821290
}
@@ -1294,7 +1302,13 @@ public void close() {
12941302
});
12951303
hash2.add(page, new GroupingAggregatorFunction.AddInput() {
12961304
@Override
1297-
public void add(int positionOffset, IntBlock groupIds) {
1305+
public void add(int positionOffset, IntArrayBlock groupIds) {
1306+
groupIds.incRef();
1307+
output2.add(new Output(positionOffset, groupIds, null));
1308+
}
1309+
1310+
@Override
1311+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
12981312
groupIds.incRef();
12991313
output2.add(new Output(positionOffset, groupIds, null));
13001314
}
@@ -1380,7 +1394,12 @@ public void testTimeSeriesBlockHash() throws Exception {
13801394
Holder<IntVector> ords1 = new Holder<>();
13811395
hash1.add(page, new GroupingAggregatorFunction.AddInput() {
13821396
@Override
1383-
public void add(int positionOffset, IntBlock groupIds) {
1397+
public void add(int positionOffset, IntArrayBlock groupIds) {
1398+
throw new AssertionError("time-series block hash should emit a vector");
1399+
}
1400+
1401+
@Override
1402+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
13841403
throw new AssertionError("time-series block hash should emit a vector");
13851404
}
13861405

@@ -1397,15 +1416,24 @@ public void close() {
13971416
});
13981417
Holder<IntVector> ords2 = new Holder<>();
13991418
hash2.add(page, new GroupingAggregatorFunction.AddInput() {
1400-
@Override
1401-
public void add(int positionOffset, IntBlock groupIds) {
1419+
private void addBlock(int positionOffset, IntBlock groupIds) {
14021420
// TODO: check why PackedValuesBlockHash doesn't emit a vector?
14031421
IntVector vector = groupIds.asVector();
14041422
assertNotNull("should emit a vector", vector);
14051423
vector.incRef();
14061424
ords2.set(vector);
14071425
}
14081426

1427+
@Override
1428+
public void add(int positionOffset, IntArrayBlock groupIds) {
1429+
addBlock(positionOffset, groupIds);
1430+
}
1431+
1432+
@Override
1433+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
1434+
addBlock(positionOffset, groupIds);
1435+
}
1436+
14091437
@Override
14101438
public void add(int positionOffset, IntVector groupIds) {
14111439
groupIds.incRef();
@@ -1498,8 +1526,7 @@ private BlockHash buildBlockHash(int emitBatchSize, Block... values) {
14981526

14991527
static void hash(boolean collectKeys, BlockHash blockHash, Consumer<OrdsAndKeys> callback, Block... values) {
15001528
blockHash.add(new Page(values), new GroupingAggregatorFunction.AddInput() {
1501-
@Override
1502-
public void add(int positionOffset, IntBlock groupIds) {
1529+
private void addBlock(int positionOffset, IntBlock groupIds) {
15031530
OrdsAndKeys result = new OrdsAndKeys(
15041531
blockHash.toString(),
15051532
positionOffset,
@@ -1532,9 +1559,19 @@ public void add(int positionOffset, IntBlock groupIds) {
15321559
}
15331560
}
15341561

1562+
@Override
1563+
public void add(int positionOffset, IntArrayBlock groupIds) {
1564+
addBlock(positionOffset, groupIds);
1565+
}
1566+
1567+
@Override
1568+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
1569+
addBlock(positionOffset, groupIds);
1570+
}
1571+
15351572
@Override
15361573
public void add(int positionOffset, IntVector groupIds) {
1537-
add(positionOffset, groupIds.asBlock());
1574+
addBlock(positionOffset, groupIds.asBlock());
15381575
}
15391576

15401577
@Override

‎x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/aggregation/blockhash/CategorizeBlockHashTests.java‎

Lines changed: 74 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import org.elasticsearch.compute.data.BytesRefBlock;
2626
import org.elasticsearch.compute.data.BytesRefVector;
2727
import org.elasticsearch.compute.data.ElementType;
28+
import org.elasticsearch.compute.data.IntArrayBlock;
29+
import org.elasticsearch.compute.data.IntBigArrayBlock;
2830
import org.elasticsearch.compute.data.IntBlock;
2931
import org.elasticsearch.compute.data.IntVector;
3032
import org.elasticsearch.compute.data.LongBlock;
@@ -99,8 +101,7 @@ public void testCategorizeRaw() {
99101
try (var hash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.SINGLE, analysisRegistry)) {
100102
for (int i = randomInt(2); i < 3; i++) {
101103
hash.add(page, new GroupingAggregatorFunction.AddInput() {
102-
@Override
103-
public void add(int positionOffset, IntBlock groupIds) {
104+
private void addBlock(int positionOffset, IntBlock groupIds) {
104105
assertEquals(groupIds.getPositionCount(), positions);
105106

106107
assertEquals(1, groupIds.getInt(0));
@@ -115,9 +116,19 @@ public void add(int positionOffset, IntBlock groupIds) {
115116
}
116117
}
117118

119+
@Override
120+
public void add(int positionOffset, IntArrayBlock groupIds) {
121+
addBlock(positionOffset, groupIds);
122+
}
123+
124+
@Override
125+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
126+
addBlock(positionOffset, groupIds);
127+
}
128+
118129
@Override
119130
public void add(int positionOffset, IntVector groupIds) {
120-
add(positionOffset, groupIds.asBlock());
131+
addBlock(positionOffset, groupIds.asBlock());
121132
}
122133

123134
@Override
@@ -162,8 +173,7 @@ public void testCategorizeRawMultivalue() {
162173
try (var hash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.SINGLE, analysisRegistry)) {
163174
for (int i = randomInt(2); i < 3; i++) {
164175
hash.add(page, new GroupingAggregatorFunction.AddInput() {
165-
@Override
166-
public void add(int positionOffset, IntBlock groupIds) {
176+
private void addBlock(int positionOffset, IntBlock groupIds) {
167177
assertEquals(groupIds.getPositionCount(), positions);
168178

169179
assertThat(groupIds.getFirstValueIndex(0), equalTo(0));
@@ -185,9 +195,19 @@ public void add(int positionOffset, IntBlock groupIds) {
185195
}
186196
}
187197

198+
@Override
199+
public void add(int positionOffset, IntArrayBlock groupIds) {
200+
addBlock(positionOffset, groupIds);
201+
}
202+
203+
@Override
204+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
205+
addBlock(positionOffset, groupIds);
206+
}
207+
188208
@Override
189209
public void add(int positionOffset, IntVector groupIds) {
190-
add(positionOffset, groupIds.asBlock());
210+
addBlock(positionOffset, groupIds.asBlock());
191211
}
192212

193213
@Override
@@ -243,8 +263,7 @@ public void testCategorizeIntermediate() {
243263
BlockHash rawHash2 = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.INITIAL, analysisRegistry);
244264
) {
245265
rawHash1.add(page1, new GroupingAggregatorFunction.AddInput() {
246-
@Override
247-
public void add(int positionOffset, IntBlock groupIds) {
266+
private void addBlock(int positionOffset, IntBlock groupIds) {
248267
assertEquals(groupIds.getPositionCount(), positions1);
249268
assertEquals(1, groupIds.getInt(0));
250269
assertEquals(2, groupIds.getInt(1));
@@ -258,9 +277,19 @@ public void add(int positionOffset, IntBlock groupIds) {
258277
}
259278
}
260279

280+
@Override
281+
public void add(int positionOffset, IntArrayBlock groupIds) {
282+
addBlock(positionOffset, groupIds);
283+
}
284+
285+
@Override
286+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
287+
addBlock(positionOffset, groupIds);
288+
}
289+
261290
@Override
262291
public void add(int positionOffset, IntVector groupIds) {
263-
add(positionOffset, groupIds.asBlock());
292+
addBlock(positionOffset, groupIds.asBlock());
264293
}
265294

266295
@Override
@@ -271,8 +300,7 @@ public void close() {
271300
intermediatePage1 = new Page(rawHash1.getKeys()[0]);
272301

273302
rawHash2.add(page2, new GroupingAggregatorFunction.AddInput() {
274-
@Override
275-
public void add(int positionOffset, IntBlock groupIds) {
303+
private void addBlock(int positionOffset, IntBlock groupIds) {
276304
assertEquals(groupIds.getPositionCount(), positions2);
277305
assertEquals(1, groupIds.getInt(0));
278306
assertEquals(2, groupIds.getInt(1));
@@ -281,9 +309,19 @@ public void add(int positionOffset, IntBlock groupIds) {
281309
assertEquals(3, groupIds.getInt(4));
282310
}
283311

312+
@Override
313+
public void add(int positionOffset, IntArrayBlock groupIds) {
314+
addBlock(positionOffset, groupIds);
315+
}
316+
317+
@Override
318+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
319+
addBlock(positionOffset, groupIds);
320+
}
321+
284322
@Override
285323
public void add(int positionOffset, IntVector groupIds) {
286-
add(positionOffset, groupIds.asBlock());
324+
addBlock(positionOffset, groupIds.asBlock());
287325
}
288326

289327
@Override
@@ -299,8 +337,7 @@ public void close() {
299337

300338
try (var intermediateHash = new CategorizeBlockHash(blockFactory, 0, AggregatorMode.FINAL, null)) {
301339
intermediateHash.add(intermediatePage1, new GroupingAggregatorFunction.AddInput() {
302-
@Override
303-
public void add(int positionOffset, IntBlock groupIds) {
340+
private void addBlock(int positionOffset, IntBlock groupIds) {
304341
List<Integer> values = IntStream.range(0, groupIds.getPositionCount())
305342
.map(groupIds::getInt)
306343
.boxed()
@@ -312,9 +349,19 @@ public void add(int positionOffset, IntBlock groupIds) {
312349
}
313350
}
314351

352+
@Override
353+
public void add(int positionOffset, IntArrayBlock groupIds) {
354+
addBlock(positionOffset, groupIds);
355+
}
356+
357+
@Override
358+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
359+
addBlock(positionOffset, groupIds);
360+
}
361+
315362
@Override
316363
public void add(int positionOffset, IntVector groupIds) {
317-
add(positionOffset, groupIds.asBlock());
364+
addBlock(positionOffset, groupIds.asBlock());
318365
}
319366

320367
@Override
@@ -325,8 +372,7 @@ public void close() {
325372

326373
for (int i = randomInt(2); i < 3; i++) {
327374
intermediateHash.add(intermediatePage2, new GroupingAggregatorFunction.AddInput() {
328-
@Override
329-
public void add(int positionOffset, IntBlock groupIds) {
375+
private void addBlock(int positionOffset, IntBlock groupIds) {
330376
List<Integer> values = IntStream.range(0, groupIds.getPositionCount())
331377
.map(groupIds::getInt)
332378
.boxed()
@@ -336,9 +382,19 @@ public void add(int positionOffset, IntBlock groupIds) {
336382
assertEquals(List.of(3, 1, 4), values);
337383
}
338384

385+
@Override
386+
public void add(int positionOffset, IntArrayBlock groupIds) {
387+
addBlock(positionOffset, groupIds);
388+
}
389+
390+
@Override
391+
public void add(int positionOffset, IntBigArrayBlock groupIds) {
392+
addBlock(positionOffset, groupIds);
393+
}
394+
339395
@Override
340396
public void add(int positionOffset, IntVector groupIds) {
341-
add(positionOffset, groupIds.asBlock());
397+
addBlock(positionOffset, groupIds.asBlock());
342398
}
343399

344400
@Override

0 commit comments

Comments
 (0)