Skip to content

Commit

Permalink
Support to pass dynamic values to timestamp Extract function (#15586)
Browse files Browse the repository at this point in the history 
Fixes #15072

Before this modification , the third parameter (timezone) require to be a Literal, it will throw a error when this parameter is column Identifier.
  • Loading branch information
@AlbericByte
AlbericByte authored Dec 21, 2023
1 parent 8a45efb commit a2e65e6
Show file tree
Hide file tree
Showing 5 changed files with 266 additions and 134 deletions.
1 change: 0 additions & 1 deletion processing/src/main/java/org/apache/druid/math/expr/NamedFunction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,6 @@ default void validationHelperCheckArgIsLiteral(Expr arg, String argName)
}
}


/**
* Helper method for implementors performing validation to check that the argument list is some expected size.
*
Expand Down
269 changes: 168 additions & 101 deletions processing/src/main/java/org/apache/druid/query/expression/TimestampExtractExprMacro.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

package org.apache.druid.query.expression;

import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.math.expr.Expr;
import org.apache.druid.math.expr.ExprEval;
Expand Down Expand Up @@ -64,6 +65,87 @@ public String name()
return FN_NAME;
}

private ExprEval getExprEval(final DateTime dateTime, final Unit unit)
{
long epoch = dateTime.getMillis() / 1000;
switch (unit) {
case EPOCH:
return ExprEval.of(epoch);
case MICROSECOND:
return ExprEval.of(epoch / 1000);
case MILLISECOND:
return ExprEval.of(dateTime.millisOfSecond().get());
case SECOND:
return ExprEval.of(dateTime.secondOfMinute().get());
case MINUTE:
return ExprEval.of(dateTime.minuteOfHour().get());
case HOUR:
return ExprEval.of(dateTime.hourOfDay().get());
case DAY:
return ExprEval.of(dateTime.dayOfMonth().get());
case DOW:
return ExprEval.of(dateTime.dayOfWeek().get());
case ISODOW:
return ExprEval.of(dateTime.dayOfWeek().get());
case DOY:
return ExprEval.of(dateTime.dayOfYear().get());
case WEEK:
return ExprEval.of(dateTime.weekOfWeekyear().get());
case MONTH:
return ExprEval.of(dateTime.monthOfYear().get());
case QUARTER:
return ExprEval.of((dateTime.monthOfYear().get() - 1) / 3 + 1);
case YEAR:
return ExprEval.of(dateTime.year().get());
case ISOYEAR:
return ExprEval.of(dateTime.year().get());
case DECADE:
// The year field divided by 10, See https://www.postgresql.org/docs/10/functions-datetime.html
return ExprEval.of(dateTime.year().get() / 10);
case CENTURY:
return ExprEval.of(Math.ceil((double) dateTime.year().get() / 100));
case MILLENNIUM:
// Years in the 1900s are in the second millennium. The third millennium started January 1, 2001.
// See https://www.postgresql.org/docs/10/functions-datetime.html
return ExprEval.of(Math.ceil((double) dateTime.year().get() / 1000));
default:
throw TimestampExtractExprMacro.this.validationFailed("unhandled unit[%s]", unit);
}
}

private static ExpressionType getOutputExpressionType(final Unit unit)
{
switch (unit) {
case CENTURY:
case MILLENNIUM:
return ExpressionType.DOUBLE;
default:
return ExpressionType.LONG;
}
}

private static String stringifyExpr(final List<Expr> args)
{
if (args.size() > 2) {
return StringUtils.format(
"%s(%s, %s, %s)",
FN_NAME,
args.get(0).stringify(),
args.get(1).stringify(),
args.get(2).stringify()
);
}
return StringUtils.format("%s(%s, %s)", FN_NAME, args.get(0).stringify(), args.get(1).stringify());
}

private static ISOChronology computeChronology(final List<Expr> args, final Expr.ObjectBinding bindings)
{
String timeZoneVal = (String) args.get(2).eval(bindings).value();
return timeZoneVal != null
? ISOChronology.getInstance(DateTimes.inferTzFromString(timeZoneVal))
: ISOChronology.getInstanceUTC();
}

@Override
public Expr apply(final List<Expr> args)
{
Expand All @@ -73,121 +155,106 @@ public Expr apply(final List<Expr> args)
throw validationFailed("unit arg must be literal");
}

if (args.size() > 2) {
validationHelperCheckArgIsLiteral(args.get(2), "timezone");
}

final Expr arg = args.get(0);
final Unit unit = Unit.valueOf(StringUtils.toUpperCase((String) args.get(1).getLiteralValue()));
final DateTimeZone timeZone;

if (args.size() > 2) {
timeZone = ExprUtils.toTimeZone(args.get(2));
} else {
timeZone = DateTimeZone.UTC;
if (args.get(2).isLiteral()) {
DateTimeZone timeZone = ExprUtils.toTimeZone(args.get(2));
ISOChronology chronology = ISOChronology.getInstance(timeZone);
return new TimestampExtractExpr(args, unit, chronology);
} else {
return new TimestampExtractDynamicExpr(args, unit);
}
}
return new TimestampExtractExpr(args, unit, ISOChronology.getInstanceUTC());
}

final ISOChronology chronology = ISOChronology.getInstance(timeZone);
public class TimestampExtractExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
private final ISOChronology chronology;
private final Unit unit;

class TimestampExtractExpr extends ExprMacroTable.BaseScalarUnivariateMacroFunctionExpr
private TimestampExtractExpr(final List<Expr> args, final Unit unit, final ISOChronology chronology)
{
private TimestampExtractExpr(Expr arg)
{
super(FN_NAME, arg);
}
super(FN_NAME, args);
this.unit = unit;
this.chronology = chronology;
}

@Nonnull
@Override
public ExprEval eval(final ObjectBinding bindings)
{
Object val = arg.eval(bindings).value();
if (val == null) {
// Return null if the argument if null.
return ExprEval.of(null);
}
final DateTime dateTime = new DateTime(val, chronology);
long epoch = dateTime.getMillis() / 1000;

switch (unit) {
case EPOCH:
return ExprEval.of(epoch);
case MICROSECOND:
return ExprEval.of(epoch / 1000);
case MILLISECOND:
return ExprEval.of(dateTime.millisOfSecond().get());
case SECOND:
return ExprEval.of(dateTime.secondOfMinute().get());
case MINUTE:
return ExprEval.of(dateTime.minuteOfHour().get());
case HOUR:
return ExprEval.of(dateTime.hourOfDay().get());
case DAY:
return ExprEval.of(dateTime.dayOfMonth().get());
case DOW:
return ExprEval.of(dateTime.dayOfWeek().get());
case ISODOW:
return ExprEval.of(dateTime.dayOfWeek().get());
case DOY:
return ExprEval.of(dateTime.dayOfYear().get());
case WEEK:
return ExprEval.of(dateTime.weekOfWeekyear().get());
case MONTH:
return ExprEval.of(dateTime.monthOfYear().get());
case QUARTER:
return ExprEval.of((dateTime.monthOfYear().get() - 1) / 3 + 1);
case YEAR:
return ExprEval.of(dateTime.year().get());
case ISOYEAR:
return ExprEval.of(dateTime.year().get());
case DECADE:
// The year field divided by 10, See https://www.postgresql.org/docs/10/functions-datetime.html
return ExprEval.of(dateTime.year().get() / 10);
case CENTURY:
return ExprEval.of(Math.ceil((double) dateTime.year().get() / 100));
case MILLENNIUM:
// Years in the 1900s are in the second millennium. The third millennium started January 1, 2001.
// See https://www.postgresql.org/docs/10/functions-datetime.html
return ExprEval.of(Math.ceil((double) dateTime.year().get() / 1000));
default:
throw TimestampExtractExprMacro.this.validationFailed("unhandled unit[%s]", unit);
}
@Nonnull
@Override
public ExprEval eval(final ObjectBinding bindings)
{
Object val = args.get(0).eval(bindings).value();
if (val == null) {
// Return null if the argument if null.
return ExprEval.of(null);
}
final DateTime dateTime = new DateTime(val, chronology);
return getExprEval(dateTime, unit);
}

@Override
public Expr visit(Shuttle shuttle)
{
return shuttle.visit(apply(shuttle.visitAll(args)));
}
@Override
public Expr visit(Shuttle shuttle)
{
return shuttle.visit(apply(shuttle.visitAll(args)));
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
switch (unit) {
case CENTURY:
case MILLENNIUM:
return ExpressionType.DOUBLE;
default:
return ExpressionType.LONG;
}
}
@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
return getOutputExpressionType(unit);
}

@Override
public String stringify()
{
return stringifyExpr(args);
}
}

public class TimestampExtractDynamicExpr extends ExprMacroTable.BaseScalarMacroFunctionExpr
{
private final Unit unit;

@Override
public String stringify()
{
if (args.size() > 2) {
return StringUtils.format(
"%s(%s, %s, %s)",
FN_NAME,
arg.stringify(),
args.get(1).stringify(),
args.get(2).stringify()
);
}
return StringUtils.format("%s(%s, %s)", FN_NAME, arg.stringify(), args.get(1).stringify());
private TimestampExtractDynamicExpr(final List<Expr> args, final Unit unit)
{
super(FN_NAME, args);
this.unit = unit;
}

@Nonnull
@Override
public ExprEval eval(final ObjectBinding bindings)
{
Object val = args.get(0).eval(bindings).value();
if (val == null) {
// Return null if the argument if null.
return ExprEval.of(null);
}
final ISOChronology chronology = computeChronology(args, bindings);
final DateTime dateTime = new DateTime(val, chronology);
return getExprEval(dateTime, unit);
}

@Override
public Expr visit(Shuttle shuttle)
{
return shuttle.visit(apply(shuttle.visitAll(args)));
}

@Nullable
@Override
public ExpressionType getOutputType(InputBindingInspector inspector)
{
return getOutputExpressionType(unit);
}

return new TimestampExtractExpr(arg);
@Override
public String stringify()
{
return stringifyExpr(args);
}
}
}
28 changes: 28 additions & 0 deletions ...essing/src/test/java/org/apache/druid/query/expression/TimestampExtractExprMacroTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,10 +20,13 @@
package org.apache.druid.query.expression;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.math.expr.Expr;
import org.apache.druid.math.expr.ExprEval;
import org.apache.druid.math.expr.ExpressionType;
import org.apache.druid.math.expr.InputBindings;
import org.apache.druid.math.expr.Parser;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
Expand Down Expand Up @@ -100,4 +103,29 @@ public void testApplyExtractMilleniumShouldBeThirdMilleniumIn2001()
));
Assert.assertEquals(3, expression.eval(InputBindings.nilBindings()).asInt());
}

@Test
public void testApplyExtractDowWithTimeZoneShouldBeFriday()
{
Expr expression = target.apply(
ImmutableList.of(
ExprEval.of("2023-12-15").toExpr(),
ExprEval.of(TimestampExtractExprMacro.Unit.DOW.toString()).toExpr(),
ExprEval.of("UTC").toExpr()
));
Assert.assertEquals(5, expression.eval(InputBindings.nilBindings()).asInt());
}

@Test
public void testApplyExtractDowWithDynamicTimeZoneShouldBeFriday()
{
Expr expression = Parser.parse("timestamp_extract(time, 'DOW', timezone)", TestExprMacroTable.INSTANCE);
Expr.ObjectBinding bindings = InputBindings.forInputSuppliers(
ImmutableMap.of(
"time", InputBindings.inputSupplier(ExpressionType.STRING, () -> "2023-12-15"),
"timezone", InputBindings.inputSupplier(ExpressionType.STRING, () -> "UTC")
)
);
Assert.assertEquals(5, expression.eval(bindings).asInt());
}
}
Loading

0 comments on commit a2e65e6

Please sign in to comment.