ARROW-7272: [C++][Java][Dataset] JNI bridge between RecordBatch and VectorSchemaRoot

cpp/src/jni/dataset/jni_util.cc

@@ -17,10 +17,13 @@
 
 #include "jni/dataset/jni_util.h"
 
-#include "arrow/util/logging.h"
-
+#include <memory>
 #include <mutex>
 
+#include "arrow/c/bridge.h"
+#include "arrow/c/helpers.h"
+#include "arrow/util/logging.h"
+
 namespace arrow {
 namespace dataset {
 namespace jni {
@@ -162,6 +165,15 @@ std::shared_ptr<ReservationListener> ReservationListenableMemoryPool::get_listen
 
 ReservationListenableMemoryPool::~ReservationListenableMemoryPool() {}
 
+Status CheckException(JNIEnv* env) {
+ if (env->ExceptionCheck()) {
+ env->ExceptionDescribe();
+ env->ExceptionClear();
+ return Status::Invalid("Error during calling Java code from native code");
+ }
+ return Status::OK();
+}
+
 jclass CreateGlobalClassReference(JNIEnv* env, const char* class_name) {
  jclass local_class = env->FindClass(class_name);
  jclass global_class = (jclass)env->NewGlobalRef(local_class);
@@ -236,6 +248,17 @@ arrow::Result<std::shared_ptr<arrow::Schema>> FromSchemaByteArray(
  env->ReleaseByteArrayElements(schemaBytes, schemaBytes_data, JNI_ABORT);
  return schema;
 }
+arrow::Status ExportRecordBatch(JNIEnv* env, const std::shared_ptr<RecordBatch>& batch,
+ jlong struct_array) {
+ return arrow::ExportRecordBatch(*batch,
+ reinterpret_cast<struct ArrowArray*>(struct_array));
+}
+
+arrow::Result<std::shared_ptr<RecordBatch>> ImportRecordBatch(
+ JNIEnv* env, const std::shared_ptr<Schema>& schema, jlong struct_array) {
+ return arrow::ImportRecordBatch(reinterpret_cast<struct ArrowArray*>(struct_array),
+ schema);
+}
 
 } // namespace jni
 } // namespace dataset

cpp/src/jni/dataset/jni_util.h

@@ -17,19 +17,21 @@
 
 #pragma once
 
+#include <jni.h>
+
 #include "arrow/array.h"
 #include "arrow/io/api.h"
 #include "arrow/ipc/api.h"
 #include "arrow/memory_pool.h"
 #include "arrow/result.h"
 #include "arrow/type.h"
 
-#include <jni.h>
-
 namespace arrow {
 namespace dataset {
 namespace jni {
 
+Status CheckException(JNIEnv* env);
+
 jclass CreateGlobalClassReference(JNIEnv* env, const char* class_name);
 
 arrow::Result<jmethodID> GetMethodID(JNIEnv* env, jclass this_class, const char* name,
@@ -48,6 +50,18 @@ arrow::Result<jbyteArray> ToSchemaByteArray(JNIEnv* env,
 arrow::Result<std::shared_ptr<arrow::Schema>> FromSchemaByteArray(JNIEnv* env,
  jbyteArray schemaBytes);
 
+/// \brief Export arrow::RecordBatch for Java (or other JVM languages) use.
+/// The exported batch is subject to C data interface specification and can be
+/// imported from Java side using provided JNI utilities.
+arrow::Status ExportRecordBatch(JNIEnv* env, const std::shared_ptr<RecordBatch>& batch,
+ jlong struct_array);
+
+/// \brief Import arrow::RecordBatch from JVM language side. The input data should
+/// ideally be exported from specific JNI utilities from JVM language side and should
+/// conform to C data interface specification.
+arrow::Result<std::shared_ptr<RecordBatch>> ImportRecordBatch(
+ JNIEnv* env, const std::shared_ptr<Schema>& schema, jlong struct_array);
+
 /// \brief Create a new shared_ptr on heap from shared_ptr t to prevent
 /// the managed object from being garbage-collected.
 ///

cpp/src/jni/dataset/jni_util_test.cc

@@ -15,11 +15,12 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include "jni/dataset/jni_util.h"
+
 #include <gtest/gtest.h>
 
 #include "arrow/memory_pool.h"
 #include "arrow/testing/gtest_util.h"
-#include "jni/dataset/jni_util.h"
 
 namespace arrow {
 namespace dataset {

cpp/src/jni/dataset/jni_wrapper.cc

@@ -24,9 +24,7 @@
 #include "arrow/filesystem/localfs.h"
 #include "arrow/ipc/api.h"
 #include "arrow/util/iterator.h"
-
 #include "jni/dataset/jni_util.h"
-
 #include "org_apache_arrow_dataset_file_JniWrapper.h"
 #include "org_apache_arrow_dataset_jni_JniWrapper.h"
 #include "org_apache_arrow_dataset_jni_NativeMemoryPool.h"
@@ -37,14 +35,8 @@ jclass illegal_access_exception_class;
 jclass illegal_argument_exception_class;
 jclass runtime_exception_class;
 
-jclass record_batch_handle_class;
-jclass record_batch_handle_field_class;
-jclass record_batch_handle_buffer_class;
 jclass java_reservation_listener_class;
 
-jmethodID record_batch_handle_constructor;
-jmethodID record_batch_handle_field_constructor;
-jmethodID record_batch_handle_buffer_constructor;
 jmethodID reserve_memory_method;
 jmethodID unreserve_memory_method;
 
@@ -100,11 +92,7 @@ class ReserveFromJava : public arrow::dataset::jni::ReservationListener {
  return arrow::Status::Invalid("JNIEnv was not attached to current thread");
  }
  env->CallObjectMethod(java_reservation_listener_, reserve_memory_method, size);
- if (env->ExceptionCheck()) {
- env->ExceptionDescribe();
- env->ExceptionClear();
- return arrow::Status::Invalid("Error calling Java side reservation listener");
- }
+ RETURN_NOT_OK(arrow::dataset::jni::CheckException(env));
  return arrow::Status::OK();
  }
 
@@ -114,11 +102,7 @@ class ReserveFromJava : public arrow::dataset::jni::ReservationListener {
  return arrow::Status::Invalid("JNIEnv was not attached to current thread");
  }
  env->CallObjectMethod(java_reservation_listener_, unreserve_memory_method, size);
- if (env->ExceptionCheck()) {
- env->ExceptionDescribe();
- env->ExceptionClear();
- return arrow::Status::Invalid("Error calling Java side reservation listener");
- }
+ RETURN_NOT_OK(arrow::dataset::jni::CheckException(env));
  return arrow::Status::OK();
  }
 
@@ -206,33 +190,10 @@ jint JNI_OnLoad(JavaVM* vm, void* reserved) {
  runtime_exception_class =
  CreateGlobalClassReference(env, "Ljava/lang/RuntimeException;");
 
- record_batch_handle_class =
- CreateGlobalClassReference(env,
- "Lorg/apache/arrow/"
- "dataset/jni/NativeRecordBatchHandle;");
- record_batch_handle_field_class =
- CreateGlobalClassReference(env,
- "Lorg/apache/arrow/"
- "dataset/jni/NativeRecordBatchHandle$Field;");
- record_batch_handle_buffer_class =
- CreateGlobalClassReference(env,
- "Lorg/apache/arrow/"
- "dataset/jni/NativeRecordBatchHandle$Buffer;");
  java_reservation_listener_class =
  CreateGlobalClassReference(env,
  "Lorg/apache/arrow/"
  "dataset/jni/ReservationListener;");
-
- record_batch_handle_constructor =
- JniGetOrThrow(GetMethodID(env, record_batch_handle_class, "<init>",
- "(J[Lorg/apache/arrow/dataset/"
- "jni/NativeRecordBatchHandle$Field;"
- "[Lorg/apache/arrow/dataset/"
- "jni/NativeRecordBatchHandle$Buffer;)V"));
- record_batch_handle_field_constructor =
- JniGetOrThrow(GetMethodID(env, record_batch_handle_field_class, "<init>", "(JJ)V"));
- record_batch_handle_buffer_constructor = JniGetOrThrow(
- GetMethodID(env, record_batch_handle_buffer_class, "<init>", "(JJJJ)V"));
  reserve_memory_method =
  JniGetOrThrow(GetMethodID(env, java_reservation_listener_class, "reserve", "(J)V"));
  unreserve_memory_method = JniGetOrThrow(
@@ -250,9 +211,6 @@ void JNI_OnUnload(JavaVM* vm, void* reserved) {
  env->DeleteGlobalRef(illegal_access_exception_class);
  env->DeleteGlobalRef(illegal_argument_exception_class);
  env->DeleteGlobalRef(runtime_exception_class);
- env->DeleteGlobalRef(record_batch_handle_class);
- env->DeleteGlobalRef(record_batch_handle_field_class);
- env->DeleteGlobalRef(record_batch_handle_buffer_class);
  env->DeleteGlobalRef(java_reservation_listener_class);
 
  default_memory_pool_id = -1L;
@@ -458,25 +416,21 @@ Java_org_apache_arrow_dataset_jni_JniWrapper_getSchemaFromScanner(JNIEnv* env, j
 /*
  * Class: org_apache_arrow_dataset_jni_JniWrapper
  * Method: nextRecordBatch
- * Signature: (J)Lorg/apache/arrow/dataset/jni/NativeRecordBatchHandle;
+ * Signature: (JJ)Z
  */
-JNIEXPORT jobject JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_nextRecordBatch(
- JNIEnv* env, jobject, jlong scanner_id) {
+JNIEXPORT jboolean JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_nextRecordBatch(
+ JNIEnv* env, jobject, jlong scanner_id, jlong struct_array) {
  JNI_METHOD_START
  std::shared_ptr<DisposableScannerAdaptor> scanner_adaptor =
  RetrieveNativeInstance<DisposableScannerAdaptor>(scanner_id);
 
  std::shared_ptr<arrow::RecordBatch> record_batch =
  JniGetOrThrow(scanner_adaptor->Next());
  if (record_batch == nullptr) {
- return nullptr; // stream ended
+ return false; // stream ended
  }
- std::shared_ptr<arrow::Schema> schema = record_batch->schema();
- jobjectArray field_array =
- env->NewObjectArray(schema->num_fields(), record_batch_handle_field_class, nullptr);
-
- std::vector<std::shared_ptr<arrow::Buffer>> buffers;
- for (int i = 0; i < schema->num_fields(); ++i) {
+ std::vector<std::shared_ptr<arrow::Array>> offset_zeroed_arrays;
+ for (int i = 0; i < record_batch->num_columns(); ++i) {
  // TODO: If the array has an offset then we need to de-offset the array
  // in order for it to be properly consumed on the Java end.
  // This forces a copy, it would be nice to avoid this if Java
@@ -485,44 +439,22 @@ JNIEXPORT jobject JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_nextRecor
  //
  // Generally a non-zero offset will occur whenever the scanner batch
  // size is smaller than the batch size of the underlying files.
- auto column = record_batch->column(i);
- if (column->offset() != 0) {
- column = JniGetOrThrow(arrow::Concatenate({column}));
- }
- auto dataArray = column->data();
- jobject field = env->NewObject(record_batch_handle_field_class,
- record_batch_handle_field_constructor,
- column->length(), column->null_count());
- env->SetObjectArrayElement(field_array, i, field);
-
- for (auto& buffer : dataArray->buffers) {
- buffers.push_back(buffer);
+ std::shared_ptr<arrow::Array> array = record_batch->column(i);
+ if (array->offset() == 0) {
+ offset_zeroed_arrays.push_back(array);
+ continue;
  }
+ std::shared_ptr<arrow::Array> offset_zeroed =
+ JniGetOrThrow(arrow::Concatenate({array}));
+ offset_zeroed_arrays.push_back(offset_zeroed);
  }
 
- jobjectArray buffer_array =
- env->NewObjectArray(buffers.size(), record_batch_handle_buffer_class, nullptr);
-
- for (size_t j = 0; j < buffers.size(); ++j) {
- auto buffer = buffers[j];
- uint8_t* data = nullptr;
- int64_t size = 0;
- int64_t capacity = 0;
- if (buffer != nullptr) {
- data = (uint8_t*)buffer->data();
- size = buffer->size();
- capacity = buffer->capacity();
- }
- jobject buffer_handle = env->NewObject(record_batch_handle_buffer_class,
- record_batch_handle_buffer_constructor,
- CreateNativeRef(buffer), data, size, capacity);
- env->SetObjectArrayElement(buffer_array, j, buffer_handle);
- }
-
- jobject ret = env->NewObject(record_batch_handle_class, record_batch_handle_constructor,
- record_batch->num_rows(), field_array, buffer_array);
- return ret;
- JNI_METHOD_END(nullptr)
+ std::shared_ptr<arrow::RecordBatch> offset_zeroed_batch = arrow::RecordBatch::Make(
+ record_batch->schema(), record_batch->num_rows(), offset_zeroed_arrays);
+ JniAssertOkOrThrow(
+ arrow::dataset::jni::ExportRecordBatch(env, offset_zeroed_batch, struct_array));
+ return true;
+ JNI_METHOD_END(false)
 }
 
 /*

java/dataset/CMakeLists.txt

@@ -33,7 +33,6 @@ message("generating headers to ${JNI_HEADERS_DIR}")
 add_jar(arrow_dataset_java
  src/main/java/org/apache/arrow/dataset/jni/JniLoader.java
  src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java
- src/main/java/org/apache/arrow/dataset/jni/NativeRecordBatchHandle.java
  src/main/java/org/apache/arrow/dataset/file/JniWrapper.java
  src/main/java/org/apache/arrow/dataset/jni/NativeMemoryPool.java
  src/main/java/org/apache/arrow/dataset/jni/ReservationListener.java

java/dataset/pom.xml

@@ -44,6 +44,12 @@
  <version>${project.version}</version>
  <scope>compile</scope>
  </dependency>
+ <dependency>
+ <groupId>org.apache.arrow</groupId>
+ <artifactId>arrow-c-data</artifactId>
+ <version>${project.version}</version>
+ <scope>compile</scope>
+ </dependency>
  <dependency>
  <groupId>org.apache.arrow</groupId>
  <artifactId>arrow-memory-netty</artifactId>
@@ -56,6 +62,12 @@
  <version>${parquet.version}</version>
  <scope>test</scope>
  </dependency>
+ <dependency>
+ <groupId>org.apache.avro</groupId>
+ <artifactId>avro</artifactId>
+ <version>${avro.version}</version>
+ <scope>test</scope>
+ </dependency>
  <dependency>
  <groupId>org.apache.parquet</groupId>
  <artifactId>parquet-hadoop</artifactId>
@@ -86,12 +98,6 @@
  </exclusion>
  </exclusions>
  </dependency>
- <dependency>
- <groupId>org.apache.avro</groupId>
- <artifactId>avro</artifactId>
- <version>${avro.version}</version>
- <scope>test</scope>
- </dependency>
  <dependency>
  <groupId>com.google.guava</groupId>
  <artifactId>guava</artifactId>

java/dataset/src/main/java/org/apache/arrow/dataset/file/FileFormat.java

@@ -24,7 +24,7 @@ public enum FileFormat {
  PARQUET(0),
  NONE(-1);
 
- private int id;
+ private final int id;
 
  FileFormat(int id) {
  this.id = id;

java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java

@@ -87,21 +87,25 @@ private JniWrapper() {
 
  /**
  * Release the Scanner by destroying its reference held by JNI wrapper.
+ *
  * @param scannerId the native pointer of the arrow::dataset::Scanner instance.
  */
  public native void closeScanner(long scannerId);
 
  /**
  * Read next record batch from the specified scanner.
+ *
  * @param scannerId the native pointer of the arrow::dataset::Scanner instance.
- * @return an instance of {@link NativeRecordBatchHandle} describing the overall layout of the native record batch.
+ * @param arrowArray pointer to an empty {@link org.apache.arrow.c.ArrowArray} struct to
+ * store C++ side record batch that conforms to C data interface.
+ * @return true if valid record batch is returned; false if stream ended.
  */
- public native NativeRecordBatchHandle nextRecordBatch(long scannerId);
+ public native boolean nextRecordBatch(long scannerId, long arrowArray);
 
  /**
  * Release the Buffer by destroying its reference held by JNI wrapper.
+ *
  * @param bufferId the native pointer of the arrow::Buffer instance.
  */
  public native void releaseBuffer(long bufferId);
-
 }