Skip to content

Commit 94553d2

Browse files
nik9000nik9000
committed
Move
1 parent 4152f19 commit 94553d2

File tree

2 files changed

+71
-69
lines changed

2 files changed

+71
-69
lines changed

‎x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/BlockTestUtils.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,23 @@
1313
import org.elasticsearch.compute.data.BlockUtils;
1414
import org.elasticsearch.compute.data.BooleanBlock;
1515
import org.elasticsearch.compute.data.BytesRefBlock;
16+
import org.elasticsearch.compute.data.BytesRefVector;
1617
import org.elasticsearch.compute.data.DocBlock;
1718
import org.elasticsearch.compute.data.DoubleBlock;
1819
import org.elasticsearch.compute.data.ElementType;
1920
import org.elasticsearch.compute.data.FloatBlock;
2021
import org.elasticsearch.compute.data.IntBlock;
2122
import org.elasticsearch.compute.data.LongBlock;
23+
import org.elasticsearch.compute.data.OrdinalBytesRefBlock;
2224
import org.elasticsearch.compute.data.Page;
25+
import org.elasticsearch.core.Releasables;
2326
import org.hamcrest.Matcher;
2427

2528
import java.util.ArrayList;
29+
import java.util.Arrays;
30+
import java.util.HashMap;
2631
import java.util.List;
32+
import java.util.Map;
2733

2834
import static org.elasticsearch.compute.data.BlockUtils.toJavaObject;
2935
import static org.elasticsearch.test.ESTestCase.between;
@@ -267,4 +273,67 @@ public static List<List<Object>> valuesAtPositions(Block block, int from, int to
267273
}
268274
return result;
269275
}
276+
277+
/**
278+
* Convert all of the {@link Block}s in a page that contain {@link BytesRef}s into
279+
* {@link OrdinalBytesRefBlock}s.
280+
*/
281+
public static Page convertBytesRefsToOrdinals(Page page) {
282+
Block[] blocks = new Block[page.getBlockCount()];
283+
try {
284+
for (int b = 0; b < page.getBlockCount(); b++) {
285+
Block block = page.getBlock(b);
286+
if (block.elementType() != ElementType.BYTES_REF) {
287+
blocks[b] = block;
288+
continue;
289+
}
290+
Map<BytesRef, Integer> dedupe = new HashMap<>();
291+
BytesRefBlock bytesRefBlock = (BytesRefBlock) block;
292+
try (
293+
IntBlock.Builder ordinals = block.blockFactory().newIntBlockBuilder(block.getPositionCount());
294+
BytesRefVector.Builder bytes = block.blockFactory().newBytesRefVectorBuilder(block.getPositionCount())
295+
) {
296+
BytesRef scratch = new BytesRef();
297+
for (int p = 0; p < block.getPositionCount(); p++) {
298+
int first = block.getFirstValueIndex(p);
299+
int count = block.getValueCount(p);
300+
if (count == 0) {
301+
ordinals.appendNull();
302+
continue;
303+
}
304+
if (count == 1) {
305+
BytesRef v = bytesRefBlock.getBytesRef(first, scratch);
306+
ordinals.appendInt(dedupe(dedupe, bytes, v));
307+
continue;
308+
}
309+
int end = first + count;
310+
ordinals.beginPositionEntry();
311+
for (int i = first; i < end; i++) {
312+
BytesRef v = bytesRefBlock.getBytesRef(i, scratch);
313+
ordinals.appendInt(dedupe(dedupe, bytes, v));
314+
}
315+
ordinals.endPositionEntry();
316+
}
317+
blocks[b] = new OrdinalBytesRefBlock(ordinals.build(), bytes.build());
318+
bytesRefBlock.decRef();
319+
}
320+
}
321+
Page p = new Page(blocks);
322+
Arrays.fill(blocks, null);
323+
return p;
324+
} finally {
325+
Releasables.close(blocks);
326+
}
327+
}
328+
329+
private static int dedupe(Map<BytesRef, Integer> dedupe, BytesRefVector.Builder bytes, BytesRef v) {
330+
Integer current = dedupe.get(v);
331+
if (current != null) {
332+
return current;
333+
}
334+
bytes.appendBytesRef(v);
335+
int o = dedupe.size();
336+
dedupe.put(v, o);
337+
return o;
338+
}
270339
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java‎

Lines changed: 2 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,12 @@
1919
import org.elasticsearch.compute.data.Block;
2020
import org.elasticsearch.compute.data.BlockFactory;
2121
import org.elasticsearch.compute.data.BlockUtils;
22-
import org.elasticsearch.compute.data.BytesRefBlock;
23-
import org.elasticsearch.compute.data.BytesRefVector;
2422
import org.elasticsearch.compute.data.ElementType;
25-
import org.elasticsearch.compute.data.IntBlock;
26-
import org.elasticsearch.compute.data.OrdinalBytesRefBlock;
2723
import org.elasticsearch.compute.data.Page;
2824
import org.elasticsearch.compute.operator.DriverContext;
2925
import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
26+
import org.elasticsearch.compute.test.BlockTestUtils;
3027
import org.elasticsearch.compute.test.TestBlockFactory;
31-
import org.elasticsearch.core.Releasables;
3228
import org.elasticsearch.indices.CrankyCircuitBreakerService;
3329
import org.elasticsearch.logging.LogManager;
3430
import org.elasticsearch.logging.Logger;
@@ -570,70 +566,7 @@ private Page maybeConvertBytesRefsToOrdinals(Page page) {
570566
break;
571567
}
572568
}
573-
return anyBytesRef && randomBoolean() ? convertBytesRefsToOrdinals(page) : page;
574-
}
575-
576-
/**
577-
* Convert all of the {@link Block}s in a page that contain {@link BytesRef}s into
578-
* {@link OrdinalBytesRefBlock}s.
579-
*/
580-
public static Page convertBytesRefsToOrdinals(Page page) {
581-
Block[] blocks = new Block[page.getBlockCount()];
582-
try {
583-
for (int b = 0; b < page.getBlockCount(); b++) {
584-
Block block = page.getBlock(b);
585-
if (block.elementType() != ElementType.BYTES_REF) {
586-
blocks[b] = block;
587-
continue;
588-
}
589-
Map<BytesRef, Integer> dedupe = new HashMap<>();
590-
BytesRefBlock bytesRefBlock = (BytesRefBlock) block;
591-
try (
592-
IntBlock.Builder ordinals = block.blockFactory().newIntBlockBuilder(block.getPositionCount());
593-
BytesRefVector.Builder bytes = block.blockFactory().newBytesRefVectorBuilder(block.getPositionCount())
594-
) {
595-
BytesRef scratch = new BytesRef();
596-
for (int p = 0; p < block.getPositionCount(); p++) {
597-
int first = block.getFirstValueIndex(p);
598-
int count = block.getValueCount(p);
599-
if (count == 0) {
600-
ordinals.appendNull();
601-
continue;
602-
}
603-
if (count == 1) {
604-
BytesRef v = bytesRefBlock.getBytesRef(first, scratch);
605-
ordinals.appendInt(dedupe(dedupe, bytes, v));
606-
continue;
607-
}
608-
int end = first + count;
609-
ordinals.beginPositionEntry();
610-
for (int i = first; i < end; i++) {
611-
BytesRef v = bytesRefBlock.getBytesRef(i, scratch);
612-
ordinals.appendInt(dedupe(dedupe, bytes, v));
613-
}
614-
ordinals.endPositionEntry();
615-
}
616-
blocks[b] = new OrdinalBytesRefBlock(ordinals.build(), bytes.build());
617-
bytesRefBlock.decRef();
618-
}
619-
}
620-
Page p = new Page(blocks);
621-
Arrays.fill(blocks, null);
622-
return p;
623-
} finally {
624-
Releasables.close(blocks);
625-
}
626-
}
627-
628-
private static int dedupe(Map<BytesRef, Integer> dedupe, BytesRefVector.Builder bytes, BytesRef v) {
629-
Integer current = dedupe.get(v);
630-
if (current != null) {
631-
return current;
632-
}
633-
bytes.appendBytesRef(v);
634-
int o = dedupe.size();
635-
dedupe.put(v, o);
636-
return o;
569+
return anyBytesRef && randomBoolean() ? BlockTestUtils.convertBytesRefsToOrdinals(page) : page;
637570
}
638571

639572
/**

0 commit comments

Comments
 (0)