This is an automated email from the ASF dual-hosted git repository.
abhishek pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new a2e65e6a89f Support to pass dynamic values to timestamp Extract
function (#15586)
a2e65e6a89f is described below
commit a2e65e6a89fd7b85b84f89b7fd7e7f5721bb05f7
Author: AlbericByte <[email protected]>
AuthorDate: Wed Dec 20 22:27:52 2023 -0800
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.
---
.../org/apache/druid/math/expr/NamedFunction.java | 1 -
.../expression/TimestampExtractExprMacro.java | 269 +++++++++++++--------
.../expression/TimestampExtractExprMacroTest.java | 28 +++
.../builtin/TimeExtractOperatorConversion.java | 42 +++-
.../sql/calcite/expression/ExpressionsTest.java | 60 +++--
5 files changed, 266 insertions(+), 134 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<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)
{
@@ -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<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);
+ }
}
}
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<String, Object> BINDINGS = ImmutableMap.<String,
Object>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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]