From a2e65e6a89fd7b85b84f89b7fd7e7f5721bb05f7 Mon Sep 17 00:00:00 2001 From: AlbericByte <132106917+AlbericByte@users.noreply.github.com> Date: Wed, 20 Dec 2023 22:27:52 -0800 Subject: [PATCH] Support to pass dynamic values to timestamp Extract function (#15586) 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. --- .../apache/druid/math/expr/NamedFunction.java | 1 - .../expression/TimestampExtractExprMacro.java | 281 +++++++++++------- .../TimestampExtractExprMacroTest.java | 28 ++ .../TimeExtractOperatorConversion.java | 42 ++- .../calcite/expression/ExpressionsTest.java | 60 ++-- 5 files changed, 272 insertions(+), 140 deletions(-) diff --git a/processing/src/main/java/org/apache/druid/math/expr/NamedFunction.java b/processing/src/main/java/org/apache/druid/math/expr/NamedFunction.java index 574535ac684..45902834b34 100644 --- a/processing/src/main/java/org/apache/druid/math/expr/NamedFunction.java +++ b/processing/src/main/java/org/apache/druid/math/expr/NamedFunction.java @@ -124,7 +124,6 @@ public interface NamedFunction } } - /** * Helper method for implementors performing validation to check that the argument list is some expected size. * diff --git a/processing/src/main/java/org/apache/druid/query/expression/TimestampExtractExprMacro.java b/processing/src/main/java/org/apache/druid/query/expression/TimestampExtractExprMacro.java index ca8b90f3676..2876d831401 100644 --- a/processing/src/main/java/org/apache/druid/query/expression/TimestampExtractExprMacro.java +++ b/processing/src/main/java/org/apache/druid/query/expression/TimestampExtractExprMacro.java @@ -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; @@ -64,6 +65,87 @@ public class TimestampExtractExprMacro implements ExprMacroTable.ExprMacro 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 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 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 args) { @@ -73,121 +155,106 @@ public class TimestampExtractExprMacro implements ExprMacroTable.ExprMacro 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 args, final Unit unit, final ISOChronology chronology) { - private TimestampExtractExpr(Expr arg) - { - super(FN_NAME, arg); - } - - @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); - } - } - - @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; - } - } - - @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()); - } + super(FN_NAME, args); + this.unit = unit; + this.chronology = chronology; } - return new TimestampExtractExpr(arg); + @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))); + } + + @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; + + private TimestampExtractDynamicExpr(final List 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); + } + + @Override + public String stringify() + { + return stringifyExpr(args); + } } } diff --git a/processing/src/test/java/org/apache/druid/query/expression/TimestampExtractExprMacroTest.java b/processing/src/test/java/org/apache/druid/query/expression/TimestampExtractExprMacroTest.java index c5f3e0fdf6e..b06608ae228 100644 --- a/processing/src/test/java/org/apache/druid/query/expression/TimestampExtractExprMacroTest.java +++ b/processing/src/test/java/org/apache/druid/query/expression/TimestampExtractExprMacroTest.java @@ -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; @@ -100,4 +103,29 @@ public class TimestampExtractExprMacroTest )); 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()); + } } diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/TimeExtractOperatorConversion.java b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/TimeExtractOperatorConversion.java index 6a36bf4dad0..721969cd60c 100644 --- a/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/TimeExtractOperatorConversion.java +++ b/sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/TimeExtractOperatorConversion.java @@ -21,6 +21,7 @@ package org.apache.druid.sql.calcite.expression.builtin; import com.google.common.collect.ImmutableList; import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexInputRef; import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlFunction; @@ -65,6 +66,23 @@ public class TimeExtractOperatorConversion implements SqlOperatorConversion ); } + public static DruidExpression applyTimeExtract( + final DruidExpression timeExpression, + final TimestampExtractExprMacro.Unit unit, + final DruidExpression timeZoneExpression + ) + { + return DruidExpression.ofFunctionCall( + timeExpression.getDruidType(), + "timestamp_extract", + ImmutableList.of( + timeExpression, + DruidExpression.ofStringLiteral(unit.name()), + timeZoneExpression + ) + ); + } + @Override public SqlFunction calciteOperator() { @@ -89,13 +107,23 @@ public class TimeExtractOperatorConversion implements SqlOperatorConversion StringUtils.toUpperCase(RexLiteral.stringValue(call.getOperands().get(1))) ); - final DateTimeZone timeZone = OperatorConversions.getOperandWithDefault( - call.getOperands(), - 2, - operand -> DateTimes.inferTzFromString(RexLiteral.stringValue(operand)), - plannerContext.getTimeZone() - ); + if (call.getOperands().size() > 2 && call.getOperands().get(2) instanceof RexInputRef) { + final RexNode timeZoneArg = call.getOperands().get(2); + final DruidExpression timeZoneExpression = Expressions.toDruidExpression( + plannerContext, + rowSignature, + timeZoneArg + ); + return applyTimeExtract(timeExpression, unit, timeZoneExpression); + } else { + final DateTimeZone timeZone = OperatorConversions.getOperandWithDefault( + call.getOperands(), + 2, + operand -> DateTimes.inferTzFromString(RexLiteral.stringValue(operand)), + plannerContext.getTimeZone() + ); - return applyTimeExtract(timeExpression, unit, timeZone); + return applyTimeExtract(timeExpression, unit, timeZone); + } } } diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionsTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionsTest.java index 2ca9323b902..0273e7d2aa9 100644 --- a/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionsTest.java +++ b/sql/src/test/java/org/apache/druid/sql/calcite/expression/ExpressionsTest.java @@ -99,34 +99,33 @@ public class ExpressionsTest extends CalciteTestBase .add("newliney", ColumnType.STRING) .add("tstr", ColumnType.STRING) .add("dstr", ColumnType.STRING) + .add("timezone", ColumnType.STRING) .build(); private static final Map BINDINGS = ImmutableMap.builder() - .put( - "t", - DateTimes.of("2000-02-03T04:05:06").getMillis() - ) - .put("a", 10) - .put("b", 25) - .put("p", 3) - .put("x", 2.25) - .put("y", 3.0) - .put("z", -2.25) - .put("o", 0) - .put("nan", Double.NaN) - .put("inf", Double.POSITIVE_INFINITY) - .put("-inf", Double.NEGATIVE_INFINITY) - .put("fnan", Float.NaN) - .put("finf", Float.POSITIVE_INFINITY) - .put("-finf", Float.NEGATIVE_INFINITY) - .put("s", "foo") - .put("hexstr", "EF") - .put("intstr", "-100") - .put("spacey", " hey there ") - .put("newliney", "beep\nboop") - .put("tstr", "2000-02-03 04:05:06") - .put("dstr", "2000-02-03") - .build(); + .put("t", DateTimes.of("2000-02-03T04:05:06").getMillis()) + .put("a", 10) + .put("b", 25) + .put("p", 3) + .put("x", 2.25) + .put("y", 3.0) + .put("z", -2.25) + .put("o", 0) + .put("nan", Double.NaN) + .put("inf", Double.POSITIVE_INFINITY) + .put("-inf", Double.NEGATIVE_INFINITY) + .put("fnan", Float.NaN) + .put("finf", Float.POSITIVE_INFINITY) + .put("-finf", Float.NEGATIVE_INFINITY) + .put("s", "foo") + .put("hexstr", "EF") + .put("intstr", "-100") + .put("spacey", " hey there ") + .put("newliney", "beep\nboop") + .put("tstr", "2000-02-03 04:05:06") + .put("dstr", "2000-02-03") + .put("timezone", "America/Los_Angeles") + .build(); private ExpressionTestHelper testHelper; @@ -1844,6 +1843,17 @@ public class ExpressionsTest extends CalciteTestBase makeExpression(ColumnType.LONG, "timestamp_extract(\"t\",'DAY','America/Los_Angeles')"), 2L ); + + testHelper.testExpressionString( + new TimeExtractOperatorConversion().calciteOperator(), + ImmutableList.of( + testHelper.makeInputRef("t"), + testHelper.makeLiteral("DAY"), + testHelper.makeInputRef("timezone") + ), + makeExpression(ColumnType.LONG, "timestamp_extract(\"t\",'DAY',\"timezone\")"), + 2L + ); } @Test