Skip to content

Commit 5ea25dc

Browse files
shmuelhanochshmuelhanoch
committed
Basic support for range types in ESQL
1 parent 77b4fd5 commit 5ea25dc

File tree

40 files changed

+893
-37
lines changed

40 files changed

+893
-37
lines changed

‎server/src/main/java/org/elasticsearch/index/mapper/BlockLoader.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -496,6 +496,8 @@ interface BlockFactory {
496496
SortedSetOrdinalsBuilder sortedSetOrdinalsBuilder(SortedSetDocValues ordinals, int count);
497497

498498
AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int count);
499+
500+
DateRangeBuilder dateRangeBuilder(int count);
499501
}
500502

501503
/**
@@ -623,4 +625,10 @@ interface AggregateMetricDoubleBuilder extends Builder {
623625

624626
IntBuilder count();
625627
}
628+
629+
interface DateRangeBuilder extends Builder {
630+
LongBuilder from();
631+
632+
LongBuilder to();
633+
}
626634
}

‎server/src/main/java/org/elasticsearch/index/mapper/RangeFieldMapper.java‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99

1010
package org.elasticsearch.index.mapper;
1111

12+
import org.apache.lucene.index.BinaryDocValues;
13+
import org.apache.lucene.index.LeafReaderContext;
1214
import org.apache.lucene.search.Query;
1315
import org.apache.lucene.search.SortField;
1416
import org.apache.lucene.util.BytesRef;
@@ -349,6 +351,96 @@ public Query rangeQuery(
349351
context
350352
);
351353
}
354+
355+
public static class DateRangeDocValuesLoader extends BlockDocValuesReader.DocValuesBlockLoader {
356+
private final String fieldName;
357+
358+
public DateRangeDocValuesLoader(String fieldName) {
359+
this.fieldName = fieldName;
360+
}
361+
362+
@Override
363+
public Builder builder(BlockFactory factory, int expectedCount) {
364+
return factory.longsFromDocValues(expectedCount);
365+
}
366+
367+
@Override
368+
public AllReader reader(LeafReaderContext context) throws IOException {
369+
var docValues = context.reader().getBinaryDocValues(fieldName);
370+
return new DateRangeDocValuesReader(docValues);
371+
}
372+
}
373+
374+
@Override
375+
public BlockLoader blockLoader(BlockLoaderContext blContext) {
376+
if (hasDocValues()) {
377+
return new DateRangeDocValuesLoader(name());
378+
}
379+
throw new IllegalStateException("Cannot load blocks without doc values");
380+
}
381+
}
382+
383+
public static class DateRangeDocValuesReader extends BlockDocValuesReader {
384+
private final BytesRef spare = new BytesRef();
385+
386+
private final BinaryDocValues numericDocValues;
387+
388+
public DateRangeDocValuesReader(BinaryDocValues numericDocValues) {
389+
this.numericDocValues = numericDocValues;
390+
}
391+
392+
private int docId = -1;
393+
394+
@Override
395+
protected int docId() {
396+
return docId;
397+
}
398+
399+
@Override
400+
public String toString() {
401+
return "BlockDocValuesReader.DateRangeDocValuesReader";
402+
}
403+
404+
@Override
405+
public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered)
406+
throws IOException {
407+
try (BlockLoader.DateRangeBuilder builder = factory.dateRangeBuilder(docs.count() - offset)) {
408+
int lastDoc = -1;
409+
for (int i = offset; i < docs.count(); i++) {
410+
int doc = docs.get(i);
411+
if (doc < lastDoc) {
412+
throw new IllegalStateException("docs within same block must be in order");
413+
}
414+
if (false == numericDocValues.advanceExact(doc)) {
415+
builder.appendNull();
416+
} else {
417+
BytesRef ref = numericDocValues.binaryValue();
418+
var ranges = BinaryRangeUtil.decodeLongRanges(ref);
419+
for (var range : ranges) {
420+
lastDoc = doc;
421+
this.docId = doc;
422+
builder.from().appendLong((long) range.from);
423+
builder.to().appendLong((long) range.to);
424+
}
425+
}
426+
}
427+
return builder.build();
428+
}
429+
}
430+
431+
@Override
432+
public void read(int doc, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException {
433+
var blockBuilder = (BlockLoader.DateRangeBuilder) builder;
434+
this.docId = doc;
435+
if (false == numericDocValues.advanceExact(doc)) {
436+
blockBuilder.appendNull();
437+
} else {
438+
var range = BinaryRangeUtil.decodeLongRanges(numericDocValues.binaryValue());
439+
assert range.size() == 1 : "stored fields should only have a single range";
440+
blockBuilder.from().appendLong((long) range.getFirst().from);
441+
blockBuilder.to().appendLong((long) range.getFirst().to);
442+
}
443+
}
352444
}
353445

354446
private final RangeType type;

‎server/src/main/resources/transport/definitions/referable/esql_date_range_type.csv‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
9177000

‎server/src/main/resources/transport/upper_bounds/9.2.csv‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
roles_security_stats,9176000
1+
esql_date_range_type,9177000

‎test/framework/src/main/java/org/elasticsearch/index/mapper/TestBlock.java‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -414,9 +414,15 @@ public SortedSetOrdinalBuilder appendOrd(int value) {
414414
return new SortedSetOrdinalBuilder();
415415
}
416416

417+
@Override
417418
public BlockLoader.AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int expectedSize) {
418419
return new AggregateMetricDoubleBlockBuilder(expectedSize);
419420
}
421+
422+
@Override
423+
public BlockLoader.DateRangeBuilder dateRangeBuilder(int expectedSize) {
424+
return new DateRangeBuilder(expectedSize);
425+
}
420426
};
421427
}
422428

@@ -624,4 +630,68 @@ public void close() {
624630

625631
}
626632
}
633+
634+
public static class DateRangeBuilder implements BlockLoader.DateRangeBuilder {
635+
private final LongBuilder from;
636+
private final LongBuilder to;
637+
638+
DateRangeBuilder(int expectedSize) {
639+
from = new LongBuilder(expectedSize);
640+
to = new LongBuilder(expectedSize);
641+
}
642+
643+
@Override
644+
public BlockLoader.LongBuilder from() {
645+
return from;
646+
}
647+
648+
@Override
649+
public BlockLoader.LongBuilder to() {
650+
return to;
651+
}
652+
653+
@Override
654+
public BlockLoader.Block build() {
655+
var fromBlock = from.build();
656+
var toBlock = to.build();
657+
assert fromBlock.size() == toBlock.size();
658+
var values = new ArrayList<>(fromBlock.size());
659+
for (int i = 0; i < fromBlock.size(); i++) {
660+
values.add(List.of(fromBlock.values.get(i), toBlock.values.get(i)));
661+
}
662+
return new TestBlock(values);
663+
}
664+
665+
@Override
666+
public BlockLoader.Builder appendNull() {
667+
throw new UnsupportedOperationException();
668+
}
669+
670+
@Override
671+
public BlockLoader.Builder beginPositionEntry() {
672+
throw new UnsupportedOperationException();
673+
}
674+
675+
@Override
676+
public BlockLoader.Builder endPositionEntry() {
677+
throw new UnsupportedOperationException();
678+
}
679+
680+
@Override
681+
public void close() {
682+
683+
}
684+
685+
private static class LongBuilder extends TestBlock.Builder implements BlockLoader.LongBuilder {
686+
private LongBuilder(int expectedSize) {
687+
super(expectedSize);
688+
}
689+
690+
@Override
691+
public BlockLoader.LongBuilder appendLong(long value) {
692+
add(value);
693+
return this;
694+
}
695+
}
696+
};
627697
}

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/plugin/EsqlCorePlugin.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,5 @@ public class EsqlCorePlugin extends Plugin implements ExtensiblePlugin {
1515

1616
public static final FeatureFlag AGGREGATE_METRIC_DOUBLE_FEATURE_FLAG = new FeatureFlag("esql_aggregate_metric_double");
1717
public static final FeatureFlag DENSE_VECTOR_FEATURE_FLAG = new FeatureFlag("esql_dense_vector");
18+
public static final FeatureFlag DATE_RANGE_FEATURE_FLAG = new FeatureFlag("esql_date_range");
1819
}

‎x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/type/DataType.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,13 @@ public enum DataType implements Writeable {
329329
// Version created just *after* we committed support for dense_vector
330330
ML_INFERENCE_SAGEMAKER_CHAT_COMPLETION
331331
)
332-
);
332+
),
333+
334+
/**
335+
* Fields with this type are used to represent a range of dates.
336+
* This is an under-construction type, and is not yet fully supported.
337+
*/
338+
DATE_RANGE(builder().esType("date_range").typeName("DATE_RANGE").estimatedSize(2 * Long.BYTES));
333339

334340
/**
335341
* Types that are actively being built. These types are

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,22 @@ public final AggregateMetricDoubleBlock newAggregateMetricDoubleBlock(
481481
return new AggregateMetricDoubleArrayBlock(min, max, sum, count);
482482
}
483483

484+
public DateRangeBlockBuilder newDateRangeBlockBuilder(int estimatedSize) {
485+
return new DateRangeBlockBuilder(estimatedSize, this);
486+
}
487+
488+
public DateRangeBlock newConstantDateRangeBlock(DateRangeBlockBuilder.DateRangeLiteral value, int positions) {
489+
var fromBuilder = newConstantLongBlockWith(value.from(), positions);
490+
var toBuilder = newConstantLongBlockWith(value.to(), positions);
491+
return new DateRangeArrayBlock(fromBuilder, toBuilder);
492+
}
493+
494+
public DateRangeBlock newDateRangeBlock(long[] fromValues, long[] toValues, int positions) {
495+
var from = newLongArrayVector(fromValues, positions).asBlock();
496+
var to = newLongArrayVector(toValues, positions).asBlock();
497+
return new DateRangeArrayBlock(from, to);
498+
}
499+
484500
/**
485501
* Returns the maximum number of bytes that a Block should be backed by a primitive array before switching to using BigArrays.
486502
*/

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ public static void appendValue(Block.Builder builder, Object val, ElementType ty
217217
case FLOAT -> ((FloatBlock.Builder) builder).appendFloat((Float) val);
218218
case DOUBLE -> ((DoubleBlock.Builder) builder).appendDouble((Double) val);
219219
case BOOLEAN -> ((BooleanBlock.Builder) builder).appendBoolean((Boolean) val);
220+
case DATE_RANGE -> ((DateRangeBlockBuilder) builder).appendDateRange((DateRangeBlockBuilder.DateRangeLiteral) val);
220221
default -> throw new UnsupportedOperationException("unsupported element type [" + type + "]");
221222
}
222223
}
@@ -246,6 +247,7 @@ private static Block constantBlock(BlockFactory blockFactory, ElementType type,
246247
case BOOLEAN -> blockFactory.newConstantBooleanBlockWith((boolean) val, size);
247248
case AGGREGATE_METRIC_DOUBLE -> blockFactory.newConstantAggregateMetricDoubleBlock((AggregateMetricDoubleLiteral) val, size);
248249
case FLOAT -> blockFactory.newConstantFloatBlockWith((float) val, size);
250+
case DATE_RANGE -> blockFactory.newConstantDateRangeBlock((DateRangeBlockBuilder.DateRangeLiteral) val, size);
249251
default -> throw new UnsupportedOperationException("unsupported element type [" + type + "]");
250252
};
251253
}
@@ -299,6 +301,12 @@ yield new AggregateMetricDoubleLiteral(
299301
aggBlock.countBlock().getInt(offset)
300302
);
301303
}
304+
case DATE_RANGE -> {
305+
DateRangeBlock b = (DateRangeBlock) block;
306+
LongBlock fromBlock = b.getFromBlock();
307+
LongBlock toBlock = b.getToBlock();
308+
yield new DateRangeBlockBuilder.DateRangeLiteral(fromBlock.getLong(offset), toBlock.getLong(offset));
309+
}
302310
case UNKNOWN -> throw new IllegalArgumentException("can't read values from [" + block + "]");
303311
};
304312
}

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

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ public final class ConstantNullBlock extends AbstractNonThreadSafeRefCounted
2626
FloatBlock,
2727
DoubleBlock,
2828
BytesRefBlock,
29-
AggregateMetricDoubleBlock {
29+
AggregateMetricDoubleBlock,
30+
DateRangeBlock {
3031

3132
private static final long BASE_RAM_BYTES_USED = RamUsageEstimator.shallowSizeOfInstance(ConstantNullBlock.class);
3233
private final int positionCount;
@@ -118,6 +119,16 @@ public ConstantNullBlock expand() {
118119
return this;
119120
}
120121

122+
@Override
123+
public LongBlock getFromBlock() {
124+
return this;
125+
}
126+
127+
@Override
128+
public LongBlock getToBlock() {
129+
return this;
130+
}
131+
121132
@Override
122133
public long ramBytesUsed() {
123134
return BASE_RAM_BYTES_USED;

0 commit comments

Comments
 (0)