feat(java): support list versions and checkout version in Dataset (#3945)

close #3944

---------

Co-authored-by: majin.nathan <majin.nathan@bytedance.com>

Author: majin1102

java/core/lance-jni/src/blocking_dataset.rs

@@ -34,7 +34,8 @@ use lance::dataset::builder::DatasetBuilder;
 use lance::dataset::statistics::{DataStatistics, DatasetStatisticsExt};
 use lance::dataset::transaction::Operation;
 use lance::dataset::{
-    ColumnAlteration, Dataset, NewColumnTransform, ProjectionRequest, ReadParams, WriteParams,
+    ColumnAlteration, Dataset, NewColumnTransform, ProjectionRequest, ReadParams, Version,
+    WriteParams,
 };
 use lance::io::{ObjectStore, ObjectStoreParams};
 use lance::table::format::Fragment;
@@ -152,6 +153,25 @@ impl BlockingDataset {
         Ok(version)
     }
 
+    pub fn list_versions(&self) -> Result<Vec<Version>> {
+        let versions = RT.block_on(self.inner.versions())?;
+        Ok(versions)
+    }
+
+    pub fn version(&self) -> Result<Version> {
+        Ok(self.inner.version())
+    }
+
+    pub fn checkout_version(&mut self, version: u64) -> Result<Self> {
+        let inner = RT.block_on(self.inner.checkout_version(version))?;
+        Ok(Self { inner })
+    }
+
+    pub fn checkout_latest(&mut self) -> Result<()> {
+        RT.block_on(self.inner.checkout_latest())?;
+        Ok(())
+    }
+
     pub fn count_rows(&self, filter: Option<String>) -> Result<usize> {
         let rows = RT.block_on(self.inner.count_rows(filter))?;
         Ok(rows)
@@ -326,6 +346,41 @@ impl IntoJava for BlockingDataset {
     }
 }
 
+impl IntoJava for Version {
+    fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result<JObject<'a>> {
+        let timestamp_str = self.timestamp.to_rfc3339();
+        let jtimestamp = env.new_string(timestamp_str)?;
+        let zdt = env
+            .call_static_method(
+                "java/time/ZonedDateTime",
+                "parse",
+                "(Ljava/lang/CharSequence;)Ljava/time/ZonedDateTime;",
+                &[JValue::Object(&jtimestamp)],
+            )?
+            .l()?;
+
+        let jmap = env.new_object("java/util/TreeMap", "()V", &[])?;
+        let map = JMap::from_env(env, &jmap)?;
+
+        for (k, v) in self.metadata {
+            let jkey = env.new_string(k)?;
+            let jval = env.new_string(v)?;
+            map.put(env, &jkey, &jval).expect("ERROR: calling jmap.put");
+        }
+
+        let java_version = env.new_object(
+            "com/lancedb/lance/Version",
+            "(JLjava/time/ZonedDateTime;Ljava/util/TreeMap;)V",
+            &[
+                JValue::Long(self.version as i64),
+                JValue::Object(&zdt),
+                JValue::Object(&jmap),
+            ],
+        )?;
+        Ok(java_version)
+    }
+}
+
 fn attach_native_dataset<'local>(
     env: &mut JNIEnv<'local>,
     dataset: BlockingDataset,
@@ -658,31 +713,87 @@ fn inner_uri<'local>(env: &mut JNIEnv<'local>, java_dataset: JObject) -> Result<
 }
 
 #[no_mangle]
-pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeVersion(
+pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeListVersions<'local>(
+    mut env: JNIEnv<'local>,
+    java_dataset: JObject,
+) -> JObject<'local> {
+    ok_or_throw!(env, inner_list_versions(&mut env, java_dataset))
+}
+
+fn inner_list_versions<'local>(
+    env: &mut JNIEnv<'local>,
+    java_dataset: JObject,
+) -> Result<JObject<'local>> {
+    let versions = {
+        let dataset_guard =
+            unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
+        dataset_guard.list_versions()?
+    };
+    let array_list = env.new_object("java/util/ArrayList", "()V", &[])?;
+
+    versions
+        .into_iter()
+        .map(|inner_ver| inner_ver.into_java(env))
+        .collect::<Result<Vec<_>>>()?
+        .into_iter()
+        .try_for_each(|java_ver| -> Result<()> {
+            env.call_method(
+                &array_list,
+                "add",
+                "(Ljava/lang/Object;)Z",
+                &[JValue::Object(&java_ver)],
+            )?;
+            Ok(())
+        })?;
+    Ok(array_list)
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeGetVersion<'local>(
+    mut env: JNIEnv<'local>,
+    java_dataset: JObject,
+) -> JObject<'local> {
+    ok_or_throw!(env, inner_get_version(&mut env, java_dataset))
+}
+
+fn inner_get_version<'local>(
+    env: &mut JNIEnv<'local>,
+    java_dataset: JObject,
+) -> Result<JObject<'local>> {
+    let version = {
+        let dataset_guard =
+            unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
+        dataset_guard.version()?
+    };
+    version.into_java(env)
+}
+
+#[no_mangle]
+pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeGetLatestVersionId(
     mut env: JNIEnv,
     java_dataset: JObject,
 ) -> jlong {
-    ok_or_throw_with_return!(env, inner_version(&mut env, java_dataset), -1) as jlong
+    ok_or_throw_with_return!(env, inner_latest_version_id(&mut env, java_dataset), -1) as jlong
 }
 
-fn inner_version(env: &mut JNIEnv, java_dataset: JObject) -> Result<u64> {
+fn inner_latest_version_id(env: &mut JNIEnv, java_dataset: JObject) -> Result<u64> {
     let dataset_guard =
         unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
-    Ok(dataset_guard.inner.version().version)
+    dataset_guard.latest_version()
 }
 
 #[no_mangle]
-pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeLatestVersion(
+pub extern "system" fn Java_com_lancedb_lance_Dataset_nativeCheckoutLatest(
     mut env: JNIEnv,
     java_dataset: JObject,
-) -> jlong {
-    ok_or_throw_with_return!(env, inner_latest_version(&mut env, java_dataset), -1) as jlong
+) {
+    ok_or_throw_without_return!(env, inner_checkout_latest(&mut env, java_dataset));
 }
 
-fn inner_latest_version(env: &mut JNIEnv, java_dataset: JObject) -> Result<u64> {
-    let dataset_guard =
+fn inner_checkout_latest(env: &mut JNIEnv, java_dataset: JObject) -> Result<()> {
+    let mut dataset_guard =
         unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
-    dataset_guard.latest_version()
+    dataset_guard.checkout_latest()
 }
 
 #[no_mangle]
@@ -700,15 +811,12 @@ fn inner_checkout_version<'local>(
     version: jlong,
 ) -> Result<JObject<'local>> {
     let new_dataset = {
-        let dataset_guard =
+        let mut dataset_guard =
             unsafe { env.get_rust_field::<_, _, BlockingDataset>(java_dataset, NATIVE_DATASET) }?;
-
-        let version_u64 = version as u64;
-        RT.block_on(dataset_guard.inner.checkout_version(version_u64))?
+        dataset_guard.checkout_version(version as u64)?
     };
 
-    let blocking_dataset = BlockingDataset { inner: new_dataset };
-    blocking_dataset.into_java(env)
+    new_dataset.into_java(env)
 }
 
 #[no_mangle]

java/core/src/main/java/com/lancedb/lance/Dataset.java

@@ -448,29 +448,62 @@ public String uri() {
 
   private native String nativeUri();
 
+  /**
+   * Get the currently checked out version id of the dataset
+   *
+   * @return the version id of the dataset
+   */
+  public long version() {
+    return getVersion().getId();
+  }
+
   /**
    * Gets the currently checked out version of the dataset.
    *
    * @return the version of the dataset
    */
-  public long version() {
+  public Version getVersion() {
     try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
       Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
-      return nativeVersion();
+      return nativeGetVersion();
     }
   }
 
-  private native long nativeVersion();
+  private native Version nativeGetVersion();
+
+  /**
+   * Get the version history of the dataset.
+   *
+   * @return the version history of the dataset
+   */
+  public List<Version> listVersions() {
+    try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
+      Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
+      return nativeListVersions();
+    }
+  }
+
+  private native List<Version> nativeListVersions();
 
   /** @return the latest version of the dataset. */
   public long latestVersion() {
-    try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
+    try (LockManager.WriteLock writeLock = lockManager.acquireWriteLock()) {
+      Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
+      return nativeGetLatestVersionId();
+    }
+  }
+
+  private native long nativeGetLatestVersionId();
+
+  /** Checkout the dataset to the latest version. */
+  public void checkoutLatest() {
+    try (LockManager.WriteLock writeLock = lockManager.acquireWriteLock()) {
       Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
-      return nativeLatestVersion();
+      nativeCheckoutLatest();
     }
   }
 
-  private native long nativeLatestVersion();
+  private native void nativeCheckoutLatest();
 
   /**
    * Checks out a specific version of the dataset. If the version is already checked out, it returns
@@ -480,17 +513,9 @@ public long latestVersion() {
    * @return a new Dataset instance with the specified version checked out
    */
   public Dataset checkoutVersion(long version) {
-    if (version < 1) {
-      throw new IllegalArgumentException("Version must be greater than 0");
-    }
-
+    Preconditions.checkArgument(version > 0, "version number must be greater than 0");
     try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
       Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
-
-      if (this.version() == version) {
-        return this;
-      }
-
       return nativeCheckoutVersion(version);
     }
   }

java/core/src/main/java/com/lancedb/lance/Version.java

@@ -0,0 +1,66 @@
+/*
+ * 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 com.lancedb.lance;
+
+import java.time.ZonedDateTime;
+import java.util.Objects;
+import java.util.SortedMap;
+import java.util.TreeMap;
+
+public class Version {
+  private final long id;
+  private final ZonedDateTime dataTime;
+  private final SortedMap<String, String> metadata;
+
+  public Version(long id, ZonedDateTime dataTime, TreeMap<String, String> metadata) {
+    this.id = id;
+    this.dataTime = dataTime;
+    this.metadata = metadata;
+  }
+
+  public ZonedDateTime getDataTime() {
+    return dataTime;
+  }
+
+  public SortedMap<String, String> getMetadata() {
+    return metadata;
+  }
+
+  public long getId() {
+    return id;
+  }
+
+  public String toString() {
+    return "Version: " + id + ", dataTime: " + dataTime + ", metadata: " + metadata;
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    Version version = (Version) o;
+    return id == version.id
+        && Objects.equals(dataTime, version.dataTime)
+        && Objects.equals(metadata, version.metadata);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(id, dataTime, metadata);
+  }
+}