From f52e7ec60e84902b1434c9acefad5a1fc5e29208 Mon Sep 17 00:00:00 2001
From: Almog Gavra <almog.gavra@gmail.com>
Date: Wed, 28 Sep 2022 07:54:11 -0700
Subject: [PATCH 1/3] add parsing for AT TIME ZONE

This commit adds the custom Calcite extension for
parsing AT TIME ZONE expressions, but does not yet
support executing these queries.
---
 pinot-common/src/main/codegen/config.fmpp     |  5 ++
 .../src/main/codegen/includes/parserImpls.ftl | 17 +++++
 .../sql/parsers/parser/SqlAtTimeZone.java     | 66 +++++++++++++++++++
 .../sql/parsers/CalciteSqlCompilerTest.java   | 57 ++++++++++++++++
 .../pinot/queries/TimestampQueriesTest.java   | 12 ++++
 5 files changed, 157 insertions(+)
 create mode 100644 pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java

diff --git a/pinot-common/src/main/codegen/config.fmpp b/pinot-common/src/main/codegen/config.fmpp
index 0b1f558e8985..e6955766bc70 100644
--- a/pinot-common/src/main/codegen/config.fmpp
+++ b/pinot-common/src/main/codegen/config.fmpp
@@ -528,6 +528,11 @@ data: {
       "SqlInsertFromFile()"
     ]
 
+    # List of custom function syntax to add
+    extraBinaryExpressions: [
+      "SqlAtTimeZone"
+    ]
+
     # List of files in @includes directory that have parser method
     implementationFiles: [
       "parserImpls.ftl"
diff --git a/pinot-common/src/main/codegen/includes/parserImpls.ftl b/pinot-common/src/main/codegen/includes/parserImpls.ftl
index 3b85f19aff78..396e256e7550 100644
--- a/pinot-common/src/main/codegen/includes/parserImpls.ftl
+++ b/pinot-common/src/main/codegen/includes/parserImpls.ftl
@@ -102,3 +102,20 @@ SqlNodeList SqlStmtsEof() :
         return stmts;
     }
 }
+
+void SqlAtTimeZone(List<Object> list, ExprContext exprContext, Span s) :
+{
+    List<Object> list2;
+    SqlOperator op;
+}
+{
+    {
+        checkNonQueryExpression(exprContext);
+        s.clear().add(this);
+    }
+    <AT> <TIME> <ZONE> { op = SqlAtTimeZone.INSTANCE; }
+    list2 = Expression2(ExprContext.ACCEPT_SUB_QUERY) {
+        list.add(new SqlParserUtil.ToTreeListItem(op, s.pos()));
+        list.addAll(list2);
+    }
+}
\ No newline at end of file
diff --git a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java
new file mode 100644
index 000000000000..38bfc277a104
--- /dev/null
+++ b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java
@@ -0,0 +1,66 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.sql.parsers.parser;
+
+import java.util.Arrays;
+import org.apache.calcite.sql.SqlCall;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.SqlSpecialOperator;
+import org.apache.calcite.sql.parser.SqlParserPos;
+import org.apache.calcite.sql.parser.SqlParserUtil;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.sql.type.SqlTypeName;
+
+import static java.util.Objects.requireNonNull;
+
+
+public class SqlAtTimeZone extends SqlSpecialOperator {
+
+  public static final SqlAtTimeZone INSTANCE = new SqlAtTimeZone();
+
+  private SqlAtTimeZone() {
+    super(
+        "AT_TIME_ZONE",
+        SqlKind.OTHER_FUNCTION,
+        32, // put this at the same precedence as LIKE (see SqlLikeOperator)
+        false,
+        opBinding -> opBinding.getTypeFactory().createSqlType(SqlTypeName.TIMESTAMP),
+        null,
+        OperandTypes.family(SqlTypeFamily.TIME)
+    );
+  }
+
+  @Override
+  public ReduceResult reduceExpr(int ordinal, TokenSequence list) {
+    SqlNode left = list.node(ordinal - 1);
+    SqlNode right = list.node(ordinal + 1);
+    return new ReduceResult(ordinal - 1,
+        ordinal + 2,
+        createCall(
+            SqlParserPos.sum(
+                Arrays.asList(requireNonNull(left, "left").getParserPosition(),
+                    requireNonNull(right, "right").getParserPosition(),
+                    list.pos(ordinal))),
+            left,
+            right));
+  }
+}
diff --git a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
index 5f69fd700106..c662af046866 100644
--- a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
@@ -2681,6 +2681,63 @@ public void testParserExtensionImpl() {
     Assert.assertEquals(sqlNodeAndOptions.getSqlType(), PinotSqlType.DML);
   }
 
+
+  @Test
+  public void shouldParseBasicAtTimeZoneExtension() {
+    // Given:
+    String sql = "SELECT ts AT TIME ZONE 'pst' FROM myTable;";
+
+    // When:
+    PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+
+    // Then:
+    Assert.assertEquals(pinotQuery.getSelectListSize(), 1);
+    Function fun = pinotQuery.getSelectList().get(0).getFunctionCall();
+    Assert.assertEquals(fun.operator, "attimezone");
+    Assert.assertEquals(fun.operands.size(), 2);
+    Assert.assertEquals(fun.operands.get(0).getIdentifier().name, "ts");
+    Assert.assertEquals(fun.operands.get(1).getLiteral().getStringValue(), "pst");
+  }
+
+  @Test
+  public void shouldParseNestedTimeExprAtTimeZoneExtension() {
+    // Given:
+    String sql = "SELECT ts + 123 AT TIME ZONE 'pst' FROM myTable;";
+
+    // When:
+    PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+
+    // Then:
+    Assert.assertEquals(pinotQuery.getSelectListSize(), 1);
+    Function fun = pinotQuery.getSelectList().get(0).getFunctionCall();
+    Assert.assertEquals(fun.operator, "attimezone");
+    Assert.assertEquals(fun.operands.size(), 2);
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operator, "plus");
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operands.size(), 2);
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operands.get(0).getIdentifier().getName(), "ts");
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operands.get(1).getLiteral().getLongValue(), 123L);
+    Assert.assertEquals(fun.operands.get(1).getLiteral().getStringValue(), "pst");
+  }
+
+  @Test
+  public void shouldParseOutsideExprAtTimeZoneExtension() {
+    // Given:
+    String sql = "SELECT ts AT TIME ZONE 'pst' > 123 FROM myTable;";
+
+    // When:
+    PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+
+    // Then:
+    Assert.assertEquals(pinotQuery.getSelectListSize(), 1);
+    Function fun = pinotQuery.getSelectList().get(0).getFunctionCall();
+    Assert.assertEquals(fun.operator, "GREATER_THAN");
+    Assert.assertEquals(fun.operands.size(), 2);
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operator, "attimezone");
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operands.size(), 2);
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operands.get(0).getIdentifier().getName(), "ts");
+    Assert.assertEquals(fun.operands.get(0).getFunctionCall().operands.get(1).getLiteral().getStringValue(), "pst");
+  }
+
   private static SqlNodeAndOptions testSqlWithCustomSqlParser(String sqlString) {
     try (StringReader inStream = new StringReader(sqlString)) {
       SqlParserImpl sqlParser = CalciteSqlParser.newSqlParser(inStream);
diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java
index 0bded079b372..c96c993ed2ca 100644
--- a/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java
+++ b/pinot-core/src/test/java/org/apache/pinot/queries/TimestampQueriesTest.java
@@ -40,6 +40,7 @@
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.apache.pinot.spi.data.Schema;
 import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.exception.BadQueryRequestException;
 import org.apache.pinot.spi.utils.ReadMode;
 import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
 import org.testng.annotations.AfterClass;
@@ -222,6 +223,17 @@ public void testQueries() {
     }
   }
 
+  @Test(
+      expectedExceptions = BadQueryRequestException.class,
+      expectedExceptionsMessageRegExp = ".*attimezone not found.*"
+  )
+  public void shouldThrowOnAtTimeZone() {
+    // this isn't yet implemented but the syntax is supported, make sure the
+    // degradation experience is clean
+    String query = "SELECT timestampColumn AT TIME ZONE 'pst' FROM testTable";
+    getBrokerResponse(query);
+  }
+
   @AfterClass
   public void tearDown()
       throws IOException {

From 370ae53f7ac77e7016a155278344e50d0400a19e Mon Sep 17 00:00:00 2001
From: Almog Gavra <almog.gavra@gmail.com>
Date: Wed, 28 Sep 2022 08:07:32 -0700
Subject: [PATCH 2/3] fix checkstyle

---
 .../org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java     | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java
index 38bfc277a104..fbeb01b7caae 100644
--- a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java
+++ b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/parser/SqlAtTimeZone.java
@@ -19,13 +19,10 @@
 package org.apache.pinot.sql.parsers.parser;
 
 import java.util.Arrays;
-import org.apache.calcite.sql.SqlCall;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlNode;
-import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.SqlSpecialOperator;
 import org.apache.calcite.sql.parser.SqlParserPos;
-import org.apache.calcite.sql.parser.SqlParserUtil;
 import org.apache.calcite.sql.type.OperandTypes;
 import org.apache.calcite.sql.type.SqlTypeFamily;
 import org.apache.calcite.sql.type.SqlTypeName;

From ac61ce70c6ffab4cb4b6b533ec93123dec815348 Mon Sep 17 00:00:00 2001
From: Almog Gavra <almog.gavra@gmail.com>
Date: Thu, 29 Sep 2022 09:13:10 -0700
Subject: [PATCH 3/3] add test and fix newline

---
 pinot-common/src/main/codegen/includes/parserImpls.ftl          | 2 +-
 .../test/java/org/apache/pinot/query/QueryCompilationTest.java  | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/pinot-common/src/main/codegen/includes/parserImpls.ftl b/pinot-common/src/main/codegen/includes/parserImpls.ftl
index 396e256e7550..449d8ab3b9a2 100644
--- a/pinot-common/src/main/codegen/includes/parserImpls.ftl
+++ b/pinot-common/src/main/codegen/includes/parserImpls.ftl
@@ -118,4 +118,4 @@ void SqlAtTimeZone(List<Object> list, ExprContext exprContext, Span s) :
         list.add(new SqlParserUtil.ToTreeListItem(op, s.pos()));
         list.addAll(list2);
     }
-}
\ No newline at end of file
+}
diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
index 0531cfaba856..9c220b443897 100644
--- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
+++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryCompilationTest.java
@@ -257,6 +257,8 @@ private Object[][] provideQueriesWithException() {
         new Object[]{"SELECT a.col1, SUM(a.col3) FROM a", "'a.col1' is not being grouped"},
         // empty IN clause fails compilation
         new Object[]{"SELECT a.col1 FROM a WHERE a.col1 IN ()", "Encountered \"\" at line"},
+        // AT TIME ZONE should fail
+        new Object[]{"SELECT a.col1 AT TIME ZONE 'PST' FROM a", "No match found for function signature AT_TIME_ZONE"},
     };
   }