Skip to content

Commit 2f8e144

Browse files
shmuelhanochshmuelhanoch
committed
Basic support for range types in ESQL
1 parent 870a155 commit 2f8e144

File tree

30 files changed

+736
-62
lines changed

30 files changed

+736
-62
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
@@ -483,6 +483,8 @@ interface BlockFactory {
483483
SortedSetOrdinalsBuilder sortedSetOrdinalsBuilder(SortedSetDocValues ordinals, int count);
484484

485485
AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int count);
486+
487+
DateRangeBuilder dateRangeBuilder(int count);
486488
}
487489

488490
/**
@@ -603,4 +605,10 @@ interface AggregateMetricDoubleBuilder extends Builder {
603605

604606
IntBuilder count();
605607
}
608+
609+
interface DateRangeBuilder extends Builder {
610+
LongBuilder from();
611+
612+
LongBuilder to();
613+
}
606614
}

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

Lines changed: 136 additions & 4 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;
@@ -420,10 +514,48 @@ protected void parseCreateField(DocumentParserContext context) throws IOExceptio
420514

421515
private Range parseRange(XContentParser parser) throws IOException {
422516
final XContentParser.Token start = parser.currentToken();
423-
if (fieldType().rangeType == RangeType.IP && start == XContentParser.Token.VALUE_STRING) {
424-
return parseIpRangeFromCidr(parser);
425-
}
517+
// Ranges representation is somewhat compilicated:
518+
// 1) "Classic" ranges from search comes as objects, while new ESQL ranges have a string representation.
519+
// 2) This ESQL string representation looks like "from..to" for all range types except IP range, which has its own special
520+
// representation.
521+
if (start == XContentParser.Token.START_OBJECT) return parseRangeObject(parser, start);
522+
return switch (fieldType().rangeType) {
523+
case RangeType.IP -> parseIpRangeFromCidr(parser);
524+
case RangeType.DATE -> parseDateRange(parser.text());
525+
case RangeType.INTEGER -> parseIntegerRange(parser.text());
526+
case RangeType.DOUBLE -> parseDoubleRange(parser.text());
527+
default -> throw new IllegalArgumentException("Cannot parse range of type " + fieldType().rangeType);
528+
};
529+
}
530+
531+
Range parseDoubleRange(String s) {
532+
var ss = s.split("\\.\\.");
533+
assert ss.length == 2 : "can't parse range: " + s;
534+
var from = Double.parseDouble(ss[0]);
535+
var to = Double.parseDouble(ss[1]);
536+
return new Range(RangeType.INTEGER, from, to, true, false);
537+
}
538+
539+
Range parseIntegerRange(String s) {
540+
var ss = s.split("\\.\\.");
541+
assert ss.length == 2 : "can't parse range: " + s;
542+
var from = Integer.parseInt(ss[0]);
543+
var to = Integer.parseInt(ss[1]);
544+
return new Range(RangeType.INTEGER, from, to, true, false);
545+
}
546+
547+
// Copied from EsqlDataTypeConverter since it's impossible to share code with this module.
548+
// We use split("..") for lightway lexing before the actual date parsers come into play.
549+
// It's both unsafe, and will not work for double ranges if we'll add them. But it's working for now
550+
Range parseDateRange(String s) {
551+
var ss = s.split("\\.\\.");
552+
assert ss.length == 2 : "can't parse range: " + s;
553+
var from = DEFAULT_DATE_TIME_FORMATTER.parseMillis(ss[0]);
554+
var to = DEFAULT_DATE_TIME_FORMATTER.parseMillis(ss[1]);
555+
return new Range(RangeType.DATE, from, to, true, false);
556+
}
426557

558+
private Range parseRangeObject(XContentParser parser, XContentParser.Token start) throws IOException {
427559
if (start != XContentParser.Token.START_OBJECT) {
428560
throw new DocumentParsingException(
429561
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;
@@ -372,9 +373,15 @@ public SortedSetOrdinalBuilder appendOrd(int value) {
372373
return new SortedSetOrdinalBuilder();
373374
}
374375

376+
@Override
375377
public BlockLoader.AggregateMetricDoubleBuilder aggregateMetricDoubleBuilder(int expectedSize) {
376378
return new AggregateMetricDoubleBlockBuilder(expectedSize);
377379
}
380+
381+
@Override
382+
public BlockLoader.DateRangeBuilder dateRangeBuilder(int expectedSize) {
383+
return new DateRangeBuilder(expectedSize);
384+
}
378385
};
379386
}
380387

@@ -582,4 +589,68 @@ public void close() {
582589

583590
}
584591
}
592+
593+
public static class DateRangeBuilder implements BlockLoader.DateRangeBuilder {
594+
private final LongBuilder from;
595+
private final LongBuilder to;
596+
597+
DateRangeBuilder(int expectedSize) {
598+
from = new LongBuilder(expectedSize);
599+
to = new LongBuilder(expectedSize);
600+
}
601+
602+
@Override
603+
public BlockLoader.LongBuilder from() {
604+
return from;
605+
}
606+
607+
@Override
608+
public BlockLoader.LongBuilder to() {
609+
return to;
610+
}
611+
612+
@Override
613+
public BlockLoader.Block build() {
614+
var fromBlock = from.build();
615+
var toBlock = to.build();
616+
assert fromBlock.size() == toBlock.size();
617+
var values = new ArrayList<>(fromBlock.size());
618+
for (int i = 0; i < fromBlock.size(); i++) {
619+
values.add(List.of(fromBlock.values.get(i), toBlock.values.get(i)));
620+
}
621+
return new TestBlock(values);
622+
}
623+
624+
@Override
625+
public BlockLoader.Builder appendNull() {
626+
return null;
627+
}
628+
629+
@Override
630+
public BlockLoader.Builder beginPositionEntry() {
631+
return null;
632+
}
633+
634+
@Override
635+
public BlockLoader.Builder endPositionEntry() {
636+
return null;
637+
}
638+
639+
@Override
640+
public void close() {
641+
Releasables.close(from, to);
642+
}
643+
644+
private static class LongBuilder extends TestBlock.Builder implements BlockLoader.LongBuilder {
645+
private LongBuilder(int expectedSize) {
646+
super(expectedSize);
647+
}
648+
649+
@Override
650+
public BlockLoader.LongBuilder appendLong(long value) {
651+
add(value);
652+
return this;
653+
}
654+
}
655+
};
585656
}

‎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").unknownSize());
314+
DENSE_VECTOR(builder().esType("dense_vector").unknownSize()),
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/module-info.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
requires org.elasticsearch.geo;
2222
requires org.elasticsearch.xcore;
2323
requires hppc;
24+
requires com.ibm.icu;
2425

2526
exports org.elasticsearch.compute;
2627
exports org.elasticsearch.compute.aggregation;

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -481,6 +481,10 @@ 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+
484488
/**
485489
* Returns the maximum number of bytes that a Block should be backed by a primitive array before switching to using BigArrays.
486490
*/

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,12 @@ yield new AggregateMetricDoubleLiteral(
299299
aggBlock.countBlock().getInt(offset)
300300
);
301301
}
302+
case DATE_RANGE -> {
303+
DateRangeBlock b = (DateRangeBlock) block;
304+
LongBlock fromBlock = b.getFromBlock();
305+
LongBlock toBlock = b.getToBlock();
306+
yield new DateRangeBlockBuilder.DateRangeLiteral(fromBlock.getLong(offset), toBlock.getLong(offset));
307+
}
302308
case UNKNOWN -> throw new IllegalArgumentException("can't read values from [" + block + "]");
303309
};
304310
}

0 commit comments

Comments
 (0)