diff --git a/docs/src/format/table/versioning.md b/docs/src/format/table/versioning.md index d90b3782c20..aa28edeafa1 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -36,3 +36,64 @@ they should return an "unsupported" error on any read or write operation. Flags with bit values 512 and above are unknown and will cause implementations to reject the dataset with an "unsupported" error. The paired mixed-version reader and writer bits must either both be set or both be clear; a half-set manifest is invalid. + +## Mixed V2 Data File Versions + +The manifest data storage version is the default for operations that do not +select an exact output version. It is a fallback, not a summary, minimum, +maximum, or profile of the data files referenced by the snapshot. Once +`FLAG_MIXED_DATA_FILE_VERSIONS` is enabled, each base data file and data overlay +file is decoded according to its own normalized version identity. + +Mixed snapshots have the following invariants: + +- Only exact V2.0, V2.1, V2.2, and V2.3 data file versions may be mixed. +- V1 and V2 data files may not appear in the same snapshot. +- A commit that first produces a mixed snapshot derives and sets both the reader + and writer capability bits from its final manifest. The bits remain set on all + later snapshots, even if a later compaction makes the files homogeneous again. +- A snapshot without the capability may only reference files matching its + manifest fallback. The only repair exception is an unambiguous, homogeneous + historical V2 snapshot whose legacy manifest metadata is stale. +- An operation-level `data_storage_version` selects the exact output version + for that operation. Omitting it uses the manifest fallback. Neither case + changes the fallback. + +For example, a dataset whose fallback is V2.1 can append V2.2 files by setting +`data_storage_version="2.2"`. The same commit adds both mixed-version capability +bits. Reads then dispatch V2.1 files to the V2.1 decoder and V2.2 files to the +V2.2 decoder. Compaction can deliberately rewrite selected fragments to any +supported exact V2 target; binary copy is only valid when every selected input +file already has that exact target version. + +### Compatibility Matrix + +| Dataset state | Mixed-aware client | Client without bit 256 support | +| --- | --- | --- | +| Historical homogeneous V1 | Reads and writes through legacy paths | Unchanged | +| Historical homogeneous V2 | Reads and writes; legacy metadata repair remains uniform-only | Unchanged | +| New homogeneous V2 without bit 256 | Reads and writes using the manifest fallback | Unchanged | +| Mixed V2.0-V2.3 with both bits set | Reads and writes by exact per-file identity | Rejects before reading or writing | +| Mixed V2 without both bits | Rejects as a per-file capability mismatch | Not a valid dataset state | +| V1/V2 mixture | Rejects | Not a valid dataset state | + +### Error Categories + +Implementations distinguish these failures in their error messages so operators +can identify the violated boundary: + +- unsupported reader or writer feature bit; +- half-set mixed-version capability corruption; +- unknown or malformed data file version identity; +- V1/V2 mixture; +- a non-fallback file without mixed-version capability; and +- binary-copy target mismatch, including the target, actual version, and path. + +### Rollout Gate + +Before the first mixed-version commit, deploy mixed-aware readers and writers +everywhere that can access the dataset. Then drain or fence writers that opened +the dataset with an older client. Only after both steps may a writer select a +different exact V2 output version. The capability bit makes clients that open +the resulting snapshot fail closed, but it cannot retroactively fence an old +writer that read an earlier manifest. diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 889c3e572eb..29b3fb86557 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -3340,6 +3340,14 @@ fn convert_java_compaction_options_to_rust( &[], )? .l()?; + let data_storage_version = env + .call_method( + &java_options, + "getDataStorageVersion", + "()Ljava/util/Optional;", + &[], + )? + .l()?; build_compaction_options( env, @@ -3357,6 +3365,7 @@ fn convert_java_compaction_options_to_rust( &max_source_rows, &max_source_bytes, &excluded_fragment_ids, + &data_storage_version, config, ) } diff --git a/java/lance-jni/src/merge_insert.rs b/java/lance-jni/src/merge_insert.rs index 2e7f869757a..751f6d59e8c 100644 --- a/java/lance-jni/src/merge_insert.rs +++ b/java/lance-jni/src/merge_insert.rs @@ -15,7 +15,9 @@ use lance::dataset::{ MergeInsertBuilder, MergeStats, WhenMatched, WhenNotMatched, WhenNotMatchedBySource, }; use lance_core::datatypes::Schema; +use lance_file::version::LanceFileVersion; use lance_index::mem_wal::CompactedSsTable; +use std::str::FromStr; use std::sync::Arc; use std::time::Duration; use uuid::Uuid; @@ -53,6 +55,7 @@ fn inner_merge_insert<'local>( let skip_auto_cleanup = extract_skip_auto_cleanup(env, &jparam)?; let use_index = extract_use_index(env, &jparam)?; let compacted_sstables = extract_compacted_sstables(env, &jparam)?; + let data_storage_version = extract_data_storage_version(env, &jparam)?; let (new_ds, merge_stats) = unsafe { let dataset = env.get_rust_field::<_, _, BlockingDataset>(jdataset, NATIVE_DATASET)?; @@ -63,7 +66,11 @@ fn inner_merge_insert<'local>( when_not_matched_by_source_delete_expr, )?; - let merge_insert_job = MergeInsertBuilder::try_new(Arc::new(dataset.clone().inner), on)? + let mut builder = MergeInsertBuilder::try_new(Arc::new(dataset.clone().inner), on)?; + if let Some(version) = data_storage_version { + builder.data_storage_version(LanceFileVersion::from_str(&version)?); + } + let merge_insert_job = builder .when_matched(when_matched) .when_not_matched(when_not_matched) .when_not_matched_by_source(when_not_matched_by_source) @@ -241,6 +248,16 @@ fn extract_use_index<'local>(env: &mut JNIEnv<'local>, jparam: &JObject) -> Resu Ok(use_index) } +fn extract_data_storage_version<'local>( + env: &mut JNIEnv<'local>, + jparam: &JObject, +) -> Result> { + let version = env + .call_method(jparam, "dataStorageVersion", "()Ljava/util/Optional;", &[])? + .l()?; + env.get_string_opt(&version) +} + fn extract_compacted_sstables<'local>( env: &mut JNIEnv<'local>, jparam: &JObject, diff --git a/java/lance-jni/src/optimize.rs b/java/lance-jni/src/optimize.rs index 1b5695ab6f9..5b63e8a39bf 100644 --- a/java/lance-jni/src/optimize.rs +++ b/java/lance-jni/src/optimize.rs @@ -17,7 +17,7 @@ use lance::dataset::{ }; use crate::{ - block_on, + JNIEnvExt, block_on, blocking_dataset::{BlockingDataset, NATIVE_DATASET}, traits::{ FromJObjectWithEnv, IntoJava, export_vec, import_vec_from_method, import_vec_to_rust, @@ -49,6 +49,7 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction max_source_rows: JObject, // Optional max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List + data_storage_version: JObject, // Optional ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -68,7 +69,8 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction max_source_fragments, max_source_rows, max_source_bytes, - excluded_fragment_ids + excluded_fragment_ids, + data_storage_version ), JObject::null() ) @@ -92,6 +94,7 @@ fn inner_plan_compaction<'local>( max_source_rows: JObject, // Optional max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List + data_storage_version: JObject, // Optional ) -> Result> { let config = { let dataset = @@ -114,6 +117,7 @@ fn inner_plan_compaction<'local>( &max_source_rows, &max_source_bytes, &excluded_fragment_ids, + &data_storage_version, &config, )?; @@ -145,6 +149,7 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativeCommitCompacti max_source_rows: JObject, // Optional max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List + data_storage_version: JObject, // Optional ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -166,6 +171,7 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativeCommitCompacti max_source_rows, max_source_bytes, excluded_fragment_ids, + data_storage_version, ), JObject::null() ) @@ -190,6 +196,7 @@ fn inner_commit_compaction<'local>( max_source_rows: JObject, // Optional max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List + data_storage_version: JObject, // Optional ) -> Result> { let config = { let dataset = @@ -212,6 +219,7 @@ fn inner_commit_compaction<'local>( &max_source_rows, &max_source_bytes, &excluded_fragment_ids, + &data_storage_version, &config, )?; let completed_tasks = import_vec_to_rust(env, &rewrite_results, |env, rewrite_result| { @@ -252,6 +260,7 @@ pub extern "system" fn Java_org_lance_compaction_CompactionTask_nativeExecute<'l max_source_rows: JObject, // Optional max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List + data_storage_version: JObject, // Optional ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -273,7 +282,8 @@ pub extern "system" fn Java_org_lance_compaction_CompactionTask_nativeExecute<'l max_source_fragments, max_source_rows, max_source_bytes, - excluded_fragment_ids + excluded_fragment_ids, + data_storage_version ), JObject::null() ) @@ -299,6 +309,7 @@ fn inner_execute_task<'local>( max_source_rows: JObject, // Optional max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List + data_storage_version: JObject, // Optional ) -> Result> { let task_data: TaskData = task_data.extract_object(env)?; let config = { @@ -322,6 +333,7 @@ fn inner_execute_task<'local>( &max_source_rows, &max_source_bytes, &excluded_fragment_ids, + &data_storage_version, &config, )?; let compaction_task = CompactionTask { @@ -345,11 +357,10 @@ const COMPACTION_PLAN_CLASS: &str = "org/lance/compaction/CompactionPlan"; const COMPACTION_PLAN_CONSTRUCTOR_SIG: &str = "(Ljava/util/List;JLorg/lance/compaction/CompactionOptions;)V"; const REWRITE_RESULT_CLASS: &str = "org/lance/compaction/RewriteResult"; -const REWRITE_RESULT_CONSTRUCTOR_SIG: &str = - "(Lorg/lance/compaction/CompactionMetrics;Ljava/util/List;Ljava/util/List;J[B)V"; +const REWRITE_RESULT_CONSTRUCTOR_SIG: &str = "(Lorg/lance/compaction/CompactionMetrics;Ljava/util/List;Ljava/util/List;J[BLjava/lang/String;)V"; const COMPACTION_OPTIONS_CLASS: &str = "org/lance/compaction/CompactionOptions"; const COMPACTION_MODE_CLASS: &str = "org/lance/compaction/CompactionMode"; -const COMPACTION_OPTIONS_CONSTRUCTOR_SIG: &str = "(Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;)V"; +const COMPACTION_OPTIONS_CONSTRUCTOR_SIG: &str = "(Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;Ljava/util/Optional;)V"; impl IntoJava for &TaskData { fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result> { @@ -431,6 +442,11 @@ impl IntoJava for &CompactionOptions { .map(|fragment_id| to_java_long_obj(env, Some(*fragment_id as i64))) .collect::>>()?; let excluded_fragment_ids = to_java_list(env, &excluded_fragment_ids)?; + let data_storage_version = match self.data_storage_version { + Some(version) => env.new_string(version.to_string())?.into(), + None => JObject::null(), + }; + let data_storage_version_opt = to_java_optional(env, data_storage_version)?; Ok(env.new_object( COMPACTION_OPTIONS_CLASS, @@ -450,6 +466,7 @@ impl IntoJava for &CompactionOptions { JValueGen::Object(&max_source_rows_opt), JValueGen::Object(&max_source_bytes_opt), JValueGen::Object(&excluded_fragment_ids), + JValueGen::Object(&data_storage_version_opt), ], )?) } @@ -481,6 +498,7 @@ impl IntoJava for &RewriteResult { } else { JObject::null() }; + let write_version: JObject<'_> = env.new_string(&self.write_version)?.into(); Ok(env.new_object( REWRITE_RESULT_CLASS, REWRITE_RESULT_CONSTRUCTOR_SIG, @@ -490,6 +508,7 @@ impl IntoJava for &RewriteResult { JValueGen::Object(&original_fragments), JValueGen::Long(self.read_version as i64), JValueGen::Object(&row_addrs), + JValueGen::Object(&write_version), ], )?) } @@ -554,13 +573,17 @@ impl FromJObjectWithEnv for JObject<'_> { } else { Some(env.convert_byte_array(row_addrs_obj)?) }; + let write_version_obj = env + .call_method(self, "getWriteVersion", "()Ljava/util/Optional;", &[])? + .l()?; + let write_version = env.get_string_opt(&write_version_obj)?.unwrap_or_default(); Ok(RewriteResult { metrics, new_fragments, read_version, original_fragments, row_addrs, - write_version: String::new(), + write_version, }) } } diff --git a/java/lance-jni/src/update.rs b/java/lance-jni/src/update.rs index 3a6f9360062..1d2db45aec7 100644 --- a/java/lance-jni/src/update.rs +++ b/java/lance-jni/src/update.rs @@ -9,6 +9,8 @@ use crate::{JNIEnvExt, block_on}; use jni::JNIEnv; use jni::objects::{JMap, JObject, JValueGen}; use lance::dataset::UpdateBuilder; +use lance_file::version::LanceFileVersion; +use std::str::FromStr; use std::sync::Arc; use std::time::Duration; @@ -30,6 +32,7 @@ fn inner_update<'local>( let where_clause = extract_where(env, &jparam)?; let conflict_retries = extract_conflict_retries(env, &jparam)?; let retry_timeout_ms = extract_retry_timeout_ms(env, &jparam)?; + let data_storage_version = extract_data_storage_version(env, &jparam)?; // Clone the inner Dataset out of the `get_rust_field` guard and drop the // guard before running the long-lived async update. Otherwise the guard @@ -44,6 +47,10 @@ fn inner_update<'local>( .conflict_retries(conflict_retries) .retry_timeout(Duration::from_millis(retry_timeout_ms)); + if let Some(version) = data_storage_version { + builder = builder.data_storage_version(LanceFileVersion::from_str(&version)?); + } + if let Some(predicate) = where_clause { builder = builder.update_where(&predicate)?; } @@ -96,6 +103,16 @@ fn extract_retry_timeout_ms<'local>(env: &mut JNIEnv<'local>, jparam: &JObject) Ok(timeout_ms) } +fn extract_data_storage_version<'local>( + env: &mut JNIEnv<'local>, + jparam: &JObject, +) -> Result> { + let version = env + .call_method(jparam, "dataStorageVersion", "()Ljava/util/Optional;", &[])? + .l()?; + env.get_string_opt(&version) +} + const UPDATE_RESULT_CLASS: &str = "org/lance/update/UpdateResult"; const UPDATE_RESULT_CONSTRUCTOR_SIG: &str = "(Lorg/lance/Dataset;J)V"; diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index c9d2d2005b1..414c0c935fd 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -194,6 +194,7 @@ pub fn build_compaction_options( max_source_rows: &JObject, // Optional max_source_bytes: &JObject, // Optional excluded_fragment_ids: &JObject, // List + data_storage_version: &JObject, // Optional config: &std::collections::HashMap, ) -> Result { let mut compaction_options = CompactionOptions::from_dataset_config(config)?; @@ -256,6 +257,9 @@ pub fn build_compaction_options( }) }) .collect::>>()?; + if let Some(version) = env.get_string_opt(data_storage_version)? { + compaction_options.data_storage_version = Some(LanceFileVersion::from_str(&version)?); + } Ok(compaction_options) } diff --git a/java/src/main/java/org/lance/compaction/Compaction.java b/java/src/main/java/org/lance/compaction/Compaction.java index 580b7ea75b7..96062383b94 100644 --- a/java/src/main/java/org/lance/compaction/Compaction.java +++ b/java/src/main/java/org/lance/compaction/Compaction.java @@ -47,7 +47,8 @@ public static CompactionPlan planCompaction( compactionOptions.getMaxSourceFragments(), compactionOptions.getMaxSourceRows(), compactionOptions.getMaxSourceBytes(), - compactionOptions.getExcludedFragmentIds()); + compactionOptions.getExcludedFragmentIds(), + compactionOptions.getDataStorageVersion()); } public static CompactionMetrics commitCompaction( @@ -71,7 +72,8 @@ public static CompactionMetrics commitCompaction( compactionOptions.getMaxSourceFragments(), compactionOptions.getMaxSourceRows(), compactionOptions.getMaxSourceBytes(), - compactionOptions.getExcludedFragmentIds()); + compactionOptions.getExcludedFragmentIds(), + compactionOptions.getDataStorageVersion()); } public static native CompactionMetrics nativeCommitCompaction( @@ -90,7 +92,8 @@ public static native CompactionMetrics nativeCommitCompaction( Optional maxSourceFragments, Optional maxSourceRows, Optional maxSourceBytes, - List excludedFragmentIds); + List excludedFragmentIds, + Optional dataStorageVersion); private static native CompactionPlan nativePlanCompaction( Dataset dataset, @@ -107,5 +110,6 @@ private static native CompactionPlan nativePlanCompaction( Optional maxSourceFragments, Optional maxSourceRows, Optional maxSourceBytes, - List excludedFragmentIds); + List excludedFragmentIds, + Optional dataStorageVersion); } diff --git a/java/src/main/java/org/lance/compaction/CompactionOptions.java b/java/src/main/java/org/lance/compaction/CompactionOptions.java index df1fd8628a4..ef6b517c307 100644 --- a/java/src/main/java/org/lance/compaction/CompactionOptions.java +++ b/java/src/main/java/org/lance/compaction/CompactionOptions.java @@ -51,6 +51,7 @@ public class CompactionOptions implements Serializable { private Optional maxSourceRows; private Optional maxSourceBytes; private List excludedFragmentIds; + private Optional dataStorageVersion; private CompactionOptions( Optional targetRowsPerFragment, @@ -66,7 +67,8 @@ private CompactionOptions( Optional maxSourceFragments, Optional maxSourceRows, Optional maxSourceBytes, - List excludedFragmentIds) { + List excludedFragmentIds, + Optional dataStorageVersion) { this.targetRowsPerFragment = targetRowsPerFragment; this.maxRowsPerGroup = maxRowsPerGroup; this.maxBytesPerFile = maxBytesPerFile; @@ -81,6 +83,7 @@ private CompactionOptions( this.maxSourceRows = maxSourceRows; this.maxSourceBytes = maxSourceBytes; this.excludedFragmentIds = List.copyOf(excludedFragmentIds); + this.dataStorageVersion = dataStorageVersion; } public Optional getDeferIndexRemap() { @@ -112,6 +115,10 @@ public List getExcludedFragmentIds() { return excludedFragmentIds; } + public Optional getDataStorageVersion() { + return dataStorageVersion; + } + public Optional getMaterializeDeletions() { return materializeDeletions; } @@ -161,6 +168,7 @@ public String toString() { .add("maxSourceRows", maxSourceRows.orElse(null)) .add("maxSourceBytes", maxSourceBytes.orElse(null)) .add("excludedFragmentIds", excludedFragmentIds) + .add("dataStorageVersion", dataStorageVersion.orElse(null)) .toString(); } @@ -179,6 +187,7 @@ private void writeObject(ObjectOutputStream output) throws IOException { output.writeObject(maxSourceRows.orElse(null)); output.writeObject(maxSourceBytes.orElse(null)); output.writeObject(excludedFragmentIds); + output.writeObject(dataStorageVersion.orElse(null)); } private void readObject(ObjectInputStream input) throws IOException, ClassNotFoundException { @@ -205,6 +214,7 @@ private void readObject(ObjectInputStream input) throws IOException, ClassNotFou this.maxSourceRows = readTrailingLong(input); this.maxSourceBytes = readTrailingLong(input); this.excludedFragmentIds = readTrailingLongList(input); + this.dataStorageVersion = readTrailingString(input); } /** @@ -238,6 +248,18 @@ private static List readTrailingLongList(ObjectInputStream input) } } + private static Optional readTrailingString(ObjectInputStream input) + throws IOException, ClassNotFoundException { + try { + return Optional.ofNullable((String) input.readObject()); + } catch (OptionalDataException e) { + if (!e.eof) { + throw e; + } + return Optional.empty(); + } + } + /** Builder for CompactionOptions. */ public static class Builder { private Optional targetRowsPerFragment = Optional.empty(); @@ -254,6 +276,7 @@ public static class Builder { private Optional maxSourceRows = Optional.empty(); private Optional maxSourceBytes = Optional.empty(); private List excludedFragmentIds = Collections.emptyList(); + private Optional dataStorageVersion = Optional.empty(); private Builder() {} @@ -345,6 +368,12 @@ public Builder withMaxSourceBytes(long maxSourceBytes) { return this; } + /** Set the exact data storage version for rewritten files. */ + public Builder withDataStorageVersion(String version) { + this.dataStorageVersion = Optional.of(java.util.Objects.requireNonNull(version)); + return this; + } + /** * Fragment IDs to exclude from compaction planning. Excluded fragments remain unchanged and act * as boundaries, so fragments on opposite sides are not combined into the same task. Duplicate @@ -393,7 +422,8 @@ public CompactionOptions build() { maxSourceFragments, maxSourceRows, maxSourceBytes, - excludedFragmentIds); + excludedFragmentIds, + dataStorageVersion); } } } diff --git a/java/src/main/java/org/lance/compaction/CompactionTask.java b/java/src/main/java/org/lance/compaction/CompactionTask.java index 50cd6f57b64..275aefaf092 100644 --- a/java/src/main/java/org/lance/compaction/CompactionTask.java +++ b/java/src/main/java/org/lance/compaction/CompactionTask.java @@ -60,7 +60,8 @@ public RewriteResult execute(Dataset dataset) { compactionOptions.getMaxSourceFragments(), compactionOptions.getMaxSourceRows(), compactionOptions.getMaxSourceBytes(), - compactionOptions.getExcludedFragmentIds()); + compactionOptions.getExcludedFragmentIds(), + compactionOptions.getDataStorageVersion()); } private native RewriteResult nativeExecute( @@ -80,7 +81,8 @@ private native RewriteResult nativeExecute( Optional maxSourceFragments, Optional maxSourceRows, Optional maxSourceBytes, - List excludedFragmentIds); + List excludedFragmentIds, + Optional dataStorageVersion); public CompactionOptions getCompactionOptions() { return compactionOptions; diff --git a/java/src/main/java/org/lance/compaction/RewriteResult.java b/java/src/main/java/org/lance/compaction/RewriteResult.java index c4c5d816c2c..3c61470303b 100644 --- a/java/src/main/java/org/lance/compaction/RewriteResult.java +++ b/java/src/main/java/org/lance/compaction/RewriteResult.java @@ -19,12 +19,15 @@ import java.io.Serializable; import java.util.List; +import java.util.Optional; /** * Rewrite Result of a single compaction task. It will be passed across different workers and be * committed later. */ public class RewriteResult implements Serializable { + private static final long serialVersionUID = 4501818269828675274L; + private final CompactionMetrics metrics; private final List newFragments; private final List originalFragments; @@ -33,18 +36,21 @@ public class RewriteResult implements Serializable { // Serialized RoaringTreemap of row addresses read from the original fragments. // null for stable row IDs. @Nullable private final byte[] rowAddrs; + @Nullable private final String writeVersion; public RewriteResult( CompactionMetrics metrics, List newFragments, List originalFragments, long readVersion, - byte[] rowAddrs) { + byte[] rowAddrs, + String writeVersion) { this.metrics = metrics; this.newFragments = newFragments; this.originalFragments = originalFragments; this.readVersion = readVersion; this.rowAddrs = rowAddrs; + this.writeVersion = writeVersion; } public long getReadVersion() { @@ -67,4 +73,9 @@ public List getNewFragments() { public List getOriginalFragments() { return originalFragments; } + + /** Returns the exact storage version used for rewritten files, if recorded. */ + public Optional getWriteVersion() { + return Optional.ofNullable(writeVersion); + } } diff --git a/java/src/main/java/org/lance/merge/MergeInsertParams.java b/java/src/main/java/org/lance/merge/MergeInsertParams.java index a9a696ead73..0e34c165893 100644 --- a/java/src/main/java/org/lance/merge/MergeInsertParams.java +++ b/java/src/main/java/org/lance/merge/MergeInsertParams.java @@ -39,6 +39,7 @@ public class MergeInsertParams { private long retryTimeoutMs = 30 * 1000; private boolean skipAutoCleanup = false; private boolean useIndex = true; + private Optional dataStorageVersion = Optional.empty(); private List compactedSstables = Collections.emptyList(); public MergeInsertParams(List on) { @@ -244,6 +245,16 @@ public MergeInsertParams withUseIndex(boolean useIndex) { return this; } + /** + * Set the exact data storage version for files written by this operation. + * + *

If omitted, the manifest fallback is used. This does not change the manifest fallback. + */ + public MergeInsertParams withDataStorageVersion(String version) { + this.dataStorageVersion = Optional.of(Preconditions.checkNotNull(version)); + return this; + } + /** * Mark MemWAL SSTables as compacted into the base table. * @@ -325,6 +336,10 @@ public boolean useIndex() { return useIndex; } + public Optional dataStorageVersion() { + return dataStorageVersion; + } + @Override public String toString() { return MoreObjects.toStringHelper(this) @@ -343,6 +358,7 @@ public String toString() { .add("retryTimeoutMs", retryTimeoutMs) .add("skipAutoCleanup", skipAutoCleanup) .add("useIndex", useIndex) + .add("dataStorageVersion", dataStorageVersion.orElse(null)) .toString(); } diff --git a/java/src/main/java/org/lance/update/UpdateParams.java b/java/src/main/java/org/lance/update/UpdateParams.java index de0224ceb1f..e260e3476ef 100644 --- a/java/src/main/java/org/lance/update/UpdateParams.java +++ b/java/src/main/java/org/lance/update/UpdateParams.java @@ -47,6 +47,7 @@ public class UpdateParams { private final Map updates; private Optional whereClause = Optional.empty(); + private Optional dataStorageVersion = Optional.empty(); private int conflictRetries = DEFAULT_CONFLICT_RETRIES; private long retryTimeoutMs = DEFAULT_RETRY_TIMEOUT_MS; @@ -105,6 +106,17 @@ public UpdateParams withRetryTimeoutMs(long timeoutMs) { return this; } + /** + * Set the exact data storage version for files written by this operation. + * + *

If omitted, the manifest fallback is used. When another V2 version produces a mixed + * snapshot, the commit derives the required capability without changing the manifest fallback. + */ + public UpdateParams withDataStorageVersion(String version) { + this.dataStorageVersion = Optional.of(Preconditions.checkNotNull(version)); + return this; + } + /** Returns an unmodifiable view of the update expressions. */ public Map updates() { return Collections.unmodifiableMap(updates); @@ -122,6 +134,10 @@ public long retryTimeoutMs() { return retryTimeoutMs; } + public Optional dataStorageVersion() { + return dataStorageVersion; + } + @Override public String toString() { return MoreObjects.toStringHelper(this) @@ -129,6 +145,7 @@ public String toString() { .add("whereClause", whereClause.orElse(null)) .add("conflictRetries", conflictRetries) .add("retryTimeoutMs", retryTimeoutMs) + .add("dataStorageVersion", dataStorageVersion.orElse(null)) .toString(); } } diff --git a/java/src/test/java/org/lance/CompactionTest.java b/java/src/test/java/org/lance/CompactionTest.java index a0d58612f79..375b73184d3 100644 --- a/java/src/test/java/org/lance/CompactionTest.java +++ b/java/src/test/java/org/lance/CompactionTest.java @@ -60,6 +60,7 @@ public void testBasicCompaction(@TempDir Path tempDir) throws Exception { .withNumThreads(1) .withMaxSourceRows(1000) .withMaxSourceBytes(10L * 1024 * 1024) + .withDataStorageVersion("2.2") .build(); CompactionPlan compactionPlan = Compaction.planCompaction(dataset, compactionOptions); @@ -69,6 +70,8 @@ public void testBasicCompaction(@TempDir Path tempDir) throws Exception { assertEquals( Optional.of(10L * 1024 * 1024), compactionPlan.getCompactionOptions().getMaxSourceBytes()); + assertEquals( + Optional.of("2.2"), compactionPlan.getCompactionOptions().getDataStorageVersion()); // will plan to compact two fragments into one. assertEquals(1, compactionPlan.getCompactionTasks().size()); @@ -80,6 +83,7 @@ public void testBasicCompaction(@TempDir Path tempDir) throws Exception { // mock network transferring task = serializeAndDeserialize(task); RewriteResult result = task.execute(dataset); + assertEquals(Optional.of("2.2"), result.getWriteVersion()); CompactionMetrics metrics = result.getMetrics(); // remove previous fragments and add new single fragment assertEquals(2, metrics.getFragmentsRemoved()); diff --git a/java/src/test/java/org/lance/MergeInsertTest.java b/java/src/test/java/org/lance/MergeInsertTest.java index b738ef8852d..8a1663e0245 100644 --- a/java/src/test/java/org/lance/MergeInsertTest.java +++ b/java/src/test/java/org/lance/MergeInsertTest.java @@ -72,7 +72,15 @@ public void testWhenNotMatchedInsertAll() throws Exception { try (ArrowArrayStream sourceStream = convertToStream(source, allocator)) { MergeInsertResult result = dataset.mergeInsert( - new MergeInsertParams(Collections.singletonList("id")), sourceStream); + new MergeInsertParams(Collections.singletonList("id")) + .withDataStorageVersion("2.2"), + sourceStream); + + Assertions.assertTrue( + result.dataset().getFragments().stream() + .flatMap(fragment -> fragment.metadata().getFiles().stream()) + .anyMatch( + file -> file.getFileMajorVersion() == 2 && file.getFileMinorVersion() == 2)); Assertions.assertEquals( "{0=Person 0, 1=Person 1, 2=Person 2, 3=Person 3, 4=Person 4, 7=Source 7, 8=Source 8, 9=Source 9}", diff --git a/java/src/test/java/org/lance/UpdateTest.java b/java/src/test/java/org/lance/UpdateTest.java index 7b69b1e6271..22993657edf 100644 --- a/java/src/test/java/org/lance/UpdateTest.java +++ b/java/src/test/java/org/lance/UpdateTest.java @@ -80,10 +80,17 @@ public void tearDown() { @Test public void testUpdateAllRows() { - UpdateResult result = dataset.update(new UpdateParams(ImmutableMap.of("name", "'updated'"))); + UpdateResult result = + dataset.update( + new UpdateParams(ImmutableMap.of("name", "'updated'")).withDataStorageVersion("2.2")); Assertions.assertEquals(ROW_COUNT, result.getNumRowsUpdated()); try (Dataset newDataset = result.getDataset()) { + Assertions.assertTrue( + newDataset.getFragments().stream() + .flatMap(fragment -> fragment.metadata().getFiles().stream()) + .allMatch( + file -> file.getFileMajorVersion() == 2 && file.getFileMinorVersion() == 2)); List names = readNames(newDataset); Assertions.assertEquals(ROW_COUNT, names.size()); for (String name : names) { diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index e386e6cb854..8bcc7847b13 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -524,6 +524,16 @@ def retry_timeout(self, timeout: timedelta) -> "MergeInsertBuilder": """ return super(MergeInsertBuilder, self).retry_timeout(timeout) + def data_storage_version(self, version: str) -> "MergeInsertBuilder": + """Set the exact data storage version for files written by this operation. + + If omitted, the dataset manifest's fallback data storage version is used. + Selecting a different V2 version causes the commit to derive mixed + data-file-version support from its final manifest. The manifest fallback + itself is not changed. + """ + return super(MergeInsertBuilder, self).data_storage_version(version) + def use_index(self, use_index: bool) -> "MergeInsertBuilder": """ Controls whether to use indices for the merge operation. @@ -2929,6 +2939,7 @@ def update( where: Optional[str] = None, conflict_retries: int = 10, retry_timeout: timedelta = timedelta(seconds=30), + data_storage_version: Optional[str] = None, ) -> UpdateResult: """ Update column values for rows matching where. @@ -2947,6 +2958,10 @@ def update( the operation before giving up. At least one attempt will be made, regardless of how long it takes to complete. Subsequent attempts will be cancelled once this timeout is reached. Default is 30 seconds. + data_storage_version : str, optional + Exact data storage version for files written by this operation. If + omitted, the manifest fallback is used. This does not change the + manifest fallback. Returns ------- @@ -2969,7 +2984,13 @@ def update( """ if isinstance(where, pa.compute.Expression): where = str(where) - return self._ds.update(updates, where, conflict_retries, retry_timeout) + return self._ds.update( + updates, + where, + conflict_retries, + retry_timeout, + data_storage_version, + ) def versions(self): """ diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index fffacc82712..27d1a95ed53 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -484,6 +484,9 @@ class _Dataset: self, updates: Dict[str, str], predicate: Optional[str] = None, + conflict_retries: Optional[int] = None, + retry_timeout: Optional[timedelta] = None, + data_storage_version: Optional[str] = None, ) -> UpdateResult: ... def count_deleted_rows(self) -> int: ... def versions(self) -> List[Version]: ... @@ -672,6 +675,7 @@ class _MergeInsertBuilder: def when_matched_fail(self) -> Self: ... def when_not_matched_insert_all(self) -> Self: ... def when_not_matched_by_source_delete(self, expr: Optional[str] = None) -> Self: ... + def data_storage_version(self, version: str) -> Self: ... def target_bases(self, bases: list[str]) -> Self: ... def target_all_bases(self, include_primary: bool = True) -> Self: ... def execute(self, new_data: pa.RecordBatchReader) -> ExecuteResult: ... diff --git a/python/python/lance/lance/optimize.pyi b/python/python/lance/lance/optimize.pyi index c4b6b6546e6..dcdb6364022 100644 --- a/python/python/lance/lance/optimize.pyi +++ b/python/python/lance/lance/optimize.pyi @@ -26,6 +26,7 @@ class CompactionMetrics: class RewriteResult: read_version: int + write_version: str metrics: CompactionMetrics old_fragments: List["FragmentMetadata"] new_fragments: List["FragmentMetadata"] diff --git a/python/python/lance/optimize.py b/python/python/lance/optimize.py index 4d45171b71f..5f223f2b7c7 100644 --- a/python/python/lance/optimize.py +++ b/python/python/lance/optimize.py @@ -121,3 +121,8 @@ class CompactionOptions(TypedDict): are not combined into the same task. Duplicate and unknown IDs are ignored. (default: None) """ + data_storage_version: Optional[str] + """ + Exact data storage version for rewritten files. If omitted, compaction + uses the manifest fallback. This does not change the manifest fallback. + """ diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index d5d137980b6..8eff854a853 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -3931,6 +3931,46 @@ def test_update_with_retry_parameters(tmp_path: Path): check_update_stats(update_dict, (10,)) +def test_update_exact_data_storage_version(tmp_path: Path): + dataset = lance.write_dataset( + pa.table({"id": [1, 2], "value": [10, 20]}), + tmp_path / "dataset", + data_storage_version="2.1", + ) + + dataset.update({"value": "value + 1"}, data_storage_version="2.2") + + files = [ + file for fragment in dataset.get_fragments() for file in fragment.metadata.files + ] + assert {(file.file_major_version, file.file_minor_version) for file in files} == { + (2, 2) + } + + +def test_merge_insert_exact_data_storage_version(tmp_path: Path): + dataset = lance.write_dataset( + pa.table({"id": [1, 2], "value": [10, 20]}), + tmp_path / "dataset", + data_storage_version="2.1", + ) + + ( + dataset.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .data_storage_version("2.2") + .execute(pa.table({"id": [2, 3], "value": [21, 30]})) + ) + + versions = { + (file.file_major_version, file.file_minor_version) + for fragment in dataset.get_fragments() + for file in fragment.metadata.files + } + assert versions == {(2, 1), (2, 2)} + + def test_scan_with_batch_size(tmp_path: Path): base_dir = tmp_path / "dataset" df = pd.DataFrame({"a": range(10000), "b": range(10000)}) diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index 661945880bf..3506d97eaaf 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -592,12 +592,22 @@ def test_dataset_distributed_optimize(tmp_path: Path): base_dir = tmp_path / "dataset" data = pa.table({"a": range(800), "b": range(800)}) - dataset = lance.write_dataset(data, base_dir, max_rows_per_file=200) + dataset = lance.write_dataset( + data, + base_dir, + max_rows_per_file=200, + data_storage_version="2.1", + ) fragments = dataset.get_fragments() assert len(fragments) == 4 plan = Compaction.plan( - dataset, options=dict(target_rows_per_fragment=400, num_threads=1) + dataset, + options=dict( + target_rows_per_fragment=400, + num_threads=1, + data_storage_version="2.2", + ), ) assert plan.read_version == 1 assert plan.num_tasks() == 2 @@ -625,6 +635,7 @@ def test_dataset_distributed_optimize(tmp_path: Path): target_rows_per_fragment=400, excluded_fragment_ids=None, num_threads=1, + data_storage_version="2.2", ), ) assert none_plan == plan @@ -637,6 +648,7 @@ def test_dataset_distributed_optimize(tmp_path: Path): assert task == plan.tasks[0] result1 = plan.tasks[0].execute(dataset) + assert result1.write_version == "2.2" result1.metrics.fragments_removed == 2 result1.metrics.fragments_added == 1 @@ -644,6 +656,7 @@ def test_dataset_distributed_optimize(tmp_path: Path): result = pickle.loads(pickled_result) assert isinstance(result, RewriteResult) assert result == result1 + assert result.write_version == "2.2" assert re.match( r"RewriteResult\(read_version=1, new_fragments=\[.+\], old_fragments=\[.+\]\)", repr(result), diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 1191309a1f7..2f91fa9f6f6 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -423,6 +423,15 @@ impl MergeInsertBuilder { Ok(slf) } + pub fn data_storage_version<'a>( + mut slf: PyRefMut<'a, Self>, + version: &str, + ) -> PyResult> { + slf.builder + .data_storage_version(version.parse().infer_error()?); + Ok(slf) + } + pub fn use_index(mut slf: PyRefMut<'_, Self>, use_index: bool) -> PyResult> { slf.builder.use_index(use_index); Ok(slf) @@ -1941,13 +1950,14 @@ impl Dataset { Ok(dict.into()) } - #[pyo3(signature=(updates, predicate=None, conflict_retries=None, retry_timeout=None))] + #[pyo3(signature=(updates, predicate=None, conflict_retries=None, retry_timeout=None, data_storage_version=None))] fn update( &mut self, updates: &Bound<'_, PyDict>, predicate: Option<&str>, conflict_retries: Option, retry_timeout: Option, + data_storage_version: Option<&str>, ) -> PyResult> { let mut builder = UpdateBuilder::new(self.ds.clone()); if let Some(predicate) = predicate { @@ -1964,6 +1974,10 @@ impl Dataset { builder = builder.retry_timeout(timeout); } + if let Some(version) = data_storage_version { + builder = builder.data_storage_version(version.parse().infer_error()?); + } + for (key, value) in updates { let column: PyBackedStr = key.cast::()?.clone().try_into()?; let expr: PyBackedStr = value.cast::()?.clone().try_into()?; diff --git a/python/src/dataset/optimize.rs b/python/src/dataset/optimize.rs index a2c1f973727..856f6225c23 100644 --- a/python/src/dataset/optimize.rs +++ b/python/src/dataset/optimize.rs @@ -86,6 +86,12 @@ fn parse_compaction_options( opts.excluded_fragment_ids = value.extract::>>()?.unwrap_or_default(); } + "data_storage_version" => { + let version: Option = value.extract()?; + opts.data_storage_version = version + .map(|version| version.parse().infer_error()) + .transpose()?; + } _ => { return Err(PyValueError::new_err(format!( "Invalid compaction option: {}", @@ -393,6 +399,12 @@ impl PyRewriteResult { self.0.read_version } + /// str : The exact data storage version used for rewritten files. + #[getter] + pub fn write_version(&self) -> &str { + &self.0.write_version + } + /// List[lance.fragment.FragmentMetadata] : The metadata for fragments that are being replaced. #[getter] pub fn original_fragments<'py>(&self, py: Python<'py>) -> PyResult>> {