JNI quarterOfYear

rapidsai · Jul 27, 2021 · fb3654a · fb3654a
1 parent 4e9f8c2
commit fb3654a
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 1 deletion.
diff --git a/java/src/main/java/ai/rapids/cudf/ColumnView.java b/java/src/main/java/ai/rapids/cudf/ColumnView.java
@@ -827,6 +827,16 @@ 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
@@ -3480,9 +3490,11 @@ 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);
+  private static native long isLeapYear(long viewHandle) throws CudfException;
 
   private static native boolean containsScalar(long columnViewHaystack, long scalarHandle) throws CudfException;
 

diff --git a/java/src/main/native/src/ColumnViewJni.cpp b/java/src/main/native/src/ColumnViewJni.cpp
@@ -752,6 +752,18 @@ 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) {

diff --git a/java/src/test/java/ai/rapids/cudf/TimestampColumnVectorTest.java b/java/src/test/java/ai/rapids/cudf/TimestampColumnVectorTest.java
@@ -311,6 +311,28 @@ 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[]{