Skip to content

Commit 6666c35

Browse files
shmuelhanochshmuelhanoch
committed
Basic support for DATE_RANGE field type in ESQL
1 parent 8c4eee6 commit 6666c35

File tree

57 files changed

+1175
-59
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

57 files changed

+1175
-59
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: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,13 @@
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;
1517
import org.elasticsearch.ElasticsearchException;
18+
import org.elasticsearch.TransportVersion;
1619
import org.elasticsearch.common.Explicit;
1720
import org.elasticsearch.common.geo.ShapeRelation;
1821
import org.elasticsearch.common.network.InetAddresses;
@@ -61,6 +64,8 @@ public class RangeFieldMapper extends FieldMapper {
6164
public static final boolean DEFAULT_INCLUDE_UPPER = true;
6265
public static final boolean DEFAULT_INCLUDE_LOWER = true;
6366

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

347442
private final RangeType type;

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

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

‎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/type/DataType.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import java.util.stream.Collectors;
3434

3535
import static java.util.stream.Collectors.toMap;
36+
import static org.elasticsearch.index.mapper.RangeFieldMapper.ESQL_DATE_RANGE_CREATED_VERSION;
3637

3738
/**
3839
* This enum represents data types the ES|QL query processing layer is able to
@@ -270,6 +271,7 @@ public enum DataType implements Writeable {
270271
* Nanosecond precision date, stored as a 64-bit signed number.
271272
*/
272273
DATE_NANOS(builder().esType("date_nanos").estimatedSize(Long.BYTES).docValues().supportedOnAllNodes()),
274+
DATE_RANGE(builder().esType("date_range").estimatedSize(2 * Long.BYTES).docValues().supportedOn(ESQL_DATE_RANGE_CREATED_VERSION)),
273275
/**
274276
* IP addresses. IPv4 address are always
275277
* <a href="https://datatracker.ietf.org/doc/html/rfc4291#section-2.5.5">embedded</a>
@@ -982,5 +984,6 @@ private static class DataTypesTransportVersions {
982984
public static final TransportVersion ESQL_AGGREGATE_METRIC_DOUBLE_CREATED_VERSION = TransportVersion.fromName(
983985
"esql_aggregate_metric_double_created_version"
984986
);
987+
985988
}
986989
}

‎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
@@ -224,6 +224,7 @@ public static void appendValue(Block.Builder builder, Object val, ElementType ty
224224
case DOUBLE -> ((DoubleBlock.Builder) builder).appendDouble((Double) val);
225225
case BOOLEAN -> ((BooleanBlock.Builder) builder).appendBoolean((Boolean) val);
226226
case AGGREGATE_METRIC_DOUBLE -> ((AggregateMetricDoubleBlockBuilder) builder).appendLiteral((AggregateMetricDoubleLiteral) val);
227+
case DATE_RANGE -> ((DateRangeBlockBuilder) builder).appendDateRange((DateRangeBlockBuilder.DateRangeLiteral) val);
227228
default -> throw new UnsupportedOperationException("unsupported element type [" + type + "]");
228229
}
229230
}
@@ -253,6 +254,7 @@ private static Block constantBlock(BlockFactory blockFactory, ElementType type,
253254
case BOOLEAN -> blockFactory.newConstantBooleanBlockWith((boolean) val, size);
254255
case AGGREGATE_METRIC_DOUBLE -> blockFactory.newConstantAggregateMetricDoubleBlock((AggregateMetricDoubleLiteral) val, size);
255256
case FLOAT -> blockFactory.newConstantFloatBlockWith((float) val, size);
257+
case DATE_RANGE -> blockFactory.newConstantDateRangeBlock((DateRangeBlockBuilder.DateRangeLiteral) val, size);
256258
default -> throw new UnsupportedOperationException("unsupported element type [" + type + "]");
257259
};
258260
}
@@ -306,6 +308,12 @@ yield new AggregateMetricDoubleLiteral(
306308
aggBlock.countBlock().getInt(offset)
307309
);
308310
}
311+
case DATE_RANGE -> {
312+
DateRangeBlock b = (DateRangeBlock) block;
313+
LongBlock fromBlock = b.getFromBlock();
314+
LongBlock toBlock = b.getToBlock();
315+
yield new DateRangeBlockBuilder.DateRangeLiteral(fromBlock.getLong(offset), toBlock.getLong(offset));
316+
}
309317
case UNKNOWN -> throw new IllegalArgumentException("can't read values from [" + block + "]");
310318
};
311319
}

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