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]

Reply via email to