feat(java): expose uri method for Dataset instance (#3231)

lancedb · Dec 12, 2024 · d5afc0a · d5afc0a
1 parent 1c8d406
commit d5afc0a
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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();