diff --git a/java/dataset/src/main/cpp/jni_wrapper.cc b/java/dataset/src/main/cpp/jni_wrapper.cc index 871a2e95b94..5640bc43496 100644 --- a/java/dataset/src/main/cpp/jni_wrapper.cc +++ b/java/dataset/src/main/cpp/jni_wrapper.cc @@ -27,6 +27,7 @@ #include "arrow/dataset/file_base.h" #include "arrow/filesystem/localfs.h" #include "arrow/filesystem/path_util.h" +#include "arrow/filesystem/s3fs.h" #include "arrow/engine/substrait/util.h" #include "arrow/ipc/api.h" #include "arrow/util/iterator.h" @@ -569,6 +570,18 @@ JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_releaseBuffe JNI_METHOD_END() } +/* + * Class: org_apache_arrow_dataset_jni_JniWrapper + * Method: ensureS3Finalized + * Signature: (J)V + */ +JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_jni_JniWrapper_ensureS3Finalized( + JNIEnv* env, jobject) { + JNI_METHOD_START + JniAssertOkOrThrow(arrow::fs::EnsureS3Finalized()); + JNI_METHOD_END() +} + /* * Class: org_apache_arrow_dataset_file_JniWrapper * Method: makeFileSystemDatasetFactory diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java b/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java index 7ada21c0582..a3b31c73e85 100644 --- a/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java +++ b/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniLoader.java @@ -59,6 +59,7 @@ public void ensureLoaded() { return; } loadRemaining(); + ensureS3FinalizedOnShutdown(); } private synchronized void loadRemaining() { @@ -109,4 +110,8 @@ private String getNormalizedArch() { } return arch; } + + private void ensureS3FinalizedOnShutdown() { + Runtime.getRuntime().addShutdownHook(new Thread(() -> { JniWrapper.get().ensureS3Finalized(); })); + } } diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java b/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java index 1a9d4188c16..93cc5d7a370 100644 --- a/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java +++ b/java/dataset/src/main/java/org/apache/arrow/dataset/jni/JniWrapper.java @@ -108,4 +108,10 @@ private JniWrapper() { * @param bufferId the native pointer of the arrow::Buffer instance. */ public native void releaseBuffer(long bufferId); + + /** + * Ensure the S3 APIs are shutdown, but only if not already done. If the S3 APIs are unintialized, + * then this is a noop. + */ + public native void ensureS3Finalized(); }