Skip to content

Commit 0591d5b

Browse files
shmuelhanochshmuelhanoch
committed
Basic support for range types in ESQL
1 parent 0722a33 commit 0591d5b

File tree

37 files changed

+899
-37
lines changed

37 files changed

+899
-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: 93 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,97 @@ 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+
if (docValues != null) {
373+
NumericDocValues singleton = DocValues.unwrapSingleton(docValues);
374+
if (singleton != null) {
375+
return new BlockDocValuesReader.SingletonLongs(singleton);
376+
}
377+
return new BlockDocValuesReader.Longs(docValues);
378+
}
379+
NumericDocValues singleton = context.reader().getNumericDocValues(fieldName);
380+
if (singleton != null) {
381+
return new BlockDocValuesReader.SingletonLongs(singleton);
382+
}
383+
return new ConstantNullsReader(); */
384+
}
385+
}
386+
387+
@Override
388+
public BlockLoader blockLoader(BlockLoaderContext blContext) {
389+
if (hasDocValues()) {
390+
return new DateRangeDocValuesLoader(name());
391+
}
392+
throw new IllegalStateException("Cannot load blocks without doc values");
393+
}
394+
}
395+
396+
public static class DateRangeDocValuesReader extends BlockDocValuesReader {
397+
private final BytesRef spare = new BytesRef();
398+
399+
private final BinaryDocValues numericDocValues;
400+
401+
public DateRangeDocValuesReader(BinaryDocValues numericDocValues) {
402+
this.numericDocValues = numericDocValues;
403+
}
404+
405+
@Override
406+
protected int docId() {
407+
return 0;
408+
}
409+
410+
@Override
411+
public String toString() {
412+
return "";
413+
}
414+
415+
@Override
416+
public BlockLoader.Block read(BlockLoader.BlockFactory factory, BlockLoader.Docs docs, int offset, boolean nullsFiltered)
417+
throws IOException {
418+
try (BlockLoader.DateRangeBuilder builder = factory.dateRangeBuilder(docs.count() - offset)) {
419+
for (int i = offset; i < docs.count(); i++) {
420+
int doc = docs.get(i);
421+
read(doc, builder);
422+
}
423+
return builder.build();
424+
}
425+
}
426+
427+
@Override
428+
public void read(int docId, BlockLoader.StoredFields storedFields, BlockLoader.Builder builder) throws IOException {
429+
read(docId, (BlockLoader.DateRangeBuilder) builder);
430+
}
431+
432+
private void read(int doc, BlockLoader.DateRangeBuilder builder) throws IOException {
433+
if (false == numericDocValues.advanceExact(doc)) {
434+
builder.appendNull();
435+
return;
436+
}
437+
438+
BytesRef ref = numericDocValues.binaryValue();
439+
var ranges = BinaryRangeUtil.decodeLongRanges(ref);
440+
for (var range : ranges) {
441+
builder.from().appendLong((long) range.from);
442+
builder.to().appendLong((long) range.to);
443+
}
444+
}
352445
}
353446

354447
private final RangeType type;

‎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: 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 DateRangeArrayBlock(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
}

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