Skip to content

Commit

Permalink
Datetime aggregation fixes. (#1061) (#1312)
Browse files Browse the repository at this point in the history 
* Update aggregation to support datetime types.

Signed-off-by: Yury-Fridlyand <[email protected]>
Signed-off-by: Yury-Fridlyand <[email protected]>
Co-authored-by: MaxKsyunz <[email protected]>
(cherry picked from commit 5220a98)

Co-authored-by: Yury-Fridlyand <[email protected]>
  • Loading branch information
@opensearch-trigger-bot @Yury-Fridlyand
opensearch-trigger-bot[bot] and Yury-Fridlyand authored Feb 1, 2023
1 parent 3939050 commit d6a00a8
Show file tree
Hide file tree
Showing 10 changed files with 689 additions and 18 deletions.
1 change: 0 additions & 1 deletion core/src/main/java/org/opensearch/sql/data/model/ExprValueUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
* SPDX-License-Identifier: Apache-2.0
*/


package org.opensearch.sql.data.model;

import java.time.Instant;
Expand Down
8 changes: 8 additions & 0 deletions core/src/main/java/org/opensearch/sql/expression/aggregation/AggregatorFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,14 @@ private static DefaultFunctionResolver avg() {
new ImmutableMap.Builder<FunctionSignature, FunctionBuilder>()
.put(new FunctionSignature(functionName, Collections.singletonList(DOUBLE)),
(functionProperties, arguments) -> new AvgAggregator(arguments, DOUBLE))
.put(new FunctionSignature(functionName, Collections.singletonList(DATE)),
(functionProperties, arguments) -> new AvgAggregator(arguments, DATE))
.put(new FunctionSignature(functionName, Collections.singletonList(DATETIME)),
(functionProperties, arguments) -> new AvgAggregator(arguments, DATETIME))
.put(new FunctionSignature(functionName, Collections.singletonList(TIME)),
(functionProperties, arguments) -> new AvgAggregator(arguments, TIME))
.put(new FunctionSignature(functionName, Collections.singletonList(TIMESTAMP)),
(functionProperties, arguments) -> new AvgAggregator(arguments, TIMESTAMP))
.build()
);
}
Expand Down
149 changes: 138 additions & 11 deletions core/src/main/java/org/opensearch/sql/expression/aggregation/AvgAggregator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,23 @@

package org.opensearch.sql.expression.aggregation;

import static java.time.temporal.ChronoUnit.MILLIS;
import static org.opensearch.sql.utils.ExpressionUtils.format;

import java.time.Instant;
import java.time.LocalTime;
import java.util.List;
import java.util.Locale;
import org.opensearch.sql.data.model.ExprDateValue;
import org.opensearch.sql.data.model.ExprDatetimeValue;
import org.opensearch.sql.data.model.ExprDoubleValue;
import org.opensearch.sql.data.model.ExprIntegerValue;
import org.opensearch.sql.data.model.ExprNullValue;
import org.opensearch.sql.data.model.ExprTimeValue;
import org.opensearch.sql.data.model.ExprTimestampValue;
import org.opensearch.sql.data.model.ExprValue;
import org.opensearch.sql.data.model.ExprValueUtils;
import org.opensearch.sql.data.type.ExprCoreType;
import org.opensearch.sql.expression.DSL;
import org.opensearch.sql.expression.Expression;
import org.opensearch.sql.expression.function.BuiltinFunctionName;

Expand All @@ -23,20 +32,39 @@
*/
public class AvgAggregator extends Aggregator<AvgAggregator.AvgState> {

/**
* To process by different ways different data types, we need to store the type.
* Input data has the same type as the result.
*/
private final ExprCoreType dataType;

public AvgAggregator(List<Expression> arguments, ExprCoreType returnType) {
super(BuiltinFunctionName.AVG.getName(), arguments, returnType);
dataType = returnType;
}

@Override
public AvgState create() {
return new AvgState();
switch (dataType) {
case DATE:
return new DateAvgState();
case DATETIME:
return new DateTimeAvgState();
case TIMESTAMP:
return new TimestampAvgState();
case TIME:
return new TimeAvgState();
case DOUBLE:
return new DoubleAvgState();
default: //unreachable code - we don't expose signatures for unsupported types
throw new IllegalArgumentException(
String.format("avg aggregation over %s type is not supported", dataType));
}
}

@Override
protected AvgState iterate(ExprValue value, AvgState state) {
state.count++;
state.total += ExprValueUtils.getDoubleValue(value);
return state;
return state.iterate(value);
}

@Override
Expand All @@ -47,18 +75,117 @@ public String toString() {
/**
* Average State.
*/
protected static class AvgState implements AggregationState {
private int count;
private double total;
protected abstract static class AvgState implements AggregationState {
protected ExprValue count;
protected ExprValue total;

AvgState() {
this.count = 0;
this.total = 0d;
this.count = new ExprIntegerValue(0);
this.total = new ExprDoubleValue(0D);
}

@Override
public abstract ExprValue result();

protected AvgState iterate(ExprValue value) {
count = DSL.add(DSL.literal(count), DSL.literal(1)).valueOf();
return this;
}
}

protected static class DoubleAvgState extends AvgState {
@Override
public ExprValue result() {
if (0 == count.integerValue()) {
return ExprNullValue.of();
}
return DSL.divide(DSL.literal(total), DSL.literal(count)).valueOf();
}

@Override
protected AvgState iterate(ExprValue value) {
total = DSL.add(DSL.literal(total), DSL.literal(value)).valueOf();
return super.iterate(value);
}
}

protected static class DateAvgState extends AvgState {
@Override
public ExprValue result() {
if (0 == count.integerValue()) {
return ExprNullValue.of();
}

return new ExprDateValue(
new ExprTimestampValue(Instant.ofEpochMilli(
DSL.divide(DSL.literal(total), DSL.literal(count)).valueOf().longValue()))
.dateValue());
}

@Override
protected AvgState iterate(ExprValue value) {
total = DSL.add(DSL.literal(total), DSL.literal(value.timestampValue().toEpochMilli()))
.valueOf();
return super.iterate(value);
}
}

protected static class DateTimeAvgState extends AvgState {
@Override
public ExprValue result() {
if (0 == count.integerValue()) {
return ExprNullValue.of();
}

return new ExprDatetimeValue(
new ExprTimestampValue(Instant.ofEpochMilli(
DSL.divide(DSL.literal(total), DSL.literal(count)).valueOf().longValue()))
.datetimeValue());
}

@Override
protected AvgState iterate(ExprValue value) {
total = DSL.add(DSL.literal(total), DSL.literal(value.timestampValue().toEpochMilli()))
.valueOf();
return super.iterate(value);
}
}

protected static class TimestampAvgState extends AvgState {
@Override
public ExprValue result() {
if (0 == count.integerValue()) {
return ExprNullValue.of();
}

return new ExprTimestampValue(Instant.ofEpochMilli(
DSL.divide(DSL.literal(total), DSL.literal(count)).valueOf().longValue()));
}

@Override
protected AvgState iterate(ExprValue value) {
total = DSL.add(DSL.literal(total), DSL.literal(value.timestampValue().toEpochMilli()))
.valueOf();
return super.iterate(value);
}
}

protected static class TimeAvgState extends AvgState {
@Override
public ExprValue result() {
return count == 0 ? ExprNullValue.of() : ExprValueUtils.doubleValue(total / count);
if (0 == count.integerValue()) {
return ExprNullValue.of();
}

return new ExprTimeValue(LocalTime.MIN.plus(
DSL.divide(DSL.literal(total), DSL.literal(count)).valueOf().longValue(), MILLIS));
}

@Override
protected AvgState iterate(ExprValue value) {
total = DSL.add(DSL.literal(total),
DSL.literal(MILLIS.between(LocalTime.MIN, value.timeValue()))).valueOf();
return super.iterate(value);
}
}
}
72 changes: 72 additions & 0 deletions core/src/test/java/org/opensearch/sql/expression/aggregation/AvgAggregatorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,18 @@
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.opensearch.sql.data.type.ExprCoreType.DATE;
import static org.opensearch.sql.data.type.ExprCoreType.DATETIME;
import static org.opensearch.sql.data.type.ExprCoreType.DOUBLE;
import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
import static org.opensearch.sql.data.type.ExprCoreType.STRING;
import static org.opensearch.sql.data.type.ExprCoreType.TIME;
import static org.opensearch.sql.data.type.ExprCoreType.TIMESTAMP;

import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.opensearch.sql.data.model.ExprValue;
import org.opensearch.sql.data.model.ExprValueUtils;
Expand Down Expand Up @@ -62,13 +71,76 @@ public void avg_with_all_missing_or_null() {
assertTrue(result.isNull());
}

@Test
public void avg_numeric_no_values() {
ExprValue result = aggregation(DSL.avg(DSL.ref("dummy", INTEGER)), List.of());
assertTrue(result.isNull());
}

@Test
public void avg_date_no_values() {
ExprValue result = aggregation(DSL.avg(DSL.ref("dummy", DATE)), List.of());
assertTrue(result.isNull());
}

@Test
public void avg_datetime_no_values() {
ExprValue result = aggregation(DSL.avg(DSL.ref("dummy", DATETIME)), List.of());
assertTrue(result.isNull());
}

@Test
public void avg_timestamp_no_values() {
ExprValue result = aggregation(DSL.avg(DSL.ref("dummy", TIMESTAMP)), List.of());
assertTrue(result.isNull());
}

@Test
public void avg_time_no_values() {
ExprValue result = aggregation(DSL.avg(DSL.ref("dummy", TIME)), List.of());
assertTrue(result.isNull());
}

@Test
public void avg_date() {
ExprValue result = aggregation(DSL.avg(DSL.date(DSL.ref("date_value", STRING))), tuples);
assertEquals(LocalDate.of(2007, 7, 2), result.dateValue());
}

@Test
public void avg_datetime() {
var result = aggregation(DSL.avg(DSL.datetime(DSL.ref("datetime_value", STRING))), tuples);
assertEquals(LocalDateTime.of(2012, 7, 2, 3, 30), result.datetimeValue());
}

@Test
public void avg_time() {
ExprValue result = aggregation(DSL.avg(DSL.time(DSL.ref("time_value", STRING))), tuples);
assertEquals(LocalTime.of(9, 30), result.timeValue());
}

@Test
public void avg_timestamp() {
var result = aggregation(DSL.avg(DSL.timestamp(DSL.ref("timestamp_value", STRING))), tuples);
assertEquals(TIMESTAMP, result.type());
assertEquals(LocalDateTime.of(2012, 7, 2, 3, 30), result.datetimeValue());
}

@Test
public void valueOf() {
ExpressionEvaluationException exception = assertThrows(ExpressionEvaluationException.class,
() -> DSL.avg(DSL.ref("double_value", DOUBLE)).valueOf(valueEnv()));
assertEquals("can't evaluate on aggregator: avg", exception.getMessage());
}

@Test
public void avg_on_unsupported_type() {
var aggregator = new AvgAggregator(List.of(DSL.ref("string", STRING)), STRING);
var exception = assertThrows(IllegalArgumentException.class,
() -> aggregator.create());
assertEquals("avg aggregation over STRING type is not supported", exception.getMessage());
}

@Test
public void test_to_string() {
Aggregator avgAggregator = DSL.avg(DSL.ref("integer_value", INTEGER));
Expand Down
8 changes: 4 additions & 4 deletions docs/user/dql/aggregations.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ SUM
Description
>>>>>>>>>>>

Usage: SUM(expr). Returns the sum of expr.
Usage: SUM(expr). Returns the sum of `expr`. `expr` could be of any of the numeric data types.

Example::

Expand All @@ -182,7 +182,7 @@ AVG
Description
>>>>>>>>>>>

Usage: AVG(expr). Returns the average value of expr.
Usage: AVG(expr). Returns the average value of `expr`. `expr` can be any numeric or datetime data type. Datetime aggregation is performed with milliseconds precision.

Example::

Expand All @@ -201,7 +201,7 @@ MAX
Description
>>>>>>>>>>>

Usage: MAX(expr). Returns the maximum value of expr.
Usage: MAX(expr). Returns the maximum value of `expr`. `expr` can be any numeric or datetime data type. Datetime aggregation is performed with milliseconds precision.

Example::

Expand All @@ -219,7 +219,7 @@ MIN
Description
>>>>>>>>>>>

Usage: MIN(expr). Returns the minimum value of expr.
Usage: MIN(expr). Returns the minimum value of `expr`. `expr` can be any numeric or datetime data type. Datetime aggregation is performed with milliseconds precision.

Example::

Expand Down
Loading

0 comments on commit d6a00a8

Please sign in to comment.