From d5afc0a46d17c9700ad20b54185d4a10528d4c51 Mon Sep 17 00:00:00 2001
From: vinoyang <yanghua1127@gmail.com>
Date: Fri, 13 Dec 2024 00:22:28 +0800
Subject: [PATCH] feat(java): expose uri method for Dataset instance (#3231)

---
 java/core/lance-jni/src/blocking_dataset.rs   | 23 +++++++++++++++++++
 .../main/java/com/lancedb/lance/Dataset.java  | 14 +++++++++++
 .../java/com/lancedb/lance/DatasetTest.java   | 12 ++++++++++
 3 files changed, 49 insertions(+)

diff --git a/java/core/lance-jni/src/blocking_dataset.rs b/java/core/lance-jni/src/blocking_dataset.rs
index 7c0ffaf717..9887cb1a76 100644
--- a/java/core/lance-jni/src/blocking_dataset.rs
+++ b/java/core/lance-jni/src/blocking_dataset.rs
@@ -579,6 +579,29 @@ fn inner_import_ffi_schema(
     Ok(())
 }
 
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeUri<'local>(
+    mut env: JNIEnv<'local>,
+    java_dataset: JObject,
+) -> JString<'local> {
+    ok_or_throw_with_return!(
+        env,
+        inner_uri(&mut env, java_dataset).map_err(|err| Error::input_error(err.to_string())),
+        JString::from(JObject::null())
+    )
+}
+
+fn inner_uri<'local>(env: &mut JNIEnv<'local>, java_dataset: JObject) -> Result<JString<'local>> {
+    let uri = {
+        let dataset_guard =
+            unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
+        dataset_guard.inner.uri().to_string()
+    };
+
+    let jstring_uri = env.new_string(uri)?;
+    Ok(jstring_uri)
+}
+
 #[no_mangle]
 pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeVersion(
     mut env: JNIEnv,
diff --git a/java/core/src/main/java/com/lancedb/lance/Dataset.java b/java/core/src/main/java/com/lancedb/lance/Dataset.java
index 73e4190418..235fbc9677 100644
--- a/java/core/src/main/java/com/lancedb/lance/Dataset.java
+++ b/java/core/src/main/java/com/lancedb/lance/Dataset.java
@@ -300,6 +300,20 @@ public LanceScanner newScan(ScanOptions options) {
     }
   }
 
+  /**
+   * Gets the URI of the dataset.
+   *
+   * @return the URI of the dataset
+   */
+  public String uri() {
+    try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
+      Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
+      return nativeUri();
+    }
+  }
+
+  private native String nativeUri();
+
   /**
    * Gets the currently checked out version of the dataset.
    *
diff --git a/java/core/src/test/java/com/lancedb/lance/DatasetTest.java b/java/core/src/test/java/com/lancedb/lance/DatasetTest.java
index 0115bb35f0..db42f58c78 100644
--- a/java/core/src/test/java/com/lancedb/lance/DatasetTest.java
+++ b/java/core/src/test/java/com/lancedb/lance/DatasetTest.java
@@ -139,6 +139,18 @@ void testDatasetVersion() {
     }
   }
 
+  @Test
+  void testDatasetUri() {
+    String datasetPath = tempDir.resolve("dataset_uri").toString();
+    try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+      TestUtils.SimpleTestDataset testDataset =
+          new TestUtils.SimpleTestDataset(allocator, datasetPath);
+      try (Dataset dataset = testDataset.createEmptyDataset()) {
+        assertEquals(datasetPath, dataset.uri());
+      }
+    }
+  }
+
   @Test
   void testOpenNonExist() throws IOException, URISyntaxException {
     String datasetPath = tempDir.resolve("non_exist").toString();