Address review comments

Signed-off-by: currantw <taylor.curran@improving.com>

Author: currantw

integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLBuiltInDateTimeFunctionITSuite.scala

@@ -396,14 +396,14 @@ class FlintSparkPPLBuiltInDateTimeFunctionITSuite
     frame = sql(s"""
                    | source=$testTable
                    | | eval last_wednesday = relative_timestamp("-1d@w3")
-                   | | eval actual_days_ago = timestampdiff(DAY, last_wednesday, now() )
+                   | | eval actual_days_ago = timestampdiff(DAY, last_wednesday, now())
                    | | eval day_of_week = day_of_week(now())
                    | | eval expected_days_ago = case(day_of_week >= 4, day_of_week - 4 else day_of_week + 3)
-                   | | eval test_result = (expected_weeks_ago = actual_weeks_ago)
+                   | | eval test_result = (expected_days_ago = actual_days_ago)
                    | | fields test_result
                    | | head 1
                    | """.stripMargin)
-    assertSameRows(Seq(Row(1)), frame)
+    assertSameRows(Seq(Row(true)), frame)
   }
 
   // TODO #957: Support earliest

ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/SerializableUdf.java

@@ -28,8 +28,8 @@
 import java.net.UnknownHostException;
 import java.sql.Timestamp;
 import java.time.Instant;
-import java.time.LocalDateTime;
 import java.time.ZoneId;
+import java.time.ZonedDateTime;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -42,7 +42,7 @@
 
 public interface SerializableUdf {
 
-    abstract class SerializableAbstractFunction1<T1,R> extends AbstractFunction1<T1,R>
+    abstract class SerializableAbstractFunction1<T1, R> extends AbstractFunction1<T1, R>
             implements Serializable {
     }
 
@@ -212,22 +212,39 @@ public BigInteger apply(String ipAddress) {
     }
 
     /**
-     * Returns the {@link Timestamp} corresponding to the given relative string and current timestamp.
+     * Returns the {@link Instant} corresponding to the given relative string, current timestamp, and current time zone ID.
      * Throws {@link RuntimeException} if the relative string is not supported.
      */
-    Function3<String, Instant, String, Timestamp> relativeTimestampFunction = new SerializableAbstractFunction3<String, Instant, String, Timestamp>() {
+    Function3<String, Object, String, Instant> relativeTimestampFunction = new SerializableAbstractFunction3<String, Object, String, Instant>() {
+
         @Override
-        public Timestamp apply(String relativeString, Instant currentInstant, String zoneId) {
-            LocalDateTime currentLocalDateTime = LocalDateTime.ofInstant(currentInstant, ZoneId.of(zoneId));
-            LocalDateTime relativeLocalDateTime = TimeUtils.getRelativeLocalDateTime(relativeString, currentLocalDateTime);
-            return Timestamp.valueOf(relativeLocalDateTime);
+        public Instant apply(String relativeString, Object currentTimestamp, String zoneIdString) {
+
+            /// If `spark.sql.datetime.java8API.enabled` is set to `true`, [org.apache.spark.sql.types.TimestampType]
+            /// is converted to [Instant] by Catalyst; otherwise, [Timestamp] is used instead.
+            Instant currentInstant =
+                    currentTimestamp instanceof Timestamp
+                            ? ((Timestamp) currentTimestamp).toInstant()
+                            : (Instant) currentTimestamp;
+
+            /// The Spark session time zone (`spark.sql.session.timeZone`)
+            /// is used, which may be different from the system time zone.
+            ZoneId zoneId = ZoneId.of(zoneIdString);
+
+            /// Relative time calculations are performed using [ZonedDateTime] because offsets (e.g. one hour ago)
+            /// need to account for changes in the time zone offset (e.g. daylight savings time), while snaps (e.g.
+            /// start of previous Wednesday) need to account for the local date time.
+            ZonedDateTime currentDateTime = ZonedDateTime.ofInstant(currentInstant, zoneId);
+            ZonedDateTime relativeDateTime = TimeUtils.getRelativeZonedDateTime(relativeString, currentDateTime);
+
+            return relativeDateTime.toInstant();
         }
     };
 
     /**
      * Get the function reference according to its name
      *
-     * @param funcName      string representing function to retrieve.
+     * @param funcName string representing function to retrieve.
      * @return relevant ScalaUDF for given function name.
      */
     static ScalaUDF visit(String funcName, List<Expression> expressions) {
@@ -270,7 +287,7 @@ static ScalaUDF visit(String funcName, List<Expression> expressions) {
                         true);
             case "ip_to_int":
                 return new ScalaUDF(geoIpUtils.ipToInt,
-                        DataTypes.createDecimalType(38,0),
+                        DataTypes.createDecimalType(38, 0),
                         seq(expressions),
                         seq(),
                         Option.empty(),

ppl-spark-integration/src/main/java/org/opensearch/sql/expression/function/TimeUtils.java

@@ -10,9 +10,9 @@
 
 import java.time.DayOfWeek;
 import java.time.Duration;
-import java.time.LocalDateTime;
 import java.time.Month;
 import java.time.Period;
+import java.time.ZonedDateTime;
 import java.time.temporal.ChronoUnit;
 import java.util.HashMap;
 import java.util.Map;
@@ -81,7 +81,7 @@ public class TimeUtils {
             Map.entry("w6", DayOfWeek.SATURDAY));
 
     /**
-     * Returns the relative {@link LocalDateTime} corresponding to the given relative string and local date time.
+     * Returns the relative {@link ZonedDateTime} corresponding to the given relative string and zoned date time.
      * <p>
      * The relative time string has syntax {@code [+|-]<offset_time_integer><offset_time_unit>@<snap_time_unit>}, and
      * is made up of two optional components:
@@ -219,12 +219,12 @@ public class TimeUtils {
      *      </body>
      * </table>
      */
-    public static LocalDateTime getRelativeLocalDateTime(String relativeString, LocalDateTime localDateTime) {
+    public static ZonedDateTime getRelativeZonedDateTime(String relativeString, ZonedDateTime zonedDateTime) {
 
-        LocalDateTime relativeLocalDateTime = localDateTime;
+        ZonedDateTime relativeZonedDateTime = zonedDateTime;
 
         if (relativeString.equalsIgnoreCase(NOW)) {
-            return localDateTime;
+            return zonedDateTime;
         }
 
         Matcher matcher = RELATIVE_PATTERN.matcher(relativeString);
@@ -235,27 +235,27 @@ public static LocalDateTime getRelativeLocalDateTime(String relativeString, Loca
 
 
         if (matcher.group("offset") != null) {
-            relativeLocalDateTime = applyOffset(
-                    relativeLocalDateTime,
+            relativeZonedDateTime = applyOffset(
+                    relativeZonedDateTime,
                     matcher.group("offsetSign"),
                     matcher.group("offsetValue"),
                     matcher.group("offsetUnit"));
         }
 
         if (matcher.group("snap") != null) {
-            relativeLocalDateTime = applySnap(
-                    relativeLocalDateTime,
+            relativeZonedDateTime = applySnap(
+                    relativeZonedDateTime,
                     matcher.group("snapUnit"));
         }
 
-        return relativeLocalDateTime;
+        return relativeZonedDateTime;
     }
 
     /**
      * Applies the offset specified by the offset sign, value,
-     * and unit to the given local date time, and returns the result.
+     * and unit to the given zoned date time, and returns the result.
      */
-    private LocalDateTime applyOffset(LocalDateTime localDateTime, String offsetSign, String offsetValue, String offsetUnit) {
+    private ZonedDateTime applyOffset(ZonedDateTime zonedDateTime, String offsetSign, String offsetValue, String offsetUnit) {
 
         int offsetValueInt = Optional.ofNullable(offsetValue).map(Integer::parseInt).orElse(1);
         if (offsetSign.equals(NEGATIVE_SIGN)) {
@@ -271,46 +271,46 @@ private LocalDateTime applyOffset(LocalDateTime localDateTime, String offsetSign
 
         if (DURATION_FOR_TIME_UNIT_MAP.containsKey(offsetUnitLowerCase)) {
             Duration offsetDuration = DURATION_FOR_TIME_UNIT_MAP.get(offsetUnitLowerCase).multipliedBy(offsetValueInt);
-            return localDateTime.plus(offsetDuration);
+            return zonedDateTime.plus(offsetDuration);
         }
 
         if (PERIOD_FOR_TIME_UNIT_MAP.containsKey(offsetUnitLowerCase)) {
             Period offsetPeriod = PERIOD_FOR_TIME_UNIT_MAP.get(offsetUnitLowerCase).multipliedBy(offsetValueInt);
-            return localDateTime.plus(offsetPeriod);
+            return zonedDateTime.plus(offsetPeriod);
         }
 
         String message = String.format("The relative date time unit '%s' is not supported.", offsetUnit);
         throw new IllegalArgumentException(message);
     }
 
     /**
-     * Snaps the given local date time to the start of the previous time
+     * Snaps the given zoned date time to the start of the previous time
      * period specified by the given snap unit, and returns the result.
      */
-    private LocalDateTime applySnap(LocalDateTime localDateTime, String snapUnit) {
+    private ZonedDateTime applySnap(ZonedDateTime zonedDateTime, String snapUnit) {
 
         // Convert to lower case to make case-insensitive.
         String snapUnitLowerCase = snapUnit.toLowerCase();
 
         if (SECOND_UNITS_SET.contains(snapUnitLowerCase)) {
-            return localDateTime.truncatedTo(ChronoUnit.SECONDS);
+            return zonedDateTime.truncatedTo(ChronoUnit.SECONDS);
         } else if (MINUTE_UNITS_SET.contains(snapUnitLowerCase)) {
-            return localDateTime.truncatedTo(ChronoUnit.MINUTES);
+            return zonedDateTime.truncatedTo(ChronoUnit.MINUTES);
         } else if (HOUR_UNITS_SET.contains(snapUnitLowerCase)) {
-            return localDateTime.truncatedTo(ChronoUnit.HOURS);
+            return zonedDateTime.truncatedTo(ChronoUnit.HOURS);
         } else if (DAY_UNITS_SET.contains(snapUnitLowerCase)) {
-            return localDateTime.truncatedTo(ChronoUnit.DAYS);
+            return zonedDateTime.truncatedTo(ChronoUnit.DAYS);
         } else if (WEEK_UNITS_SET.contains(snapUnitLowerCase)) {
-            return applySnapToDayOfWeek(localDateTime, DayOfWeek.SUNDAY);
+            return applySnapToDayOfWeek(zonedDateTime, DayOfWeek.SUNDAY);
         } else if (MONTH_UNITS_SET.contains(snapUnitLowerCase)) {
-            return localDateTime.truncatedTo(ChronoUnit.DAYS).withDayOfMonth(1);
+            return zonedDateTime.truncatedTo(ChronoUnit.DAYS).withDayOfMonth(1);
         } else if (QUARTER_UNITS_SET.contains(snapUnitLowerCase)) {
-            Month snapMonth = localDateTime.getMonth().firstMonthOfQuarter();
-            return localDateTime.truncatedTo(ChronoUnit.DAYS).withDayOfMonth(1).withMonth(snapMonth.getValue());
+            Month snapMonth = zonedDateTime.getMonth().firstMonthOfQuarter();
+            return zonedDateTime.truncatedTo(ChronoUnit.DAYS).withDayOfMonth(1).withMonth(snapMonth.getValue());
         } else if (YEAR_UNITS_SET.contains(snapUnitLowerCase)) {
-            return localDateTime.truncatedTo(ChronoUnit.DAYS).withDayOfYear(1);
+            return zonedDateTime.truncatedTo(ChronoUnit.DAYS).withDayOfYear(1);
         } else if (DAY_OF_THE_WEEK_FOR_SNAP_UNIT_MAP.containsKey(snapUnitLowerCase)) {
-            return applySnapToDayOfWeek(localDateTime, DAY_OF_THE_WEEK_FOR_SNAP_UNIT_MAP.get(snapUnitLowerCase));
+            return applySnapToDayOfWeek(zonedDateTime, DAY_OF_THE_WEEK_FOR_SNAP_UNIT_MAP.get(snapUnitLowerCase));
         }
 
         String message = String.format("The relative date time unit '%s' is not supported.", snapUnit);
@@ -321,10 +321,10 @@ private LocalDateTime applySnap(LocalDateTime localDateTime, String snapUnit) {
      * Snaps the given date time to the start of the previous
      * specified day of the week, and returns the result.
      */
-    private LocalDateTime applySnapToDayOfWeek(LocalDateTime dateTime, DayOfWeek snapDayOfWeek) {
-        LocalDateTime snappedDateTime = dateTime.truncatedTo(ChronoUnit.DAYS);
+    private ZonedDateTime applySnapToDayOfWeek(ZonedDateTime zonedDateTime, DayOfWeek snapDayOfWeek) {
+        ZonedDateTime snappedDateTime = zonedDateTime.truncatedTo(ChronoUnit.DAYS);
 
-        int daysToSnap = dateTime.getDayOfWeek().getValue() - snapDayOfWeek.getValue();
+        int daysToSnap = zonedDateTime.getDayOfWeek().getValue() - snapDayOfWeek.getValue();
         if (daysToSnap < 0) daysToSnap += DayOfWeek.values().length;
 
         return snappedDateTime.minusDays(daysToSnap);

ppl-spark-integration/src/test/java/org/opensearch/sql/expression/function/SerializableTimeUdfTest.java

@@ -13,49 +13,62 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
-import static org.opensearch.sql.expression.function.SerializableUdf.relativeTimestampFunction;
 
 public class SerializableTimeUdfTest {
 
-    // Monday, Jan 03, 2000 @ 01:01:01.100
-    private final Instant MOCK_INSTANT = Instant.parse("2000-01-03T01:01:01.100Z");
+    // Monday, Jan 03, 2000 @ 01:01:01.100 UTC
     private final ZoneId MOCK_ZONE_ID = ZoneId.of("UTC");
+    private final Instant MOCK_INSTANT = Instant.parse("2000-01-03T01:01:01.100Z");
+    private final Timestamp MOCK_TIMESTAMP = Timestamp.from(MOCK_INSTANT);
 
     @Test
     public void relativeTimestampTest() {
 
-        /* These are only basic tests of the relative date time functionality.
-          For more comprehensive tests, see {@link TimeUtilsTest}.
-         */
-
-        testValid("-60m", "2000-01-03 00:01:01.100");
-        testValid("-H", "2000-01-03 00:01:01.100");
-        testValid("+2wk", "2000-01-17 01:01:01.100");
-        testValid("-1h@W3", "1999-12-29 00:00:00");
-        testValid("@d", "2000-01-03 00:00:00");
-        testValid("now", "2000-01-03 01:01:01.100");
-
-        testInvalid("invalid", "The relative date time 'invalid' is not supported.");
-        testInvalid("INVALID", "The relative date time 'INVALID' is not supported.");
-        testInvalid("~h", "The relative date time '~h' is not supported.");
-        testInvalid("+1.1h", "The relative date time '+1.1h' is not supported.");
-        testInvalid("+ms", "The relative date time unit 'ms' is not supported.");
-        testInvalid("+1INVALID", "The relative date time unit 'INVALID' is not supported.");
-        testInvalid("@INVALID", "The relative date time unit 'INVALID' is not supported.");
-        testInvalid("@ms", "The relative date time unit 'ms' is not supported.");
-        testInvalid("@w8", "The relative date time unit 'w8' is not supported.");
+        /// These are only basic tests of the relative time functionality.
+        /// For more comprehensive tests, see [TimeUtilsTest].
+
+        testValidInstant("-60m", "2000-01-03T00:01:01.100Z");
+        testValidInstant("-H", "2000-01-03T00:01:01.100Z");
+        testValidInstant("+2wk", "2000-01-17T01:01:01.100Z");
+        testValidInstant("-1h@W3", "1999-12-29T00:00:00Z");
+        testValidInstant("@d", "2000-01-03T00:00:00Z");
+        testValidInstant("now", "2000-01-03T01:01:01.100Z");
+
+        testValidTimestamp("-60m", "2000-01-03T00:01:01.100Z");
+        testValidTimestamp("-H", "2000-01-03T00:01:01.100Z");
+        testValidTimestamp("+2wk", "2000-01-17T01:01:01.100Z");
+        testValidTimestamp("-1h@W3", "1999-12-29T00:00:00Z");
+        testValidTimestamp("@d", "2000-01-03T00:00:00Z");
+        testValidTimestamp("now", "2000-01-03T01:01:01.100Z");
+
+        testInvalidString("invalid", "The relative date time 'invalid' is not supported.");
+        testInvalidString("INVALID", "The relative date time 'INVALID' is not supported.");
+        testInvalidString("~h", "The relative date time '~h' is not supported.");
+        testInvalidString("+1.1h", "The relative date time '+1.1h' is not supported.");
+        testInvalidString("+ms", "The relative date time unit 'ms' is not supported.");
+        testInvalidString("+1INVALID", "The relative date time unit 'INVALID' is not supported.");
+        testInvalidString("@INVALID", "The relative date time unit 'INVALID' is not supported.");
+        testInvalidString("@ms", "The relative date time unit 'ms' is not supported.");
+        testInvalidString("@w8", "The relative date time unit 'w8' is not supported.");
+    }
+
+    private void testValidInstant(String relativeString, String expectedInstantString) {
+        String testMessage = String.format("\"%s\"", relativeString);
+        Instant expectedInstant = Instant.parse(expectedInstantString);
+        Instant actualTimestamp = SerializableUdf.relativeTimestampFunction.apply(relativeString, MOCK_INSTANT, MOCK_ZONE_ID.toString());
+        assertEquals(testMessage, expectedInstant, actualTimestamp);
     }
 
-    private void testValid(String relativeString, String expectedTimestampString) {
+    private void testValidTimestamp(String relativeString, String expectedInstantString) {
         String testMessage = String.format("\"%s\"", relativeString);
-        Timestamp expectedTimestamp = Timestamp.valueOf(expectedTimestampString);
-        Timestamp actualTimestamp = relativeTimestampFunction.apply(relativeString, MOCK_INSTANT, MOCK_ZONE_ID.toString());
-        assertEquals(testMessage, expectedTimestamp, actualTimestamp);
+        Instant expectedInstant = Instant.parse(expectedInstantString);
+        Instant actualTimestamp = SerializableUdf.relativeTimestampFunction.apply(relativeString, MOCK_TIMESTAMP, MOCK_ZONE_ID.toString());
+        assertEquals(testMessage, expectedInstant, actualTimestamp);
     }
 
-    private void testInvalid(String relativeString, String expectedExceptionMessage) {
+    private void testInvalidString(String relativeString, String expectedExceptionMessage) {
         String testMessage = String.format("\"%s\"", relativeString);
-        String actualExceptionMessage = assertThrows(testMessage, RuntimeException.class, () -> relativeTimestampFunction.apply(relativeString, MOCK_INSTANT, MOCK_ZONE_ID.toString())).getMessage();
+        String actualExceptionMessage = assertThrows(testMessage, RuntimeException.class, () -> SerializableUdf.relativeTimestampFunction.apply(relativeString, MOCK_INSTANT, MOCK_ZONE_ID.toString())).getMessage();
         assertEquals(expectedExceptionMessage, actualExceptionMessage);
     }
 }