From 82154635b88fc6da8945311af48b8ca9b1cc6602 Mon Sep 17 00:00:00 2001 From: "fangbo.0511" Date: Tue, 3 Feb 2026 11:06:05 +0800 Subject: [PATCH 1/2] fix(java): init allocator for new dataset when checkout branch/tag --- java/src/main/java/org/lance/Dataset.java | 24 ++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) 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; } } From 0e22da6388fe51b04e5988d23a7ff11278812714 Mon Sep 17 00:00:00 2001 From: "fangbo.0511" Date: Tue, 3 Feb 2026 14:03:01 +0800 Subject: [PATCH 2/2] add test cases for checkout out new dataset operation --- java/src/test/java/org/lance/DatasetTest.java | 3 +++ 1 file changed, 3 insertions(+) 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) }