Skip to content

Commit 6626682

Browse files
leontyevdvleontyevdv
committed
ES|QL: Address PR review comments
- Add a new test index to AbstractLogicalPlanOptimizerTests to test with TBucket - Add tbucket to bucket translation test - Add date_nano support to TBucket Closes #131068
1 parent 06cd95a commit 6626682

File tree

5 files changed

+132
-10
lines changed

5 files changed

+132
-10
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/grouping/TBucket.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ public class TBucket extends GroupingFunction.EvaluatableGroupingFunction implem
3636
private final Expression timestamp;
3737

3838
@FunctionInfo(
39-
returnType = { "date" },
39+
returnType = { "date", "date_nanos" },
4040
description = """
4141
Creates groups of values - buckets - out of a @timestamp attribute. The size of the buckets must be provided directly.""",
4242
examples = {
@@ -102,7 +102,7 @@ protected TypeResolution resolveType() {
102102

103103
@Override
104104
public DataType dataType() {
105-
return DataType.DATETIME;
105+
return timestamp.dataType();
106106
}
107107

108108
@Override

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4136,13 +4136,14 @@ public void testGroupingOverridesInInlinestats() {
41364136
public void testTBucketWithIntervalInStringInBothAggregationAndGrouping() {
41374137
LogicalPlan plan = analyze("""
41384138
FROM sample_data
4139-
| STATS min = MIN(@timestamp), max = MAX(@timestamp) BY bucket = TBUCKET("1 week")
4139+
| STATS min = MIN(@timestamp), max = MAX(@timestamp) BY bucket = TBUCKET(1 week)
41404140
| SORT min
41414141
""", "mapping-sample_data.json");
41424142

41434143
Limit limit = as(plan, Limit.class);
41444144
OrderBy orderBy = as(limit.child(), OrderBy.class);
41454145
Aggregate agg = as(orderBy.child(), Aggregate.class);
4146+
41464147
List<? extends NamedExpression> aggregates = agg.aggregates();
41474148
assertThat(aggregates, hasSize(3));
41484149
Alias a = as(aggregates.get(0), Alias.class);
@@ -4157,6 +4158,7 @@ public void testTBucketWithIntervalInStringInBothAggregationAndGrouping() {
41574158
assertEquals("@timestamp", fa.name());
41584159
ReferenceAttribute ra = as(aggregates.get(2), ReferenceAttribute.class);
41594160
assertEquals("bucket", ra.name());
4161+
41604162
List<Expression> groupings = agg.groupings();
41614163
assertEquals(1, groupings.size());
41624164
a = as(groupings.get(0), Alias.class); // reference in groupings is resolved

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/grouping/TBucketTests.java‎

Lines changed: 60 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
1212

1313
import org.elasticsearch.common.Rounding;
14+
import org.elasticsearch.common.time.DateUtils;
1415
import org.elasticsearch.index.mapper.DateFieldMapper;
1516
import org.elasticsearch.logging.LogManager;
1617
import org.elasticsearch.xpack.esql.core.expression.Expression;
@@ -19,6 +20,7 @@
1920
import org.elasticsearch.xpack.esql.expression.function.AbstractScalarFunctionTestCase;
2021
import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
2122
import org.hamcrest.Matcher;
23+
import org.hamcrest.Matchers;
2224

2325
import java.time.Duration;
2426
import java.time.Instant;
@@ -29,7 +31,6 @@
2931
import java.util.function.Supplier;
3032

3133
import static org.hamcrest.Matchers.equalTo;
32-
import static org.hamcrest.Matchers.hasSize;
3334

3435
public class TBucketTests extends AbstractScalarFunctionTestCase {
3536
public TBucketTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
@@ -55,6 +56,20 @@ public static Iterable<Object[]> parameters() {
5556
Duration.ofDays(1L),
5657
"[86400000 in Z][fixed]"
5758
);
59+
dateNanosCasesWithSpan(
60+
suppliers,
61+
"fixed date nanos with period",
62+
() -> DateUtils.toLong(Instant.parse("2023-01-01T00:00:00.00Z")),
63+
DataType.DATE_PERIOD,
64+
Period.ofYears(1)
65+
);
66+
dateNanosCasesWithSpan(
67+
suppliers,
68+
"fixed date nanos with duration",
69+
() -> DateUtils.toLong(Instant.parse("2023-02-17T09:00:00.00Z")),
70+
DataType.TIME_DURATION,
71+
Duration.ofDays(1L)
72+
);
5873
return parameterSuppliersFromTypedData(suppliers);
5974
}
6075

@@ -69,7 +84,7 @@ private static void dateCasesWithSpan(
6984
suppliers.add(new TestCaseSupplier(name, List.of(spanType, DataType.DATETIME), () -> {
7085
List<TestCaseSupplier.TypedData> args = new ArrayList<>();
7186
args.add(new TestCaseSupplier.TypedData(span, spanType, "buckets").forceLiteral());
72-
args.add(new TestCaseSupplier.TypedData(date.getAsLong(), DataType.DATETIME, "field"));
87+
args.add(new TestCaseSupplier.TypedData(date.getAsLong(), DataType.DATETIME, "@timestamp"));
7388

7489
return new TestCaseSupplier.TestCase(
7590
args,
@@ -80,7 +95,38 @@ private static void dateCasesWithSpan(
8095
}));
8196
}
8297

98+
private static void dateNanosCasesWithSpan(
99+
List<TestCaseSupplier> suppliers,
100+
String name,
101+
LongSupplier date,
102+
DataType spanType,
103+
Object span
104+
) {
105+
suppliers.add(new TestCaseSupplier(name, List.of(spanType, DataType.DATE_NANOS), () -> {
106+
List<TestCaseSupplier.TypedData> args = new ArrayList<>();
107+
args.add(new TestCaseSupplier.TypedData(span, spanType, "buckets").forceLiteral());
108+
args.add(new TestCaseSupplier.TypedData(date.getAsLong(), DataType.DATE_NANOS, "@timestamp"));
109+
return new TestCaseSupplier.TestCase(
110+
args,
111+
Matchers.startsWith("DateTruncDateNanosEvaluator[fieldVal=Attribute[channel=0], rounding=Rounding["),
112+
DataType.DATE_NANOS,
113+
resultsMatcher(args)
114+
);
115+
}));
116+
}
117+
83118
private static Matcher<Object> resultsMatcher(List<TestCaseSupplier.TypedData> typedData) {
119+
if (typedData.get(1).type() == DataType.DATE_NANOS) {
120+
long nanos = ((Number) typedData.get(1).data()).longValue();
121+
long expected = DateUtils.toNanoSeconds(
122+
Rounding.builder(Rounding.DateTimeUnit.DAY_OF_MONTH).build().prepareForUnknown().round(DateUtils.toMilliSeconds(nanos))
123+
);
124+
LogManager.getLogger(getTestClass()).info("Expected: " + DateUtils.toInstant(expected));
125+
LogManager.getLogger(getTestClass()).info("Input: " + DateUtils.toInstant(nanos));
126+
return equalTo(expected);
127+
}
128+
129+
// For DATETIME, we use the millis value directly
84130
long millis = ((Number) typedData.get(1).data()).longValue();
85131
long expected = Rounding.builder(Rounding.DateTimeUnit.DAY_OF_MONTH).build().prepareForUnknown().round(millis);
86132
LogManager.getLogger(getTestClass()).info("Expected: " + Instant.ofEpochMilli(expected));
@@ -98,9 +144,16 @@ protected boolean canSerialize() {
98144
return false;
99145
}
100146

101-
public static List<DataType> signatureTypes(List<DataType> testCaseTypes) {
102-
assertThat(testCaseTypes, hasSize(2));
103-
assertThat(testCaseTypes.get(1), equalTo(DataType.DATETIME));
104-
return List.of(testCaseTypes.get(0));
105-
}
147+
// TODO: With this method commented the V3Doc fails
148+
// public static List<DataType> signatureTypes(List<DataType> testCaseTypes) {
149+
// // DATE_PERIOD, DATETIME
150+
// // DATE_PERIOD, DATE_NANOS
151+
//
152+
// // TIME_DURATION, DATETIME
153+
// // TIME_DURATION, DATE_NANOS
154+
//
155+
// assertThat(testCaseTypes, hasSize(2));
156+
// // assertThat(testCaseTypes.get(1), anyOf(equalTo(DataType.DATE_NANOS), equalTo(DataType.DATETIME)));
157+
// return List.of(testCaseTypes.get(0), testCaseTypes.get(1));
158+
// }
106159
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/AbstractLogicalPlanOptimizerTests.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ public abstract class AbstractLogicalPlanOptimizerTests extends ESTestCase {
5151
protected static Map<String, EsField> metricMapping;
5252
protected static Analyzer metricsAnalyzer;
5353
protected static Analyzer multiIndexAnalyzer;
54+
protected static Analyzer sampleDataIndexAnalyzer;
5455

5556
protected static EnrichResolution enrichResolution;
5657

@@ -169,6 +170,21 @@ public static void init() {
169170
),
170171
TEST_VERIFIER
171172
);
173+
174+
var sampleDataMapping = loadMapping("mapping-sample_data.json");
175+
var sampleDataIndex = IndexResolution.valid(
176+
new EsIndex("sample_data", sampleDataMapping, Map.of("sample_data", IndexMode.STANDARD))
177+
);
178+
sampleDataIndexAnalyzer = new Analyzer(
179+
new AnalyzerContext(
180+
EsqlTestUtils.TEST_CFG,
181+
new EsqlFunctionRegistry(),
182+
sampleDataIndex,
183+
enrichResolution,
184+
emptyInferenceResolution()
185+
),
186+
TEST_VERIFIER
187+
);
172188
}
173189

174190
protected LogicalPlan optimizedPlan(String query) {
@@ -211,6 +227,11 @@ protected LogicalPlan planMultiIndex(String query) {
211227
return logicalOptimizer.optimize(multiIndexAnalyzer.analyze(parser.createStatement(query, EsqlTestUtils.TEST_CFG)));
212228
}
213229

230+
protected LogicalPlan planSample(String query) {
231+
var analyzed = sampleDataIndexAnalyzer.analyze(parser.createStatement(query, EsqlTestUtils.TEST_CFG));
232+
return logicalOptimizer.optimize(analyzed);
233+
}
234+
214235
@Override
215236
protected List<String> filteredWarnings() {
216237
return withDefaultLimitWarning(super.filteredWarnings());

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8122,4 +8122,50 @@ public List<Attribute> output() {
81228122
assertThat(e.getMessage(), containsString("Output has changed from"));
81238123
}
81248124

8125+
public void testTranslateDataGroupedByTBucket() {
8126+
assumeTrue("requires snapshot builds", Build.current().isSnapshot());
8127+
var query = """
8128+
FROM sample_data
8129+
| STATS min = MIN(@timestamp), max = MAX(@timestamp) BY bucket = TBUCKET(1 hour)
8130+
| SORT min
8131+
""";
8132+
8133+
var plan = planSample(query);
8134+
var topN = as(plan, TopN.class);
8135+
8136+
Aggregate aggregate = as(topN.child(), Aggregate.class);
8137+
assertThat(aggregate, not(instanceOf(TimeSeriesAggregate.class)));
8138+
8139+
assertThat(aggregate.groupings(), hasSize(1));
8140+
assertThat(aggregate.groupings().get(0), instanceOf(ReferenceAttribute.class));
8141+
assertThat(as(aggregate.groupings().getFirst(), ReferenceAttribute.class).name(), equalTo("bucket"));
8142+
8143+
assertThat(aggregate.aggregates(), hasSize(3));
8144+
List<? extends NamedExpression> aggregates = aggregate.aggregates();
8145+
assertThat(aggregates, hasSize(3));
8146+
Alias a = as(aggregates.get(0), Alias.class);
8147+
assertEquals("min", a.name());
8148+
Min min = as(a.child(), Min.class);
8149+
FieldAttribute fa = as(min.field(), FieldAttribute.class);
8150+
assertEquals("@timestamp", fa.name());
8151+
a = as(aggregates.get(1), Alias.class);
8152+
assertEquals("max", a.name());
8153+
Max max = as(a.child(), Max.class);
8154+
fa = as(max.field(), FieldAttribute.class);
8155+
assertEquals("@timestamp", fa.name());
8156+
ReferenceAttribute ra = as(aggregates.get(2), ReferenceAttribute.class);
8157+
assertEquals("bucket", ra.name());
8158+
assertThat(Expressions.attribute(aggregate.groupings().get(0)).id(), equalTo(aggregate.aggregates().get(2).id()));
8159+
Eval eval = as(aggregate.child(), Eval.class);
8160+
assertThat(eval.fields(), hasSize(1));
8161+
Alias bucketAlias = eval.fields().get(0);
8162+
assertThat(bucketAlias.child(), instanceOf(Bucket.class));
8163+
assertThat(Expressions.attribute(bucketAlias).id(), equalTo(aggregate.aggregates().get(2).id()));
8164+
Bucket bucket = as(Alias.unwrap(bucketAlias), Bucket.class);
8165+
assertThat(Expressions.attribute(bucket.field()).name(), equalTo("@timestamp"));
8166+
assertThat(bucket.children().get(0), instanceOf(FieldAttribute.class));
8167+
assertThat(((FieldAttribute)bucket.children().get(0)).name(), equalTo("@timestamp"));
8168+
assertThat(bucket.children().get(1), instanceOf(Literal.class));
8169+
assertThat(((Literal)bucket.children().get(1)).value(), equalTo(Duration.ofHours(1)));
8170+
}
81258171
}

0 commit comments

Comments
 (0)