From 839100fa4b38402bd0883e6c511e7a2dec599fe5 Mon Sep 17 00:00:00 2001 From: Jerry He Date: Mon, 17 Aug 2026 14:11:51 -0700 Subject: [PATCH 1/8] fix: advance _row_last_updated_at_version for UPDATE COLUMNS FROM (#528) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Lance dependency — required to compile * Depends on lance-format/lance#6748, which adds `Update.Builder.updatedFragmentOffsets(...)` / `Update.updatedFragmentOffsets()` and JNI FromJava + IntoJava so the driver commit passes matched offsets into Rust. Bump **`lance.version`** in `pom.xml` to a release that includes lance#6748 before building or merging this PR; otherwise build fails on the new API calls. ## Summary * Fixes https://github.com/lance-format/lance-spark/issues/418. * `UpdateColumnsWriter.processFragment`: after each `fragment.updateColumns()` call, reads `result.getUpdatedRowOffsets()` and accumulates a `Map` of fragment id → matched physical row offsets. * `TaskCommit`: carries the per-fragment offset map alongside the existing `updatedFragments` and `fieldsModified` fields. * `UpdateColumnsBackfillBatchWrite.commit()`: merges offset maps from all task commits and passes them to `Update.builder().updatedFragmentOffsets(...)`. Lance's `build_manifest` then calls the partial `_row_last_updated_at_version` refresh only for the matched rows, leaving unmatched rows and untouched fragments unchanged. * `BaseUpdateColumnsBackfillTest`: flips `testUpdateColumnsPreservesCreatedAtAndAdvancesLastUpdatedWithStableRowIds` from the "known gap" pin (assertEquals, no change) to the correct assertion (`assertTrue(after > before)`); updates Javadoc accordingly. ## Background `UPDATE COLUMNS FROM` rewrites column data in place via Lance's `Operation::Update` with `RewriteColumns` mode. Lance's `build_manifest` can partially refresh `_row_last_updated_at_version` for only the matched rows — but only when the `updated_fragment_offsets` map is non-empty on the commit. Previously `UpdateColumnsBackfillBatchWrite` never populated this map, so the partial refresh never activated and `_row_last_updated_at_version` stayed stale after every UPDATE COLUMNS commit, breaking CDF consumers. The matched row offsets are already computed inside Lance during `fragment.updateColumns()` and surfaced via `FragmentUpdateResult.getUpdatedRowOffsets()` (lance#6650). The missing piece was wiring those offsets from the executor task result through `TaskCommit` to the driver commit, and then setting them on the `Update` operation — which this PR does. ## Test plan * `BaseUpdateColumnsBackfillTest#testUpdateColumnsPreservesCreatedAtAndAdvancesLastUpdatedWithStableRowIds` — creates a stable-row-id table, runs UPDATE COLUMNS over all rows, and asserts `_row_last_updated_at_version` strictly increases for each row while `_row_created_at_version` is unchanged. --------- Co-authored-by: Jing chen He Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Daniel Rammer --- .../UpdateColumnsBackfillBatchWrite.java | 30 ++++++++++-- .../update/BaseUpdateColumnsBackfillTest.java | 47 +++++++++---------- 2 files changed, 48 insertions(+), 29 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java index 3d8cb58db..895086964 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java @@ -39,6 +39,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -129,6 +130,12 @@ public void commit(WriterCommitMessage[] messages) { .findFirst() .orElse(new long[0]); + Map mergedUpdatedFragmentOffsets = new HashMap<>(); + Arrays.stream(messages) + .map(m -> (TaskCommit) m) + .map(TaskCommit::getUpdatedFragmentOffsets) + .forEach(m -> m.forEach(mergedUpdatedFragmentOffsets::put)); + if (updatedFragments.isEmpty()) { logger.info("No updated fragments to commit."); return; @@ -145,12 +152,14 @@ public void commit(WriterCommitMessage[] messages) { .map(Fragment::metadata) .forEach(updatedFragments::add); - // Commit update operation using CommitBuilder + // Commit update operation using CommitBuilder. Pass matched physical row offsets per + // fragment so Lance can partially refresh _row_last_updated_at_version (stable row IDs). Update update = Update.builder() .updatedFragments(updatedFragments) .fieldsModified(fieldsModified) .updateMode(Optional.of(Update.UpdateMode.RewriteColumns)) + .updatedFragmentOffsets(mergedUpdatedFragmentOffsets) .build(); long version = Objects.requireNonNull( @@ -175,6 +184,7 @@ public void commit(WriterCommitMessage[] messages) { public static class UpdateColumnsWriter extends AbstractBackfillWriter { private final List updatedFragments = new ArrayList<>(); + private final Map updatedFragmentOffsets = new HashMap<>(); private long[] fieldsModified; public UpdateColumnsWriter( @@ -204,11 +214,15 @@ protected void processFragment(Fragment fragment, ArrowArrayStream stream) { LanceDataset.ROW_ADDRESS_COLUMN.name()); updatedFragments.add(result.getUpdatedFragment()); fieldsModified = result.getFieldsModified(); + byte[] rowOffsetBytes = result.getUpdatedRowOffsetBytes(); + if (rowOffsetBytes != null && rowOffsetBytes.length > 0) { + updatedFragmentOffsets.put((long) fragment.getId(), rowOffsetBytes); + } } @Override protected WriterCommitMessage buildCommitMessage() { - return new TaskCommit(updatedFragments, fieldsModified); + return new TaskCommit(updatedFragments, fieldsModified, updatedFragmentOffsets); } } @@ -275,10 +289,16 @@ public String toString() { public static class TaskCommit implements WriterCommitMessage { private final List updatedFragments; private final long[] fieldsModified; + private final Map updatedFragmentOffsets; - TaskCommit(List updatedFragments, long[] fieldsModified) { + TaskCommit( + List updatedFragments, + long[] fieldsModified, + Map updatedFragmentOffsets) { this.updatedFragments = updatedFragments; this.fieldsModified = fieldsModified; + this.updatedFragmentOffsets = + updatedFragmentOffsets != null ? updatedFragmentOffsets : Collections.emptyMap(); } List getUpdatedFragments() { @@ -288,5 +308,9 @@ List getUpdatedFragments() { long[] getFieldsModified() { return fieldsModified; } + + Map getUpdatedFragmentOffsets() { + return updatedFragmentOffsets; + } } } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseUpdateColumnsBackfillTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseUpdateColumnsBackfillTest.java index c747b796c..d00366cb5 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseUpdateColumnsBackfillTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseUpdateColumnsBackfillTest.java @@ -51,6 +51,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; /** * Base test class for UPDATE COLUMNS FROM command. @@ -328,19 +329,11 @@ public void testUpdatePreservesRowIdAndFragId() { } /** - * Pins down the version-column behavior of UPDATE COLUMNS FROM on a stable-row-id table. - * - *

UPDATE COLUMNS goes through Lance's {@code Update} operation, which (unlike ADD COLUMNS via - * {@code Merge} and unlike row-level UPDATE) does not bump {@code - * _row_last_updated_at_version}. CDF consumers therefore cannot detect column-level rewrites via - * the version columns today. - * - *

This test pins down current behavior so a future change to make UPDATE COLUMNS CDF-aware - * shows up as a deliberate test update rather than a silent regression. - * - *

Tracking upstream fix: https://github.com/lance-format/lance/issues/6734 — once that lands, - * flip the {@code _row_last_updated_at_version} assertion below from {@code assertEquals} to a - * strict-greater check (mirroring the ADD COLUMNS version test). + * UPDATE COLUMNS FROM on a stable-row-id table must preserve {@code _row_created_at_version} for + * every row, advance {@code _row_last_updated_at_version} only for the matched row, and leave + * unmatched rows' last-updated version unchanged. The connector passes matched physical row + * offsets on commit so Lance can partially refresh last-updated metadata (see + * lance-format/lance#6734 and the Java {@code Update.updatedFragmentOffsets} API). */ @Test public void testUpdateColumnsPreservesCreatedAtAndAdvancesLastUpdatedWithStableRowIds() { @@ -354,9 +347,10 @@ public void testUpdateColumnsPreservesCreatedAtAndAdvancesLastUpdatedWithStableR fullTable)) .collectAsList(); + // Only update id=2; id=1 and id=3 are unmatched and must not have their last-updated advanced. spark.sql( String.format( - "CREATE TEMPORARY VIEW tmp_view_cdf AS SELECT _rowaddr, _fragid, value * 100 AS value FROM %s", + "CREATE TEMPORARY VIEW tmp_view_cdf AS SELECT _rowaddr, _fragid, value * 100 AS value FROM %s WHERE id = 2", fullTable)); spark.sql(String.format("ALTER TABLE %s UPDATE COLUMNS value FROM tmp_view_cdf", fullTable)); @@ -372,19 +366,20 @@ public void testUpdateColumnsPreservesCreatedAtAndAdvancesLastUpdatedWithStableR for (int i = 0; i < before.size(); i++) { Row b = before.get(i); Row a = after.get(i); - assertEquals(b.getInt(0), a.getInt(0)); + int id = b.getInt(0); + assertEquals(id, a.getInt(0)); assertEquals( - b.getLong(1), - a.getLong(1), - "_row_created_at_version must be unchanged for id=" + b.getInt(0)); - // Known gap (lance-format/lance#6734): UPDATE COLUMNS does not currently advance - // last_updated. When that issue is fixed, flip this assertion to a strict-greater check. - assertEquals( - b.getLong(2), - a.getLong(2), - "_row_last_updated_at_version is currently NOT advanced by UPDATE COLUMNS (id=" - + b.getInt(0) - + ") — if this changes, update the assertion"); + b.getLong(1), a.getLong(1), "_row_created_at_version must be unchanged for id=" + id); + if (id == 2) { + assertTrue( + a.getLong(2) > b.getLong(2), + "_row_last_updated_at_version must advance for matched id=" + id); + } else { + assertEquals( + b.getLong(2), + a.getLong(2), + "_row_last_updated_at_version must not change for unmatched id=" + id); + } } } From e1ca61a257db3e588d0f876eb9f70f27a646da6a Mon Sep 17 00:00:00 2001 From: fangbo Date: Tue, 18 Aug 2026 15:25:06 +0800 Subject: [PATCH 2/8] refactor: refactor Spark options to use LanceRef reference instead of raw version number (#770) Related https://github.com/lance-format/lance-spark/issues/713 ### Description This PR refactors Spark read and write options to use LanceRef instead of passing around raw version numbers directly. The goal is to make reference handling more explicit and extensible, especially for future branch and tag support. The change updates the Spark option model and related call sites to carry reference information through the connector, while preserving the ability to resolve snapshot versions where needed. This lays the groundwork for consistent handling of main, versioned, branch, and tag references across read, scan, blob/search, and write paths. ### Summary 1. Replace raw version-based Spark options with LanceRef. 2. Update read/write option APIs and affected Spark call sites. 3. Align test utilities with the new reference model. 4. Prepare the connector for branch/tag-aware reference handling. --------- Co-authored-by: fangbo --- .../spark/write/SparkPositionDeltaWrite.java | 12 +- .../spark/write/SparkPositionDeltaWrite.java | 13 ++- .../spark/BlobV2RowLevelContextTest.java | 2 +- .../main/java/org/lance/spark/LanceRef.java | 106 ++++++++++++++++++ .../lance/spark/LanceSparkReadOptions.java | 38 +++---- .../lance/spark/LanceSparkWriteOptions.java | 32 +++--- .../spark/internal/LanceFragmentScanner.java | 2 +- .../lance/spark/read/LanceScanBuilder.java | 8 +- .../java/org/lance/spark/utils/Utils.java | 13 ++- .../write/AddColumnsBackfillBatchWrite.java | 11 +- .../lance/spark/write/LanceBatchWrite.java | 12 +- .../UpdateColumnsBackfillBatchWrite.java | 12 +- .../LanceBlobSourceContextRule.scala | 6 +- .../java/org/lance/spark/LanceRefTest.java | 71 ++++++++++++ ...anceSparkReadOptionsSerializationTest.java | 26 ++--- .../spark/LanceSparkWriteOptionsTest.java | 22 ++-- .../spark/read/LanceScanBuilderTest.java | 5 +- .../org/lance/spark/read/LanceSplitTest.java | 8 +- .../org/lance/spark/write/SparkWriteTest.java | 5 +- .../scala/org/lance/spark/BlobPlanProbe.scala | 7 +- 20 files changed, 308 insertions(+), 103 deletions(-) create mode 100644 lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java create mode 100644 lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java diff --git a/lance-spark-3.4_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java b/lance-spark-3.4_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java index 576b24cd2..e07fedd1e 100644 --- a/lance-spark-3.4_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java +++ b/lance-spark-3.4_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java @@ -22,6 +22,7 @@ import org.lance.namespace.LanceNamespace; import org.lance.operation.Update; import org.lance.spark.LanceConstant; +import org.lance.spark.LanceRef; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkWriteOptions; import org.lance.spark.function.LanceFragmentIdWithDefaultFunction; @@ -95,9 +96,8 @@ public SparkPositionDeltaWrite( List tableId) { this.sparkSchema = sparkSchema; try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { - this.writeOptions = writeOptions.withVersion(ds.version()); - logger.debug( - "Resolved dataset version for position delta write: {}", this.writeOptions.getVersion()); + this.writeOptions = writeOptions.withRef(LanceRef.ofMain(ds.version())); + logger.debug("Resolved dataset ref for position delta write: {}", this.writeOptions.getRef()); } this.initialStorageOptions = initialStorageOptions; this.namespaceImpl = namespaceImpl; @@ -173,8 +173,10 @@ public void commit(WriterCommitMessage[] messages) { long version = Objects.requireNonNull( - writeOptions.getVersion(), - "version must be set (resolved in SparkPositionDeltaWrite constructor)"); + writeOptions.getRef(), + "ref must be set (resolved in SparkPositionDeltaWrite constructor)") + .getVersionNumber() + .get(); try (Dataset dataset = Utils.openDatasetBuilder(writeOptions).build()) { // Parallel stream is safe: each deleteRows() operates on an independent // FileFragment value writing to a distinct object store path (see lance-core). diff --git a/lance-spark-3.5_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java b/lance-spark-3.5_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java index 707a1642b..c3921fd45 100644 --- a/lance-spark-3.5_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java +++ b/lance-spark-3.5_2.12/src/main/java/org/lance/spark/write/SparkPositionDeltaWrite.java @@ -23,6 +23,7 @@ import org.lance.namespace.LanceNamespace; import org.lance.operation.Update; import org.lance.spark.LanceConstant; +import org.lance.spark.LanceRef; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkCatalogConfig; import org.lance.spark.LanceSparkWriteOptions; @@ -110,11 +111,11 @@ public SparkPositionDeltaWrite( Map blobSourceContexts) { this.sparkSchema = sparkSchema; try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { - this.writeOptions = writeOptions.withVersion(ds.version()); + this.writeOptions = writeOptions.withRef(LanceRef.ofMain(ds.version())); this.hasStableRowIds = hasStableRowIds(ds, writeOptions); LOG.debug( - "Resolved dataset version for position delta write: {}, stableRowIds={}", - this.writeOptions.getVersion(), + "Resolved dataset ref for position delta write: {}, stableRowIds={}", + this.writeOptions.getRef(), this.hasStableRowIds); } this.initialStorageOptions = initialStorageOptions; @@ -197,8 +198,10 @@ public void commit(WriterCommitMessage[] messages) { long version = Objects.requireNonNull( - writeOptions.getVersion(), - "version must be set (resolved in SparkPositionDeltaWrite constructor)"); + writeOptions.getRef(), + "ref must be set (resolved in SparkPositionDeltaWrite constructor)") + .getVersionNumber() + .get(); try (Dataset dataset = Utils.openDatasetBuilder(writeOptions).build()) { List> deletionResults = aggregatedDeletions.entrySet().parallelStream() diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/BlobV2RowLevelContextTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/BlobV2RowLevelContextTest.java index 68b51fa0c..900f1ee1f 100644 --- a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/BlobV2RowLevelContextTest.java +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/BlobV2RowLevelContextTest.java @@ -55,7 +55,7 @@ public void testMergeCopyingSourceBlobAttachesPinnedSourceContexts() throws Exce Map contexts = operation.blobSourceContexts(); assertFalse(contexts.isEmpty()); - contexts.values().forEach(context -> assertNotNull(context.getReadOptions().getVersion())); + contexts.values().forEach(context -> assertNotNull(context.getReadOptions().getRef())); } finally { spark.sql("DROP TABLE IF EXISTS " + fqSrc); spark.sql("DROP TABLE IF EXISTS " + fqTgt); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java new file mode 100644 index 000000000..05dd55357 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java @@ -0,0 +1,106 @@ +/* + * 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 org.lance.spark; + +import org.lance.spark.utils.Optional; + +import com.google.common.base.Preconditions; + +import java.io.Serializable; +import java.util.Objects; + +public class LanceRef implements Serializable { + private static final long serialVersionUID = 7967911672598737764L; + + private final Optional versionNumber; + private final Optional branchName; + private final Optional tagName; + + public LanceRef( + Optional versionNumber, Optional branchName, Optional tagName) { + this.versionNumber = versionNumber; + this.branchName = branchName; + this.tagName = tagName; + } + + public Optional getVersionNumber() { + return versionNumber; + } + + public Optional getBranchName() { + return branchName; + } + + public Optional getTagName() { + return tagName; + } + + public static LanceRef ofMain(long versionNumber) { + Preconditions.checkArgument(versionNumber > 0, "versionNumber must be greater than 0"); + return new LanceRef(Optional.of(versionNumber), Optional.empty(), Optional.empty()); + } + + public static LanceRef ofMain() { + return new LanceRef(Optional.empty(), Optional.empty(), Optional.empty()); + } + + public static LanceRef ofBranch(String branchName) { + Preconditions.checkArgument( + branchName != null && !branchName.isEmpty(), "branchName must not be empty"); + return new LanceRef(Optional.empty(), Optional.of(branchName), Optional.empty()); + } + + public static LanceRef ofBranch(String branchName, long versionNumber) { + Preconditions.checkArgument( + branchName != null && !branchName.isEmpty(), "branchName must not be empty"); + Preconditions.checkArgument(versionNumber > 0, "versionNumber must be greater than 0"); + return new LanceRef(Optional.of(versionNumber), Optional.of(branchName), Optional.empty()); + } + + public static LanceRef ofTag(String tagName) { + Preconditions.checkArgument(tagName != null && !tagName.isEmpty(), "tagName must not be empty"); + return new LanceRef(Optional.empty(), Optional.empty(), Optional.of(tagName)); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + LanceRef lanceRef = (LanceRef) o; + return Objects.equals(versionNumber, lanceRef.versionNumber) + && Objects.equals(branchName, lanceRef.branchName) + && Objects.equals(tagName, lanceRef.tagName); + } + + @Override + public int hashCode() { + return Objects.hash(versionNumber, branchName, tagName); + } + + @Override + public String toString() { + return "LanceRef{" + + "versionNumber=" + + (versionNumber.isPresent() ? versionNumber.get() : null) + + ", branchName=" + + (branchName.isPresent() ? branchName.get() : null) + + ", tagName=" + + (tagName.isPresent() ? tagName.get() : null) + + '}'; + } +} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java index 83f08d784..149c846e6 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java @@ -47,7 +47,7 @@ * } */ public class LanceSparkReadOptions implements Serializable { - private static final long serialVersionUID = 3L; + private static final long serialVersionUID = 4L; public static final String CONFIG_DATASET_URI = "path"; public static final String CONFIG_PUSH_DOWN_FILTERS = "pushDownFilters"; @@ -109,7 +109,7 @@ public class LanceSparkReadOptions implements Serializable { private final String datasetName; private final boolean pushDownFilters; private final Integer blockSize; - private final Long version; + private final LanceRef ref; private final Integer indexCacheSize; private final Integer metadataCacheSize; private final int batchSize; @@ -145,7 +145,7 @@ private LanceSparkReadOptions(Builder builder) { this.datasetName = paths[1]; this.pushDownFilters = builder.pushDownFilters; this.blockSize = builder.blockSize; - this.version = builder.version; + this.ref = builder.ref; this.indexCacheSize = builder.indexCacheSize; this.metadataCacheSize = builder.metadataCacheSize; this.batchSize = builder.batchSize; @@ -249,8 +249,8 @@ public Integer getBlockSize() { return blockSize; } - public Long getVersion() { - return version; + public LanceRef getRef() { + return ref; } public Integer getIndexCacheSize() { @@ -324,19 +324,19 @@ public void setNamespace(LanceNamespace namespace) { } /** - * Creates a copy of this options with a different version. + * Creates a copy of these options with a different reference. * - *

This is used to pin the version during scan planning for snapshot isolation. + *

This is used to pin the resolved reference during scan planning for snapshot isolation. * - * @param newVersion the version to use - * @return a new LanceSparkReadOptions with the specified version + * @param newRef the reference to use + * @return new read options with the specified reference */ - public LanceSparkReadOptions withVersion(long newVersion) { + public LanceSparkReadOptions withRef(LanceRef newRef) { return builder() .datasetUri(this.datasetUri) .pushDownFilters(this.pushDownFilters) .blockSize(this.blockSize) - .version(newVersion) + .ref(newRef) .indexCacheSize(this.indexCacheSize) .metadataCacheSize(this.metadataCacheSize) .batchSize(this.batchSize) @@ -364,8 +364,8 @@ public ReadOptions toReadOptions() { if (blockSize != null) { builder.setBlockSize(blockSize); } - if (version != null) { - builder.setVersion(version); + if (ref != null) { + ref.getVersionNumber().ifPresent(builder::setVersion); } if (indexCacheSize != null) { builder.setIndexCacheSize(indexCacheSize); @@ -404,7 +404,7 @@ public boolean equals(Object o) { && FullTextQueryUtils.equals(fullTextQuery, that.fullTextQuery) && Objects.equals(datasetUri, that.datasetUri) && Objects.equals(blockSize, that.blockSize) - && Objects.equals(version, that.version) + && Objects.equals(ref, that.ref) && Objects.equals(indexCacheSize, that.indexCacheSize) && Objects.equals(metadataCacheSize, that.metadataCacheSize) && Objects.equals(storageOptions, that.storageOptions) @@ -420,7 +420,7 @@ public int hashCode() { datasetUri, pushDownFilters, blockSize, - version, + ref, indexCacheSize, metadataCacheSize, batchSize, @@ -441,7 +441,7 @@ public static class Builder { private boolean pushDownFilters = DEFAULT_PUSH_DOWN_FILTERS; private Integer blockSize; private FullTextQuery fullTextQuery; - private Long version; + private LanceRef ref; private Integer indexCacheSize; private Integer metadataCacheSize; private int batchSize = DEFAULT_BATCH_SIZE; @@ -477,8 +477,8 @@ public Builder fullTextQuery(FullTextQuery fullTextQuery) { return this; } - public Builder version(Long version) { - this.version = version; + public Builder ref(LanceRef ref) { + this.ref = ref; return this; } @@ -590,7 +590,7 @@ private void parseTypedFlags(Map opts) { this.blockSize = Integer.parseInt(opts.get(CONFIG_BLOCK_SIZE)); } if (opts.containsKey(CONFIG_VERSION)) { - this.version = Long.parseLong(opts.get(CONFIG_VERSION)); + this.ref = LanceRef.ofMain(Long.parseLong(opts.get(CONFIG_VERSION))); } if (opts.containsKey(CONFIG_INDEX_CACHE_SIZE)) { this.indexCacheSize = Integer.parseInt(opts.get(CONFIG_INDEX_CACHE_SIZE)); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java index 3fcab33f8..362b588a0 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java @@ -46,7 +46,7 @@ * } */ public class LanceSparkWriteOptions implements Serializable { - private static final long serialVersionUID = 1L; + private static final long serialVersionUID = 2L; public static final String CONFIG_DATASET_URI = "path"; public static final String CONFIG_WRITE_MODE = "write_mode"; @@ -96,8 +96,8 @@ public class LanceSparkWriteOptions implements Serializable { /** The table identifier within the namespace, used for credential refresh. */ private final List tableId; - /** Use this version to open the dataset and apply write if set. */ - private final Long version; + /** Use this reference to open the dataset and apply write if set. */ + private final LanceRef ref; /** Catalog and cache backend configuration for process-local session reuse. */ private final String catalogName; @@ -123,7 +123,7 @@ private LanceSparkWriteOptions(Builder builder) { this.storageOptions = new HashMap<>(builder.storageOptions); this.namespace = builder.namespace; this.tableId = builder.tableId; - this.version = builder.version; + this.ref = builder.ref; this.catalogName = builder.catalogName; this.indexCacheBackend = builder.indexCacheBackend; this.metadataCacheBackend = builder.metadataCacheBackend; @@ -223,8 +223,8 @@ public List getTableId() { return tableId; } - public Long getVersion() { - return version; + public LanceRef getRef() { + return ref; } public String getCatalogName() { @@ -258,15 +258,15 @@ public Builder toBuilder() { .storageOptions(storageOptions) .namespace(namespace) .tableId(tableId) - .version(version) + .ref(ref) .catalogName(catalogName) .indexCacheBackend(indexCacheBackend) .metadataCacheBackend(metadataCacheBackend); } - /** Returns a copy of these options with version set to the given version. */ - public LanceSparkWriteOptions withVersion(long version) { - return toBuilder().version(version).build(); + /** Returns a copy of these options with the given reference. */ + public LanceSparkWriteOptions withRef(LanceRef ref) { + return toBuilder().ref(ref).build(); } public boolean hasNamespace() { @@ -353,7 +353,7 @@ public boolean equals(Object o) { && Objects.equals(blobPackFileSizeThreshold, that.blobPackFileSizeThreshold) && Objects.equals(storageOptions, that.storageOptions) && Objects.equals(tableId, that.tableId) - && Objects.equals(version, that.version) + && Objects.equals(ref, that.ref) && Objects.equals(catalogName, that.catalogName) && Objects.equals(indexCacheBackend, that.indexCacheBackend) && Objects.equals(metadataCacheBackend, that.metadataCacheBackend); @@ -377,7 +377,7 @@ public int hashCode() { blobPackFileSizeThreshold, storageOptions, tableId, - version, + ref, catalogName, indexCacheBackend, metadataCacheBackend); @@ -401,7 +401,7 @@ public static class Builder { private Map storageOptions = new HashMap<>(); private LanceNamespace namespace; private List tableId; - private Long version; + private LanceRef ref; private String catalogName; private String indexCacheBackend; private String metadataCacheBackend; @@ -492,9 +492,9 @@ public Builder tableId(List tableId) { return this; } - /** Pin opens to this dataset manifest version. */ - public Builder version(Long version) { - this.version = version; + /** Pin opens to this dataset reference. */ + public Builder ref(LanceRef ref) { + this.ref = ref; return this; } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java index a235a1171..ca11ae6bb 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java @@ -101,7 +101,7 @@ public static LanceFragmentScanner create(int fragmentId, LanceInputPartition in throw new IllegalStateException( String.format( "Fragment %d not found in dataset at %s (version=%s)", - fragmentId, readOptions.getDatasetUri(), readOptions.getVersion())); + fragmentId, readOptions.getDatasetUri(), readOptions.getRef())); } ScanOptions.Builder scanOptions = new ScanOptions.Builder(); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java index dd077184a..685edb7d3 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java @@ -25,6 +25,7 @@ import org.lance.schema.LanceField; import org.lance.schema.LanceSchema; import org.lance.spark.LanceConstant; +import org.lance.spark.LanceRef; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.search.LanceSearchQuery; @@ -269,7 +270,7 @@ public Scan build() { // long-lived high-write-frequency datasets do not silently truncate to a wrong version. LanceSplit.ScanPlanResult scanPlan = LanceSplit.planScan(dataset); LanceSparkReadOptions resolvedReadOptions = - readOptions.withVersion(scanPlan.getResolvedVersion()); + readOptions.withRef(LanceRef.ofMain(scanPlan.getResolvedVersion())); Optional whereCondition = FilterPushDown.compileFiltersToSqlWhereClause(pushedPredicates); @@ -336,7 +337,10 @@ private Scan buildNamespaceFtsScan() { .topK(k) .offset(pushedOffset) .filter(whereCondition.isPresent() ? whereCondition.get() : null) - .version(readOptions.getVersion()) + .version( + readOptions.getRef() == null || readOptions.getRef().getVersionNumber().isEmpty() + ? null + : readOptions.getRef().getVersionNumber().get()) .withRowId(withRowId ? Boolean.TRUE : null) .build(); return new LanceSearchScan(schema, query); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java index 797997ba4..277f0aa2b 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java @@ -17,6 +17,7 @@ import org.lance.ReadOptions; import org.lance.Version; import org.lance.namespace.LanceNamespace; +import org.lance.spark.LanceRef; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkCatalogConfig; import org.lance.spark.LanceSparkReadOptions; @@ -74,7 +75,7 @@ public static class OpenDatasetBuilder { private final String catalogName; private final String indexCacheBackend; private final String metadataCacheBackend; - private final Long version; + private final LanceRef ref; private final Integer blockSize; private final Integer indexCacheSize; private final Integer metadataCacheSize; @@ -87,7 +88,7 @@ public static class OpenDatasetBuilder { private OpenDatasetBuilder(LanceSparkReadOptions opts) { this.uri = opts.getDatasetUri(); this.storageOptions = opts.getStorageOptions(); - this.version = opts.getVersion(); + this.ref = opts.getRef(); this.catalogName = opts.getCatalogName(); this.indexCacheBackend = opts.getIndexCacheBackend(); this.metadataCacheBackend = opts.getMetadataCacheBackend(); @@ -106,7 +107,7 @@ private OpenDatasetBuilder(LanceSparkWriteOptions opts) { this.catalogName = opts.getCatalogName(); this.indexCacheBackend = opts.getIndexCacheBackend(); this.metadataCacheBackend = opts.getMetadataCacheBackend(); - this.version = opts.getVersion(); + this.ref = opts.getRef(); this.blockSize = null; this.indexCacheSize = null; this.metadataCacheSize = null; @@ -138,8 +139,8 @@ public Dataset build() { .setStorageOptions(merged) .setSession( LanceRuntime.session(catalogName, indexCacheBackend, metadataCacheBackend)); - if (version != null) { - roBuilder.setVersion(version); + if (ref != null) { + ref.getVersionNumber().ifPresent(roBuilder::setVersion); } if (blockSize != null) { roBuilder.setBlockSize(blockSize); @@ -205,7 +206,7 @@ public static LanceSparkReadOptions createReadOptions( .catalogName(catalogName); if (versionId.isPresent()) { - builder.version(versionId.get()); + builder.ref(LanceRef.ofMain(versionId.get())); } if (tableId.isPresent()) { builder.tableId(tableId.get()); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/AddColumnsBackfillBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/AddColumnsBackfillBatchWrite.java index 398dd6cfe..870b1ee35 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/AddColumnsBackfillBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/AddColumnsBackfillBatchWrite.java @@ -21,6 +21,7 @@ import org.lance.fragment.FragmentMergeResult; import org.lance.operation.Merge; import org.lance.spark.LanceDataset; +import org.lance.spark.LanceRef; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkWriteOptions; import org.lance.spark.utils.Utils; @@ -75,8 +76,8 @@ public AddColumnsBackfillBatchWrite( List tableId) { this.schema = schema; try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { - this.writeOptions = writeOptions.withVersion(ds.version()); - logger.debug("Resolved dataset version for ADD COLUMNS: {}", this.writeOptions.getVersion()); + this.writeOptions = writeOptions.withRef(LanceRef.ofMain(ds.version())); + logger.debug("Resolved dataset ref for ADD COLUMNS: {}", this.writeOptions.getRef()); } this.newColumns = newColumns; this.initialStorageOptions = initialStorageOptions; @@ -137,8 +138,10 @@ public void commit(WriterCommitMessage[] messages) { Schema arrowSchema = LanceArrowUtils.toArrowSchema(sparkSchema, "UTC", false); long version = Objects.requireNonNull( - writeOptions.getVersion(), - "version must be set (resolved in AddColumnsBackfillBatchWrite constructor)"); + writeOptions.getRef(), + "ref must be set (resolved in AddColumnsBackfillBatchWrite constructor)") + .getVersionNumber() + .get(); // Get existing fragments try (Dataset dataset = Utils.openDatasetBuilder(writeOptions).build()) { diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java index d3792dac8..82a7b3974 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java @@ -22,6 +22,7 @@ import org.lance.operation.Append; import org.lance.operation.Operation; import org.lance.operation.Overwrite; +import org.lance.spark.LanceRef; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkWriteOptions; import org.lance.spark.utils.BlobSourceContext; @@ -128,9 +129,8 @@ public LanceBatchWrite( this.writeOptions = writeOptions; } else { try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { - this.writeOptions = writeOptions.withVersion(ds.version()); - logger.debug( - "Resolved dataset version for batch write: {}", this.writeOptions.getVersion()); + this.writeOptions = writeOptions.withRef(LanceRef.ofMain(ds.version())); + logger.debug("Resolved dataset ref for batch write: {}", this.writeOptions.getRef()); } } } @@ -192,8 +192,10 @@ public void commit(WriterCommitMessage[] messages) { // For non-staged tables, commit immediately long version = Objects.requireNonNull( - writeOptions.getVersion(), - "version must be set (resolved in LanceBatchWrite constructor)"); + writeOptions.getRef(), + "ref must be set (resolved in LanceBatchWrite constructor)") + .getVersionNumber() + .get(); try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { Operation operation; if (isOverwrite) { diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java index 895086964..366b213d8 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/UpdateColumnsBackfillBatchWrite.java @@ -21,6 +21,7 @@ import org.lance.fragment.FragmentUpdateResult; import org.lance.operation.Update; import org.lance.spark.LanceDataset; +import org.lance.spark.LanceRef; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkWriteOptions; import org.lance.spark.utils.Utils; @@ -84,9 +85,8 @@ public UpdateColumnsBackfillBatchWrite( List tableId) { this.schema = schema; try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { - this.writeOptions = writeOptions.withVersion(ds.version()); - logger.debug( - "Resolved dataset version for UPDATE COLUMNS: {}", this.writeOptions.getVersion()); + this.writeOptions = writeOptions.withRef(LanceRef.ofMain(ds.version())); + logger.debug("Resolved dataset ref for UPDATE COLUMNS: {}", this.writeOptions.getRef()); } this.updateColumns = updateColumns; this.initialStorageOptions = initialStorageOptions; @@ -163,8 +163,10 @@ public void commit(WriterCommitMessage[] messages) { .build(); long version = Objects.requireNonNull( - writeOptions.getVersion(), - "version must be set (resolved in UpdateColumnsBackfillBatchWrite constructor)"); + writeOptions.getRef(), + "ref must be set (resolved in UpdateColumnsBackfillBatchWrite constructor)") + .getVersionNumber() + .get(); CommitBuilder commitBuilder = new CommitBuilder(dataset) .writeParams( diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/optimizer/LanceBlobSourceContextRule.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/optimizer/LanceBlobSourceContextRule.scala index 2d28b657d..40595ec0c 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/optimizer/LanceBlobSourceContextRule.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/optimizer/LanceBlobSourceContextRule.scala @@ -20,7 +20,7 @@ import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation} import org.apache.spark.sql.util.{CaseInsensitiveStringMap, LanceSerializeUtil} -import org.lance.spark.{LanceConstant, LanceDataset, LanceSparkReadOptions} +import org.lance.spark.{LanceConstant, LanceDataset, LanceRef, LanceSparkReadOptions} import org.lance.spark.read.LanceScan import org.lance.spark.utils.{BlobSourceContext, BlobUtils, Utils} @@ -133,12 +133,12 @@ object LanceBlobSourceContextRule extends Logging { // Pin to the driver-visible version when time travel is not set; fall back on open failure. private def pinToCurrentVersion(ds: LanceDataset): LanceSparkReadOptions = { val opts = ds.readOptions() - if (opts.getVersion != null) { + if (opts.getRef != null) { return opts } try { val dataset = Utils.openDatasetBuilder(opts).build() - try opts.withVersion(dataset.version()) + try opts.withRef(LanceRef.ofMain(dataset.version())) finally dataset.close() } catch { case NonFatal(e) => diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java new file mode 100644 index 000000000..38e5b45c7 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java @@ -0,0 +1,71 @@ +/* + * 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 org.lance.spark; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; + +public class LanceRefTest { + + @Test + public void refsUseValueEquality() { + assertEquals(LanceRef.ofMain(), LanceRef.ofMain()); + assertEquals(LanceRef.ofMain(7), LanceRef.ofMain(7)); + assertEquals(LanceRef.ofBranch("dev"), LanceRef.ofBranch("dev")); + assertEquals(LanceRef.ofBranch("dev", 7), LanceRef.ofBranch("dev", 7)); + assertEquals(LanceRef.ofTag("release"), LanceRef.ofTag("release")); + + assertNotEquals(LanceRef.ofMain(7), LanceRef.ofMain(8)); + assertNotEquals(LanceRef.ofBranch("dev"), LanceRef.ofBranch("prod")); + assertNotEquals(LanceRef.ofTag("release"), LanceRef.ofTag("latest")); + } + + @Test + public void equalRefsHaveEqualHashCodes() { + assertEquals(LanceRef.ofMain(7).hashCode(), LanceRef.ofMain(7).hashCode()); + assertEquals(LanceRef.ofBranch("dev", 7).hashCode(), LanceRef.ofBranch("dev", 7).hashCode()); + assertEquals(LanceRef.ofTag("release").hashCode(), LanceRef.ofTag("release").hashCode()); + } + + @Test + public void optionsCompareRefsByValue() { + LanceSparkReadOptions readLeft = + LanceSparkReadOptions.builder() + .datasetUri("file:///tmp/test") + .ref(LanceRef.ofBranch("dev", 7)) + .build(); + LanceSparkReadOptions readRight = + LanceSparkReadOptions.builder() + .datasetUri("file:///tmp/test") + .ref(LanceRef.ofBranch("dev", 7)) + .build(); + assertEquals(readLeft, readRight); + assertEquals(readLeft.hashCode(), readRight.hashCode()); + + LanceSparkWriteOptions writeLeft = + LanceSparkWriteOptions.builder() + .datasetUri("file:///tmp/test") + .ref(LanceRef.ofTag("release")) + .build(); + LanceSparkWriteOptions writeRight = + LanceSparkWriteOptions.builder() + .datasetUri("file:///tmp/test") + .ref(LanceRef.ofTag("release")) + .build(); + assertEquals(writeLeft, writeRight); + assertEquals(writeLeft.hashCode(), writeRight.hashCode()); + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java index ca790e97e..c19b126a1 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java @@ -61,15 +61,15 @@ public void testFullTextQuerySerializationRoundTrip() throws IOException, ClassN } @Test - public void testWithVersionPropagatesFullTextQuery() { + public void testWithRefPropagatesFullTextQuery() { FullTextQuery fts = FullTextQuery.phrase("quick brown fox", "body", 1); LanceSparkReadOptions options = LanceSparkReadOptions.builder().datasetUri("s3://bucket/path").fullTextQuery(fts).build(); - LanceSparkReadOptions versioned = options.withVersion(42); + LanceSparkReadOptions versioned = options.withRef(LanceRef.ofMain(42)); Assertions.assertNotNull( - versioned.getFullTextQuery(), "fullTextQuery must not be dropped by withVersion()"); - Assertions.assertEquals(42, versioned.getVersion()); + versioned.getFullTextQuery(), "fullTextQuery must not be dropped by withRef()"); + Assertions.assertEquals(42, versioned.getRef().getVersionNumber().get()); Assertions.assertTrue( org.lance.spark.utils.FullTextQueryUtils.equals( options.getFullTextQuery(), versioned.getFullTextQuery())); @@ -125,15 +125,15 @@ public void testMultiMatchQuerySerializationRoundTrip() } @Test - public void testWithVersionPropagatesMultiMatchQuery() { + public void testWithRefPropagatesMultiMatchQuery() { FullTextQuery fts = FullTextQuery.multiMatch("quick fox", Arrays.asList("title", "body")); LanceSparkReadOptions options = LanceSparkReadOptions.builder().datasetUri("s3://bucket/path").fullTextQuery(fts).build(); - LanceSparkReadOptions versioned = options.withVersion(7); + LanceSparkReadOptions versioned = options.withRef(LanceRef.ofMain(7)); Assertions.assertNotNull( - versioned.getFullTextQuery(), "fullTextQuery must not be dropped by withVersion()"); - Assertions.assertEquals(7, versioned.getVersion()); + versioned.getFullTextQuery(), "fullTextQuery must not be dropped by withRef()"); + Assertions.assertEquals(7, versioned.getRef().getVersionNumber().get()); Assertions.assertTrue( org.lance.spark.utils.FullTextQueryUtils.equals( options.getFullTextQuery(), versioned.getFullTextQuery())); @@ -276,7 +276,7 @@ public void testCacheBackendConfigurationSurvivesSerialization() } @Test - public void testWithVersionPreservesCacheBackendConfiguration() { + public void testWithRefPreservesCacheBackendConfiguration() { LanceSparkReadOptions options = LanceSparkReadOptions.builder() .datasetUri("s3://bucket/path") @@ -285,7 +285,7 @@ public void testWithVersionPreservesCacheBackendConfiguration() { .metadataCacheBackend("moka://?capacity=524288") .build(); - LanceSparkReadOptions pinned = options.withVersion(7); + LanceSparkReadOptions pinned = options.withRef(LanceRef.ofMain(7)); Assertions.assertEquals("cache-catalog", pinned.getCatalogName()); Assertions.assertEquals("moka://?capacity=1048576", pinned.getIndexCacheBackend()); @@ -293,17 +293,17 @@ public void testWithVersionPreservesCacheBackendConfiguration() { } @Test - public void testExecutorCredentialRefreshPreservedByWithVersion() { + public void testExecutorCredentialRefreshPreservedByWithRef() { LanceSparkReadOptions options = LanceSparkReadOptions.builder() .datasetUri("s3://bucket/path") .executorCredentialRefresh(false) .build(); - LanceSparkReadOptions pinned = options.withVersion(7); + LanceSparkReadOptions pinned = options.withRef(LanceRef.ofMain(7)); Assertions.assertFalse( pinned.isExecutorCredentialRefresh(), - "withVersion() must propagate the executor_credential_refresh flag"); + "withRef() must propagate the executor_credential_refresh flag"); } @Test diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java index 0d1050ca6..dfd97021e 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java @@ -37,16 +37,16 @@ public class LanceSparkWriteOptionsTest { private final String TEMP_URL = "file:///tmp/test"; @Test - public void versionIsNullByDefault() { + public void refIsNullByDefault() { LanceSparkWriteOptions opts = LanceSparkWriteOptions.from(TEMP_URL); - assertNull(opts.getVersion()); + assertNull(opts.getRef()); } @Test - public void builderSetsVersion() { + public void builderSetsRef() { LanceSparkWriteOptions opts = - LanceSparkWriteOptions.builder().datasetUri(TEMP_URL).version(7L).build(); - assertEquals(7L, opts.getVersion()); + LanceSparkWriteOptions.builder().datasetUri(TEMP_URL).ref(LanceRef.ofMain(7L)).build(); + assertEquals(7L, opts.getRef().getVersionNumber().get()); } @Test @@ -67,7 +67,7 @@ public void fileFormatVersionUsesValueEquality() { } @Test - public void withVersionCopiesOptions() { + public void withRefCopiesOptions() { LanceSparkWriteOptions base = LanceSparkWriteOptions.builder() .datasetUri(TEMP_URL) @@ -75,9 +75,9 @@ public void withVersionCopiesOptions() { .indexCacheBackend("moka://?capacity=1048576") .metadataCacheBackend("moka://?capacity=524288") .build(); - LanceSparkWriteOptions pinned = base.withVersion(3L); - assertEquals(3L, pinned.getVersion()); - assertNull(base.getVersion()); + LanceSparkWriteOptions pinned = base.withRef(LanceRef.ofMain(3L)); + assertEquals(3L, pinned.getRef().getVersionNumber().get()); + assertNull(base.getRef()); assertEquals("cache-catalog", pinned.getCatalogName()); assertEquals("moka://?capacity=1048576", pinned.getIndexCacheBackend()); assertEquals("moka://?capacity=524288", pinned.getMetadataCacheBackend()); @@ -352,7 +352,7 @@ public void testRebuiltOverwriteOptionsSurviveJavaSerialization() throws Excepti .datasetUri(TEMP_URL) .fromOptions(options) .blobPackFileSizeThreshold(8192L) - .version(7L) + .ref(LanceRef.ofMain(7L)) .namespace(stubNamespace) .build() .toBuilder() @@ -375,7 +375,7 @@ public void testRebuiltOverwriteOptionsSurviveJavaSerialization() throws Excepti assertEquals(256, copy.getBatchSize()); assertEquals(4096L, copy.getMaxBatchBytes()); assertEquals(Long.valueOf(8192L), copy.getBlobPackFileSizeThreshold()); - assertEquals(7L, copy.getVersion()); + assertEquals(7L, copy.getRef().getVersionNumber().get()); assertNull( copy.getNamespace(), "namespace is transient: the non-null stub set above must not survive serialization"); diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java index 35cdcedad..850ea26e7 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java @@ -13,6 +13,7 @@ */ package org.lance.spark.read; +import org.lance.spark.LanceRef; import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.TestUtils; import org.lance.spark.utils.BlobUtils; @@ -348,9 +349,9 @@ public void testBuildPinsResolvedVersionOnReadOptions() { org.apache.spark.sql.connector.read.InputPartition[] partitions = scan.planInputPartitions(); assertTrue(partitions.length > 0); LanceInputPartition first = (LanceInputPartition) partitions[0]; - Long pinned = first.getReadOptions().getVersion(); + LanceRef pinned = first.getReadOptions().getRef(); assertNotNull(pinned, "build() must pin the resolved version onto readOptions"); - assertTrue(pinned > 0); + assertTrue(pinned.getVersionNumber().get() > 0); } @Test diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java index 92ae5fe94..a5b1aeacc 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java @@ -78,7 +78,7 @@ public void testPlanScanWithDatasetDoesNotCloseExternalDataset() { /** * Contract test: the long-typed resolved version returned by {@link LanceSplit#planScan(Dataset)} - * must round-trip through {@link org.lance.spark.LanceSparkReadOptions#withVersion(long)} without + * must round-trip through {@link org.lance.spark.LanceSparkReadOptions} as a LanceRef without * truncation. This guards against silently casting to {@code int}, which would corrupt the * snapshot-isolation guarantee for long-lived high-write-frequency datasets. */ @@ -88,7 +88,11 @@ public void testResolvedVersionRoundTripsAsLong() { long resolved = result.getResolvedVersion(); assertEquals( resolved, - TestUtils.TestTable1Config.readOptions.withVersion(resolved).getVersion().longValue()); + TestUtils.TestTable1Config.readOptions + .withRef(org.lance.spark.LanceRef.ofMain(resolved)) + .getRef() + .getVersionNumber() + .get()); } @SuppressWarnings("deprecation") diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java index 2e84bb76c..e4c8f05cc 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java @@ -17,6 +17,7 @@ import org.lance.WriteParams; import org.lance.memwal.InitializeMemWalParams; import org.lance.namespace.LanceNamespace; +import org.lance.spark.LanceRef; import org.lance.spark.LanceSparkWriteOptions; import org.lance.spark.TestUtils; @@ -158,7 +159,7 @@ public void testTruncatePreservesWriteOptionsAndOverwritesMode(TestInfo testInfo .storageOptions(storageOptions) .namespace(stubNamespace) .tableId(Arrays.asList("default", "test_table")) - .version(7L) + .ref(LanceRef.ofMain(7L)) .build(); SparkWrite.SparkWriteBuilder builder = new SparkWrite.SparkWriteBuilder( @@ -199,7 +200,7 @@ public void testTruncatePreservesWriteOptionsAndOverwritesMode(TestInfo testInfo assertEquals(storageOptions, truncatedOptions.getStorageOptions()); assertSame(stubNamespace, truncatedOptions.getNamespace()); assertEquals(Arrays.asList("default", "test_table"), truncatedOptions.getTableId()); - assertEquals(7L, truncatedOptions.getVersion()); + assertEquals(7L, truncatedOptions.getRef().getVersionNumber().get()); } // --- requiredDistribution / requiredOrdering tests --- diff --git a/lance-spark-base_2.12/src/test/scala/org/lance/spark/BlobPlanProbe.scala b/lance-spark-base_2.12/src/test/scala/org/lance/spark/BlobPlanProbe.scala index 231362690..de39c1dd7 100644 --- a/lance-spark-base_2.12/src/test/scala/org/lance/spark/BlobPlanProbe.scala +++ b/lance-spark-base_2.12/src/test/scala/org/lance/spark/BlobPlanProbe.scala @@ -35,7 +35,12 @@ object BlobPlanProbe { /** The dataset version each encoded blob source context would resolve against; null = latest. */ def blobSourceContextVersions(plan: LogicalPlan): java.util.Map[String, java.lang.Long] = { val versions = new java.util.HashMap[String, java.lang.Long]() - decodedContexts(plan).forEach((uri, ctx) => versions.put(uri, ctx.getReadOptions.getVersion)) + decodedContexts(plan).forEach { (uri, ctx) => + val ref = ctx.getReadOptions.getRef + val version = + if (ref == null || ref.getVersionNumber.isEmpty) null else ref.getVersionNumber.get + versions.put(uri, version) + } versions } From 37b56fe2e09b40c1c8cb8610e68a9f1e9f995f9f Mon Sep 17 00:00:00 2001 From: fangbo Date: Fri, 21 Aug 2026 15:28:46 +0800 Subject: [PATCH 3/8] feat: support DQL for specific tag (#773) Related https://github.com/lance-format/lance-spark/issues/713 # Motivation In some use cases, users need to run DQL against a specific tag rather than relying on the default resolution behavior. Tags are read-only. Tags could just stay on time travel. # How to query ```SQL SELECT * FROM lance.users VERSION AS OF 'specific_tag'; ``` # Changes * Add support for executing DQL against a specific tag. * Enable more explicit tag-scoped reads in query workflows. * Improve flexibility for scenarios where users need to access data from a designated tag. # Testing * Added/updated tests for DQL queries with a specific tag. * Verified that existing query behavior remains unchanged when no tag is specified. * Validated the new tag-specific query path works as expected. --------- Co-authored-by: fangbo --- docs/src/operations/dql/select.md | 42 +++- integration-tests/test_lance_spark.py | 45 +++- .../spark/LancePositionDeltaDataset.java | 1 + .../java/org/lance/spark/tag/TagDQLTest.java | 16 ++ .../spark/LancePositionDeltaDataset.java | 1 + .../java/org/lance/spark/tag/TagDQLTest.java | 16 ++ .../spark/BaseLanceNamespaceSparkCatalog.java | 56 ++--- .../java/org/lance/spark/LanceDataset.java | 19 ++ .../lance/spark/read/LanceScanBuilder.java | 22 +- .../java/org/lance/spark/read/LanceSplit.java | 45 +++- .../java/org/lance/spark/utils/Utils.java | 47 +++- .../spark/read/LanceScanBuilderTest.java | 16 ++ .../org/lance/spark/read/LanceSplitTest.java | 22 +- .../org/lance/spark/tag/BaseTagDQLTest.java | 217 ++++++++++++++++++ 14 files changed, 498 insertions(+), 67 deletions(-) create mode 100644 lance-spark-3.4_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java create mode 100644 lance-spark-3.5_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java create mode 100644 lance-spark-base_2.12/src/test/java/org/lance/spark/tag/BaseTagDQLTest.java diff --git a/docs/src/operations/dql/select.md b/docs/src/operations/dql/select.md index ad649f008..2eea0ba50 100644 --- a/docs/src/operations/dql/select.md +++ b/docs/src/operations/dql/select.md @@ -199,6 +199,46 @@ Use `VERSION AS OF` to query a specific version of the table: df.show(); ``` +### Query by Tag + +Use `VERSION AS OF` with a quoted tag name to query the snapshot referenced by a tag: + +=== "SQL" + ```sql + -- Query the snapshot referenced by the release_candidate tag + SELECT * FROM users VERSION AS OF 'release_candidate'; + + -- Query specific columns from a tagged snapshot + SELECT id, name FROM users VERSION AS OF 'v1.0'; + ``` + +=== "Python" + ```python + # Query a tag using SQL + spark.sql("SELECT * FROM users VERSION AS OF 'release_candidate'").show() + ``` + +=== "Scala" + ```scala + // Query a tag using SQL + spark.sql("SELECT * FROM users VERSION AS OF 'release_candidate'").show() + ``` + +=== "Java" + ```java + // Query a tag using SQL + spark.sql("SELECT * FROM users VERSION AS OF 'release_candidate'").show(); + ``` + +!!! note + Tags whose names consist entirely of integer digits cannot be queried. Numeric values are + interpreted as table versions, even when quoted. For example, both `VERSION AS OF 123` and + `VERSION AS OF '123'` query table version 123 rather than a tag named `123`. Use a tag name that + contains at least one non-digit character. + +Tag queries are read-only. `UPDATE`, `DELETE`, `INSERT`, `MERGE INTO`, `ADD COLUMNS`, and +`UPDATE COLUMNS` operations cannot target a tagged snapshot. + ### Query by Timestamp Use `TIMESTAMP AS OF` to query the table as it existed at a specific point in time: @@ -305,4 +345,4 @@ Spark UDT columns (e.g. MLlib's `VectorUDT`) round-trip transparently. On write .format("lance") .option("batch_size", "4096") .load("/path/to/dataset.lance") - ``` + ``` \ No newline at end of file diff --git a/integration-tests/test_lance_spark.py b/integration-tests/test_lance_spark.py index 8d84f6022..a38f5c55f 100644 --- a/integration-tests/test_lance_spark.py +++ b/integration-tests/test_lance_spark.py @@ -2702,6 +2702,49 @@ def test_version_as_of(self, spark): assert len(result) == 1 assert result[0].id == 1 + def test_tag_as_of_excludes_data_inserted_after_tag_creation(self, spark): + """Test that a tag remains on its snapshot after the main table advances.""" + spark.sql(""" + CREATE TABLE default.test_table ( + id INT, + name STRING + ) + """) + spark.sql(""" + INSERT INTO default.test_table VALUES + (1, 'before_tag_1'), + (2, 'before_tag_2') + """) + spark.sql("ALTER TABLE default.test_table CREATE TAG stable") + + spark.sql(""" + INSERT INTO default.test_table VALUES + (3, 'after_tag_1'), + (4, 'after_tag_2') + """) + + tagged = spark.sql(""" + SELECT id, name + FROM default.test_table VERSION AS OF 'stable' + ORDER BY id + """).collect() + current = spark.sql(""" + SELECT id, name + FROM default.test_table + ORDER BY id + """).collect() + + assert [(row.id, row.name) for row in tagged] == [ + (1, "before_tag_1"), + (2, "before_tag_2"), + ] + assert [(row.id, row.name) for row in current] == [ + (1, "before_tag_1"), + (2, "before_tag_2"), + (3, "after_tag_1"), + (4, "after_tag_2"), + ] + @requires_update_or_merge def test_version_as_of_after_update(self, spark): """Test VERSION AS OF returns data before an update.""" @@ -3323,4 +3366,4 @@ def test_show_functions_lists_fts(self, spark): if __name__ == "__main__": - pytest.main([__file__, "-v"]) + pytest.main([__file__, "-v"]) \ No newline at end of file diff --git a/lance-spark-3.4_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java b/lance-spark-3.4_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java index 79340cee7..5a04b5e92 100644 --- a/lance-spark-3.4_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java +++ b/lance-spark-3.4_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java @@ -74,6 +74,7 @@ public LancePositionDeltaDataset( @Override public RowLevelOperationBuilder newRowLevelOperationBuilder( RowLevelOperationInfo rowLevelOperationInfo) { + ensureWritable(); return new LanceRowLevelOperationBuilder( rowLevelOperationInfo.command(), sparkSchema, diff --git a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java new file mode 100644 index 000000000..8f38a3c66 --- /dev/null +++ b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java @@ -0,0 +1,16 @@ +/* + * 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 org.lance.spark.tag; + +public class TagDQLTest extends BaseTagDQLTest {} diff --git a/lance-spark-3.5_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java b/lance-spark-3.5_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java index 79340cee7..5a04b5e92 100644 --- a/lance-spark-3.5_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java +++ b/lance-spark-3.5_2.12/src/main/java/org/lance/spark/LancePositionDeltaDataset.java @@ -74,6 +74,7 @@ public LancePositionDeltaDataset( @Override public RowLevelOperationBuilder newRowLevelOperationBuilder( RowLevelOperationInfo rowLevelOperationInfo) { + ensureWritable(); return new LanceRowLevelOperationBuilder( rowLevelOperationInfo.command(), sparkSchema, diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java new file mode 100644 index 000000000..8f38a3c66 --- /dev/null +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/tag/TagDQLTest.java @@ -0,0 +1,16 @@ +/* + * 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 org.lance.spark.tag; + +public class TagDQLTest extends BaseTagDQLTest {} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java index a4277fe10..859a4811f 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java @@ -1653,6 +1653,14 @@ private void updateDatasetConfig( } } + private LanceRef parseVersionRef(String version) { + try { + return LanceRef.ofMain(Utils.parseVersion(version)); + } catch (NumberFormatException e) { + return LanceRef.ofTag(version); + } + } + private Table loadTableInternal( Identifier ident, Optional timestamp, Optional version) throws NoSuchTableException { @@ -1666,31 +1674,21 @@ private Table loadTableInternal( DescribeTableResponse describeResponse = resolved.describeResponse; Map initialStorageOptions = describeResponse.getStorageOptions(); - Optional versionId = Optional.empty(); + Optional ref = Optional.empty(); if (timestamp.isPresent()) { try (Dataset dataset = Utils.openDatasetBuilder(resolved.readOptions).build()) { - versionId = Optional.of(Utils.findVersion(dataset.listVersions(), timestamp.get())); + ref = + Optional.of( + LanceRef.ofMain(Utils.findVersion(dataset.listVersions(), timestamp.get()))); } catch (TableNotFoundException e) { throw new NoSuchTableException(ident); } } else if (version.isPresent()) { - versionId = Optional.of(Utils.parseVersion(version.get())); + ref = Optional.of(parseVersionRef(version.get())); } - // If time travel requested, rebuild readOptions with the resolved version - LanceSparkReadOptions readOptions; - if (versionId.isPresent()) { - readOptions = - createReadOptions( - describeResponse.getLocation(), - catalogConfig, - versionId, - Optional.of(namespace), - Optional.of(resolved.tableIdList), - name); - } else { - readOptions = resolved.readOptions; - } + LanceSparkReadOptions readOptions = + ref.isPresent() ? resolved.readOptions.withRef(ref.get()) : resolved.readOptions; // Read schema, file format version, and config from the dataset String fileFormatVersion; @@ -1749,28 +1747,24 @@ private Table loadTableFromPath( throws NoSuchTableException { String datasetUri = getDatasetUri(ident); - Optional versionId = Optional.empty(); + LanceSparkReadOptions baseReadOptions = + createReadOptions( + datasetUri, catalogConfig, Optional.empty(), Optional.empty(), Optional.empty(), name); + Optional ref = Optional.empty(); if (version.isPresent()) { - versionId = Optional.of(Utils.parseVersion(version.get())); + ref = Optional.of(parseVersionRef(version.get())); } else if (timestamp.isPresent()) { - LanceSparkReadOptions readOptions = - createReadOptions( - datasetUri, - catalogConfig, - Optional.empty(), - Optional.empty(), - Optional.empty(), - name); - try (Dataset dataset = Utils.openDatasetBuilder(readOptions).build()) { - versionId = Optional.of(Utils.findVersion(dataset.listVersions(), timestamp.get())); + try (Dataset dataset = Utils.openDatasetBuilder(baseReadOptions).build()) { + ref = + Optional.of( + LanceRef.ofMain(Utils.findVersion(dataset.listVersions(), timestamp.get()))); } catch (IllegalArgumentException e) { throw new NoSuchTableException(ident); } } LanceSparkReadOptions readOptions = - createReadOptions( - datasetUri, catalogConfig, versionId, Optional.empty(), Optional.empty(), name); + ref.isPresent() ? baseReadOptions.withRef(ref.get()) : baseReadOptions; // Read schema, file format version, and config from the dataset String fileFormatVersion; diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java index 7498f8206..71a78dc33 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java @@ -56,6 +56,9 @@ public class LanceDataset private static final Logger LOG = LoggerFactory.getLogger(LanceDataset.class); + private static final Set READ_ONLY_CAPABILITIES = + ImmutableSet.of(TableCapability.BATCH_READ); + private static final Set CAPABILITIES = ImmutableSet.of( TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE); @@ -321,6 +324,7 @@ public ScanBuilder newScanBuilder(CaseInsensitiveStringMap caseInsensitiveString .indexCacheBackend(readOptions.getIndexCacheBackend()) .metadataCacheBackend(readOptions.getMetadataCacheBackend()) .fromOptions(mergedOptions) + .ref(readOptions.getRef()) .build(); } return new LanceScanBuilder( @@ -354,11 +358,26 @@ public Map properties() { @Override public Set capabilities() { + if (isTagReference()) { + return READ_ONLY_CAPABILITIES; + } return BlobUtils.hasBlobV2Fields(sparkSchema) ? CAPABILITIES_WITH_BLOB_V2 : CAPABILITIES; } + protected void ensureWritable() { + if (isTagReference()) { + throw new UnsupportedOperationException( + "Writes are not supported for tag: " + readOptions.getRef().getTagName().get()); + } + } + + private boolean isTagReference() { + return readOptions.getRef() != null && readOptions.getRef().getTagName().isPresent(); + } + @Override public WriteBuilder newWriteBuilder(LogicalWriteInfo logicalWriteInfo) { + ensureWritable(); if (capabilities().contains(TableCapability.ACCEPT_ANY_SCHEMA)) { LanceWriteSchemaValidator.validate(sparkSchema, logicalWriteInfo.schema()); } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java index 685edb7d3..f8a8c57a5 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java @@ -176,9 +176,7 @@ public Scan build() { // partition). A full-text query without a namespace, or against a catalog-only namespace such // as Glue that does not implement queryTable, falls through to the local per-fragment scan // below. COUNT(*) is excluded because countTableRows has no full-text field. - if (readOptions.getFullTextQuery() != null - && LanceRuntime.supportsQueryTable(namespaceImpl) - && !pushedAggregation.isPresent()) { + if (shouldNamespaceFtsScan()) { return buildNamespaceFtsScan(); } @@ -268,9 +266,8 @@ public Scan build() { // the resolved version onto the read options shipped to workers, providing snapshot // isolation across all tasks of this query. The version is kept as a long end-to-end so // long-lived high-write-frequency datasets do not silently truncate to a wrong version. - LanceSplit.ScanPlanResult scanPlan = LanceSplit.planScan(dataset); - LanceSparkReadOptions resolvedReadOptions = - readOptions.withRef(LanceRef.ofMain(scanPlan.getResolvedVersion())); + LanceSplit.ScanPlanResult scanPlan = LanceSplit.planScan(dataset, readOptions); + LanceSparkReadOptions resolvedReadOptions = readOptions.withRef(scanPlan.getRef()); Optional whereCondition = FilterPushDown.compileFiltersToSqlWhereClause(pushedPredicates); @@ -298,6 +295,19 @@ public Scan build() { } } + boolean shouldNamespaceFtsScan() { + LanceRef ref = readOptions.getRef(); + boolean hasTag = ref != null && ref.getTagName().isPresent(); + if (hasTag) { + return false; + } + + return readOptions.getFullTextQuery() != null + && LanceRuntime.supportsQueryTable(namespaceImpl) + && !pushedAggregation.isPresent() + && !hasTag; + } + /** * Builds a single-partition scan that runs the full-text query server-side through the namespace * {@code queryTable} endpoint. Reuses the search-package namespace scan/reader, driven from the diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java index 18c07962d..7ca6b7f1e 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java @@ -15,6 +15,8 @@ import org.lance.Dataset; import org.lance.Fragment; +import org.lance.Tag; +import org.lance.spark.LanceRef; import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.utils.Utils; @@ -24,6 +26,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; public class LanceSplit implements Serializable { private static final long serialVersionUID = 2983749283749283749L; @@ -41,15 +44,15 @@ public List getFragments() { /** Result of scan planning containing splits, resolved version, and per-fragment row counts. */ public static class ScanPlanResult { private final List splits; - private final long resolvedVersion; + private final LanceRef ref; /** Per-fragment logical row counts (after deletions). Key is fragment ID. */ private final Map fragmentRowCounts; public ScanPlanResult( - List splits, long resolvedVersion, Map fragmentRowCounts) { + List splits, LanceRef ref, Map fragmentRowCounts) { this.splits = splits; - this.resolvedVersion = resolvedVersion; + this.ref = ref; this.fragmentRowCounts = fragmentRowCounts; } @@ -57,8 +60,8 @@ public List getSplits() { return splits; } - public long getResolvedVersion() { - return resolvedVersion; + public LanceRef getRef() { + return ref; } public Map getFragmentRowCounts() { @@ -75,7 +78,7 @@ public Map getFragmentRowCounts() { */ public static ScanPlanResult planScan(LanceSparkReadOptions readOptions) { try (Dataset dataset = Utils.openDatasetBuilder(readOptions).build()) { - return planScan(dataset); + return planScan(dataset, readOptions); } } @@ -88,7 +91,7 @@ public static ScanPlanResult planScan(LanceSparkReadOptions readOptions) { * *

The caller retains ownership of the dataset; this method does not close it. */ - public static ScanPlanResult planScan(Dataset dataset) { + public static ScanPlanResult planScan(Dataset dataset, LanceSparkReadOptions readOptions) { List fragments = dataset.getFragments(); List splits = new ArrayList<>(fragments.size()); Map fragmentRowCounts = new HashMap<>(fragments.size()); @@ -97,8 +100,32 @@ public static ScanPlanResult planScan(Dataset dataset) { splits.add(new LanceSplit(Collections.singletonList(id))); fragmentRowCounts.put(id, fragment.metadata().getNumRows()); } - long resolvedVersion = dataset.getVersion().getId(); - return new ScanPlanResult(splits, resolvedVersion, fragmentRowCounts); + + LanceRef plannedRef; + + LanceRef ref = readOptions.getRef(); + if (ref != null && ref.getBranchName().isPresent()) { + plannedRef = LanceRef.ofBranch(ref.getBranchName().get(), dataset.getVersion().getId()); + } else if (ref != null && ref.getTagName().isPresent()) { + Optional tag = + dataset.tags().list().stream() + .filter(t -> t.getName().equals(ref.getTagName().get())) + .findFirst(); + if (tag.isEmpty()) { + throw new RuntimeException("No existed tag with name:" + ref.getTagName().get()); + } + + Optional tagBranch = tag.get().getBranch(); + if (tagBranch.isEmpty()) { + plannedRef = LanceRef.ofMain(dataset.getVersion().getId()); + } else { + plannedRef = LanceRef.ofBranch(tagBranch.get(), dataset.version()); + } + } else { + plannedRef = LanceRef.ofMain(dataset.getVersion().getId()); + } + + return new ScanPlanResult(splits, plannedRef, fragmentRowCounts); } /** diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java index 277f0aa2b..afad5b9e5 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java @@ -15,6 +15,7 @@ import org.lance.Dataset; import org.lance.ReadOptions; +import org.lance.Ref; import org.lance.Version; import org.lance.namespace.LanceNamespace; import org.lance.spark.LanceRef; @@ -129,35 +130,63 @@ public OpenDatasetBuilder runtimeNamespace( } public Dataset build() { + if (ref != null && (ref.getTagName().isPresent() || ref.getBranchName().isPresent())) { + // Open specific tag or branch/version + Dataset main = openMain(null); + try { + if (ref.getTagName().isPresent()) { + return main.checkout(Ref.ofTag(ref.getTagName().get())); + } else { + return ref.getVersionNumber().isPresent() + ? main.checkout( + Ref.ofBranch(ref.getBranchName().get(), ref.getVersionNumber().get())) + : main.checkout(Ref.ofBranch(ref.getBranchName().get())); + } + } finally { + main.close(); + } + } else { + return openMain( + ref != null && ref.getVersionNumber().isPresent() + ? ref.getVersionNumber().get() + : null); + } + } + + private Dataset openMain(Long version) { LanceRuntime.enableOpenTelemetry(); Map base = storageOptions != null ? storageOptions : Collections.emptyMap(); Map merged = LanceRuntime.mergeStorageOptions(base, initialStorageOptions); - ReadOptions.Builder roBuilder = + ReadOptions.Builder builder = new ReadOptions.Builder() .setStorageOptions(merged) .setSession( LanceRuntime.session(catalogName, indexCacheBackend, metadataCacheBackend)); - if (ref != null) { - ref.getVersionNumber().ifPresent(roBuilder::setVersion); + if (version != null) { + builder.setVersion(version); } if (blockSize != null) { - roBuilder.setBlockSize(blockSize); + builder.setBlockSize(blockSize); } if (indexCacheSize != null) { - roBuilder.setIndexCacheSize(indexCacheSize); + builder.setIndexCacheSize(indexCacheSize); } if (metadataCacheSize != null) { - roBuilder.setMetadataCacheSize(metadataCacheSize); + builder.setMetadataCacheSize(metadataCacheSize); } + return open(builder.build()); + } + + private Dataset open(ReadOptions readOptions) { if (namespace != null && tableId != null) { return Dataset.open() .allocator(LanceRuntime.allocator()) .namespaceClient(namespace) .tableId(tableId) - .readOptions(roBuilder.build()) + .readOptions(readOptions) .build(); } if (runtimeNamespaceImpl != null) { @@ -169,14 +198,14 @@ public Dataset build() { .allocator(LanceRuntime.allocator()) .namespaceClient(runtimeNamespace) .tableId(effectiveTableId) - .readOptions(roBuilder.build()) + .readOptions(readOptions) .build(); } } return Dataset.open() .allocator(LanceRuntime.allocator()) .uri(uri) - .readOptions(roBuilder.build()) + .readOptions(readOptions) .build(); } } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java index 850ea26e7..ff2b57a96 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java @@ -13,6 +13,7 @@ */ package org.lance.spark.read; +import org.lance.ipc.FullTextQuery; import org.lance.spark.LanceRef; import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.TestUtils; @@ -354,6 +355,21 @@ public void testBuildPinsResolvedVersionOnReadOptions() { assertTrue(pinned.getVersionNumber().get() > 0); } + @Test + public void testTagFullTextQueryDoesNotUseNamespaceScan() { + LanceSparkReadOptions options = + LanceSparkReadOptions.builder() + .datasetUri(TestUtils.TestTable1Config.datasetUri) + .ref(LanceRef.ofTag("stable")) + .fullTextQuery(FullTextQuery.match("hello", "b")) + .build(); + LanceScanBuilder builder = + new LanceScanBuilder( + TEST_SCHEMA, options, Collections.emptyMap(), "dir", Collections.emptyMap()); + + assertFalse(builder.shouldNamespaceFtsScan()); + } + @Test public void testBuildWithCountStarReturnsLocalScan() { LanceScanBuilder builder = createBuilder(); diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java index a5b1aeacc..91cc7909b 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceSplitTest.java @@ -14,6 +14,7 @@ package org.lance.spark.read; import org.lance.Dataset; +import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.TestUtils; import org.lance.spark.utils.Utils; @@ -29,7 +30,7 @@ public class LanceSplitTest { public void testPlanScanReturnsNonEmptySplits() { LanceSplit.ScanPlanResult result = LanceSplit.planScan(TestUtils.TestTable1Config.readOptions); assertFalse(result.getSplits().isEmpty()); - assertTrue(result.getResolvedVersion() > 0); + assertTrue(result.getRef().getVersionNumber().get() > 0); } @Test @@ -58,15 +59,16 @@ public void testPlanScanReturnsFragmentRowCounts() { } /** - * Contract test: {@link LanceSplit#planScan(Dataset)} must not close the externally-owned dataset - * handle. {@link LanceScanBuilder#build()} relies on this so it can keep using the single open - * handle for both manifest/zonemap loading and split planning. + * Contract test: {@link LanceSplit#planScan(Dataset, LanceSparkReadOptions)} must not close the + * externally-owned dataset handle. {@link LanceScanBuilder#build()} relies on this so it can keep + * using the single open handle for both manifest/zonemap loading and split planning. */ @Test public void testPlanScanWithDatasetDoesNotCloseExternalDataset() { try (Dataset dataset = Utils.openDatasetBuilder(TestUtils.TestTable1Config.readOptions).build()) { - LanceSplit.ScanPlanResult result = LanceSplit.planScan(dataset); + LanceSplit.ScanPlanResult result = + LanceSplit.planScan(dataset, TestUtils.TestTable1Config.readOptions); assertFalse(result.getSplits().isEmpty()); // If planScan(Dataset) accidentally closed the handle, subsequent native calls would @@ -77,15 +79,15 @@ public void testPlanScanWithDatasetDoesNotCloseExternalDataset() { } /** - * Contract test: the long-typed resolved version returned by {@link LanceSplit#planScan(Dataset)} - * must round-trip through {@link org.lance.spark.LanceSparkReadOptions} as a LanceRef without - * truncation. This guards against silently casting to {@code int}, which would corrupt the - * snapshot-isolation guarantee for long-lived high-write-frequency datasets. + * Contract test: the long-typed resolved version returned by {@link LanceSplit#planScan(Dataset, + * LanceSparkReadOptions)} must round-trip through {@link org.lance.spark.LanceSparkReadOptions} + * as a LanceRef without truncation. This guards against silently casting to {@code int}, which + * would corrupt the snapshot-isolation guarantee for long-lived high-write-frequency datasets. */ @Test public void testResolvedVersionRoundTripsAsLong() { LanceSplit.ScanPlanResult result = LanceSplit.planScan(TestUtils.TestTable1Config.readOptions); - long resolved = result.getResolvedVersion(); + long resolved = result.getRef().getVersionNumber().get(); assertEquals( resolved, TestUtils.TestTable1Config.readOptions diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/tag/BaseTagDQLTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/tag/BaseTagDQLTest.java new file mode 100644 index 000000000..5b1092871 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/tag/BaseTagDQLTest.java @@ -0,0 +1,217 @@ +/* + * 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 org.lance.spark.tag; + +import org.lance.Ref; +import org.lance.spark.LanceDataset; +import org.lance.spark.LanceRef; + +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.Table; +import org.apache.spark.sql.connector.catalog.TableCapability; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.nio.file.FileSystems; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import java.util.UUID; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +/** Base tests for querying tagged table snapshots. */ +public abstract class BaseTagDQLTest { + private static final String CATALOG_NAME = "lance_test"; + + private SparkSession spark; + private String tableName; + private String fullTable; + private String tableDir; + + @TempDir Path tempDir; + + @BeforeEach + public void setup() throws IOException { + Path rootPath = tempDir.resolve(UUID.randomUUID().toString()); + Files.createDirectories(rootPath); + String testRoot = rootPath.toString(); + spark = + SparkSession.builder() + .appName("lance-tag-dql-test") + .master("local[4]") + .config( + "spark.sql.catalog." + CATALOG_NAME, "org.lance.spark.LanceNamespaceSparkCatalog") + .config( + "spark.sql.extensions", "org.lance.spark.extensions.LanceSparkSessionExtensions") + .config("spark.sql.catalog." + CATALOG_NAME + ".impl", "dir") + .config("spark.sql.catalog." + CATALOG_NAME + ".root", testRoot) + .config("spark.sql.catalog." + CATALOG_NAME + ".single_level_ns", "true") + .getOrCreate(); + tableName = "tag_dql_test_" + UUID.randomUUID().toString().replace("-", ""); + fullTable = CATALOG_NAME + ".default." + tableName; + tableDir = FileSystems.getDefault().getPath(testRoot, tableName + ".lance").toString(); + } + + @AfterEach + public void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + @Test + public void testTagSnapshotDoesNotIncludeDataWrittenAfterTagCreation() { + spark.sql(String.format("create table %s (id int, text string) using lance", fullTable)); + insertRange(0, 5); + String tag = "snapshot_before_new_data"; + createTag(tag); + + insertRange(5, 10); + + String taggedTable = String.format("%s version as of '%s'", fullTable, tag); + + Assertions.assertEquals(5, spark.sql("select * from " + taggedTable).count()); + Assertions.assertEquals( + 0, spark.sql("select * from " + taggedTable + " where id >= 5").count()); + + Assertions.assertEquals(10, spark.table(fullTable).count()); + Assertions.assertEquals(5, spark.sql("select * from " + fullTable + " where id >= 5").count()); + } + + @Test + public void testQueryTagCreatedFromBranchAfterMainAdvances() { + spark.sql(String.format("create table %s (id int, text string) using lance", fullTable)); + insertRange(0, 5); + + String branch = "source_branch"; + String tag = "branch_snapshot"; + spark.sql(String.format("alter table %s create branch %s", fullTable, branch)); + spark.sql( + String.format("alter table %s create tag %s as of branch %s", fullTable, tag, branch)); + + insertRange(5, 10); + + String taggedTable = String.format("%s version as of '%s'", fullTable, tag); + Assertions.assertEquals(5, spark.sql("select * from " + taggedTable).count()); + Assertions.assertEquals( + 0, spark.sql("select * from " + taggedTable + " where id >= 5").count()); + Assertions.assertEquals(10, spark.table(fullTable).count()); + Assertions.assertEquals(5, spark.sql("select * from " + fullTable + " where id >= 5").count()); + } + + @Test + public void testQueryNonexistentTagThrowsException() { + spark.sql(String.format("create table %s (id int, text string) using lance", fullTable)); + insertRange(0, 5); + + String query = String.format("select * from %s version as of 'nonexistent_tag'", fullTable); + + Assertions.assertThrows(Exception.class, () -> spark.sql(query).collectAsList()); + } + + @Test + public void testTagTableUsesTagReferenceAndIsReadOnly() throws Exception { + DatasetVersions versions = prepareDatasetWithHistory(); + TableCatalog catalog = + (TableCatalog) spark.sessionState().catalogManager().catalog(CATALOG_NAME); + + Table taggedTable = + catalog.loadTable( + Identifier.of(new String[] {"default"}, tableName), versions.firstInsertTag); + LanceRef ref = ((LanceDataset) taggedTable).readOptions().getRef(); + + Assertions.assertEquals(versions.firstInsertTag, ref.getTagName().get()); + Assertions.assertTrue(ref.getVersionNumber().isEmpty()); + Assertions.assertEquals( + Collections.singleton(TableCapability.BATCH_READ), taggedTable.capabilities()); + } + + @Test + public void testRejectWritesToTag() { + DatasetVersions versions = prepareDatasetWithHistory(); + String taggedTable = String.format("%s version as of '%s'", fullTable, versions.firstInsertTag); + spark.sql( + String.format( + "create temporary view tag_write_source as " + + "select id, text, _rowaddr, _fragid from %s", + taggedTable)); + + String[] statements = { + String.format("update %s set text = 'updated' where id = 0", taggedTable), + String.format("delete from %s where id = 0", taggedTable), + String.format("insert into %s values (10, 'inserted')", taggedTable), + String.format( + "merge into %s t using tag_write_source s on t.id = s.id " + + "when matched then update set text = s.text", + taggedTable), + String.format("alter table %s add columns copied_text from tag_write_source", taggedTable), + String.format("alter table %s update columns text from tag_write_source", taggedTable) + }; + + for (String statement : statements) { + Assertions.assertThrows( + Exception.class, () -> spark.sql(statement).collectAsList(), statement); + } + Assertions.assertEquals(10, spark.table(fullTable).count()); + Assertions.assertEquals(2, spark.table(fullTable).schema().size()); + } + + private DatasetVersions prepareDatasetWithHistory() { + spark.sql(String.format("create table %s (id int, text string) using lance", fullTable)); + insertRange(0, 5); + long firstInsertVersion = currentVersion(); + String firstInsertTag = "tag_" + firstInsertVersion; + createTag(firstInsertTag); + insertRange(5, 10); + return new DatasetVersions(firstInsertTag); + } + + private void insertRange(int startInclusive, int endExclusive) { + spark.sql( + String.format( + "insert into %s (id, text) values %s", + fullTable, + IntStream.range(startInclusive, endExclusive) + .boxed() + .map(i -> String.format("(%d, 'text_%d')", i, i)) + .collect(Collectors.joining(",")))); + } + + private long currentVersion() { + try (org.lance.Dataset dataset = org.lance.Dataset.open().uri(tableDir).build()) { + return dataset.getVersion().getId(); + } + } + + private void createTag(String tag) { + try (org.lance.Dataset dataset = org.lance.Dataset.open().uri(tableDir).build()) { + dataset.tags().create(tag, Ref.ofMain()); + } + } + + private static final class DatasetVersions { + private final String firstInsertTag; + + private DatasetVersions(String firstInsertTag) { + this.firstInsertTag = firstInsertTag; + } + } +} From 4ce67ca6ab48449476de0635629542df0e76fa87 Mon Sep 17 00:00:00 2001 From: Drew Gallardo Date: Fri, 21 Aug 2026 14:45:50 -0700 Subject: [PATCH 4/8] feat: support branch reads (#776) --- docs/src/config.md | 11 + docs/src/operations/dql/select.md | 51 ++- integration-tests/test_lance_spark.py | 97 ++++- .../spark/BaseLanceNamespaceSparkCatalog.java | 79 +++- .../java/org/lance/spark/LanceDataset.java | 99 +++-- .../main/java/org/lance/spark/LanceRef.java | 61 ++- .../lance/spark/LanceSparkReadOptions.java | 8 +- .../lance/spark/LanceSparkWriteOptions.java | 3 + .../lance/spark/read/LanceScanBuilder.java | 6 +- .../java/org/lance/spark/read/LanceSplit.java | 29 +- .../java/org/lance/spark/utils/Utils.java | 23 ++ .../v2/AddColumnsBackfillExec.scala | 7 +- .../datasources/v2/AddIndexExec.scala | 5 +- .../datasources/v2/DropIndexExec.scala | 6 +- .../v2/LanceCreateBranchExec.scala | 5 +- .../datasources/v2/LanceCreateTagExec.scala | 5 +- .../datasources/v2/LanceDropBranchExec.scala | 5 +- .../datasources/v2/LanceDropTagExec.scala | 5 +- .../datasources/v2/OptimizeExec.scala | 6 +- .../v2/SetUnenforcedPrimaryKeyExec.scala | 8 +- .../v2/UpdateColumnsBackfillExec.scala | 8 +- .../execution/datasources/v2/VacuumExec.scala | 6 +- .../java/org/lance/spark/LanceRefTest.java | 19 + ...anceSparkReadOptionsSerializationTest.java | 24 ++ .../spark/LanceSparkWriteOptionsTest.java | 15 + .../lance/spark/branch/BaseBranchDDLTest.java | 354 ++++++++++++++++++ .../spark/read/LanceScanBuilderTest.java | 15 + 27 files changed, 823 insertions(+), 137 deletions(-) diff --git a/docs/src/config.md b/docs/src/config.md index 6cd8eed64..84dcfd75b 100644 --- a/docs/src/config.md +++ b/docs/src/config.md @@ -579,6 +579,17 @@ The parent configuration effectively "anchors" your Spark catalog at a specific hierarchy, making the extra levels transparent to Spark users while maintaining compatibility with the underlying namespace implementation. +## Branch Read Option + +Set `branch` to read the current head of a named branch. Do not set `version` on the same read. + +```python +df = spark.read \ + .format("lance") \ + .option("branch", "audit") \ + .load("/path/to/dataset.lance") +``` + ## Memory Configuration Lance Spark uses Arrow for data transfer between native code and Spark, and maintains caches for improved performance. diff --git a/docs/src/operations/dql/select.md b/docs/src/operations/dql/select.md index 2eea0ba50..ff4d064e5 100644 --- a/docs/src/operations/dql/select.md +++ b/docs/src/operations/dql/select.md @@ -239,6 +239,54 @@ Use `VERSION AS OF` with a quoted tag name to query the snapshot referenced by a Tag queries are read-only. `UPDATE`, `DELETE`, `INSERT`, `MERGE INTO`, `ADD COLUMNS`, and `UPDATE COLUMNS` operations cannot target a tagged snapshot. +### Query by Branch + +Read the current head of a named branch. Do not set `branch` and `version` on the same read. + +=== "SQL" + ```sql + SELECT * FROM catalog.db.users.branch_audit; + ``` + +=== "Python" + ```python + audit = spark.read.option("branch", "audit").table("catalog.db.users") + + audit_path = spark.read \ + .format("lance") \ + .option("branch", "audit") \ + .load("/path/to/dataset.lance") + ``` + +=== "Scala" + ```scala + val audit = spark.read.option("branch", "audit").table("catalog.db.users") + + val auditPath = spark.read + .format("lance") + .option("branch", "audit") + .load("/path/to/dataset.lance") + ``` + +=== "Java" + ```java + Dataset audit = spark.read() + .option("branch", "audit") + .table("catalog.db.users"); + + Dataset auditPath = spark.read() + .format("lance") + .option("branch", "audit") + .load("/path/to/dataset.lance"); + ``` + +If a table named `users.branch_audit` already exists, Spark reads that table. Otherwise the last +segment is a branch on the parent table. Do not add `VERSION AS OF` or `TIMESTAMP AS OF`. Branch +identifiers are read-only; mutating commands are rejected. + +`.option("branch").table(...)` applies the branch at scan time. Spark still analyzes with the table +schema. Use `table.branch_name` when the branch schema differs from the table. + ### Query by Timestamp Use `TIMESTAMP AS OF` to query the table as it existed at a specific point in time: @@ -279,6 +327,7 @@ These options control how data is read from Lance datasets. They can be set usin | `batch_size` | Integer | `8192` | Number of rows to read per batch during scanning. Larger values may improve throughput but increase memory usage. | | `use_scalar_index` | Boolean | `true` | Whether to use scalar indices (e.g. btree) for filter acceleration during scanning. | | `version` | Integer | Latest | Specific dataset version to read. If not specified, reads the latest version. | +| `branch` | String | Main | Named branch to read. Reads the current head. Do not set with `version`. | | `block_size` | Integer | - | Block size in bytes for reading data. | | `index_cache_size` | Integer | - | Size of the index cache in number of entries. | | `metadata_cache_size` | Integer | - | Size of the metadata cache in number of entries. | @@ -345,4 +394,4 @@ Spark UDT columns (e.g. MLlib's `VectorUDT`) round-trip transparently. On write .format("lance") .option("batch_size", "4096") .load("/path/to/dataset.lance") - ``` \ No newline at end of file + ``` diff --git a/integration-tests/test_lance_spark.py b/integration-tests/test_lance_spark.py index a38f5c55f..ff1d84bc2 100644 --- a/integration-tests/test_lance_spark.py +++ b/integration-tests/test_lance_spark.py @@ -2804,6 +2804,101 @@ def test_version_as_of_after_delete(self, spark): assert len(result) == 3 +class TestDQLBranchRead: + def test_branch_reference_matches_iceberg_read_surfaces(self, spark): + spark.sql("CREATE TABLE default.test_table (id INT, name STRING)") + spark.sql( + "INSERT INTO default.test_table VALUES (1, 'a'), (2, 'b')" + ) + expected = [(1, "a"), (2, "b")] + spark.sql( + "ALTER TABLE default.test_table CREATE BRANCH test_branch" + ) + spark.sql( + "INSERT INTO default.test_table VALUES (3, 'c'), (4, 'd')" + ) + + identifier = spark.sql( + "SELECT * FROM default.test_table.branch_test_branch ORDER BY id" + ).collect() + option_table = ( + spark.read.option("branch", "test_branch") + .table("default.test_table") + .orderBy("id") + .collect() + ) + option_path = ( + spark.read.format("lance") + .option("branch", "test_branch") + .load(_table_location(spark, "default.test_table")) + .orderBy("id") + .collect() + ) + main = spark.sql( + "SELECT * FROM default.test_table ORDER BY id" + ).collect() + + assert [(row.id, row.name) for row in identifier] == expected + assert [(row.id, row.name) for row in option_table] == expected + assert [(row.id, row.name) for row in option_path] == expected + assert [(row.id, row.name) for row in main] == expected + [(3, "c"), (4, "d")] + + def test_branch_identifier_rejects_as_of_and_conflicting_options(self, spark): + spark.sql("CREATE TABLE default.test_table (id INT, name STRING)") + spark.sql("INSERT INTO default.test_table VALUES (1, 'main')") + spark.sql("ALTER TABLE default.test_table CREATE BRANCH audit") + + with pytest.raises(Exception, match="Cannot combine"): + spark.sql( + "SELECT * FROM default.test_table.branch_audit VERSION AS OF 1" + ).collect() + with pytest.raises(Exception, match="Cannot combine"): + spark.sql( + "SELECT * FROM default.test_table.branch_audit TIMESTAMP AS OF now()" + ).collect() + with pytest.raises(Exception): + spark.read.option("branch", "audit").option("version", "1").table( + "default.test_table" + ).collect() + with pytest.raises(Exception, match="no_such_branch"): + spark.read.option("branch", "no_such_branch").table( + "default.test_table" + ).collect() + + def test_branch_identifier_is_read_only(self, spark): + spark.sql("CREATE TABLE default.test_table (id INT, name STRING)") + spark.sql("INSERT INTO default.test_table VALUES (1, 'main')") + spark.sql("ALTER TABLE default.test_table CREATE BRANCH audit") + + with pytest.raises(Exception): + spark.sql( + "INSERT INTO default.test_table.branch_audit VALUES (2, 'branch')" + ).collect() + + assert spark.table("default.test_table").count() == 1 + assert spark.table("default.test_table.branch_audit").count() == 1 + + def test_existing_table_wins_over_branch_identifier(self, spark): + spark.sql("CREATE TABLE default.test_table (id INT, name STRING)") + spark.sql("INSERT INTO default.test_table VALUES (1, 'branch_row')") + spark.sql("ALTER TABLE default.test_table CREATE BRANCH audit") + spark.sql( + "CREATE TABLE default.test_table.branch_audit (id INT, name STRING)" + ) + spark.sql( + "INSERT INTO default.test_table.branch_audit VALUES (99, 'literal')" + ) + + rows = spark.table("default.test_table.branch_audit").collect() + assert [(row.id, row.name) for row in rows] == [(99, "literal")] + assert [ + row.id + for row in spark.read.option("branch", "audit") + .table("default.test_table") + .collect() + ] == [1] + + @requires_update_or_merge class TestDMLMergeDelete: """Test MERGE INTO with WHEN MATCHED THEN DELETE.""" @@ -3366,4 +3461,4 @@ def test_show_functions_lists_fts(self, spark): if __name__ == "__main__": - pytest.main([__file__, "-v"]) \ No newline at end of file + pytest.main([__file__, "-v"]) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java index 859a4811f..ada3b7cf1 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java @@ -79,6 +79,7 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.regex.Matcher; import java.util.regex.Pattern; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -117,6 +118,8 @@ public abstract class BaseLanceNamespaceSparkCatalog private static final String CONFIG_PARENT_DELIMITER = "parent_delimiter"; private static final String CONFIG_PARENT_DELIMITER_DEFAULT = "."; + private static final Pattern BRANCH_SUFFIX = Pattern.compile("branch_(.+)"); + private boolean pathBasedOnly = false; private LanceNamespace namespace; @@ -596,17 +599,60 @@ private boolean tableExistsAtPath(Identifier ident) { @Override public Table loadTable(Identifier ident) throws NoSuchTableException { - return loadTableInternal(ident, Optional.empty(), Optional.empty()); + return loadTableOrBranchSelector(ident, Optional.empty(), Optional.empty()); } @Override public Table loadTable(Identifier ident, String version) throws NoSuchTableException { - return loadTableInternal(ident, Optional.empty(), Optional.of(version)); + return loadTableOrBranchSelector(ident, Optional.empty(), Optional.of(version)); } @Override public Table loadTable(Identifier ident, long timestamp) throws NoSuchTableException { - return loadTableInternal(ident, Optional.of(timestamp), Optional.empty()); + return loadTableOrBranchSelector(ident, Optional.of(timestamp), Optional.empty()); + } + + private Table loadTableOrBranchSelector( + Identifier ident, Optional timestamp, Optional version) + throws NoSuchTableException { + Optional branch = branchSelector(ident); + try { + return loadTableInternal(ident, timestamp, version, Optional.empty()); + } catch (NoSuchTableException e) { + if (!branch.isPresent()) { + throw e; + } + return loadBranchSelector(ident, timestamp, version, branch.get()); + } catch (LanceNamespaceException e) { + if (!branch.isPresent() || e.getErrorCode() != ErrorCode.NAMESPACE_NOT_FOUND) { + throw e; + } + return loadBranchSelector(ident, timestamp, version, branch.get()); + } + } + + private Table loadBranchSelector( + Identifier ident, Optional timestamp, Optional version, LanceRef branch) + throws NoSuchTableException { + if (timestamp.isPresent() || version.isPresent()) { + throw new IllegalArgumentException( + "Cannot combine a branch identifier with VERSION AS OF or TIMESTAMP AS OF"); + } + return loadTableInternal(parentIdent(ident), timestamp, version, Optional.of(branch)); + } + + private static Optional branchSelector(Identifier ident) { + if (ident.namespace().length == 0) { + return Optional.empty(); + } + Matcher matcher = BRANCH_SUFFIX.matcher(ident.name()); + return matcher.matches() ? Optional.of(LanceRef.ofBranch(matcher.group(1))) : Optional.empty(); + } + + private static Identifier parentIdent(Identifier ident) { + String[] namespace = ident.namespace(); + return Identifier.of( + Arrays.copyOf(namespace, namespace.length - 1), namespace[namespace.length - 1]); } @Override @@ -1662,19 +1708,22 @@ private LanceRef parseVersionRef(String version) { } private Table loadTableInternal( - Identifier ident, Optional timestamp, Optional version) + Identifier ident, + Optional timestamp, + Optional version, + Optional branchRef) throws NoSuchTableException { // Handle path-based access if (isPathBasedIdentifier(ident)) { - return loadTableFromPath(ident, timestamp, version); + return loadTableFromPath(ident, timestamp, version, branchRef); } ResolvedTable resolved = resolveIdentifier(ident); DescribeTableResponse describeResponse = resolved.describeResponse; Map initialStorageOptions = describeResponse.getStorageOptions(); - Optional ref = Optional.empty(); + Optional ref = branchRef; if (timestamp.isPresent()) { try (Dataset dataset = Utils.openDatasetBuilder(resolved.readOptions).build()) { ref = @@ -1698,6 +1747,7 @@ private Table loadTableInternal( schema = LanceArrowUtils.fromArrowSchema(dataset.getSchema()); fileFormatVersion = dataset.getLanceFileFormatVersion(); tableProperties = dataset.getConfig(); + readOptions = pinLoadedBranch(readOptions, dataset); } // Create read options with namespace support @@ -1743,14 +1793,17 @@ private DescribeTableResponse describeTableOrThrow(DescribeTableRequest request, * spark.read.format("lance").load(path). */ private Table loadTableFromPath( - Identifier ident, Optional timestamp, Optional version) + Identifier ident, + Optional timestamp, + Optional version, + Optional branchRef) throws NoSuchTableException { String datasetUri = getDatasetUri(ident); LanceSparkReadOptions baseReadOptions = createReadOptions( datasetUri, catalogConfig, Optional.empty(), Optional.empty(), Optional.empty(), name); - Optional ref = Optional.empty(); + Optional ref = branchRef; if (version.isPresent()) { ref = Optional.of(parseVersionRef(version.get())); } else if (timestamp.isPresent()) { @@ -1774,6 +1827,7 @@ private Table loadTableFromPath( schema = LanceArrowUtils.fromArrowSchema(dataset.getSchema()); fileFormatVersion = dataset.getLanceFileFormatVersion(); tableProperties = dataset.getConfig(); + readOptions = pinLoadedBranch(readOptions, dataset); } catch (IllegalArgumentException e) { throw new NoSuchTableException(ident); } @@ -1782,6 +1836,15 @@ private Table loadTableFromPath( readOptions, schema, null, null, null, false, fileFormatVersion, tableProperties, null); } + private static LanceSparkReadOptions pinLoadedBranch( + LanceSparkReadOptions readOptions, Dataset dataset) { + LanceRef ref = readOptions.getRef(); + if (ref == null || !ref.isBranch()) { + return readOptions; + } + return readOptions.withRef(Utils.pinOpenedRef(dataset, ref)); + } + public abstract LanceDataset createDataset( LanceSparkReadOptions readOptions, StructType sparkSchema, diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java index 71a78dc33..239dd9c08 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java @@ -13,22 +13,26 @@ */ package org.lance.spark; +import org.lance.Dataset; import org.lance.memwal.ShardingSpec; import org.lance.spark.read.LanceScanBuilder; import org.lance.spark.utils.BlobSourceContext; import org.lance.spark.utils.BlobUtils; +import org.lance.spark.utils.Utils; import org.lance.spark.write.AddColumnsBackfillWrite; import org.lance.spark.write.LanceWriteSchemaValidator; import org.lance.spark.write.SparkWrite; import org.lance.spark.write.StagedCommit; import org.lance.spark.write.UpdateColumnsBackfillWrite; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableSet; import org.apache.spark.sql.connector.catalog.MetadataColumn; import org.apache.spark.sql.connector.catalog.StagedTable; import org.apache.spark.sql.connector.catalog.SupportsMetadataColumns; import org.apache.spark.sql.connector.catalog.SupportsRead; import org.apache.spark.sql.connector.catalog.SupportsWrite; +import org.apache.spark.sql.connector.catalog.Table; import org.apache.spark.sql.connector.catalog.TableCapability; import org.apache.spark.sql.connector.read.ScanBuilder; import org.apache.spark.sql.connector.write.LogicalWriteInfo; @@ -38,6 +42,7 @@ import org.apache.spark.sql.types.StructField; import org.apache.spark.sql.types.StructType; import org.apache.spark.sql.util.CaseInsensitiveStringMap; +import org.apache.spark.sql.util.LanceArrowUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -309,26 +314,22 @@ public String getFileFormatVersion() { } @Override - public ScanBuilder newScanBuilder(CaseInsensitiveStringMap caseInsensitiveStringMap) { - // Merge scan-time options with the existing read options - LanceSparkReadOptions scanOptions = readOptions; - if (!caseInsensitiveStringMap.isEmpty()) { - Map mergedOptions = new HashMap<>(readOptions.getStorageOptions()); - mergedOptions.putAll(caseInsensitiveStringMap.asCaseSensitiveMap()); - scanOptions = - LanceSparkReadOptions.builder() - .datasetUri(readOptions.getDatasetUri()) - .namespace(readOptions.getNamespace()) - .tableId(readOptions.getTableId()) - .catalogName(readOptions.getCatalogName()) - .indexCacheBackend(readOptions.getIndexCacheBackend()) - .metadataCacheBackend(readOptions.getMetadataCacheBackend()) - .fromOptions(mergedOptions) - .ref(readOptions.getRef()) - .build(); + public ScanBuilder newScanBuilder(CaseInsensitiveStringMap scanMap) { + LanceSparkReadOptions scanOptions = mergeScanOptions(scanMap); + StructType scanSchema = sparkSchema; + LanceRef scanRef = scanOptions.getRef(); + if (scanRef != null && scanRef.isBranch() && scanRef.getVersionNumber().isEmpty()) { + try (Dataset dataset = + Utils.openDatasetBuilder(scanOptions) + .initialStorageOptions(initialStorageOptions) + .runtimeNamespace(namespaceImpl, namespaceProperties, readOptions.getTableId()) + .build()) { + scanOptions = scanOptions.withRef(Utils.pinOpenedRef(dataset, scanRef)); + scanSchema = LanceArrowUtils.fromArrowSchema(dataset.getSchema()); + } } return new LanceScanBuilder( - sparkSchema, + scanSchema, scanOptions, initialStorageOptions, namespaceImpl, @@ -336,6 +337,49 @@ public ScanBuilder newScanBuilder(CaseInsensitiveStringMap caseInsensitiveString shardingSpec); } + private LanceSparkReadOptions mergeScanOptions(CaseInsensitiveStringMap scanMap) { + if (scanMap.isEmpty()) { + return readOptions; + } + LanceRef tableRef = readOptions.getRef(); + boolean scanSetsBranchOrVersion = + scanMap.containsKey(LanceSparkReadOptions.CONFIG_VERSION) + || scanMap.containsKey(LanceSparkReadOptions.CONFIG_BRANCH); + Map merged = new HashMap<>(readOptions.getStorageOptions()); + // Stored options can still hold version/branch keys. Strip them so fromOptions cannot + // overwrite the table ref with a leftover key. + merged.remove(LanceSparkReadOptions.CONFIG_VERSION); + merged.remove(LanceSparkReadOptions.CONFIG_BRANCH); + merged.putAll(scanMap.asCaseSensitiveMap()); + LanceSparkReadOptions scanOptions = + LanceSparkReadOptions.builder() + .datasetUri(readOptions.getDatasetUri()) + .namespace(readOptions.getNamespace()) + .tableId(readOptions.getTableId()) + .catalogName(readOptions.getCatalogName()) + .indexCacheBackend(readOptions.getIndexCacheBackend()) + .metadataCacheBackend(readOptions.getMetadataCacheBackend()) + .ref(tableRef) + .fromOptions(merged) + .build(); + if (tableRef != null && scanSetsBranchOrVersion) { + LanceRef scanRef = scanOptions.getRef(); + if (sameNamedBranch(tableRef, scanRef)) { + return scanOptions.withRef(tableRef); + } + Preconditions.checkArgument( + tableRef.equals(scanRef), "Cannot combine %s with %s", tableRef, scanRef); + } + return scanOptions; + } + + private static boolean sameNamedBranch(LanceRef tableRef, LanceRef scanRef) { + return tableRef.isBranch() + && scanRef != null + && scanRef.isBranch() + && tableRef.getBranchName().equals(scanRef.getBranchName()); + } + @Override public String name() { return this.readOptions.getDatasetName(); @@ -358,21 +402,30 @@ public Map properties() { @Override public Set capabilities() { - if (isTagReference()) { + if (isBranchOrTag()) { return READ_ONLY_CAPABILITIES; } return BlobUtils.hasBlobV2Fields(sparkSchema) ? CAPABILITIES_WITH_BLOB_V2 : CAPABILITIES; } + public static LanceDataset requireWritable(Table table, String command) { + if (!(table instanceof LanceDataset)) { + throw new UnsupportedOperationException(command + " only supports LanceDataset"); + } + LanceDataset dataset = (LanceDataset) table; + dataset.ensureWritable(); + return dataset; + } + protected void ensureWritable() { - if (isTagReference()) { + if (isBranchOrTag()) { throw new UnsupportedOperationException( - "Writes are not supported for tag: " + readOptions.getRef().getTagName().get()); + "Writes are not supported for " + readOptions.getRef()); } } - private boolean isTagReference() { - return readOptions.getRef() != null && readOptions.getRef().getTagName().isPresent(); + private boolean isBranchOrTag() { + return readOptions.getRef() != null && readOptions.getRef().isBranchOrTag(); } @Override diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java index 05dd55357..1b97440f1 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceRef.java @@ -34,18 +34,6 @@ public LanceRef( this.tagName = tagName; } - public Optional getVersionNumber() { - return versionNumber; - } - - public Optional getBranchName() { - return branchName; - } - - public Optional getTagName() { - return tagName; - } - public static LanceRef ofMain(long versionNumber) { Preconditions.checkArgument(versionNumber > 0, "versionNumber must be greater than 0"); return new LanceRef(Optional.of(versionNumber), Optional.empty(), Optional.empty()); @@ -73,6 +61,34 @@ public static LanceRef ofTag(String tagName) { return new LanceRef(Optional.empty(), Optional.empty(), Optional.of(tagName)); } + public Optional getVersionNumber() { + return versionNumber; + } + + public Optional getBranchName() { + return branchName; + } + + public Optional getTagName() { + return tagName; + } + + public boolean isMain() { + return branchName.isEmpty() && tagName.isEmpty(); + } + + public boolean isBranch() { + return branchName.isPresent(); + } + + public boolean isTag() { + return tagName.isPresent(); + } + + public boolean isBranchOrTag() { + return isBranch() || isTag(); + } + @Override public boolean equals(Object o) { if (this == o) { @@ -94,13 +110,18 @@ public int hashCode() { @Override public String toString() { - return "LanceRef{" - + "versionNumber=" - + (versionNumber.isPresent() ? versionNumber.get() : null) - + ", branchName=" - + (branchName.isPresent() ? branchName.get() : null) - + ", tagName=" - + (tagName.isPresent() ? tagName.get() : null) - + '}'; + if (isTag()) { + return "tag " + tagName.get(); + } + if (isBranch() && versionNumber.isPresent()) { + return "branch " + branchName.get() + " version " + versionNumber.get(); + } + if (isBranch()) { + return "branch " + branchName.get(); + } + if (versionNumber.isPresent()) { + return "version " + versionNumber.get(); + } + return "main"; } } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java index 149c846e6..79a7da945 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkReadOptions.java @@ -53,6 +53,7 @@ public class LanceSparkReadOptions implements Serializable { public static final String CONFIG_PUSH_DOWN_FILTERS = "pushDownFilters"; public static final String CONFIG_BLOCK_SIZE = "block_size"; public static final String CONFIG_VERSION = "version"; + public static final String CONFIG_BRANCH = "branch"; public static final String CONFIG_INDEX_CACHE_SIZE = "index_cache_size"; public static final String CONFIG_METADATA_CACHE_SIZE = "metadata_cache_size"; public static final String CONFIG_BATCH_SIZE = "batch_size"; @@ -364,7 +365,7 @@ public ReadOptions toReadOptions() { if (blockSize != null) { builder.setBlockSize(blockSize); } - if (ref != null) { + if (ref != null && ref.isMain()) { ref.getVersionNumber().ifPresent(builder::setVersion); } if (indexCacheSize != null) { @@ -583,6 +584,9 @@ private void parseTypedFlags(Map opts) { Preconditions.checkArgument( !opts.containsKey(CONFIG_NEAREST), "The nearest read option is no longer supported; use VECTOR_SEARCH table function"); + Preconditions.checkArgument( + !(opts.containsKey(CONFIG_VERSION) && opts.containsKey(CONFIG_BRANCH)), + "Specify only one of branch or version"); if (opts.containsKey(CONFIG_PUSH_DOWN_FILTERS)) { this.pushDownFilters = Boolean.parseBoolean(opts.get(CONFIG_PUSH_DOWN_FILTERS)); } @@ -591,6 +595,8 @@ private void parseTypedFlags(Map opts) { } if (opts.containsKey(CONFIG_VERSION)) { this.ref = LanceRef.ofMain(Long.parseLong(opts.get(CONFIG_VERSION))); + } else if (opts.containsKey(CONFIG_BRANCH)) { + this.ref = LanceRef.ofBranch(opts.get(CONFIG_BRANCH)); } if (opts.containsKey(CONFIG_INDEX_CACHE_SIZE)) { this.indexCacheSize = Integer.parseInt(opts.get(CONFIG_INDEX_CACHE_SIZE)); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java index 362b588a0..b564af76a 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java @@ -520,6 +520,9 @@ public Builder metadataCacheBackend(String metadataCacheBackend) { * @return this builder */ public Builder fromOptions(Map options) { + Preconditions.checkArgument( + !options.containsKey(LanceSparkReadOptions.CONFIG_BRANCH), + "The branch option is read-only"); this.storageOptions = new HashMap<>(options); if (options.containsKey(CONFIG_WRITE_MODE)) { this.writeMode = WriteMode.valueOf(options.get(CONFIG_WRITE_MODE).toUpperCase()); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java index f8a8c57a5..d33a7f38e 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceScanBuilder.java @@ -297,15 +297,13 @@ public Scan build() { boolean shouldNamespaceFtsScan() { LanceRef ref = readOptions.getRef(); - boolean hasTag = ref != null && ref.getTagName().isPresent(); - if (hasTag) { + if (ref != null && ref.isBranchOrTag()) { return false; } return readOptions.getFullTextQuery() != null && LanceRuntime.supportsQueryTable(namespaceImpl) - && !pushedAggregation.isPresent() - && !hasTag; + && !pushedAggregation.isPresent(); } /** diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java index 7ca6b7f1e..506a145a4 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/read/LanceSplit.java @@ -15,7 +15,6 @@ import org.lance.Dataset; import org.lance.Fragment; -import org.lance.Tag; import org.lance.spark.LanceRef; import org.lance.spark.LanceSparkReadOptions; import org.lance.spark.utils.Utils; @@ -26,7 +25,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; public class LanceSplit implements Serializable { private static final long serialVersionUID = 2983749283749283749L; @@ -101,31 +99,8 @@ public static ScanPlanResult planScan(Dataset dataset, LanceSparkReadOptions rea fragmentRowCounts.put(id, fragment.metadata().getNumRows()); } - LanceRef plannedRef; - - LanceRef ref = readOptions.getRef(); - if (ref != null && ref.getBranchName().isPresent()) { - plannedRef = LanceRef.ofBranch(ref.getBranchName().get(), dataset.getVersion().getId()); - } else if (ref != null && ref.getTagName().isPresent()) { - Optional tag = - dataset.tags().list().stream() - .filter(t -> t.getName().equals(ref.getTagName().get())) - .findFirst(); - if (tag.isEmpty()) { - throw new RuntimeException("No existed tag with name:" + ref.getTagName().get()); - } - - Optional tagBranch = tag.get().getBranch(); - if (tagBranch.isEmpty()) { - plannedRef = LanceRef.ofMain(dataset.getVersion().getId()); - } else { - plannedRef = LanceRef.ofBranch(tagBranch.get(), dataset.version()); - } - } else { - plannedRef = LanceRef.ofMain(dataset.getVersion().getId()); - } - - return new ScanPlanResult(splits, plannedRef, fragmentRowCounts); + LanceRef ref = Utils.pinOpenedRef(dataset, readOptions.getRef()); + return new ScanPlanResult(splits, ref, fragmentRowCounts); } /** diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java index afad5b9e5..e71f2fcde 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/Utils.java @@ -16,6 +16,7 @@ import org.lance.Dataset; import org.lance.ReadOptions; import org.lance.Ref; +import org.lance.Tag; import org.lance.Version; import org.lance.namespace.LanceNamespace; import org.lance.spark.LanceRef; @@ -35,6 +36,28 @@ public static long parseVersion(String version) { return Long.parseUnsignedLong(version); } + // Tag names and branch heads move. Store the opened version, and the branch if a tag points at + // one. + public static LanceRef pinOpenedRef(Dataset dataset, LanceRef requested) { + long version = dataset.getVersion().getId(); + if (requested != null && requested.isTag()) { + String tagName = requested.getTagName().get(); + for (Tag tag : dataset.tags().list()) { + if (tag.getName().equals(tagName)) { + if (tag.getBranch().isPresent()) { + return LanceRef.ofBranch(tag.getBranch().get(), version); + } + return LanceRef.ofMain(version); + } + } + throw new RuntimeException("Tag not found: " + tagName); + } + if (requested != null && requested.isBranch()) { + return LanceRef.ofBranch(requested.getBranchName().get(), version); + } + return LanceRef.ofMain(version); + } + public static long findVersion(List versions, long timestamp) { long versionID = -1; Instant instant = instantFromTimestamp(timestamp); diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddColumnsBackfillExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddColumnsBackfillExec.scala index 4a250d4f8..db8d20578 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddColumnsBackfillExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddColumnsBackfillExec.scala @@ -29,11 +29,8 @@ case class AddColumnsBackfillExec( override def output: Seq[Attribute] = Seq.empty override protected def run(): Seq[InternalRow] = { - val originalTable = catalog.loadTable(ident) match { - case lanceTable: LanceDataset => lanceTable - case _ => - throw new UnsupportedOperationException("AddColumnsBackfill only supports for LanceDataset") - } + val originalTable = + LanceDataset.requireWritable(catalog.loadTable(ident), "AddColumnsBackfill") // Check the added columns must not exist val originalFields = originalTable.schema().fieldNames.toSet diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala index 030825b66..b365c1c2b 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala @@ -80,10 +80,7 @@ case class AddIndexExec( override def output: Seq[Attribute] = AddIndexOutputType.SCHEMA override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case d: LanceDataset => d - case _ => throw new UnsupportedOperationException("AddIndex only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "AddIndex") val readOptions = lanceDataset.readOptions() val indexType = IndexUtils.buildIndexType(method) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DropIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DropIndexExec.scala index 63f3c1c58..5e81f1137 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DropIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DropIndexExec.scala @@ -35,11 +35,7 @@ case class LanceDropIndexExec( override def output: Seq[Attribute] = LanceDropIndexOutputType.SCHEMA override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case ds: LanceDataset => ds - case _ => - throw new UnsupportedOperationException("DropIndex only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "DropIndex") val readOptions = lanceDataset.readOptions() diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateBranchExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateBranchExec.scala index 52877bc69..948def0ef 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateBranchExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateBranchExec.scala @@ -33,10 +33,7 @@ case class LanceCreateBranchExec( override def output: Seq[Attribute] = LanceCreateBranchOutputType.SCHEMA override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case d: LanceDataset => d - case _ => throw new UnsupportedOperationException("CreateBranch only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "CreateBranch") val dataset = Utils.openDatasetBuilder(lanceDataset.readOptions()) .initialStorageOptions(lanceDataset.getInitialStorageOptions) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateTagExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateTagExec.scala index 3c6114e44..bc1f8cb50 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateTagExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceCreateTagExec.scala @@ -33,10 +33,7 @@ case class LanceCreateTagExec( override def output: Seq[Attribute] = LanceCreateTagOutputType.SCHEMA override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case d: LanceDataset => d - case _ => throw new UnsupportedOperationException("CreateTag only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "CreateTag") val dataset = Utils.openDatasetBuilder(lanceDataset.readOptions()) .initialStorageOptions(lanceDataset.getInitialStorageOptions) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropBranchExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropBranchExec.scala index 95885e535..be30b1a64 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropBranchExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropBranchExec.scala @@ -32,10 +32,7 @@ case class LanceDropBranchExec( override def output: Seq[Attribute] = LanceDropBranchOutputType.SCHEMA override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case d: LanceDataset => d - case _ => throw new UnsupportedOperationException("DropBranch only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "DropBranch") val dataset = Utils.openDatasetBuilder(lanceDataset.readOptions()) .initialStorageOptions(lanceDataset.getInitialStorageOptions) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropTagExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropTagExec.scala index 13c16b39f..cf2d23f3a 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropTagExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDropTagExec.scala @@ -32,10 +32,7 @@ case class LanceDropTagExec( override def output: Seq[Attribute] = LanceDropTagOutputType.SCHEMA override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case d: LanceDataset => d - case _ => throw new UnsupportedOperationException("DropTag only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "DropTag") val dataset = Utils.openDatasetBuilder(lanceDataset.readOptions()) .initialStorageOptions(lanceDataset.getInitialStorageOptions) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeExec.scala index f40dd5363..1618d60e3 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/OptimizeExec.scala @@ -56,11 +56,7 @@ case class OptimizeExec( } override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case lanceDataset: LanceDataset => lanceDataset - case _ => - throw new UnsupportedOperationException("Optimize only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "Optimize") // Build compaction options from arguments val options = buildOptions() diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/SetUnenforcedPrimaryKeyExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/SetUnenforcedPrimaryKeyExec.scala index 9d4af6689..a6f6ee2f3 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/SetUnenforcedPrimaryKeyExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/SetUnenforcedPrimaryKeyExec.scala @@ -39,12 +39,8 @@ case class SetUnenforcedPrimaryKeyExec( override def output: Seq[Attribute] = SetUnenforcedPrimaryKeyOutputType.SCHEMA override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case ds: LanceDataset => ds - case _ => - throw new UnsupportedOperationException( - "SET UNENFORCED PRIMARY KEY only supports LanceDataset") - } + val lanceDataset = + LanceDataset.requireWritable(catalog.loadTable(ident), "SET UNENFORCED PRIMARY KEY") val readOptions = lanceDataset.readOptions() val dataset = Utils.openDatasetBuilder(readOptions).build() diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/UpdateColumnsBackfillExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/UpdateColumnsBackfillExec.scala index a44606cac..d6184cb56 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/UpdateColumnsBackfillExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/UpdateColumnsBackfillExec.scala @@ -37,12 +37,8 @@ case class UpdateColumnsBackfillExec( override def output: Seq[Attribute] = Seq.empty override protected def run(): Seq[InternalRow] = { - val originalTable = catalog.loadTable(ident) match { - case lanceTable: LanceDataset => lanceTable - case _ => - throw new UnsupportedOperationException( - "UpdateColumnsBackfill only supports LanceDataset") - } + val originalTable = + LanceDataset.requireWritable(catalog.loadTable(ident), "UpdateColumnsBackfill") // Check the updated columns must exist val originalFields = originalTable.schema().fieldNames.toSet diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/VacuumExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/VacuumExec.scala index 31d212712..ca2267d55 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/VacuumExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/VacuumExec.scala @@ -48,11 +48,7 @@ case class VacuumExec( } override protected def run(): Seq[InternalRow] = { - val lanceDataset = catalog.loadTable(ident) match { - case lanceDataset: LanceDataset => lanceDataset - case _ => - throw new UnsupportedOperationException("Vacuum only supports LanceDataset") - } + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "Vacuum") val policy = buildPolicy() val readOptions = lanceDataset.readOptions() diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java index 38e5b45c7..905a7b7eb 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceRefTest.java @@ -16,7 +16,9 @@ import org.junit.jupiter.api.Test; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; public class LanceRefTest { @@ -33,6 +35,23 @@ public void refsUseValueEquality() { assertNotEquals(LanceRef.ofTag("release"), LanceRef.ofTag("latest")); } + @Test + public void testRefKind() { + assertTrue(LanceRef.ofMain().isMain()); + assertFalse(LanceRef.ofMain().isBranch()); + assertTrue(LanceRef.ofBranch("dev").isBranch()); + assertFalse(LanceRef.ofBranch("dev").isMain()); + assertTrue(LanceRef.ofTag("release").isTag()); + assertTrue(LanceRef.ofBranch("dev").isBranchOrTag()); + assertTrue(LanceRef.ofTag("release").isBranchOrTag()); + assertFalse(LanceRef.ofMain().isBranchOrTag()); + assertEquals("branch audit", LanceRef.ofBranch("audit").toString()); + assertEquals("branch audit version 2", LanceRef.ofBranch("audit", 2).toString()); + assertEquals("tag release", LanceRef.ofTag("release").toString()); + assertEquals("version 7", LanceRef.ofMain(7).toString()); + assertEquals("main", LanceRef.ofMain().toString()); + } + @Test public void equalRefsHaveEqualHashCodes() { assertEquals(LanceRef.ofMain(7).hashCode(), LanceRef.ofMain(7).hashCode()); diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java index c19b126a1..19970c6bf 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkReadOptionsSerializationTest.java @@ -75,6 +75,30 @@ public void testWithRefPropagatesFullTextQuery() { options.getFullTextQuery(), versioned.getFullTextQuery())); } + @Test + public void testFromOptionsParsesBranch() { + LanceSparkReadOptions options = + LanceSparkReadOptions.from( + Collections.singletonMap(LanceSparkReadOptions.CONFIG_BRANCH, "audit"), + "s3://bucket/path"); + + Assertions.assertEquals(LanceRef.ofBranch("audit"), options.getRef()); + } + + @Test + public void testFromOptionsRejectsBranchAndVersion() { + Map options = new HashMap<>(); + options.put(LanceSparkReadOptions.CONFIG_BRANCH, "audit"); + options.put(LanceSparkReadOptions.CONFIG_VERSION, "2"); + + IllegalArgumentException exception = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> LanceSparkReadOptions.from(options, "s3://bucket/path")); + Assertions.assertTrue(exception.getMessage().contains("branch")); + Assertions.assertTrue(exception.getMessage().contains("version")); + } + @Test public void testFromOptionsParsesFtsSubtype() { FullTextQuery original = FullTextQuery.phrase("hello world", "body", 0); diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java index dfd97021e..632eddcf3 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java @@ -29,6 +29,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; /** Tests for {@link LanceSparkWriteOptions}. */ @@ -49,6 +50,20 @@ public void builderSetsRef() { assertEquals(7L, opts.getRef().getVersionNumber().get()); } + @Test + public void testWriteOptionsRejectBranch() { + Map options = new HashMap<>(); + options.put(LanceSparkReadOptions.CONFIG_BRANCH, "audit"); + + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + LanceSparkWriteOptions.builder().datasetUri(TEMP_URL).fromOptions(options).build()); + + assertTrue(exception.getMessage().contains("read-only")); + } + @Test public void fileFormatVersionUsesValueEquality() { LanceSparkWriteOptions left = diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java index a8cdeb44d..22d202a21 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java @@ -14,10 +14,18 @@ package org.lance.spark.branch; import org.lance.Ref; +import org.lance.spark.LanceDataset; +import org.lance.spark.LanceRef; +import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.read.LanceInputPartition; +import org.lance.spark.read.LanceSplit; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.apache.spark.sql.util.CaseInsensitiveStringMap; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; @@ -28,11 +36,14 @@ import java.nio.file.FileSystems; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Arrays; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.UUID; import java.util.stream.Collectors; import java.util.stream.IntStream; +import java.util.stream.Stream; /** Base tests for BRANCH DDL commands. */ public abstract class BaseBranchDDLTest { @@ -381,6 +392,349 @@ public void testCreateBranchWithBacktickQuotedName() { branches.containsKey(branchName), "Expected backtick-quoted branch to be returned"); } + @Test + public void testBranchReadUsesBranchSchemaAfterMainAddsColumn() throws Exception { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + spark.sql( + String.format( + "create temporary view extra_cols as select _rowaddr, _fragid, 1 as extra from %s", + fullTable)); + spark.sql(String.format("alter table %s add columns extra from extra_cols", fullTable)); + + TableCatalog catalog = + (TableCatalog) spark.sessionState().catalogManager().catalog(catalogName); + List expectedBranchSchema = + columnNames( + spark.sql( + "select * from " + fullTable + " version as of " + versions.firstInsertVersion)); + List catalogSchema = + Arrays.asList( + catalog + .loadTable(Identifier.of(new String[] {"default", tableName}, "branch_audit")) + .schema() + .fieldNames()); + List identifierSchema = columnNames(spark.table(fullTable + ".branch_audit")); + + Assertions.assertEquals(expectedBranchSchema, catalogSchema); + Assertions.assertEquals(expectedBranchSchema, identifierSchema); + } + + @Test + public void testBranchIdentifierPinsVersionOnLoad() throws Exception { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + + TableCatalog catalog = + (TableCatalog) spark.sessionState().catalogManager().catalog(catalogName); + LanceDataset table = + (LanceDataset) + catalog.loadTable(Identifier.of(new String[] {"default", tableName}, "branch_audit")); + LanceRef ref = table.readOptions().getRef(); + + Assertions.assertEquals("audit", ref.getBranchName().get()); + Assertions.assertEquals(versions.firstInsertVersion, ref.getVersionNumber().get()); + } + + @Test + public void testBranchOptionPinsVersionOnScan() throws Exception { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + + TableCatalog catalog = + (TableCatalog) spark.sessionState().catalogManager().catalog(catalogName); + LanceDataset table = + (LanceDataset) catalog.loadTable(Identifier.of(new String[] {"default"}, tableName)); + Map options = new HashMap<>(); + options.put("branch", "audit"); + LanceInputPartition partition = + (LanceInputPartition) + table + .newScanBuilder(new CaseInsensitiveStringMap(options)) + .build() + .toBatch() + .planInputPartitions()[0]; + LanceRef ref = partition.getReadOptions().getRef(); + + Assertions.assertEquals("audit", ref.getBranchName().get()); + Assertions.assertEquals(versions.firstInsertVersion, ref.getVersionNumber().get()); + } + + @Test + public void testTagScanPinsVersionNotName() { + DatasetVersions versions = prepareDatasetWithHistory(); + LanceSparkReadOptions readOptions = + LanceSparkReadOptions.builder() + .datasetUri(tableDir) + .ref(LanceRef.ofTag(versions.firstInsertTag)) + .build(); + + LanceRef plannedRef = LanceSplit.planScan(readOptions).getRef(); + + Assertions.assertTrue(plannedRef.getTagName().isEmpty()); + Assertions.assertEquals(versions.firstInsertVersion, plannedRef.getVersionNumber().get()); + } + + @Test + public void testReadBranchByOptionPathAndIdentifier() { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + insertRange(10, 15); + + Assertions.assertEquals(15, spark.table(fullTable).count()); + Assertions.assertEquals(5, spark.read().option("branch", "audit").table(fullTable).count()); + Assertions.assertEquals( + 5, spark.read().format("lance").option("branch", "audit").load(tableDir).count()); + Assertions.assertEquals(5, spark.table(fullTable + ".branch_audit").count()); + Assertions.assertEquals( + 5, spark.read().option("branch", "audit").table(fullTable + ".branch_audit").count()); + } + + @Test + public void testExistingTableWinsOverBranchIdentifier() throws Exception { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + + String literalTable = fullTable + ".branch_audit"; + spark.sql("create table " + literalTable + " (id int, text string) using lance"); + spark.sql("insert into " + literalTable + " values (99, 'literal')"); + + Assertions.assertEquals(1, spark.table(literalTable).count()); + Assertions.assertEquals(99, spark.table(literalTable).collectAsList().get(0).getInt(0)); + Assertions.assertEquals(5, spark.read().option("branch", "audit").table(fullTable).count()); + + TableCatalog catalog = + (TableCatalog) spark.sessionState().catalogManager().catalog(catalogName); + LanceDataset table = + (LanceDataset) + catalog.loadTable(Identifier.of(new String[] {"default", tableName}, "branch_audit")); + LanceRef ref = table.readOptions().getRef(); + Assertions.assertTrue(ref == null || ref.isMain()); + } + + @Test + public void testBranchIdentifierKeepsRowsWhenBatchSizeIsSet() { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + insertRange(10, 15); + + Assertions.assertEquals( + 5, spark.read().option("batch_size", "1024").table(fullTable + ".branch_audit").count()); + } + + @Test + public void testBranchScanPinsBranchAndVersion() { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + LanceSparkReadOptions readOptions = + LanceSparkReadOptions.builder() + .datasetUri(tableDir) + .ref(LanceRef.ofBranch("audit")) + .build(); + + LanceRef plannedRef = LanceSplit.planScan(readOptions).getRef(); + + Assertions.assertEquals("audit", plannedRef.getBranchName().get()); + Assertions.assertTrue(plannedRef.getVersionNumber().isPresent()); + Assertions.assertEquals(versions.firstInsertVersion, plannedRef.getVersionNumber().get()); + } + + @Test + public void testMissingBranchFails() { + prepareDatasetWithHistory(); + + Exception missing = + Assertions.assertThrows( + Exception.class, + () -> spark.read().option("branch", "no_such_branch").table(fullTable).collectAsList()); + Assertions.assertTrue(exceptionChainMessages(missing).contains("no_such_branch")); + } + + @Test + public void testBranchAndVersionOptionsFail() { + prepareDatasetWithHistory(); + + Exception conflicting = + Assertions.assertThrows( + Exception.class, + () -> + spark + .read() + .option("branch", "audit") + .option("version", "1") + .table(fullTable) + .collectAsList()); + String conflictMessages = exceptionChainMessages(conflicting); + Assertions.assertTrue(conflictMessages.contains("branch")); + Assertions.assertTrue(conflictMessages.contains("version")); + } + + @Test + public void testBranchIdentifierRejectsVersionAsOf() { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + + Exception asOf = + Assertions.assertThrows( + Exception.class, + () -> + spark + .sql("select * from " + fullTable + ".branch_audit version as of 1") + .collectAsList()); + Assertions.assertTrue(exceptionChainMessages(asOf).contains("Cannot combine")); + + Exception timestampAsOf = + Assertions.assertThrows( + Exception.class, + () -> + spark + .sql("select * from " + fullTable + ".branch_audit timestamp as of now()") + .collectAsList()); + Assertions.assertTrue(exceptionChainMessages(timestampAsOf).contains("Cannot combine")); + } + + @Test + public void testBranchIdentifierRejectsDifferentBranchOption() { + DatasetVersions versions = prepareDatasetWithHistory(); + spark.sql( + String.format( + "alter table %s create branch audit as of version %d", + fullTable, versions.firstInsertVersion)); + + Exception mismatch = + Assertions.assertThrows( + Exception.class, + () -> + spark + .read() + .option("branch", "other") + .table(fullTable + ".branch_audit") + .collectAsList()); + String mismatchMessages = exceptionChainMessages(mismatch); + Assertions.assertTrue(mismatchMessages.contains("audit")); + Assertions.assertTrue(mismatchMessages.contains("other")); + } + + @Test + public void testInsertIntoBranchIdentifierFails() { + prepareDatasetWithHistory(); + spark.sql(String.format("alter table %s create branch audit", fullTable)); + + Assertions.assertThrows( + Exception.class, + () -> + spark + .sql("insert into " + fullTable + ".branch_audit values (99, 'branch')") + .collectAsList()); + Assertions.assertThrows( + Exception.class, + () -> + spark + .sql("select 99 as id, 'branch' as text") + .write() + .format("lance") + .option("branch", "audit") + .mode("append") + .save(tableDir)); + Assertions.assertEquals(10, spark.table(fullTable).count()); + Assertions.assertEquals(10, spark.table(fullTable + ".branch_audit").count()); + } + + @Test + public void testBranchIdentifierRejectsCreateIndexAndVacuum() throws Exception { + prepareDatasetWithHistory(); + spark.sql(String.format("alter table %s create branch audit", fullTable)); + + TableCatalog catalog = + (TableCatalog) spark.sessionState().catalogManager().catalog(catalogName); + Identifier branchIdentifier = + Identifier.of(new String[] {"default", tableName}, "branch_audit"); + long branchVersionBefore = + ((LanceDataset) catalog.loadTable(branchIdentifier)) + .readOptions() + .getRef() + .getVersionNumber() + .get(); + long fileCountBefore = datasetFileCount(); + + Exception createIndex = + Assertions.assertThrows( + Exception.class, + () -> + spark + .sql( + "alter table " + + fullTable + + ".branch_audit create index id_idx using zonemap (id) " + + "with (train=false)") + .collectAsList()); + Assertions.assertTrue(exceptionChainMessages(createIndex).contains("Writes are not supported")); + + Exception vacuum = + Assertions.assertThrows( + Exception.class, + () -> + spark + .sql("vacuum " + fullTable + ".branch_audit with (before_version=1000000)") + .collectAsList()); + Assertions.assertTrue(exceptionChainMessages(vacuum).contains("Writes are not supported")); + + long branchVersionAfter = + ((LanceDataset) catalog.loadTable(branchIdentifier)) + .readOptions() + .getRef() + .getVersionNumber() + .get(); + Assertions.assertEquals(branchVersionBefore, branchVersionAfter); + Assertions.assertEquals(fileCountBefore, datasetFileCount()); + Assertions.assertEquals(10, spark.table(fullTable).count()); + Assertions.assertEquals(10, spark.table(fullTable + ".branch_audit").count()); + } + + private long datasetFileCount() throws IOException { + try (Stream files = Files.walk(FileSystems.getDefault().getPath(tableDir))) { + return files.filter(Files::isRegularFile).count(); + } + } + + private static String exceptionChainMessages(Throwable throwable) { + StringBuilder messages = new StringBuilder(); + for (Throwable current = throwable; current != null; current = current.getCause()) { + if (current.getMessage() != null) { + messages.append(current.getMessage()).append('\n'); + } + } + return messages.toString(); + } + + private static List columnNames(Dataset rows) { + return Arrays.asList(rows.columns()); + } + private DatasetVersions prepareDatasetWithHistory() { spark.sql(String.format("create table %s (id int, text string) using lance;", fullTable)); insertRange(0, 5); diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java index ff2b57a96..a15adc3ff 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/read/LanceScanBuilderTest.java @@ -370,6 +370,21 @@ public void testTagFullTextQueryDoesNotUseNamespaceScan() { assertFalse(builder.shouldNamespaceFtsScan()); } + @Test + public void testBranchFullTextQueryDoesNotUseNamespaceScan() { + LanceSparkReadOptions options = + LanceSparkReadOptions.builder() + .datasetUri(TestUtils.TestTable1Config.datasetUri) + .ref(LanceRef.ofBranch("audit")) + .fullTextQuery(FullTextQuery.match("hello", "b")) + .build(); + LanceScanBuilder builder = + new LanceScanBuilder( + TEST_SCHEMA, options, Collections.emptyMap(), "dir", Collections.emptyMap()); + + assertFalse(builder.shouldNamespaceFtsScan()); + } + @Test public void testBuildWithCountStarReturnsLocalScan() { LanceScanBuilder builder = createBuilder(); From 27c8a381b25b02eb6a80871b86d66bd98ce55b1c Mon Sep 17 00:00:00 2001 From: Drew Gallardo Date: Fri, 21 Aug 2026 17:39:08 -0700 Subject: [PATCH 5/8] fix: resolve branch identifiers through the parent table (#777) Branch identifiers were failing on Glue because we first tried to resolve the full branch shaped identifier as a literal table. Glue only accepts database.table, so that lookup failed before we could fall back to the parent table and branch. This changes branch identifier resolution so the literal table lookup is just a probe. If it does not resolve as a real table, we load the parent table and apply the branch instead. Literal tables still win on catalogs that support them. The collision test is skipped on Glue since Glue cannot represent that table shape. --- integration-tests/test_lance_spark.py | 4 ++- .../spark/BaseLanceNamespaceSparkCatalog.java | 36 ++++++++++++++----- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/integration-tests/test_lance_spark.py b/integration-tests/test_lance_spark.py index ff1d84bc2..7b54f5797 100644 --- a/integration-tests/test_lance_spark.py +++ b/integration-tests/test_lance_spark.py @@ -2805,7 +2805,7 @@ def test_version_as_of_after_delete(self, spark): class TestDQLBranchRead: - def test_branch_reference_matches_iceberg_read_surfaces(self, spark): + def test_branch_identifier_matches_option_and_path(self, spark): spark.sql("CREATE TABLE default.test_table (id INT, name STRING)") spark.sql( "INSERT INTO default.test_table VALUES (1, 'a'), (2, 'b')" @@ -2879,6 +2879,8 @@ def test_branch_identifier_is_read_only(self, spark): assert spark.table("default.test_table.branch_audit").count() == 1 def test_existing_table_wins_over_branch_identifier(self, spark): + if getattr(spark, "_lance_backend", None) == "glue": + pytest.skip("Glue table identifiers are database.table") spark.sql("CREATE TABLE default.test_table (id INT, name STRING)") spark.sql("INSERT INTO default.test_table VALUES (1, 'branch_row')") spark.sql("ALTER TABLE default.test_table CREATE BRANCH audit") diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java index ada3b7cf1..e879dd966 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/BaseLanceNamespaceSparkCatalog.java @@ -616,18 +616,27 @@ private Table loadTableOrBranchSelector( Identifier ident, Optional timestamp, Optional version) throws NoSuchTableException { Optional branch = branchSelector(ident); - try { + if (!branch.isPresent()) { return loadTableInternal(ident, timestamp, version, Optional.empty()); + } + + try { + ResolvedTable literal = resolveLiteralTable(ident); + return loadResolvedTable(ident, literal, timestamp, version, Optional.empty()); } catch (NoSuchTableException e) { - if (!branch.isPresent()) { - throw e; - } return loadBranchSelector(ident, timestamp, version, branch.get()); + } + } + + private ResolvedTable resolveLiteralTable(Identifier ident) throws NoSuchTableException { + try { + return resolveIdentifier(ident); } catch (LanceNamespaceException e) { - if (!branch.isPresent() || e.getErrorCode() != ErrorCode.NAMESPACE_NOT_FOUND) { - throw e; + ErrorCode code = e.getErrorCode(); + if (code == ErrorCode.NAMESPACE_NOT_FOUND || code == ErrorCode.INVALID_INPUT) { + throw new NoSuchTableException(ident); } - return loadBranchSelector(ident, timestamp, version, branch.get()); + throw e; } } @@ -642,7 +651,7 @@ private Table loadBranchSelector( } private static Optional branchSelector(Identifier ident) { - if (ident.namespace().length == 0) { + if (ident.namespace().length == 0 || isPathBasedIdentifier(ident)) { return Optional.empty(); } Matcher matcher = BRANCH_SUFFIX.matcher(ident.name()); @@ -1719,7 +1728,16 @@ private Table loadTableInternal( return loadTableFromPath(ident, timestamp, version, branchRef); } - ResolvedTable resolved = resolveIdentifier(ident); + return loadResolvedTable(ident, resolveIdentifier(ident), timestamp, version, branchRef); + } + + private Table loadResolvedTable( + Identifier ident, + ResolvedTable resolved, + Optional timestamp, + Optional version, + Optional branchRef) + throws NoSuchTableException { DescribeTableResponse describeResponse = resolved.describeResponse; Map initialStorageOptions = describeResponse.getStorageOptions(); From 6114a1b40cdd18cbf99f12303802e0efb18f283d Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Sat, 22 Aug 2026 01:12:21 +0000 Subject: [PATCH 6/8] chore: bump version to 0.8.0-beta.1 --- .bumpversion.toml | 2 +- lance-spark-3.4_2.12/pom.xml | 2 +- lance-spark-3.4_2.13/pom.xml | 2 +- lance-spark-3.5_2.12/pom.xml | 2 +- lance-spark-3.5_2.13/pom.xml | 2 +- lance-spark-4.0_2.13/pom.xml | 2 +- lance-spark-4.1_2.13/pom.xml | 2 +- lance-spark-4.2_2.13/pom.xml | 2 +- lance-spark-base_2.12/pom.xml | 2 +- lance-spark-base_2.13/pom.xml | 2 +- lance-spark-bundle-3.4_2.12/pom.xml | 2 +- lance-spark-bundle-3.4_2.13/pom.xml | 2 +- lance-spark-bundle-3.5_2.12/pom.xml | 2 +- lance-spark-bundle-3.5_2.13/pom.xml | 2 +- lance-spark-bundle-4.0_2.13/pom.xml | 2 +- lance-spark-bundle-4.1_2.13/pom.xml | 2 +- lance-spark-bundle-4.2_2.13/pom.xml | 2 +- pom.xml | 4 ++-- 18 files changed, 19 insertions(+), 19 deletions(-) diff --git a/.bumpversion.toml b/.bumpversion.toml index 29c07b8a3..82e0d1485 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.7.1" +current_version = "0.8.0-beta.1" parse = "(?P\\d+)\\.(?P\\d+)\\.(?P\\d+)(-(?Palpha|beta|rc)\\.(?P\\d+))?" serialize = [ "{major}.{minor}.{patch}-{pre_label}.{pre_n}", diff --git a/lance-spark-3.4_2.12/pom.xml b/lance-spark-3.4_2.12/pom.xml index cc45d24fb..528bb7206 100644 --- a/lance-spark-3.4_2.12/pom.xml +++ b/lance-spark-3.4_2.12/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-3.4_2.13/pom.xml b/lance-spark-3.4_2.13/pom.xml index 5883d86b4..e60e18164 100644 --- a/lance-spark-3.4_2.13/pom.xml +++ b/lance-spark-3.4_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-3.5_2.12/pom.xml b/lance-spark-3.5_2.12/pom.xml index 70172fa72..653506cb3 100644 --- a/lance-spark-3.5_2.12/pom.xml +++ b/lance-spark-3.5_2.12/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-3.5_2.13/pom.xml b/lance-spark-3.5_2.13/pom.xml index d2de3e58c..581205e5e 100644 --- a/lance-spark-3.5_2.13/pom.xml +++ b/lance-spark-3.5_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-4.0_2.13/pom.xml b/lance-spark-4.0_2.13/pom.xml index 9824ab191..f29a7863c 100644 --- a/lance-spark-4.0_2.13/pom.xml +++ b/lance-spark-4.0_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-4.1_2.13/pom.xml b/lance-spark-4.1_2.13/pom.xml index 7869688f0..c422b4a25 100644 --- a/lance-spark-4.1_2.13/pom.xml +++ b/lance-spark-4.1_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-4.2_2.13/pom.xml b/lance-spark-4.2_2.13/pom.xml index 566dd9e3a..c63c5a92e 100644 --- a/lance-spark-4.2_2.13/pom.xml +++ b/lance-spark-4.2_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-base_2.12/pom.xml b/lance-spark-base_2.12/pom.xml index ec4cf1e23..4d9c8d714 100644 --- a/lance-spark-base_2.12/pom.xml +++ b/lance-spark-base_2.12/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-base_2.13/pom.xml b/lance-spark-base_2.13/pom.xml index 64b4e4308..0a727fd09 100644 --- a/lance-spark-base_2.13/pom.xml +++ b/lance-spark-base_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-bundle-3.4_2.12/pom.xml b/lance-spark-bundle-3.4_2.12/pom.xml index a68bb4837..620d21605 100644 --- a/lance-spark-bundle-3.4_2.12/pom.xml +++ b/lance-spark-bundle-3.4_2.12/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-bundle-3.4_2.13/pom.xml b/lance-spark-bundle-3.4_2.13/pom.xml index b70f60c18..b409f99e4 100644 --- a/lance-spark-bundle-3.4_2.13/pom.xml +++ b/lance-spark-bundle-3.4_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-bundle-3.5_2.12/pom.xml b/lance-spark-bundle-3.5_2.12/pom.xml index ddf9df3ac..f686a24f5 100644 --- a/lance-spark-bundle-3.5_2.12/pom.xml +++ b/lance-spark-bundle-3.5_2.12/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-bundle-3.5_2.13/pom.xml b/lance-spark-bundle-3.5_2.13/pom.xml index 002fa24a5..6923b4f1c 100644 --- a/lance-spark-bundle-3.5_2.13/pom.xml +++ b/lance-spark-bundle-3.5_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-bundle-4.0_2.13/pom.xml b/lance-spark-bundle-4.0_2.13/pom.xml index e82d9dcd8..d1dc04d30 100644 --- a/lance-spark-bundle-4.0_2.13/pom.xml +++ b/lance-spark-bundle-4.0_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-bundle-4.1_2.13/pom.xml b/lance-spark-bundle-4.1_2.13/pom.xml index 74e6cef15..afbe1aad2 100644 --- a/lance-spark-bundle-4.1_2.13/pom.xml +++ b/lance-spark-bundle-4.1_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/lance-spark-bundle-4.2_2.13/pom.xml b/lance-spark-bundle-4.2_2.13/pom.xml index 2acf70271..08dd92d60 100644 --- a/lance-spark-bundle-4.2_2.13/pom.xml +++ b/lance-spark-bundle-4.2_2.13/pom.xml @@ -5,7 +5,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 ../pom.xml diff --git a/pom.xml b/pom.xml index 572c3ce07..488425ef2 100644 --- a/pom.xml +++ b/pom.xml @@ -4,7 +4,7 @@ org.lance lance-spark-root - 0.7.1 + 0.8.0-beta.1 pom ${project.artifactId} @@ -50,7 +50,7 @@ - 0.7.1 + 0.8.0-beta.1 11.0.0-beta.10 0.8.6 0.4.0 From 9e49c48effd34763c6aa048842e58024ba1abf74 Mon Sep 17 00:00:00 2001 From: Eunjin Song Date: Sun, 23 Aug 2026 09:08:39 -0700 Subject: [PATCH 7/8] build(deps): bump lance-core to 11.0.0-beta.21 The Arrow C stream export path added in the following commit calls LanceScanner#exportArrowStream(long), which is first available in lance-core 11.0.0-beta.21 (lance-format/lance#7259). Co-Authored-By: Claude Opus 4.8 --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 488425ef2..91d9e4fee 100644 --- a/pom.xml +++ b/pom.xml @@ -51,7 +51,7 @@ 0.8.0-beta.1 - 11.0.0-beta.10 + 11.0.0-beta.21 0.8.6 0.4.0 From d94b5c6e34da2c88fc76b1cefb82e497d5d72269 Mon Sep 17 00:00:00 2001 From: Eunjin Song Date: Sun, 23 Aug 2026 09:08:49 -0700 Subject: [PATCH 8/8] feat(gluten): export a Lance fragment scan as an Arrow C stream Add LanceArrowStreamScanner, which plans a fragment scan via the existing LanceFragmentScanner and exports it as an Arrow C Data Interface stream (ArrowArrayStream) for native consumers such as Apache Gluten / Velox. Only the ArrowArrayStream C-struct address crosses the JVM/native boundary, so the consumer's Arrow build and classloader do not need to match lance-spark's (Gluten builds Arrow 15 to match Spark 3.5 / Velox, while the Lance Java SDK is on Arrow 18). The Lance native core populates the caller-owned stream directly through LanceScanner#exportArrowStream(long), so no Arrow data is materialized on the JVM heap on this path. All scan planning (column projection, filter pushdown, limit/offset, row-id / row-address, batch size) is delegated to LanceFragmentScanner, so the exported stream yields exactly the same rows in the same order as the Spark columnar reader. LanceArrowStream owns the exported stream and the scan behind it; closing it releases the native scan (via the stream's release callback) and then the scanner and dataset handles. Co-Authored-By: Claude Opus 4.8 --- .../internal/LanceArrowStreamScanner.java | 151 ++++++++++++++++++ .../spark/internal/LanceFragmentScanner.java | 13 ++ .../internal/LanceArrowStreamScannerTest.java | 71 ++++++++ 3 files changed, 235 insertions(+) create mode 100644 lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceArrowStreamScanner.java create mode 100644 lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceArrowStreamScannerTest.java diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceArrowStreamScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceArrowStreamScanner.java new file mode 100644 index 000000000..fbe91391f --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceArrowStreamScanner.java @@ -0,0 +1,151 @@ +/* + * 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 org.lance.spark.internal; + +import org.lance.spark.LanceRuntime; +import org.lance.spark.read.LanceInputPartition; + +import org.apache.arrow.c.ArrowArrayStream; + +import java.io.IOException; + +/** + * Exports a Lance fragment scan as an Arrow C Data Interface stream ({@link ArrowArrayStream}) for + * native consumers such as Apache Gluten / Velox. + * + *

Only the {@link ArrowArrayStream} C-struct address ({@link LanceArrowStream#streamAddress()}) + * crosses the JVM/native boundary, so the consumer's Arrow build and classloader do not need to + * match lance-spark's. All scan planning — column projection, filter pushdown, limit/offset, row-id + * / row-address, batch size — is delegated to {@link LanceFragmentScanner}, so this path produces + * exactly the same rows in the same order as the Spark columnar reader. + * + *

The Lance native core writes batches into the caller-owned stream on demand as the consumer + * pulls them, so no Arrow data is materialized on the JVM heap on this path. + */ +public final class LanceArrowStreamScanner { + + private LanceArrowStreamScanner() {} + + /** + * Plans a fragment scan and exports it as an Arrow C stream. + * + *

The returned {@link LanceArrowStream} owns the exported stream and the scan backing it. Hand + * {@link LanceArrowStream#streamAddress()} to the native consumer, let it drain the stream to + * exhaustion, then {@link LanceArrowStream#close() close} the handle. Closing before the consumer + * has finished reading is a use-after-free on caller-owned native memory. + * + * @param fragmentId the Lance fragment to scan + * @param inputPartition the planned partition (schema, filter, limit/offset, storage options) + * @return an open Arrow C stream handle over the fragment scan + */ + public static LanceArrowStream export(int fragmentId, LanceInputPartition inputPartition) { + LanceFragmentScanner fragmentScanner = LanceFragmentScanner.create(fragmentId, inputPartition); + ArrowArrayStream stream = ArrowArrayStream.allocateNew(LanceRuntime.allocator()); + try { + // The Lance native core populates the caller-owned stream directly from the planned scan, so + // no Arrow batch is ever materialized on the JVM heap here — the consumer pulls batches over + // the C Data Interface. The stream's release callback routes back to the native side, so + // releasing the stream (by the consumer, or by LanceArrowStream#close) tears down the scan. + fragmentScanner.exportArrowStream(stream.memoryAddress()); + } catch (Throwable t) { + closeQuietly(stream); + closeQuietly(fragmentScanner); + if (t instanceof RuntimeException) { + throw (RuntimeException) t; + } + if (t instanceof Error) { + throw (Error) t; + } + throw new RuntimeException(t); + } + return new LanceArrowStream(stream, fragmentScanner); + } + + private static void closeQuietly(AutoCloseable closeable) { + if (closeable != null) { + try { + closeable.close(); + } catch (Exception ignore) { + // Best effort on the construction error path. + } + } + } + + /** + * Owns an exported {@link ArrowArrayStream} together with the fragment scan behind it. + * + *

{@link #close()} releases, in order, the exported stream (whose release callback tears down + * the native scan, freeing its buffers) and then the Lance scanner and dataset handles. These are + * distinct resources: the stream drives the native scan, while the scanner owns the open dataset. + */ + public static final class LanceArrowStream implements AutoCloseable { + private final ArrowArrayStream stream; + private final LanceFragmentScanner fragmentScanner; + + LanceArrowStream(ArrowArrayStream stream, LanceFragmentScanner fragmentScanner) { + this.stream = stream; + this.fragmentScanner = fragmentScanner; + } + + /** The Arrow C Data Interface stream backing this scan. */ + public ArrowArrayStream stream() { + return stream; + } + + /** + * The C-struct address to hand to a native consumer (e.g. a Velox Arrow-stream source). Valid + * until {@link #close()}. + */ + public long streamAddress() { + return stream.memoryAddress(); + } + + @Override + public void close() throws IOException { + Throwable primary = null; + // Closing the stream runs its release callback, tearing down the native scan and its buffers; + // then release the scanner and the dataset it holds open. + primary = closeAndAccumulate(stream, primary); + primary = closeAndAccumulate(fragmentScanner, primary); + if (primary != null) { + if (primary instanceof IOException) { + throw (IOException) primary; + } + if (primary instanceof RuntimeException) { + throw (RuntimeException) primary; + } + if (primary instanceof Error) { + throw (Error) primary; + } + throw new IOException(primary); + } + } + + private static Throwable closeAndAccumulate(AutoCloseable closeable, Throwable primary) { + if (closeable == null) { + return primary; + } + try { + closeable.close(); + return primary; + } catch (Throwable t) { + if (primary != null) { + primary.addSuppressed(t); + return primary; + } + return t; + } + } + } +} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java index ca11ae6bb..3268c48b3 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/internal/LanceFragmentScanner.java @@ -186,6 +186,19 @@ public ArrowReader getArrowReader() { return scanner.scanBatches(); } + /** + * Exports this fragment scan into a caller-owned Arrow C Data Interface stream. The Lance native + * side populates the {@code ArrowArrayStream} at {@code streamAddress} directly, so only the + * C-struct address crosses the JVM/native boundary. The caller owns the stream and must close it + * (which releases the native scan via the stream's release callback); the scanner and dataset + * held by this object are released separately by {@link #close()}. + * + * @param streamAddress the memory address of a freshly-allocated, empty {@code ArrowArrayStream} + */ + public void exportArrowStream(long streamAddress) throws IOException { + scanner.exportArrowStream(streamAddress); + } + @Override public void close() throws IOException { Throwable primary = null; diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceArrowStreamScannerTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceArrowStreamScannerTest.java new file mode 100644 index 000000000..12f6e2e55 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/internal/LanceArrowStreamScannerTest.java @@ -0,0 +1,71 @@ +/* + * 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 org.lance.spark.internal; + +import org.lance.spark.LanceRuntime; +import org.lance.spark.TestUtils; + +import org.apache.arrow.c.Data; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +public class LanceArrowStreamScannerTest { + + /** + * Exports each fragment of the bundled test table as an Arrow C Data Interface stream, re-imports + * it on the JVM (standing in for a native consumer), and asserts the rows match what the Spark + * columnar reader produces. Closing the imported reader and then the {@link + * LanceArrowStreamScanner.LanceArrowStream} handle under the leak-checking allocator also + * verifies the export/reader/scanner lifecycle releases cleanly. + */ + @Test + public void exportsFragmentAsArrowCStream() throws Exception { + List> expectedValues = TestUtils.TestTable1Config.expectedValues; + int rowIndex = 0; + for (int fragmentId = 0; fragmentId <= 1; fragmentId++) { + try (LanceArrowStreamScanner.LanceArrowStream handle = + LanceArrowStreamScanner.export( + fragmentId, TestUtils.TestTable1Config.inputPartition); + ArrowReader reader = Data.importArrayStream(LanceRuntime.allocator(), handle.stream())) { + + // Schema is available before the first batch: x, y, b, c. + assertEquals(4, reader.getVectorSchemaRoot().getSchema().getFields().size()); + + while (reader.loadNextBatch()) { + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + int columns = root.getFieldVectors().size(); + for (int r = 0; r < root.getRowCount(); r++) { + List expectedRow = expectedValues.get(rowIndex); + for (int col = 0; col < columns; col++) { + Object actual = root.getVector(col).getObject(r); + assertNotNull(actual, "Null at row " + rowIndex + " column " + col); + assertEquals( + expectedRow.get(col).longValue(), + ((Number) actual).longValue(), + "Mismatch at row " + rowIndex + " column " + col); + } + rowIndex++; + } + } + } + } + assertEquals(4, rowIndex); + } +}