Skip to content

Commit edcc5ca

Browse files
nik9000dakrone
authored andcommitted
ESQL: Method to convert BooleanBlock to a "mask" (elastic#112253)
This adds a method, `BooleanBlock#toMask` to convert `BooleanBlock`s into a "mask" for use with `keepMask`.
1 parent 68c5b9d commit edcc5ca

File tree

15 files changed

+306
-4
lines changed

15 files changed

+306
-4
lines changed

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanArrayBlock.java

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBigArrayBlock.java

Lines changed: 21 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanBlock.java

Lines changed: 7 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/compute/src/main/generated-src/org/elasticsearch/compute/data/BooleanVectorBlock.java

Lines changed: 6 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/ConstantNullBlock.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,11 @@ public OrdinalBytesRefBlock asOrdinals() {
4848
return null;
4949
}
5050

51+
@Override
52+
public ToMask toMask() {
53+
return new ToMask(blockFactory.newConstantBooleanVector(false, positionCount), false);
54+
}
55+
5156
@Override
5257
public boolean isNull(int position) {
5358
return true;
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.compute.data;
9+
10+
import org.elasticsearch.core.Releasable;
11+
12+
/**
13+
* Result from calling {@link BooleanBlock#toMask}. {@link #close closing} this will
14+
* close the contained {@link #mask()}. If you want to keep a reference to it then you'll
15+
* have to {@link Block#incRef()} it.
16+
*/
17+
public record ToMask(BooleanVector mask, boolean hadMultivaluedFields) implements Releasable {
18+
@Override
19+
public void close() {
20+
mask.close();
21+
}
22+
}

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-ArrayBlock.java.st

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,28 @@ $if(BytesRef)$
101101
public OrdinalBytesRefBlock asOrdinals() {
102102
return null;
103103
}
104+
105+
$elseif(boolean)$
106+
@Override
107+
public ToMask toMask() {
108+
if (getPositionCount() == 0) {
109+
return new ToMask(blockFactory().newConstantBooleanVector(false, 0), false);
110+
}
111+
try (BooleanVector.FixedBuilder builder = blockFactory().newBooleanVectorFixedBuilder(getPositionCount())) {
112+
boolean hasMv = false;
113+
for (int p = 0; p < getPositionCount(); p++) {
114+
builder.appendBoolean(switch (getValueCount(p)) {
115+
case 0 -> false;
116+
case 1 -> getBoolean(getFirstValueIndex(p));
117+
default -> {
118+
hasMv = true;
119+
yield false;
120+
}
121+
});
122+
}
123+
return new ToMask(builder.build(), hasMv);
124+
}
125+
}
104126
$endif$
105127

106128
@Override

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-BigArrayBlock.java.st

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,29 @@ public final class $Type$BigArrayBlock extends AbstractArrayBlock implements $Ty
8686
return null;
8787
}
8888

89+
$if(boolean)$
90+
@Override
91+
public ToMask toMask() {
92+
if (getPositionCount() == 0) {
93+
return new ToMask(blockFactory().newConstantBooleanVector(false, 0), false);
94+
}
95+
try (BooleanVector.FixedBuilder builder = blockFactory().newBooleanVectorFixedBuilder(getPositionCount())) {
96+
boolean hasMv = false;
97+
for (int p = 0; p < getPositionCount(); p++) {
98+
builder.appendBoolean(switch (getValueCount(p)) {
99+
case 0 -> false;
100+
case 1 -> getBoolean(getFirstValueIndex(p));
101+
default -> {
102+
hasMv = true;
103+
yield false;
104+
}
105+
});
106+
}
107+
return new ToMask(builder.build(), hasMv);
108+
}
109+
}
110+
$endif$
111+
89112
@Override
90113
public $type$ get$Type$(int valueIndex) {
91114
return vector.get$Type$(valueIndex);

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Block.java.st

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,16 @@ $if(BytesRef)$
6363
* returns null. Callers must not release the returned block as no extra reference is retained by this method.
6464
*/
6565
OrdinalBytesRefBlock asOrdinals();
66-
$endif$
6766

67+
$elseif(boolean)$
68+
/**
69+
* Convert this to a {@link BooleanVector "mask"} that's appropriate for
70+
* passing to {@link #keepMask}. Null and multivalued positions will be
71+
* converted to {@code false}.
72+
*/
73+
ToMask toMask();
74+
75+
$endif$
6876
@Override
6977
$Type$Block filter(int... positions);
7078

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/data/X-Vector.java.st

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,8 +51,8 @@ $if(BytesRef)$
5151
* returns null. Callers must not release the returned vector as no extra reference is retained by this method.
5252
*/
5353
OrdinalBytesRefVector asOrdinals();
54-
$endif$
5554

55+
$endif$
5656
@Override
5757
$Type$Vector filter(int... positions);
5858

0 commit comments

Comments
 (0)