Add JNI for extract_quarter, add_calendrical_months, and is_leap_year

java/src/main/java/ai/rapids/cudf/ColumnView.java

@@ -827,6 +827,34 @@ public final ColumnVector dayOfYear() {
     return new ColumnVector(dayOfYear(getNativeView()));
   }
 
+  /**
+   * Get the quarter of the year from a timestamp.
+   * @return A new INT16 vector allocated on the GPU. It will be a value from {1, 2, 3, 4}
+   * corresponding to the quarter of the year.
+   */
+  public final ColumnVector quarterOfYear() {
+    assert type.isTimestampType();
+    return new ColumnVector(quarterOfYear(getNativeView()));
+  }
+
+  /**
+   * Add the specified number of months to the timestamp.
+   * @param months must be a INT16 column indicating the number of months to add. A negative number
+   *               of months works too.
+   * @return the updated timestamp
+   */
+  public final ColumnVector addCalendricalMonths(ColumnView months) {
+    return new ColumnVector(addCalendricalMonths(getNativeView(), months.getNativeView()));
+  }
+
+  /**
+   * Check to see if the year for this timestamp is a leap year or not.
+   * @return BOOL8 vector of results
+   */
+  public final ColumnVector isLeapYear() {
+    return new ColumnVector(isLeapYear(getNativeView()));
+  }
+
   /**
    * Rounds all the values in a column to the specified number of decimal places.
    *
@@ -3462,6 +3490,12 @@ private static native long scan(long viewHandle, long aggregation,
 
   private static native long dayOfYear(long viewHandle) throws CudfException;
 
+  private static native long quarterOfYear(long viewHandle) throws CudfException;
+
+  private static native long addCalendricalMonths(long tsViewHandle, long monthsViewHandle);
+
+  private static native long isLeapYear(long viewHandle) throws CudfException;
+
   private static native boolean containsScalar(long columnViewHaystack, long scalarHandle) throws CudfException;
 
   private static native long containsVector(long columnViewHaystack, long columnViewNeedles) throws CudfException;

java/src/main/native/src/ColumnViewJni.cpp

@@ -752,6 +752,45 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_dayOfYear(JNIEnv *env, jc
   CATCH_STD(env, 0);
 }
 
+JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_quarterOfYear(JNIEnv *env, jclass,
+                                                                     jlong input_ptr) {
+  JNI_NULL_CHECK(env, input_ptr, "input is null", 0);
+  try {
+    cudf::jni::auto_set_device(env);
+    const cudf::column_view *input = reinterpret_cast<cudf::column_view *>(input_ptr);
+    std::unique_ptr<cudf::column> output = cudf::datetime::extract_quarter(*input);
+    return reinterpret_cast<jlong>(output.release());
+  }
+  CATCH_STD(env, 0);
+}
+
+JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_addCalendricalMonths(JNIEnv *env, jclass,
+                                                                            jlong ts_ptr,
+                                                                            jlong months_ptr) {
+  JNI_NULL_CHECK(env, ts_ptr, "ts is null", 0);
+  JNI_NULL_CHECK(env, months_ptr, "months is null", 0);
+  try {
+    cudf::jni::auto_set_device(env);
+    const cudf::column_view *ts = reinterpret_cast<cudf::column_view *>(ts_ptr);
+    const cudf::column_view *months = reinterpret_cast<cudf::column_view *>(months_ptr);
+    std::unique_ptr<cudf::column> output = cudf::datetime::add_calendrical_months(*ts, *months);
+    return reinterpret_cast<jlong>(output.release());
+  }
+  CATCH_STD(env, 0);
+}
+
+JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_isLeapYear(JNIEnv *env, jclass,
+                                                                  jlong input_ptr) {
+  JNI_NULL_CHECK(env, input_ptr, "input is null", 0);
+  try {
+    cudf::jni::auto_set_device(env);
+    const cudf::column_view *input = reinterpret_cast<cudf::column_view *>(input_ptr);
+    std::unique_ptr<cudf::column> output = cudf::datetime::is_leap_year(*input);
+    return reinterpret_cast<jlong>(output.release());
+  }
+  CATCH_STD(env, 0);
+}
+
 JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_ColumnView_castTo(JNIEnv *env, jclass, jlong handle,
                                                               jint type, jint scale) {
   JNI_NULL_CHECK(env, handle, "native handle is null", 0);

java/src/test/java/ai/rapids/cudf/TimestampColumnVectorTest.java

@@ -311,6 +311,78 @@ public void testDayOfYear() {
     }
   }
 
+  @Test
+  public void testQuarterOfYear() {
+    short[] EXPECTED = new short[]{4, 3, 1, 4, 3};
+    try (ColumnVector timestampColumnVector = ColumnVector.timestampMilliSecondsFromLongs(TIMES_MS);
+         ColumnVector result = timestampColumnVector.quarterOfYear();
+         ColumnVector expected = ColumnVector.fromShorts(EXPECTED)) {
+      assertColumnsAreEqual(expected, result);
+    }
+
+    try (ColumnVector timestampColumnVector = ColumnVector.timestampSecondsFromLongs(TIMES_S);
+         ColumnVector result = timestampColumnVector.quarterOfYear();
+         ColumnVector expected = ColumnVector.fromShorts(EXPECTED)) {
+      assertColumnsAreEqual(expected, result);
+    }
+
+    try (ColumnVector timestampColumnVector = ColumnVector.daysFromInts(TIMES_DAY);
+         ColumnVector result = timestampColumnVector.quarterOfYear();
+         ColumnVector expected = ColumnVector.fromShorts(EXPECTED)) {
+      assertColumnsAreEqual(expected, result);
+    }
+  }
+
+  @Test
+  public void testAddMonths() {
+    long[] EXPECTED = new long[]{
+        -131968727762L,   //'1965-10-26 14:01:12.238' Tuesday
+        1533384000115L,   //'2018-08-04 12:00:00.115' Saturday
+        1679729532929L,   //'2023-03-25 07:32:12.929' Saturday
+        -124019927762L,   //'1966-01-26 14:01:12.238' Wednesday
+        1520164800115L};  //'2018-03-04 12:00:00.115' Sunday
+    try (ColumnVector timestampColumnVector = ColumnVector.timestampMilliSecondsFromLongs(TIMES_MS);
+         ColumnVector months = ColumnVector.fromShorts(
+             (short)0, (short)1, (short)2, (short)3, (short)-4);
+         ColumnVector result = timestampColumnVector.addCalendricalMonths(months);
+         ColumnVector expected = ColumnVector.timestampMilliSecondsFromLongs(EXPECTED)) {
+      assertColumnsAreEqual(expected, result);
+    }
+  }
+
+  @Test
+  public void testIsLeapYear() {
+    Boolean[] EXPECTED = new Boolean[]{false, false, false, false, false};
+    try (ColumnVector timestampColumnVector = ColumnVector.timestampMilliSecondsFromLongs(TIMES_MS);
+         ColumnVector result = timestampColumnVector.isLeapYear();
+         ColumnVector expected = ColumnVector.fromBoxedBooleans(EXPECTED)) {
+      assertColumnsAreEqual(expected, result);
+    }
+
+    try (ColumnVector timestampColumnVector = ColumnVector.timestampSecondsFromLongs(TIMES_S);
+         ColumnVector result = timestampColumnVector.isLeapYear();
+         ColumnVector expected = ColumnVector.fromBoxedBooleans(EXPECTED)) {
+      assertColumnsAreEqual(expected, result);
+    }
+
+    try (ColumnVector timestampColumnVector = ColumnVector.daysFromInts(TIMES_DAY);
+         ColumnVector result = timestampColumnVector.isLeapYear();
+         ColumnVector expected = ColumnVector.fromBoxedBooleans(EXPECTED)) {
+      assertColumnsAreEqual(expected, result);
+    }
+
+    final long[] LEAP_TIMES_S = {1073865600L,   // Monday, January 12, 2004 0:00:00
+        947635200L, // Wednesday, January 12, 2000 0:00:00
+        -2208038400L // Friday, January 12, 1900 0:00:00
+    };
+
+    try (ColumnVector timestampColumnVector = ColumnVector.timestampSecondsFromLongs(LEAP_TIMES_S);
+         ColumnVector result = timestampColumnVector.isLeapYear();
+         ColumnVector expected = ColumnVector.fromBoxedBooleans(true, true, false)) {
+      assertColumnsAreEqual(expected, result);
+    }
+  }
+
   @Test
   public void testCastToTimestamp() {
     try (ColumnVector timestampMillis = ColumnVector.timestampMilliSecondsFromLongs(TIMES_MS);