diff --git a/java/src/main/java/org/lance/Dataset.java b/java/src/main/java/org/lance/Dataset.java
index 23b7c051934..52785a70262 100644
--- a/java/src/main/java/org/lance/Dataset.java
+++ b/java/src/main/java/org/lance/Dataset.java
@@ -811,7 +811,13 @@ public Dataset checkoutVersion(long version) {
     Preconditions.checkArgument(version > 0, "version number must be greater than 0");
     try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
       Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
-      return nativeCheckoutVersion(version);
+      Dataset newDataset = nativeCheckoutVersion(version);
+      if (selfManagedAllocator) {
+        newDataset.allocator = new RootAllocator(Long.MAX_VALUE);
+      } else {
+        newDataset.allocator = allocator;
+      }
+      return newDataset;
     }
   }
 
@@ -828,7 +834,13 @@ public Dataset checkoutTag(String tag) {
     Preconditions.checkArgument(tag != null, "Tag can not be null");
     try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
       Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
-      return nativeCheckoutTag(tag);
+      Dataset newDataset = nativeCheckoutTag(tag);
+      if (selfManagedAllocator) {
+        newDataset.allocator = new RootAllocator(Long.MAX_VALUE);
+      } else {
+        newDataset.allocator = allocator;
+      }
+      return newDataset;
     }
   }
 
@@ -1380,7 +1392,13 @@ public Dataset checkout(Ref ref) {
     Preconditions.checkNotNull(ref);
     try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) {
       Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed");
-      return nativeCheckout(ref);
+      Dataset newDataset = nativeCheckout(ref);
+      if (selfManagedAllocator) {
+        newDataset.allocator = new RootAllocator(Long.MAX_VALUE);
+      } else {
+        newDataset.allocator = allocator;
+      }
+      return newDataset;
     }
   }
 
diff --git a/java/src/test/java/org/lance/DatasetTest.java b/java/src/test/java/org/lance/DatasetTest.java
index 156be38ad69..6ea4533accb 100644
--- a/java/src/test/java/org/lance/DatasetTest.java
+++ b/java/src/test/java/org/lance/DatasetTest.java
@@ -257,6 +257,7 @@ void testDatasetCheckoutVersion(@TempDir Path tempDir) {
 
         // checkout the dataset at version 1
         try (Dataset checkoutV1 = dataset2.checkoutVersion(1)) {
+          assertNotNull(checkoutV1.getSchema());
           assertEquals(1, checkoutV1.version());
           assertEquals(2, checkoutV1.latestVersion());
           assertEquals(0, checkoutV1.countRows());
@@ -305,6 +306,7 @@ void testTags(@TempDir Path tempDir) {
 
         // checkout the dataset at version 1
         try (Dataset checkoutV1 = dataset2.checkoutTag("tag1")) {
+          assertNotNull(checkoutV1.getSchema());
           assertEquals(1, checkoutV1.version());
           assertEquals(2, checkoutV1.latestVersion());
           assertEquals(0, checkoutV1.countRows());
@@ -1696,6 +1698,7 @@ void testBranches(@TempDir Path tempDir) {
 
                   // Step 6. use checkout_branch to checkout branch2
                   try (Dataset branch2V4New = mainV2.checkout(Ref.ofBranch("branch2"))) {
+                    assertNotNull(branch2V4New.getSchema());
                     assertEquals(4, branch2V4New.version());
                     assertEquals(10, branch2V4New.countRows()); // A(5) + B(3) + C(2)
                   }