Skip to content

Commit cba120d

Browse files
shmuelhanochshmuelhanoch
committed
Basic support for range types in ESQL
1 parent fae3656 commit cba120d

File tree

35 files changed

+779
-43
lines changed

35 files changed

+779
-43
lines changed

‎docs/reference/query-languages/esql/images/functions/to_dense_vector.svg‎

Lines changed: 1 addition & 1 deletion
Loading

‎docs/reference/query-languages/esql/kibana/definition/commands/inlinestats.json‎

Lines changed: 0 additions & 5 deletions
This file was deleted.

‎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: 95 additions & 2 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;
@@ -49,6 +51,7 @@
4951
import java.util.Objects;
5052
import java.util.Set;
5153

54+
import static org.elasticsearch.index.mapper.DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER;
5255
import static org.elasticsearch.index.query.RangeQueryBuilder.GTE_FIELD;
5356
import static org.elasticsearch.index.query.RangeQueryBuilder.GT_FIELD;
5457
import static org.elasticsearch.index.query.RangeQueryBuilder.LTE_FIELD;
@@ -62,7 +65,7 @@ public class RangeFieldMapper extends FieldMapper {
6265
public static final boolean DEFAULT_INCLUDE_LOWER = true;
6366

6467
public static class Defaults {
65-
public static final DateFormatter DATE_FORMATTER = DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER;
68+
public static final DateFormatter DATE_FORMATTER = DEFAULT_DATE_TIME_FORMATTER;
6669
public static final Locale LOCALE = DateFieldMapper.DEFAULT_LOCALE;
6770
}
6871

@@ -349,6 +352,97 @@ public Query rangeQuery(
349352
context
350353
);
351354
}
355+
356+
public static class DateRangeDocValuesLoader extends BlockDocValuesReader.DocValuesBlockLoader {
357+
private final String fieldName;
358+
359+
public DateRangeDocValuesLoader(String fieldName) {
360+
this.fieldName = fieldName;
361+
}
362+
363+
@Override
364+
public Builder builder(BlockFactory factory, int expectedCount) {
365+
return factory.longsFromDocValues(expectedCount);
366+
}
367+
368+
@Override
369+
public AllReader reader(LeafReaderContext context) throws IOException {
370+
var docValues = context.reader().getBinaryDocValues(fieldName);
371+
return new DateRangeDocValuesReader(docValues);
372+
/*
373+
if (docValues != null) {
374+
NumericDocValues singleton = DocValues.unwrapSingleton(docValues);
375+
if (singleton != null) {
376+
return new BlockDocValuesReader.SingletonLongs(singleton);
377+
}
378+
return new BlockDocValuesReader.Longs(docValues);
379+
}
380+
NumericDocValues singleton = context.reader().getNumericDocValues(fieldName);
381+
if (singleton != null) {
382+
return new BlockDocValuesReader.SingletonLongs(singleton);
383+
}
384+
return new ConstantNullsReader(); */
385+
}
386+
}
387+
388+
@Override
389+
public BlockLoader blockLoader(BlockLoaderContext blContext) {
390+
if (hasDocValues()) {
391+
return new DateRangeDocValuesLoader(name());
392+
}
393+
throw new IllegalStateException("Cannot load blocks without doc values");
394+
}
395+
}
396+
397+
public static class DateRangeDocValuesReader extends BlockDocValuesReader {
398+
private final BytesRef spare = new BytesRef();
399+
400+
private final BinaryDocValues numericDocValues;
401+
402+
public DateRangeDocValuesReader(BinaryDocValues numericDocValues) {
403+
this.numericDocValues = numericDocValues;
404+
}
405+
406+
@Override
407+
protected int docId() {
408+
return 0;
409+
}
410+
411+
@Override
412+
public String toString() {
413+
return "";
414+
}
415+
416+
@Override
417+
public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered)
418+
throws IOException {
419+
try (BlockLoader.DateRangeBuilder builder = factory.dateRangeBuilder(docs.count() - offset)) {
420+
for (int i = offset; i < docs.count(); i++) {
421+
int doc = docs.get(i);
422+
read(doc, builder);
423+
}
424+
return builder.build();
425+
}
426+
}
427+
428+
@Override
429+
public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException {
430+
read(docId, (BlockLoader.DateRangeBuilder) builder);
431+
}
432+
433+
private void read(int doc, BlockLoader.DateRangeBuilder builder) throws IOException {
434+
if (false == numericDocValues.advanceExact(doc)) {
435+
builder.appendNull();
436+
return;
437+
}
438+
439+
BytesRef ref = numericDocValues.binaryValue();
440+
var ranges = BinaryRangeUtil.decodeLongRanges(ref);
441+
for (var range : ranges) {
442+
builder.from().appendLong((long) range.from);
443+
builder.to().appendLong((long) range.to);
444+
}
445+
}
352446
}
353447

354448
private final RangeType type;
@@ -423,7 +517,6 @@ private Range parseRange(XContentParser parser) throws IOException {
423517
if (fieldType().rangeType == RangeType.IP && start == XContentParser.Token.VALUE_STRING) {
424518
return parseIpRangeFromCidr(parser);
425519
}
426-
427520
if (start != XContentParser.Token.START_OBJECT) {
428521
throw new DocumentParsingException(
429522
parser.getTokenLocation(),

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

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.apache.lucene.index.SortedDocValues;
1414
import org.apache.lucene.index.SortedSetDocValues;
1515
import org.apache.lucene.util.BytesRef;
16+
import org.elasticsearch.core.Releasables;
1617
import org.hamcrest.Matcher;
1718

1819
import java.io.IOException;
@@ -414,9 +415,15 @@ public SortedSetOrdinalBuilder appendOrd(int value) {
414415
return new SortedSetOrdinalBuilder();
415416
}
416417

418+
@Override
417419
public BlockLoader.AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int expectedSize) {
418420
return new AggregateMetricDoubleBlockBuilder(expectedSize);
419421
}
422+
423+
@Override
424+
public BlockLoader.DateRangeBuilder dateRangeBuilder(int expectedSize) {
425+
return new DateRangeBuilder(expectedSize);
426+
}
420427
};
421428
}
422429

@@ -624,4 +631,68 @@ public void close() {
624631

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

‎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: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,13 @@ public enum DataType {
311311
/**
312312
* Fields with this type are dense vectors, represented as an array of double values.
313313
*/
314-
DENSE_VECTOR(builder().esType("dense_vector").estimatedSize(4096));
314+
DENSE_VECTOR(builder().esType("dense_vector").estimatedSize(4096)),
315+
316+
/**
317+
* Fields with this type are used to represent a range of dates.
318+
* This is an under-construction type, and is not yet fully supported.
319+
*/
320+
DATE_RANGE(builder().esType("date_range").typeName("DATE_RANGE").estimatedSize(2 * Long.BYTES));
315321

316322
/**
317323
* Types that are actively being built. These types are
@@ -332,7 +338,8 @@ public enum DataType {
332338
*/
333339
public static final Map<DataType, FeatureFlag> UNDER_CONSTRUCTION = Map.ofEntries(
334340
Map.entry(AGGREGATE_METRIC_DOUBLE, EsqlCorePlugin.AGGREGATE_METRIC_DOUBLE_FEATURE_FLAG),
335-
Map.entry(DENSE_VECTOR, EsqlCorePlugin.DENSE_VECTOR_FEATURE_FLAG)
341+
Map.entry(DENSE_VECTOR, EsqlCorePlugin.DENSE_VECTOR_FEATURE_FLAG),
342+
Map.entry(DATE_RANGE, EsqlCorePlugin.DATE_RANGE_FEATURE_FLAG)
336343
);
337344

338345
private final String typeName;

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,32 @@ 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+
try (DateRangeBlockBuilder builder = newDateRangeBlockBuilder(positions)) {
490+
for (int i = 0; i < positions; i++) {
491+
if (value.from() != null) {
492+
builder.from().appendLong(value.from());
493+
} else {
494+
builder.appendNull();
495+
}
496+
if (value.to() != null) {
497+
builder.to().appendLong(value.to());
498+
}
499+
}
500+
return builder.build();
501+
}
502+
}
503+
504+
public DateRangeBlock newDateRangeBlock(long[] fromValues, long[] toValues, int positions) {
505+
var from = newLongArrayVector(fromValues, positions).asBlock();
506+
var to = newLongArrayVector(toValues, positions).asBlock();
507+
return new DateRangeBlock(from, to);
508+
}
509+
484510
/**
485511
* Returns the maximum number of bytes that a Block should be backed by a primitive array before switching to using BigArrays.
486512
*/

‎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
}

0 commit comments

Comments
 (0)