Add support for EXTRACT syntax and converts it to appropriate Pinot expression

pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java

@@ -96,6 +96,8 @@ public enum TransformFunctionType {
   LOOKUP("lookUp"),
   GROOVY("groovy"),
 
+  EXTRACT("extract"),
+
   // Regexp functions
   REGEXP_EXTRACT("regexpExtract"),
 

pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java

@@ -669,6 +669,8 @@ private static Expression toExpression(SqlNode node) {
           return RequestUtils.getIdentifierExpression(((SqlIdentifier) node).getSimple());
         }
         return RequestUtils.getIdentifierExpression(node.toString());
+      case INTERVAL_QUALIFIER:
+        return RequestUtils.getLiteralExpression(node.toString());
       case LITERAL:
         return RequestUtils.getLiteralExpression((SqlLiteral) node);
       case AS:

pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java

@@ -184,6 +184,38 @@ public void testQuotedStrings() {
             .getIdentifier().getName(), "Martha''s Vineyard");
   }
 
+  @Test
+  public void testExtract() {
+    {
+      // Case 1 -- Year and date format ('2017-06-15')
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT EXTRACT(YEAR FROM '2017-06-15')");
+      Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
+      Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(), "YEAR");
+      Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(), "2017-06-15");
+    }
+    {
+      // Case 2 -- date format ('2017-06-15 09:34:21')
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT EXTRACT(YEAR FROM '2017-06-15 09:34:21')");
+      Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
+      Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(), "YEAR");
+      Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(), "2017-06-15 09:34:21");
+    }
+    {
+      // Case 3 -- Month
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT EXTRACT(MONTH FROM '2017-06-15')");
+      Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
+      Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(), "MONTH");
+      Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(), "2017-06-15");
+    }
+    {
+      // Case 4 -- Day
+      PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery("SELECT EXTRACT(DAY FROM '2017-06-15')");
+      Function function = pinotQuery.getSelectList().get(0).getFunctionCall();
+      Assert.assertEquals(function.getOperands().get(0).getLiteral().getStringValue(), "DAY");
+      Assert.assertEquals(function.getOperands().get(1).getLiteral().getStringValue(), "2017-06-15");
+    }
+  }
+
   @Test
   public void testFilterClauses() {
     {
@@ -274,7 +306,7 @@ public void testFilterClauses() {
 
     {
       PinotQuery pinotQuery =
-          CalciteSqlParser.compileToPinotQuery("select * from vegetable where startsWith(g, " + "'str')");
+          CalciteSqlParser.compileToPinotQuery("select * from vegetable where startsWith(g, 'str')");
       Function func = pinotQuery.getFilterExpression().getFunctionCall();
       Assert.assertEquals(func.getOperator(), FilterKind.EQUALS.name());
       Assert.assertEquals(func.getOperands().get(0).getFunctionCall().getOperator(), "startswith");
@@ -283,7 +315,7 @@ public void testFilterClauses() {
 
     {
       PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(
-          "select * from vegetable where startsWith(g, " + "'str')=true and startsWith(f, 'str')");
+          "select * from vegetable where startsWith(g, 'str')=true and startsWith(f, 'str')");
       Function func = pinotQuery.getFilterExpression().getFunctionCall();
       Assert.assertEquals(func.getOperator(), FilterKind.AND.name());
       List<Expression> operands = func.getOperands();
@@ -302,7 +334,7 @@ public void testFilterClauses() {
 
     {
       PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(
-          "select * from vegetable where (startsWith(g, " + "'str')=true and startsWith(f, 'str')) AND (e and d=true)");
+          "select * from vegetable where (startsWith(g, 'str')=true and startsWith(f, 'str')) AND (e and d=true)");
       Function func = pinotQuery.getFilterExpression().getFunctionCall();
       Assert.assertEquals(func.getOperator(), FilterKind.AND.name());
       List<Expression> operands = func.getOperands();
@@ -609,7 +641,7 @@ public void testGroupbys() {
       pinotQuery = CalciteSqlParser.compileToPinotQuery(
           "select concat(upper(playerName), lower(teamID), '-') playerTeam, "
               + "upper(league) leagueUpper, count(playerName) cnt from baseballStats group by playerTeam, lower"
-              + "(teamID), leagueUpper " + "having cnt > 1 order by cnt desc limit 10");
+              + "(teamID), leagueUpper having cnt > 1 order by cnt desc limit 10");
     } catch (SqlCompilationException e) {
       throw e;
     }
@@ -724,8 +756,7 @@ public void testQueryOptions() {
       Assert.assertTrue(e.getCause().getMessage().contains("OPTION"));
     }
     try {
-      CalciteSqlParser.compileToPinotQuery(
-          "select * from vegetables where name <> 'Brussels OPTION (delicious=yes)");
+      CalciteSqlParser.compileToPinotQuery("select * from vegetables where name <> 'Brussels OPTION (delicious=yes)");
     } catch (SqlCompilationException e) {
       Assert.assertTrue(e.getCause() instanceof ParseException);
     }
@@ -770,24 +801,25 @@ public void testQuerySetOptions() {
 
     // test invalid options
     try {
-      CalciteSqlParser.compileToPinotQuery("select * from vegetables SET delicious='yes', foo='1234' "
-          + "where name <> 'Brussels sprouts'");
+      CalciteSqlParser.compileToPinotQuery(
+          "select * from vegetables SET delicious='yes', foo='1234' where name <> 'Brussels sprouts'");
       Assert.fail("SQL should not be compiled");
     } catch (SqlCompilationException sce) {
       // expected.
     }
 
     try {
-      CalciteSqlParser.compileToPinotQuery("select * from vegetables where name <> 'Brussels sprouts'; "
-          + "SET (delicious='yes', foo=1234)");
+      CalciteSqlParser.compileToPinotQuery(
+          "select * from vegetables where name <> 'Brussels sprouts'; SET (delicious='yes', foo=1234)");
       Assert.fail("SQL should not be compiled");
     } catch (SqlCompilationException sce) {
       // expected.
     }
 
     try {
-      CalciteSqlParser.compileToPinotQuery("select * from vegetables where name <> 'Brussels sprouts'; "
-          + "SET (delicious='yes', foo=1234); select * from meat");
+      CalciteSqlParser.compileToPinotQuery(
+          "select * from vegetables where name <> 'Brussels sprouts'; SET (delicious='yes', foo=1234); select * from "
+              + "meat");
       Assert.fail("SQL should not be compiled");
     } catch (SqlCompilationException sce) {
       // expected.
@@ -2522,7 +2554,7 @@ public void testInvalidQueryWithSemicolon() {
 
     // Query having multiple SQL statements
     Assert.expectThrows(SqlCompilationException.class, () -> CalciteSqlParser.compileToPinotQuery(
-        "SELECT col1, count(*) FROM foo GROUP BY col1; SELECT col2," + "count(*) FROM foo GROUP BY col2"));
+        "SELECT col1, count(*) FROM foo GROUP BY col1; SELECT col2, count(*) FROM foo GROUP BY col2"));
 
     // Query having multiple SQL statements with trailing and leading whitespaces
     Assert.expectThrows(SqlCompilationException.class, () -> CalciteSqlParser.compileToPinotQuery(

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunction.java

@@ -0,0 +1,111 @@
+/**
+ * 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.core.operator.transform.function;
+
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.core.operator.blocks.ProjectionBlock;
+import org.apache.pinot.core.operator.transform.TransformResultMetadata;
+import org.apache.pinot.segment.spi.datasource.DataSource;
+import org.joda.time.Chronology;
+import org.joda.time.DateTimeField;
+import org.joda.time.chrono.ISOChronology;
+
+
+public class ExtractTransformFunction extends BaseTransformFunction {
+  public static final String FUNCTION_NAME = "extract";
+  private TransformFunction _mainTransformFunction;
+  protected Field _field;
+  protected Chronology _chronology = ISOChronology.getInstanceUTC();
+
+  private enum Field {
+    YEAR, MONTH, DAY, HOUR, MINUTE, SECOND
+  }
+
+  @Override
+  public String getName() {
+    return FUNCTION_NAME;
+  }
+
+  @Override
+  public void init(List<TransformFunction> arguments, Map<String, DataSource> dataSourceMap) {
+    if (arguments.size() != 2) {
+      throw new IllegalArgumentException("Exactly 2 arguments are required for EXTRACT transform function");
+    }
+
+    _field = Field.valueOf(((LiteralTransformFunction) arguments.get(0)).getLiteral());
+
+    _mainTransformFunction = arguments.get(1);
+  }
+
+  @Override
+  public TransformResultMetadata getResultMetadata() {
+    return INT_SV_NO_DICTIONARY_METADATA;
+  }
+
+  @Override
+  public int[] transformToIntValuesSV(ProjectionBlock projectionBlock) {
+    int numDocs = projectionBlock.getNumDocs();
+
+    if (_intValuesSV == null || _intValuesSV.length < numDocs) {
+      _intValuesSV = new int[numDocs];
+    }
+
+    long[] timestamps = _mainTransformFunction.transformToLongValuesSV(projectionBlock);
+
+    convert(timestamps, numDocs, _intValuesSV);
+
+    return _intValuesSV;
+  }
+
+  private void convert(long[] timestamps, int numDocs, int[] output) {
+    for (int i = 0; i < numDocs; i++) {
+      DateTimeField accessor;
+
+      switch (_field) {
+        case YEAR:
+          accessor = _chronology.year();
+          output[i] = accessor.get(timestamps[i]);
+          break;
+        case MONTH:
+          accessor = _chronology.monthOfYear();
+          output[i] = accessor.get(timestamps[i]);
+          break;
+        case DAY:
+          accessor = _chronology.dayOfMonth();
+          output[i] = accessor.get(timestamps[i]);
+          break;
+        case HOUR:
+          accessor = _chronology.hourOfDay();
+          output[i] = accessor.get(timestamps[i]);
+          break;
+        case MINUTE:
+          accessor = _chronology.minuteOfHour();
+          output[i] = accessor.get(timestamps[i]);
+          break;
+        case SECOND:
+          accessor = _chronology.secondOfMinute();
+          output[i] = accessor.get(timestamps[i]);
+          break;
+        default:
+          throw new IllegalArgumentException("Unsupported FIELD type");
+      }
+    }
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java

@@ -132,6 +132,8 @@ private static Map<String, Class<? extends TransformFunction>> createRegistry()
     typeToImplementation.put(TransformFunctionType.INIDSET, InIdSetTransformFunction.class);
     typeToImplementation.put(TransformFunctionType.LOOKUP, LookupTransformFunction.class);
 
+    typeToImplementation.put(TransformFunctionType.EXTRACT, ExtractTransformFunction.class);
+
     // Regexp functions
     typeToImplementation.put(TransformFunctionType.REGEXP_EXTRACT,
         RegexpExtractTransformFunction.class);

pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/ExtractTransformFunctionTest.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.core.operator.transform.function;
+
+import java.util.function.LongToIntFunction;
+import org.apache.pinot.common.function.scalar.DateTimeFunctions;
+import org.apache.pinot.common.request.context.ExpressionContext;
+import org.apache.pinot.common.request.context.RequestContextUtils;
+import org.testng.Assert;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+
+public class ExtractTransformFunctionTest extends BaseTransformFunctionTest {
+
+  @DataProvider
+  public static Object[][] testCases() {
+    return new Object[][]{
+        //@formatter:off
+        {"year", (LongToIntFunction) DateTimeFunctions::year},
+        {"month", (LongToIntFunction) DateTimeFunctions::month},
+        {"day", (LongToIntFunction) DateTimeFunctions::dayOfMonth},
+        {"hour", (LongToIntFunction) DateTimeFunctions::hour},
+        {"minute", (LongToIntFunction) DateTimeFunctions::minute},
+        {"second", (LongToIntFunction) DateTimeFunctions::second},
+        // TODO: Need to add timezone_hour and timezone_minute
+//      "timezone_hour",
+//      "timezone_minute",
+      //@formatter:on
+    };
+  }
+
+  @Test(dataProvider = "testCases")
+  public void testExtractTransformFunction(String field, LongToIntFunction expected) {
+    // NOTE: functionality of ExtractTransformFunction is covered in ExtractTransformFunctionTest
+    // SELECT EXTRACT(YEAR FROM '2017-10-10')
+
+    ExpressionContext expression =
+        RequestContextUtils.getExpression(String.format("extract(%s FROM %s)", field, TIMESTAMP_COLUMN));
+
+    TransformFunction transformFunction = TransformFunctionFactory.get(expression, _dataSourceMap);
+    Assert.assertTrue(transformFunction instanceof ExtractTransformFunction);
+    int[] value = transformFunction.transformToIntValuesSV(_projectionBlock);
+    for (int i = 0; i < _projectionBlock.getNumDocs(); i++) {
+      assertEquals(value[i], expected.applyAsInt(_timeValues[i]));
+    }
+  }
+}