Expose OutOfBoundsPolicy in JNI for Table.gather

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

@@ -0,0 +1,39 @@
+/*
+ *
+ *  Copyright (c) 2021, NVIDIA CORPORATION.
+ *
+ *  Licensed 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 ai.rapids.cudf;
+
+/**
+ * Policy to account for possible out-of-bounds indices
+ *
+ * `NULLIFY` means to nullify output values corresponding to out-of-bounds gather map values.
+ *
+ * `DONT_CHECK` means do not check whether the indices are out-of-bounds, for better
+ *   performance. Use `DONT_CHECK` carefully, as it can result in a CUDA exception if
+ *   the gather map values are actually out of range.
+ *
+ * @note This enum doesn't have a nativeId because the C++ out_of_bounds_policy is a
+ *        a boolean enum. It is just added for clarity in the Java API.
+ */
+public enum OutOfBoundsPolicy {
+  /* Output values corresponding to out-of-bounds indices are null */
+  NULLIFY,  
+
+  /* No bounds checking is performed, better performance */
+  DONT_CHECK
+}

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

@@ -2016,7 +2016,7 @@ public ColumnVector rowBitCount() {
    * @return the resulting Table.
    */
   public Table gather(ColumnView gatherMap) {
-    return gather(gatherMap, true);
+    return gather(gatherMap, OutOfBoundsPolicy.NULLIFY);
   }
 
   /**
@@ -2027,16 +2027,36 @@ public Table gather(ColumnView gatherMap) {
    *
    * A negative value `i` in the `gatherMap` is interpreted as `i+n`, where
    * `n` is the number of rows in this table.
-
+   *
+   * @deprecated Use {@link #gather(ColumnView, OutOfBoundsPolicy)}
    * @param gatherMap the map of indexes.  Must be non-nullable and integral type.
    * @param checkBounds if true bounds checking is performed on the value. Be very careful
    *                    when setting this to false.
    * @return the resulting Table.
    */
+  @Deprecated
   public Table gather(ColumnView gatherMap, boolean checkBounds) {
     return new Table(gather(nativeHandle, gatherMap.getNativeView(), checkBounds));
   }
 
+  /**
+   * Gathers the rows of this table according to `gatherMap` such that row "i"
+   * in the resulting table's columns will contain row "gatherMap[i]" from this table.
+   * The number of rows in the result table will be equal to the number of elements in
+   * `gatherMap`.
+   *
+   * A negative value `i` in the `gatherMap` is interpreted as `i+n`, where
+   * `n` is the number of rows in this table.
+   *
+   * @param gatherMap the map of indexes.  Must be non-nullable and integral type.
+   * @param outOfBoundsPolicy policy to use when an out-of-range value is in `gatherMap`
+   * @return the resulting Table.
+   */
+  public Table gather(ColumnView gatherMap, OutOfBoundsPolicy outOfBoundsPolicy) {
+    boolean checkBounds = outOfBoundsPolicy == OutOfBoundsPolicy.NULLIFY;
+    return new Table(gather(nativeHandle, gatherMap.getNativeView(), checkBounds));
+  }
+
   private GatherMap[] buildJoinGatherMaps(long[] gatherMapData) {
     long bufferSize = gatherMapData[0];
     long leftAddr = gatherMapData[1];

java/src/main/native/src/TableJni.cpp

@@ -143,6 +143,13 @@ class jni_writer_data_sink final : public cudf::io::data_sink {
     stream.synchronize();
   }
 
+  std::future<void> device_write_async(void const *gpu_data, size_t size,
+                                       rmm::cuda_stream_view stream) override {
+    // Call the sync version until figuring out how to write asynchronously.
+    device_write(gpu_data, size, stream);
+    return std::async(std::launch::deferred, [] {});
+  }
+
   void flush() override {
     if (current_buffer_written > 0) {
       JNIEnv *env = cudf::jni::get_jni_env(jvm);