Datetime transform functions

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

@@ -18,7 +18,10 @@
  */
 package org.apache.pinot.common.function;
 
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -69,6 +72,18 @@ public enum TransformFunctionType {
   TIMECONVERT("timeConvert"),
   DATETIMECONVERT("dateTimeConvert"),
   DATETRUNC("dateTrunc"),
+  YEAR("year"),
+  YEAR_OF_WEEK("yearOfWeek", "yow"),
+  QUARTER("quarter"),
+  MONTH_OF_YEAR("monthOfYear", "month"),
+  WEEK_OF_YEAR("weekOfYear", "week"),
+  DAY_OF_YEAR("dayOfYear", "doy"),
+  DAY_OF_MONTH("dayOfMonth", "day"),
+  DAY_OF_WEEK("dayOfWeek", "dow"),
+  HOUR("hour"),
+  MINUTE("minute"),
+  SECOND("second"),
+  MILLISECOND("millisecond"),
   ARRAYLENGTH("arrayLength"),
   ARRAYAVERAGE("arrayAverage"),
   ARRAYMIN("arrayMin"),
@@ -112,14 +127,24 @@ public enum TransformFunctionType {
   GEOTOH3("geoToH3");
 
   private static final Set<String> NAMES = Arrays.stream(values()).flatMap(
-      func -> Stream.of(func.getName(), StringUtils.remove(func.getName(), '_').toUpperCase(),
-          StringUtils.remove(func.getName(), '_').toLowerCase(), func.getName().toUpperCase(),
-          func.getName().toLowerCase(), func.name(), func.name().toLowerCase())).collect(Collectors.toSet());
+      func -> func.getAliases().stream().flatMap(name -> Stream.of(name, StringUtils.remove(name, '_').toUpperCase(),
+          StringUtils.remove(name, '_').toLowerCase(), name.toUpperCase(),
+          name.toLowerCase()))).collect(Collectors.toSet());
 
   private final String _name;
+  private final List<String> _aliases;
 
-  TransformFunctionType(String name) {
+  TransformFunctionType(String name, String... aliases) {
     _name = name;
+    List<String> all = new ArrayList<>(aliases.length + 2);
+    all.add(name);
+    all.add(name());
+    all.addAll(Arrays.asList(aliases));
+    _aliases = Collections.unmodifiableList(all);
+  }
+
+  public List<String> getAliases() {
+    return _aliases;
   }
 
   public static boolean isTransformFunction(String functionName) {

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

@@ -0,0 +1,249 @@
+/**
+ * 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 com.google.common.base.Preconditions;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.function.TransformFunctionType;
+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.apache.pinot.spi.data.FieldSpec;
+import org.joda.time.DateTime;
+import org.joda.time.DateTimeZone;
+
+
+public abstract class DateTimeTransformFunction extends BaseTransformFunction {
+
+  private static final TransformResultMetadata METADATA =
+      new TransformResultMetadata(FieldSpec.DataType.INT, true, false);
+  private final String _name;
+  private TransformFunction _timestampsFunction;
+  protected DateTimeZone _zoneId;
+
+  protected DateTimeTransformFunction(String name) {
+    _name = name;
+  }
+
+  @Override
+  public void init(List<TransformFunction> arguments, Map<String, DataSource> dataSourceMap) {
+    Preconditions.checkArgument(!arguments.isEmpty() && arguments.size() <= 2, "%s takes one or two arguments", _name);
+    _timestampsFunction = arguments.get(0);
+    if (arguments.size() == 2) {
+      Preconditions.checkArgument(arguments.get(1) instanceof LiteralTransformFunction,
+          "zoneId parameter %s must be a literal",
+          _name);
+      _zoneId = DateTimeZone.forID(((LiteralTransformFunction) arguments.get(1)).getLiteral());
+    } else {
+      _zoneId = DateTimeZone.UTC;
+    }
+  }
+
+  @Override
+  public String getName() {
+    return _name;
+  }
+
+  @Override
+  public TransformResultMetadata getResultMetadata() {
+    return METADATA;
+  }
+
+  @Override
+  public int[] transformToIntValuesSV(ProjectionBlock projectionBlock) {
+    int numDocs = projectionBlock.getNumDocs();
+    if (_intValuesSV == null || _intValuesSV.length < numDocs) {
+      _intValuesSV = new int[numDocs];
+    }
+    long[] timestamps = _timestampsFunction.transformToLongValuesSV(projectionBlock);
+    convert(timestamps, numDocs, _intValuesSV);
+    return _intValuesSV;
+  }
+
+  protected abstract void convert(long[] timestamps, int numDocs, int[] output);
+
+  public static final class Year extends DateTimeTransformFunction {
+
+    public Year() {
+      super(TransformFunctionType.YEAR.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getYear();
+      }
+    }
+  }
+
+  public static final class YearOfWeek extends DateTimeTransformFunction {
+
+    public YearOfWeek() {
+      super(TransformFunctionType.YEAR_OF_WEEK.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getWeekyear();
+      }
+    }
+  }
+
+  public static final class Month extends DateTimeTransformFunction {
+
+    public Month() {
+      super(TransformFunctionType.MONTH_OF_YEAR.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getMonthOfYear();
+      }
+    }
+  }
+
+  public static final class WeekOfYear extends DateTimeTransformFunction {
+
+    public WeekOfYear() {
+      super(TransformFunctionType.WEEK_OF_YEAR.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getWeekOfWeekyear();
+      }
+    }
+  }
+
+  public static final class DayOfYear extends DateTimeTransformFunction {
+
+    public DayOfYear() {
+      super(TransformFunctionType.DAY_OF_YEAR.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getDayOfYear();
+      }
+    }
+  }
+
+  public static final class DayOfMonth extends DateTimeTransformFunction {
+
+    public DayOfMonth() {
+      super(TransformFunctionType.DAY_OF_MONTH.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getDayOfMonth();
+      }
+    }
+  }
+
+  public static final class DayOfWeek extends DateTimeTransformFunction {
+
+    public DayOfWeek() {
+      super(TransformFunctionType.DAY_OF_WEEK.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getDayOfWeek();
+      }
+    }
+  }
+
+  public static final class Hour extends DateTimeTransformFunction {
+
+    public Hour() {
+      super(TransformFunctionType.HOUR.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getHourOfDay();
+      }
+    }
+  }
+
+  public static final class Minute extends DateTimeTransformFunction {
+
+    public Minute() {
+      super(TransformFunctionType.MINUTE.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getMinuteOfHour();
+      }
+    }
+  }
+
+  public static final class Second extends DateTimeTransformFunction {
+
+    public Second() {
+      super(TransformFunctionType.SECOND.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getSecondOfMinute();
+      }
+    }
+  }
+
+  public static final class Millisecond extends DateTimeTransformFunction {
+
+    public Millisecond() {
+      super(TransformFunctionType.MILLISECOND.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = new DateTime(timestamps[i], _zoneId).getMillisOfSecond();
+      }
+    }
+  }
+
+  public static final class Quarter extends DateTimeTransformFunction {
+
+    public Quarter() {
+      super(TransformFunctionType.QUARTER.getName());
+    }
+
+    @Override
+    protected void convert(long[] timestamps, int numDocs, int[] output) {
+      for (int i = 0; i < numDocs; i++) {
+        output[i] = (new DateTime(timestamps[i], _zoneId).getMonthOfYear() - 1) / 3 + 1;
+      }
+    }
+  }
+}