From 08bebbe56e37a2da39181abffdf6a2926303b370 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 20 Aug 2026 17:44:24 +0800 Subject: [PATCH 01/24] feat: stabilize field IDs across schema evolution --- docs/src/format/table/index.md | 6 + docs/src/format/table/schema.md | 56 ++- docs/src/format/table/versioning.md | 3 +- java/lance-jni/src/transaction.rs | 327 ++++++++++++++---- .../org/lance/operation/SchemaOperation.java | 7 +- .../java/org/lance/operation/MergeTest.java | 12 +- protos/table.proto | 11 + python/python/tests/test_dataset.py | 13 + python/python/tests/test_schema.py | 10 + python/src/schema.rs | 10 +- python/src/transaction.rs | 12 +- rust/lance-core/src/datatypes/field.rs | 21 ++ rust/lance-core/src/datatypes/schema.rs | 82 ++++- rust/lance-table/src/feature_flags.rs | 105 +++++- rust/lance-table/src/format/manifest.rs | 120 ++++++- rust/lance-table/src/transaction.rs | 5 +- .../src/transaction/manifest_build.rs | 117 ++++++- .../src/transaction/test_support.rs | 1 + rust/lance-table/src/transaction/validate.rs | 322 +++++++++++++++++ rust/lance/src/blob.rs | 31 ++ rust/lance/src/dataset.rs | 69 +++- rust/lance/src/dataset/fragment.rs | 2 +- rust/lance/src/dataset/schema_evolution.rs | 55 ++- rust/lance/src/dataset/tests/dataset_io.rs | 104 +++++- .../src/dataset/tests/dataset_migrations.rs | 205 ++++++++++- rust/lance/src/dataset/updater.rs | 3 +- rust/lance/src/dataset/write.rs | 7 + rust/lance/src/dataset/write/commit.rs | 13 + rust/lance/src/io/commit.rs | 130 ++++++- 29 files changed, 1716 insertions(+), 143 deletions(-) diff --git a/docs/src/format/table/index.md b/docs/src/format/table/index.md index f9da132cf3b..a0037dd3f50 100644 --- a/docs/src/format/table/index.md +++ b/docs/src/format/table/index.md @@ -28,6 +28,12 @@ A manifest describes a single version of the dataset. It contains the complete schema definition including nested fields, the list of data fragments comprising this version, a monotonically increasing version number, and an optional reference to the index section that describes a list of index metadata. +`max_allocated_field_id` is optional allocator state. Its presence activates the stable field-ID +contract for this branch ancestry. It is a high-water mark: every non-negative field ID in the +canonical manifest schema, base data-file mappings, and overlay mappings must be less than or equal +to it, and a successor manifest must never lower it. Schema-preserving rewrites retain the value; +new identities allocate above it with checked arithmetic. See [Field IDs](schema.md#field-ids). +
Manifest protobuf message diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index 15e8946c708..2ad0d25f5fb 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -221,16 +221,44 @@ Assigned IDs with parent relationships: Note: A `parent_id` of -1 indicates a top-level field. For nested fields, `parent_id` references the ID of the parent field. Child fields reference their parent via `parent_id` rather than being stored as separate "children" arrays in the protobuf message (though the Rust in-memory representation maintains a children vector for convenience). **New field assignment (incremental):** -When fields are added later (e.g., through schema evolution), they receive the next available ID -incrementally. This preserves the history of field additions. +When fields are added later (e.g., through schema evolution), they receive incrementally assigned +IDs. On a dataset with the stable field-ID contract activated, allocation starts immediately after +the manifest's persistent `max_allocated_field_id`; IDs retired by drop or replacement are never +reused. + +Field-ID metadata supplied on an incoming Arrow schema is not an allocation authority. When that +schema is merged into an activated dataset, IDs for newly introduced logical fields are cleared and +assigned by the dataset allocator; callers cannot select or reserve IDs through Arrow metadata. + +The presence of `Manifest.max_allocated_field_id` is the activation marker. If it is absent, the +dataset has legacy allocation semantics and an implementation may derive the next ID from fields +still referenced by the current snapshot. Activation initializes the high-water mark from the +maximum non-negative ID referenced by the canonical manifest schema, base data files, and overlay +files. Activation provides a forward guarantee only; it cannot reconstruct identities that were +dropped or reused in older snapshots. + +New datasets activate this contract in their initial manifest and require both readers and writers +to understand it. Existing legacy datasets remain unchanged until an explicit migration commit; +controlled deployments may use a writer-only migration only after every pre-gate writer has been +retired. ### Field ID Properties -- **Immutable**: Once assigned, a field's ID never changes +- **Immutable after activation**: An identity keeps its ID throughout the activated branch ancestry - **Unique**: Each field within a table has a unique ID -- **Stable**: IDs are preserved across schema evolution operations +- **Never reused after activation**: Dropped and replaced identities permanently retire their IDs +- **Monotonic**: New identities are allocated densely above the persistent high-water mark - **Sparse**: Field IDs may not form a contiguous sequence after schema evolution +The guarantee is scoped to one dataset branch ancestry. Field IDs are not globally unique across +datasets or independently evolving branches. A persistent cross-dataset or cross-branch reference +must carry the corresponding dataset and ancestry identity. + +Two branches may allocate the same integer after their common ancestor. A future branch-merge +operation must reject the merge when both sides created post-ancestor identities whose bindings +differ; it must not silently choose one binding, renumber already-persisted fields, or merge them by +name. Branch-local collision avoidance requires a separate allocator design. + ### Using Field IDs When referencing fields internally within the format, use the field ids rather than field names or positions. @@ -296,13 +324,29 @@ The complete schema is represented as a collection of top-level fields plus meta Field IDs enable efficient schema evolution: - **Add Column**: Assign a new field ID and add to schema -- **Drop Column**: Remove field from schema; its ID may be reused in some systems +- **Drop Column**: Remove the field and permanently retire its ID after activation - **Rename Column**: Change field name; ID remains the same - **Reorder Columns**: Change field order in schema; IDs remain the same -- **Type Evolution**: Data type can be changed. This might require rewriting the column in the data, depending on how the type was changed. +- **Metadata or Nullability Change**: Preserve the field ID +- **Type Replacement**: Allocate a new field ID and retire the old identity +- **Overwrite**: On an activated dataset, replace all field identities with fresh IDs above the + current high-water mark The use of field IDs ensures that data files can be correctly interpreted even as the schema changes over time. +### Blob Identity Namespace + +A Blob column's canonical logical fields in the manifest schema participate in stable field-ID +allocation. The top-level Blob field is the public binding identity; logical children already +present in the manifest schema are also allocated and retired normally. + +Blob writer-prepared fields and stored descriptor fields are representation details. Synthetic +children such as `kind`, `blob_id`, `blob_size`, `position`, `size`, and `blob_uri` do not enter the +dataset field-ID namespace unless they are part of the canonical manifest schema or a +`DataFile.fields` mapping. Their IDs may remain `-1` or use a file-local namespace, and they do not +advance `max_allocated_field_id`. The `blob_id` value identifies a sidecar object and is unrelated +to schema field IDs. + ## Example Schemas The examples below use a simplified representation of the field structure. In the actual protobuf format, `type` refers to the field type enum (PARENT/REPEATED/LEAF) and `logical_type` contains the data type string representation. diff --git a/docs/src/format/table/versioning.md b/docs/src/format/table/versioning.md index e0b7d480702..4ce218db314 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -31,7 +31,8 @@ they should return an "unsupported" error on any read or write operation. | 32 | `FLAG_DISABLE_TRANSACTION_FILE` | No | Yes | Transactions are recorded in the manifest rather than in a separate transaction file. | | 64 | `FLAG_UNSTABLE_DATA_OVERLAY_FILES` | Yes | Yes | Fragments may carry data overlay files. Unstable: release builds reject it unless explicitly opted in. | | 128 | `FLAG_MEM_WAL_INDEX_CATCHUP` | Yes | Yes | `index_catchup` is maintained on this table, so an index absent from it is *not* caught up. See [MemWAL](mem_wal.md). | +| 256 | `FLAG_STABLE_FIELD_IDS` | Conditional | Yes | The manifest carries a persistent field-ID high-water mark. The safe activation mode also requires readers; a controlled rollout may require writers only. See [Field IDs](schema.md#field-ids). | -Flags with bit values 256 and above are unknown and will cause implementations to reject the dataset with an "unsupported" error. +Flags with bit values 512 and above are unknown and will cause implementations to reject the dataset with an "unsupported" error. diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index 5a4a295b0f2..53457fd2113 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -921,7 +921,8 @@ fn convert_schema_from_operation( java_allocator: &JObject, dataset: Option<&mut BlockingDataset>, read_version: u64, -) -> Result { + replaces_all_identities: bool, +) -> Result<(LanceSchema, HashMap)> { let schema_ptr = env .call_method( java_operation, @@ -933,9 +934,15 @@ fn convert_schema_from_operation( let c_schema_ptr = schema_ptr as *mut FFI_ArrowSchema; let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) }; - if let Some(dataset) = dataset { - let arrow_schema = Schema::try_from(&c_schema)?; + let arrow_schema = Schema::try_from(&c_schema)?; + let original_schema = LanceSchema::try_from(&arrow_schema).map_err(|e| { + Error::input_error(format!( + "Failed to convert Arrow schema to Lance schema: {}", + e + )) + })?; + let schema = if let Some(dataset) = dataset { // Derive field ids based on the transaction read dataset schema. let read_schema = { if dataset.inner.version().version == read_version { @@ -947,17 +954,42 @@ fn convert_schema_from_operation( }; let max_field_id = dataset.inner.manifest().max_field_id(); - let schema = - LanceSchema::from_arrow_schema(&arrow_schema, Some(read_schema), Some(max_field_id))?; - Ok(schema) + let stable_field_ids = dataset.inner.manifest().uses_stable_field_ids(); + LanceSchema::from_arrow_schema( + &arrow_schema, + Some(read_schema), + Some(max_field_id), + stable_field_ids, + replaces_all_identities, + )? } else { - let schema = Schema::try_from(&c_schema)?; - LanceSchema::try_from(&schema).map_err(|e| { - Error::input_error(format!( - "Failed to convert Arrow schema to Lance schema: {}", - e - )) + // New datasets use stable field IDs by default, so Arrow metadata is + // not an allocation authority even before a manifest exists. + LanceSchema::from_arrow_schema(&arrow_schema, None, None, true, true)? + }; + + let field_id_remap = original_schema + .fields_pre_order() + .zip(schema.fields_pre_order()) + .filter_map(|(original, canonical)| { + (original.id >= 0 && original.id != canonical.id).then_some((original.id, canonical.id)) }) + .collect(); + Ok((schema, field_id_remap)) +} + +fn remap_fragment_field_ids(fragments: &mut [Fragment], field_id_remap: &HashMap) { + if field_id_remap.is_empty() { + return; + } + for fragment in fragments { + for file in fragment.referenced_lance_files_mut() { + for field_id in Arc::make_mut(&mut file.fields) { + if let Some(canonical_id) = field_id_remap.get(field_id) { + *field_id = *canonical_id; + } + } + } } } @@ -982,14 +1014,17 @@ trait SchemaExt { max_existing_id: Option, ) -> Result<()>; - /// Create schema from `arrow_schema`, with field id priority below: - /// 1. arrow metadata field id. - /// 2. field id from `base_schema`. - /// 3. field id from `max_existing_id`. + /// Create a schema from `arrow_schema`. + /// + /// For an existing dataset, Arrow field-ID metadata is not an allocation + /// authority: matching identities inherit IDs from `base_schema` and every + /// new identity is allocated above `max_existing_id`. fn from_arrow_schema( arrow_schema: &Schema, base_schema: Option, max_existing_id: Option, + dataset_allocator_is_authoritative: bool, + replaces_all_identities: bool, ) -> Result; } @@ -999,7 +1034,6 @@ impl SchemaExt for LanceSchema { base_schema: Option, max_existing_id: Option, ) -> Result<()> { - // Set id from base_schema if let Some(base_schema) = &base_schema { for field in self.fields.iter_mut() { if let Some(base_field) = base_schema.field(&field.name) { @@ -1013,7 +1047,7 @@ impl SchemaExt for LanceSchema { .map(|s| s.max_field_id().unwrap_or(-1)) .unwrap_or(-1); let max_id = max_id.max(max_existing_id.unwrap_or(-1)); - self.set_field_id(Some(max_id)); + self.try_set_field_id(Some(max_id))?; Ok(()) } @@ -1021,6 +1055,8 @@ impl SchemaExt for LanceSchema { arrow_schema: &Schema, base_schema: Option, max_existing_id: Option, + dataset_allocator_is_authoritative: bool, + replaces_all_identities: bool, ) -> Result { let mut schema = Self { fields: arrow_schema @@ -1030,7 +1066,25 @@ impl SchemaExt for LanceSchema { .collect::>()?, metadata: arrow_schema.metadata.clone(), }; - schema.set_field_id_from_schema(base_schema, max_existing_id)?; + if dataset_allocator_is_authoritative { + if replaces_all_identities || base_schema.is_none() { + schema.try_reassign_field_ids(max_existing_id)?; + } else if let Some(base_schema) = &base_schema { + for field in &mut schema.fields { + let base_field = if field.id >= 0 { + base_schema + .field_by_id(field.id) + .or_else(|| base_schema.field(&field.name)) + } else { + base_schema.field(&field.name) + }; + field.set_stable_field_id_from_field(-1, base_schema, base_field)?; + } + schema.try_set_field_id(max_existing_id)?; + } + } else { + schema.set_field_id_from_schema(base_schema, max_existing_id)?; + } schema.validate()?; schema.verify_primary_key()?; @@ -1039,6 +1093,15 @@ impl SchemaExt for LanceSchema { } trait FieldExt { + fn clear_field_ids(&mut self); + + fn set_stable_field_id_from_field( + &mut self, + parent_id: i32, + base_schema: &LanceSchema, + base_field: Option<&Field>, + ) -> lance_core::Result<()>; + /// Recursively set field ID and parent ID for this field and all its children. fn set_field_id_from_field( &mut self, @@ -1048,6 +1111,51 @@ trait FieldExt { } impl FieldExt for Field { + fn clear_field_ids(&mut self) { + self.id = -1; + self.parent_id = -1; + for child in &mut self.children { + child.clear_field_ids(); + } + } + + fn set_stable_field_id_from_field( + &mut self, + parent_id: i32, + base_schema: &LanceSchema, + base_field: Option<&Field>, + ) -> lance_core::Result<()> { + let Some(base_field) = base_field.filter(|base| base.logical_type == self.logical_type) + else { + self.clear_field_ids(); + self.parent_id = parent_id; + return Ok(()); + }; + + self.id = base_field.id; + self.parent_id = parent_id; + for child in &mut self.children { + let base_child = if child.id >= 0 { + base_schema + .field_by_id(child.id) + .filter(|base| base.parent_id == base_field.id) + .or_else(|| { + base_field + .children + .iter() + .find(|base| base.name == child.name) + }) + } else { + base_field + .children + .iter() + .find(|base| base.name == child.name) + }; + child.set_stable_field_id_from_field(self.id, base_schema, base_child)?; + } + Ok(()) + } + fn set_field_id_from_field( &mut self, parent_id: i32, @@ -1092,10 +1200,8 @@ fn convert_to_rust_operation( ) -> Result { let op_name = env.get_string_from_method(java_operation, "name")?; let op = match op_name.as_str() { - "Project" => Operation::Project { - preserves_nullability: env - .get_boolean_from_method(java_operation, "preservesNullability")?, - schema: convert_schema_from_operation( + "Project" => { + let (schema, _) = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1105,8 +1211,14 @@ fn convert_to_rust_operation( })?, dataset, read_version, - )?, - }, + false, + )?; + Operation::Project { + preserves_nullability: env + .get_boolean_from_method(java_operation, "preservesNullability")?, + schema, + } + } "UpdateConfig" => { let config_updates_obj = env .call_method( @@ -1213,7 +1325,7 @@ fn convert_to_rust_operation( } } "Overwrite" => { - let fragments: Vec = + let mut fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; @@ -1226,10 +1338,7 @@ fn convert_to_rust_operation( to_rust_map(env, &config_upsert_values) }, )?; - // Pass None for dataset so that the new schema is not validated - // against the old schema. Overwrite replaces the entire dataset, - // so fields with the same name but different types are allowed. - let schema = convert_schema_from_operation( + let (schema, field_id_remap) = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1237,9 +1346,11 @@ fn convert_to_rust_operation( "BufferAllocator is required for Overwrite operations".to_string(), ) })?, - None, + dataset, read_version, + true, )?; + remap_fragment_field_ids(&mut fragments, &field_id_remap); Operation::Overwrite { fragments, schema, @@ -1381,25 +1492,28 @@ fn convert_to_rust_operation( Operation::DataReplacement { replacements } } "Merge" => { - let fragments: Vec = + let mut fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; + let (schema, field_id_remap) = convert_schema_from_operation( + env, + java_operation, + allocator.ok_or_else(|| { + Error::input_error( + "BufferAllocator is required for Merge operations".to_string(), + ) + })?, + dataset, + read_version, + false, + )?; + remap_fragment_field_ids(&mut fragments, &field_id_remap); Operation::Merge { fragments, preserves_nullability: env .get_boolean_from_method(java_operation, "preservesNullability")?, - schema: convert_schema_from_operation( - env, - java_operation, - allocator.ok_or_else(|| { - Error::input_error( - "BufferAllocator is required for Merge operations".to_string(), - ) - })?, - dataset, - read_version, - )?, + schema, } } "Restore" => { @@ -1716,6 +1830,68 @@ mod tests { pub const LANCE_FIELD_ID_KEY: &str = "lance:field_id"; + #[test] + fn java_schema_conversion_rejects_field_id_exhaustion() { + let arrow_schema = ArrowSchema::new(vec![ArrowField::new( + "new_field", + ArrowDataType::Int32, + false, + )]); + + let err = LanceSchema::from_arrow_schema(&arrow_schema, None, Some(i32::MAX), true, true) + .unwrap_err(); + + assert!(err.to_string().contains("IDs are exhausted"), "{err}"); + } + + #[test] + fn legacy_java_schema_conversion_preserves_arrow_field_ids() { + let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); + base.id = 0; + let base_schema = LanceSchema { + fields: vec![base], + metadata: HashMap::new(), + }; + let arrow_schema = ArrowSchema::new(vec![ + ArrowField::new("a", ArrowDataType::Int32, false).with_metadata(HashMap::from([( + LANCE_FIELD_ID_KEY.to_string(), + "5".to_string(), + )])), + ArrowField::new("b", ArrowDataType::Int32, false).with_metadata(HashMap::from([( + LANCE_FIELD_ID_KEY.to_string(), + "9".to_string(), + )])), + ]); + + let schema = + LanceSchema::from_arrow_schema(&arrow_schema, Some(base_schema), Some(0), false, false) + .unwrap(); + + assert_eq!(schema.field("a").unwrap().id, 5); + assert_eq!(schema.field("b").unwrap().id, 9); + } + + #[test] + fn stable_java_schema_conversion_preserves_rename_identity() { + let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); + base.id = 10; + let base_schema = LanceSchema { + fields: vec![base], + metadata: HashMap::new(), + }; + let arrow_schema = ArrowSchema::new(vec![ + ArrowField::new("renamed", ArrowDataType::Int32, false).with_metadata(HashMap::from([ + (LANCE_FIELD_ID_KEY.to_string(), "10".to_string()), + ])), + ]); + + let schema = + LanceSchema::from_arrow_schema(&arrow_schema, Some(base_schema), Some(10), true, false) + .unwrap(); + + assert_eq!(schema.field("renamed").unwrap().id, 10); + } + #[test] fn test_create_schema_from_arrow() { // base_schema has an existing field id @@ -1802,22 +1978,19 @@ mod tests { metadata: HashMap::from([("base_schema_k".to_string(), "base_schema_v".to_string())]), }; - // new_schema specifies: - // - field a: manual field id + // new_schema specifies field IDs in Arrow metadata to prove they are + // ignored for an existing dataset: + // - field a: inherits from base_schema // - field b: no id -> should inherit from base_schema // - field c: new field -> should be assigned based on max_field_id - // - struct s: parent+child(x) manual, child(y) inherit, child(z) max_field_id - // - list l: parent manual, child(item) inherit - // - list l2: parent manual, child(item) max_field_id - // - map m: parent manual, child(entries/key/value) inherit - // - map m2: parent manual, child(entries/key/value) max_field_id + // - existing nested fields inherit; new nested fields use the allocator let mut a_meta = HashMap::new(); a_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "5".to_string()); let arrow_a = ArrowField::new("a", ArrowDataType::Int32, false).with_metadata(a_meta); let arrow_b = ArrowField::new("b", ArrowDataType::Int32, false); let arrow_c = ArrowField::new("c", ArrowDataType::Int32, false); - // struct s: manual parent + manual child x + // struct s: metadata IDs on the parent and child x are ignored. let mut s_meta = HashMap::new(); s_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "50".to_string()); let mut x_meta = HashMap::new(); @@ -1833,7 +2006,7 @@ mod tests { ) .with_metadata(s_meta); - // list l: parent manual, item inherit + // list l: the metadata parent ID is ignored; both identities inherit. let mut l_meta = HashMap::new(); l_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "60".to_string()); let arrow_l = ArrowField::new( @@ -1847,7 +2020,7 @@ mod tests { ) .with_metadata(l_meta); - // list l2: parent manual, item max_field_id (no base match) + // list l2: the metadata parent ID is ignored; both identities are new. let mut l2_meta = HashMap::new(); l2_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "61".to_string()); let arrow_l2 = ArrowField::new( @@ -1861,7 +2034,7 @@ mod tests { ) .with_metadata(l2_meta); - // map m: parent manual, entries/key/value inherit + // map m: the metadata parent ID is ignored; all identities inherit. let map_entries = ArrowField::new( "entries", ArrowDataType::Struct(ArrowFields::from(vec![ @@ -1875,7 +2048,7 @@ mod tests { let arrow_m = ArrowField::new("m", ArrowDataType::Map(Arc::new(map_entries), false), true) .with_metadata(m_meta); - // map m2: parent manual, entries/key/value max_field_id (no base match) + // map m2: metadata IDs are ignored; the full new subtree is allocated. let map_entries = ArrowField::new( "entries", ArrowDataType::Struct(ArrowFields::from(vec![ @@ -1897,12 +2070,18 @@ mod tests { HashMap::from([("new_schema_k".to_string(), "new_schema_v".to_string())]), ); - let schema = - LanceSchema::from_arrow_schema(&arrow_schema, Some(base_schema), Some(100)).unwrap(); + let schema = LanceSchema::from_arrow_schema( + &arrow_schema, + Some(base_schema), + Some(100), + true, + false, + ) + .unwrap(); - // 1. Manually specified field id + // 1. Arrow metadata cannot override an existing identity. let got_a = schema.field("a").unwrap(); - assert_eq!(got_a.id, 5); + assert_eq!(got_a.id, 10); assert!(!got_a.metadata.contains_key(LANCE_FIELD_ID_KEY)); // 2. Inherit field id + metadata from base_schema (field b) @@ -1913,31 +2092,31 @@ mod tests { let got_c = schema.field("c").unwrap(); assert_eq!(got_c.id, 101); - // 4. struct: parent+child(x) manual, child(y) inherit, child(z) max_field_id + // 4. struct: existing parent and children inherit; z is allocated. let got_s = schema.field("s").unwrap(); - assert_eq!(got_s.id, 50); + assert_eq!(got_s.id, 20); let got_sx = schema.field("s.x").unwrap(); - assert_eq!(got_sx.id, 51); + assert_eq!(got_sx.id, 21); let got_sy = schema.field("s.y").unwrap(); assert_eq!(got_sy.id, 22); let got_sz = schema.field("s.z").unwrap(); assert_eq!(got_sz.id, 102); - // 5. list l: parent manual, item inherit + // 5. list l: parent and item inherit. let got_l = schema.field("l").unwrap(); - assert_eq!(got_l.id, 60); + assert_eq!(got_l.id, 30); let got_li = schema.field("l.item").unwrap(); assert_eq!(got_li.id, 31); - // 6. list l2: parent manual, item max_field_id + // 6. list l2: parent and item are newly allocated. let got_l2 = schema.field("l2").unwrap(); - assert_eq!(got_l2.id, 61); + assert_eq!(got_l2.id, 103); let got_l2i = schema.field("l2.item").unwrap(); - assert_eq!(got_l2i.id, 103); + assert_eq!(got_l2i.id, 104); - // 7. map m: parent manual, entries/key/value inherit + // 7. map m: parent, entries, key, and value inherit. let got_m = schema.field("m").unwrap(); - assert_eq!(got_m.id, 70); + assert_eq!(got_m.id, 40); let got_me = schema.field("m.entries").unwrap(); assert_eq!(got_me.id, 41); let got_mk = schema.field("m.entries.key").unwrap(); @@ -1945,15 +2124,15 @@ mod tests { let got_mv = schema.field("m.entries.value").unwrap(); assert_eq!(got_mv.id, 43); - // 8. map m2: parent manual, entries/key/value max_field_id + // 8. map m2: the full new subtree is allocated in pre-order. let got_m2 = schema.field("m2").unwrap(); - assert_eq!(got_m2.id, 71); + assert_eq!(got_m2.id, 105); let got_m2e = schema.field("m2.entries").unwrap(); - assert_eq!(got_m2e.id, 104); + assert_eq!(got_m2e.id, 106); let got_m2k = schema.field("m2.entries.key").unwrap(); - assert_eq!(got_m2k.id, 105); + assert_eq!(got_m2k.id, 107); let got_m2v = schema.field("m2.entries.value").unwrap(); - assert_eq!(got_m2v.id, 106); + assert_eq!(got_m2v.id, 108); // 9. Schema metadata: when new_schema.metadata is non-empty, use new_schema metadata assert_eq!( diff --git a/java/src/main/java/org/lance/operation/SchemaOperation.java b/java/src/main/java/org/lance/operation/SchemaOperation.java index 509492c852d..2fbcc289001 100644 --- a/java/src/main/java/org/lance/operation/SchemaOperation.java +++ b/java/src/main/java/org/lance/operation/SchemaOperation.java @@ -23,13 +23,18 @@ /** * Schema related base operation. * - *

Each field will be assigned a field id when transaction commits, in the following order: + *

For legacy datasets, each field is assigned a field id in the following order: * *

    *
  1. Parse from field metadata with key {@code lance:field_id}. *
  2. Otherwise, set field id from txn read version dataset's schema field (with the same name). *
  3. Otherwise, allocate based on the max field id of the dataset. *
+ * + *

New datasets use stable field IDs. On those datasets, metadata may identify an existing field + * for rename or reorder, but it cannot allocate a new identity. New IDs are assigned by the dataset + * allocator, and field mappings in fragments committed by the same operation are remapped to the + * canonical IDs. */ public abstract class SchemaOperation implements Operation { private final Schema schema; diff --git a/java/src/test/java/org/lance/operation/MergeTest.java b/java/src/test/java/org/lance/operation/MergeTest.java index 1e037285217..9f69eeaa738 100644 --- a/java/src/test/java/org/lance/operation/MergeTest.java +++ b/java/src/test/java/org/lance/operation/MergeTest.java @@ -192,6 +192,10 @@ void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Ex int addressFieldId = maxFieldId + 20; int cityFieldId = maxFieldId + 30; int countryFieldId = maxFieldId + 40; + int canonicalAgeFieldId = maxFieldId + 1; + int canonicalAddressFieldId = maxFieldId + 2; + int canonicalCityFieldId = maxFieldId + 3; + int canonicalCountryFieldId = maxFieldId + 4; Field idField = new Field( @@ -301,18 +305,18 @@ void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Ex // Verify field id. LanceField evolvedAgeField = findField(evolvedDataset.getLanceSchema().fields(), "age"); - Assertions.assertEquals(ageFieldId, evolvedAgeField.getId()); + Assertions.assertEquals(canonicalAgeFieldId, evolvedAgeField.getId()); LanceField evolvedAddressField = findField(evolvedDataset.getLanceSchema().fields(), "address"); - Assertions.assertEquals(addressFieldId, evolvedAddressField.getId()); + Assertions.assertEquals(canonicalAddressFieldId, evolvedAddressField.getId()); LanceField evolvedCityField = findField(evolvedAddressField.getChildren(), "city"); - Assertions.assertEquals(cityFieldId, evolvedCityField.getId()); + Assertions.assertEquals(canonicalCityFieldId, evolvedCityField.getId()); LanceField evolvedCountryField = findField(evolvedAddressField.getChildren(), "country"); - Assertions.assertEquals(countryFieldId, evolvedCountryField.getId()); + Assertions.assertEquals(canonicalCountryFieldId, evolvedCountryField.getId()); // Verify merged data try (LanceScanner scanner = evolvedDataset.newScan()) { diff --git a/protos/table.proto b/protos/table.proto index 0f8a17d0ff6..9df1e6f0850 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -125,6 +125,9 @@ message Manifest { // MemWAL SSTables still hold. Writers must refuse it too: one that does not // maintain index_catchup can change an index without withdrawing the // position recorded for it. Setting it is one-way. + // * 1 << 8: field IDs use a persistent allocation high-water mark and are + // never reused. This bit is set for readers only by fail-closed activation; + // writers always require the corresponding writer bit. uint64 reader_feature_flags = 9; // Feature flags for writers. @@ -217,6 +220,14 @@ message Manifest { // The branch of the dataset. None means main branch. optional string branch = 20; + + // The highest field ID allocated since stable field identity was activated. + // + // Presence activates the stable-field-ID contract. The value may name a field + // that is absent from the current schema and files because dropped IDs are + // permanently retired. Absence means this is a legacy dataset whose field-ID + // allocator is still derived from currently referenced fields. + optional int32 max_allocated_field_id = 22; } // Manifest // external dataset base path diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index ae49b8cab3d..9cddbf83dac 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5578,6 +5578,19 @@ def test_schema_project_drop_column(tmp_path: Path): assert tbl == expected +def test_schema_project_raw_arrow_cannot_allocate_field_id(tmp_path: Path): + dataset = lance.write_dataset( + pa.table({"a": pa.array([1, 2], pa.int32())}), tmp_path + ) + schema = pa.schema([pa.field("a", pa.int32(), metadata={b"lance:field_id": b"42"})]) + + project = lance.LanceOperation.Project(schema) + dataset = lance.LanceDataset.commit(dataset, project, read_version=dataset.version) + + assert dataset.lance_schema.fields()[0].id() == 0 + assert dataset.to_table() == pa.table({"a": pa.array([1, 2], pa.int32())}) + + def test_schema_project_rename_column(tmp_path: Path): table = pa.Table.from_pydict({"a": range(100, 200), "b": range(300, 400)}) base_dir = tmp_path / "test" diff --git a/python/python/tests/test_schema.py b/python/python/tests/test_schema.py index c384466082f..34022461ff6 100644 --- a/python/python/tests/test_schema.py +++ b/python/python/tests/test_schema.py @@ -73,3 +73,13 @@ def test_lance_schema_from_protos_rejects_missing_parent(): match="Field 'child' \\(id=7\\) references parent id 42", ): LanceSchema._from_protos("{}", field_proto) + + +def test_lance_schema_from_pyarrow_ignores_field_id_metadata(): + arrow_schema = pa.schema( + [pa.field("x", pa.int32(), metadata={b"lance:field_id": b"42"})] + ) + + schema = LanceSchema.from_pyarrow(arrow_schema) + + assert schema.fields()[0].id() == 0 diff --git a/python/src/schema.rs b/python/src/schema.rs index 8cdc2115cd1..926dc039ef9 100644 --- a/python/src/schema.rs +++ b/python/src/schema.rs @@ -129,12 +129,16 @@ impl LanceSchema { /// Create a Lance schema from a PyArrow schema. /// - /// This will assign field ids in depth-first order. Be aware this may not - /// match the correct schema for a particular table. + /// This assigns field ids in depth-first order. Arrow field-ID metadata is + /// descriptive only and cannot select identities for a dataset. Be aware + /// this assignment may not match the correct schema for a particular table. #[staticmethod] pub fn from_pyarrow(schema: PyArrowType) -> PyResult { - let schema = Schema::try_from(&schema.0) + let mut schema = Schema::try_from(&schema.0) .map_err(|err| PyValueError::new_err(format!("Failed to convert schema: {}", err)))?; + schema + .try_reassign_field_ids(None) + .map_err(|err| PyValueError::new_err(format!("Failed to assign field ids: {err}")))?; Ok(Self(schema)) } diff --git a/python/src/transaction.rs b/python/src/transaction.rs index bf70cf3dd6d..ae48fbc7dfd 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -996,11 +996,17 @@ fn extract_schema(schema: &Bound<'_, PyAny>) -> PyResult { } fn convert_schema(arrow_schema: &ArrowSchema) -> PyResult { - // Note: the field ids here are wrong. - Schema::try_from(arrow_schema).map_err(|e| { + let mut schema = Schema::try_from(arrow_schema).map_err(|e| { PyValueError::new_err(format!( "Failed to convert Arrow schema to Lance schema: {}", e )) - }) + })?; + // A raw Arrow schema is not an allocation authority. Assign a canonical + // standalone sequence; dataset commits will either match existing + // identities or reject it in favor of a dataset-derived LanceSchema. + schema + .try_reassign_field_ids(None) + .map_err(|e| PyValueError::new_err(format!("Failed to assign field ids: {e}")))?; + Ok(schema) } diff --git a/rust/lance-core/src/datatypes/field.rs b/rust/lance-core/src/datatypes/field.rs index d5eb89dccb0..138db5d8174 100644 --- a/rust/lance-core/src/datatypes/field.rs +++ b/rust/lance-core/src/datatypes/field.rs @@ -1029,6 +1029,27 @@ impl Field { .for_each(|f| f.set_id(self.id, id_seed)); } + /// Recursively assign missing field IDs without overflowing the `i32` ID space. + /// + /// Callers that persist the result should prefer this method over [`Self::set_id`]. + /// The `i64` seed represents the next candidate ID so `i32::MAX + 1` can be + /// reported as an error instead of wrapping. + pub fn try_set_id(&mut self, parent_id: i32, id_seed: &mut i64) -> Result<()> { + self.parent_id = parent_id; + if self.id < 0 { + self.id = i32::try_from(*id_seed).map_err(|_| { + Error::invalid_input( + "No further field ID can be allocated because IDs are exhausted", + ) + })?; + *id_seed += 1; + } + for child in &mut self.children { + child.try_set_id(self.id, id_seed)?; + } + Ok(()) + } + /// Recursively reset field ID for this field and all its children. pub(super) fn reset_id(&mut self) { self.id = -1; diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index 2e328c88d24..883a2ce5f07 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -693,6 +693,42 @@ impl Schema { .for_each(|f| f.set_id(-1, &mut current_id)); } + /// Assign IDs to every unassigned field using checked arithmetic. + /// + /// Existing IDs are preserved. New IDs start after both this schema's + /// maximum ID and `max_existing_id`. + pub fn try_set_field_id(&mut self, max_existing_id: Option) -> Result<()> { + let schema_max_id = self.max_field_id().unwrap_or(-1); + let max_existing_id = max_existing_id.unwrap_or(-1); + let mut current_id = i64::from(schema_max_id.max(max_existing_id)) + 1; + let unassigned_count = self.fields_pre_order().filter(|field| field.id < 0).count() as i64; + if unassigned_count > 0 && current_id + unassigned_count - 1 > i64::from(i32::MAX) { + return Err(Error::invalid_input( + "No further field ID can be allocated because IDs are exhausted", + )); + } + for field in &mut self.fields { + field.try_set_id(-1, &mut current_id)?; + } + Ok(()) + } + + /// Replace every field ID with a fresh checked allocation. + /// + /// The first assigned ID is one greater than `max_existing_id`. This is used + /// when an operation replaces the complete schema identity, such as overwrite. + pub fn try_reassign_field_ids(&mut self, max_existing_id: Option) -> Result<()> { + let field_count = self.fields_pre_order().count() as i64; + let first_id = i64::from(max_existing_id.unwrap_or(-1)) + 1; + if field_count > 0 && first_id + field_count - 1 > i64::from(i32::MAX) { + return Err(Error::invalid_input( + "No further field ID can be allocated because IDs are exhausted", + )); + } + self.reset_id(); + self.try_set_field_id(max_existing_id) + } + fn reset_id(&mut self) { self.fields.iter_mut().for_each(|f| f.reset_id()); } @@ -888,7 +924,7 @@ impl TryFrom<&ArrowSchema> for Schema { .collect::>()?, metadata: schema.metadata.clone(), }; - schema.set_field_id(None); + schema.try_set_field_id(None)?; schema.validate()?; schema.verify_primary_key()?; @@ -1748,11 +1784,45 @@ pub fn escape_field_path_for_project(name: &str) -> String { #[cfg(test)] mod tests { + use crate::datatypes::field::LANCE_FIELD_ID_KEY; use arrow_schema::{DataType as ArrowDataType, Fields as ArrowFields}; use std::{collections::HashMap, sync::Arc}; use super::*; + #[test] + fn checked_field_id_allocation_is_atomic_on_exhaustion() { + let mut schema = Schema { + fields: vec![ + Field::try_from(&ArrowField::new("a", ArrowDataType::Int32, false)).unwrap(), + Field::try_from(&ArrowField::new("b", ArrowDataType::Int32, false)).unwrap(), + ], + metadata: HashMap::new(), + }; + + let err = schema.try_set_field_id(Some(i32::MAX - 1)).unwrap_err(); + + assert!(err.to_string().contains("IDs are exhausted"), "{err}"); + assert!(schema.fields.iter().all(|field| field.id == -1)); + } + + #[test] + fn checked_field_id_reassignment_is_atomic_on_exhaustion() { + let arrow_schema = ArrowSchema::new(vec![ + ArrowField::new("a", ArrowDataType::Int32, false), + ArrowField::new("b", ArrowDataType::Int32, false), + ]); + let mut schema = Schema::try_from(&arrow_schema).unwrap(); + let original = schema.clone(); + + let err = schema + .try_reassign_field_ids(Some(i32::MAX - 1)) + .unwrap_err(); + + assert!(err.to_string().contains("IDs are exhausted"), "{err}"); + assert_eq!(schema, original); + } + #[test] fn test_resolve_with_quoted_fields() { // Create a schema with fields containing dots @@ -2399,8 +2469,14 @@ mod tests { assert_eq!(schema.max_field_id(), Some(5)); let to_merged_arrow_schema = ArrowSchema::new(vec![ - ArrowField::new("d", DataType::Int32, false), - ArrowField::new("e", DataType::Binary, false), + ArrowField::new("d", DataType::Int32, false).with_metadata(HashMap::from([( + LANCE_FIELD_ID_KEY.to_string(), + "100".to_string(), + )])), + ArrowField::new("e", DataType::Binary, false).with_metadata(HashMap::from([( + LANCE_FIELD_ID_KEY.to_string(), + "101".to_string(), + )])), ]); let mut merged = schema.merge(&to_merged_arrow_schema).unwrap(); merged.set_field_id(None); diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 9c4d626829d..84d35911974 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -39,12 +39,18 @@ pub const FLAG_UNSTABLE_DATA_OVERLAY_FILES: u64 = 64; /// invalidating the catch-up position recorded for that index, leaving a stale /// position behind. Both must refuse the table. pub const FLAG_MEM_WAL_INDEX_CATCHUP: u64 = 128; +/// Field IDs are allocated from a persistent high-water mark and are never reused. +/// +/// This is always a writer requirement. It is also a reader requirement when an +/// activation chooses fail-closed compatibility with pre-gate binaries. +pub const FLAG_STABLE_FIELD_IDS: u64 = 256; /// The first bit that is unknown as a feature flag -pub const FLAG_UNKNOWN: u64 = 256; +pub const FLAG_UNKNOWN: u64 = 512; // This build only understands flags below the unknown boundary, so a bit // allocated at or above it would be refused by the very readers meant to use it. const _: () = assert!(FLAG_MEM_WAL_INDEX_CATCHUP < FLAG_UNKNOWN); +const _: () = assert!(FLAG_STABLE_FIELD_IDS < FLAG_UNKNOWN); /// Environment variable that opts a release build into reading and writing data /// overlay files before the feature is generally released. @@ -75,6 +81,13 @@ pub fn apply_feature_flags( } else { 0 }; + let stable_field_ids_reader = if manifest.max_allocated_field_id.is_some() + && manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0 + { + FLAG_STABLE_FIELD_IDS + } else { + 0 + }; // Reset flags manifest.reader_feature_flags = 0; @@ -134,6 +147,11 @@ pub fn apply_feature_flags( manifest.writer_feature_flags |= FLAG_DISABLE_TRANSACTION_FILE; } + if manifest.max_allocated_field_id.is_some() { + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + manifest.reader_feature_flags |= stable_field_ids_reader; + } + manifest.reader_feature_flags |= mem_wal_index_catchup; manifest.writer_feature_flags |= mem_wal_index_catchup; @@ -143,11 +161,10 @@ pub fn apply_feature_flags( /// Carry [`FLAG_MEM_WAL_INDEX_CATCHUP`] from the manifest a new one is derived /// from. /// -/// [`apply_feature_flags`] carries this bit across its own reset, but it only -/// ever sees one manifest. It cannot help where a *new* manifest is derived from -/// an existing one -- `Manifest::new_from_previous` and `shallow_clone` both -/// zero the feature words -- because the destination starts with nothing to -/// carry. That transition is this function's job. +/// [`apply_feature_flags`] carries this bit across its own reset, while this +/// helper validates and explicitly carries the transition between source and +/// destination manifests. Constructors preserve the bit as defense in depth, +/// but callers must not depend on a constructor choice for protocol safety. /// /// A half-set state is refused rather than normalized: one bit set means a /// legacy reader or a legacy writer is still permitted, which is neither mode. @@ -230,6 +247,24 @@ pub fn validate_mem_wal_index_catchup_flags(manifest: &Manifest) -> Result<()> { Ok(()) } +/// Refuse a manifest whose stable-field-ID marker and required flags disagree. +/// +/// The high-water mark is the activation marker and the writer bit is what +/// keeps pre-feature writers away. Either one without the other is an unsafe, +/// undefined state. The reader bit is optional migration policy, but cannot be +/// set on a legacy manifest. +pub fn validate_stable_field_id_flags(manifest: &Manifest) -> Result<()> { + let activated = manifest.max_allocated_field_id.is_some(); + let reader = manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0; + let writer = manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS != 0; + if activated != writer || (reader && !activated) { + return Err(Error::invalid_input( + "Manifest stable-field-ID high-water mark and reader/writer feature flags disagree", + )); + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -244,6 +279,7 @@ mod tests { assert!(can_read_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_read_dataset(super::FLAG_BASE_PATHS)); assert!(can_read_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + assert!(can_read_dataset(super::FLAG_STABLE_FIELD_IDS)); // Overlay support is gated on the build profile / env opt-in, so the // flag is readable exactly when overlays are enabled (see // test_data_overlay_flag_release_gating for the full policy). @@ -320,6 +356,7 @@ mod tests { assert!(can_write_dataset(super::FLAG_TABLE_CONFIG)); assert!(can_write_dataset(super::FLAG_BASE_PATHS)); assert!(can_write_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); + assert!(can_write_dataset(super::FLAG_STABLE_FIELD_IDS)); // Overlay support is gated on the build profile / env opt-in, so the // flag is writable exactly when overlays are enabled (see // test_data_overlay_flag_release_gating for the full policy). @@ -399,7 +436,7 @@ mod tests { let mut source = empty_manifest(); source.reader_feature_flags = FLAG_MEM_WAL_INDEX_CATCHUP; source.writer_feature_flags = FLAG_MEM_WAL_INDEX_CATCHUP; - // What `Manifest::new_from_previous` hands us: both words zeroed. + // Model a destination assembled without protocol-specific inheritance. let mut destination = empty_manifest(); inherit_mem_wal_index_catchup(&mut destination, &source).unwrap(); @@ -470,6 +507,34 @@ mod tests { ); } + #[test] + fn apply_feature_flags_sets_writer_gate_for_stable_field_ids() { + let mut manifest = empty_manifest(); + manifest.activate_stable_field_ids(); + + apply_feature_flags(&mut manifest, false, false).unwrap(); + + assert_eq!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } + + #[test] + fn apply_feature_flags_preserves_fail_closed_stable_field_id_gate() { + let mut manifest = empty_manifest(); + manifest.activate_stable_field_ids(); + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + + // This runs more than once while committing a manifest. The reader + // requirement is migration policy, not something schema contents can + // derive, so it must survive every recomputation. + apply_feature_flags(&mut manifest, false, false).unwrap(); + apply_feature_flags(&mut manifest, false, false).unwrap(); + + assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } + fn empty_manifest() -> Manifest { use crate::format::DataStorageFormat; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; @@ -496,4 +561,30 @@ mod tests { // silently widen what this build claims to understand. assert!(!can_read_dataset(FLAG_UNKNOWN)); } + + #[test] + fn the_stable_field_id_bit_is_below_the_unknown_boundary() { + assert!(can_read_dataset(FLAG_STABLE_FIELD_IDS)); + assert!(can_write_dataset(FLAG_STABLE_FIELD_IDS)); + assert!(!can_write_dataset(FLAG_UNKNOWN)); + } + + #[test] + fn stable_field_id_marker_and_writer_gate_must_agree() { + let mut activated_without_gate = empty_manifest(); + activated_without_gate.activate_stable_field_ids(); + assert!(validate_stable_field_id_flags(&activated_without_gate).is_err()); + + let mut gate_without_marker = empty_manifest(); + gate_without_marker.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + assert!(validate_stable_field_id_flags(&gate_without_marker).is_err()); + + let mut writer_only = empty_manifest(); + writer_only.activate_stable_field_ids(); + writer_only.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + validate_stable_field_id_flags(&writer_only).unwrap(); + + writer_only.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + validate_stable_field_id_flags(&writer_only).unwrap(); + } } diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index c305d9f4e4a..3d48f57e340 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -18,7 +18,7 @@ use std::ops::Range; use std::sync::Arc; use super::Fragment; -use crate::feature_flags::FLAG_MEM_WAL_INDEX_CATCHUP; +use crate::feature_flags::{FLAG_MEM_WAL_INDEX_CATCHUP, FLAG_STABLE_FIELD_IDS}; use crate::feature_flags::{FLAG_STABLE_ROW_IDS, has_deprecated_v2_feature_flag}; use crate::format::fragment::DataFileFieldInterner; use crate::format::pb; @@ -76,6 +76,10 @@ pub struct Manifest { /// None means never set, Some(0) means max ID used so far is 0 pub max_fragment_id: Option, + /// The highest field ID allocated since stable field identity was activated. + /// `None` means the dataset still uses the legacy live-reference allocator. + pub max_allocated_field_id: Option, + /// The path to the transaction file, relative to the root of the dataset pub transaction_file: Option, @@ -191,6 +195,7 @@ impl Manifest { reader_feature_flags: 0, writer_feature_flags: 0, max_fragment_id: None, + max_allocated_field_id: None, transaction_file: None, transaction_section: None, fragment_offsets, @@ -219,9 +224,12 @@ impl Manifest { index_section: None, // Caller should update index if they want to keep them. timestamp_nanos: 0, // This will be set on commit tag: None, - reader_feature_flags: 0, // These will be set on commit - writer_feature_flags: 0, // These will be set on commit + reader_feature_flags: previous.reader_feature_flags + & (FLAG_MEM_WAL_INDEX_CATCHUP | FLAG_STABLE_FIELD_IDS), + writer_feature_flags: previous.writer_feature_flags + & (FLAG_MEM_WAL_INDEX_CATCHUP | FLAG_STABLE_FIELD_IDS), max_fragment_id: previous.max_fragment_id, + max_allocated_field_id: previous.max_allocated_field_id, transaction_file: None, transaction_section: None, fragment_offsets, @@ -276,12 +284,14 @@ impl Manifest { index_section: None, // These will be set on commit timestamp_nanos: self.timestamp_nanos, tag: None, - // Not derivable from the manifest, so it would be lost like any - // other zeroed word -- and a clone of a table that requires index - // catch-up would silently come back as legacy. - reader_feature_flags: self.reader_feature_flags & FLAG_MEM_WAL_INDEX_CATCHUP, - writer_feature_flags: self.writer_feature_flags & FLAG_MEM_WAL_INDEX_CATCHUP, + // Preserve protocol mode bits that are not derivable from fragments. + // Dropping either would silently downgrade the clone's semantics. + reader_feature_flags: self.reader_feature_flags + & (FLAG_MEM_WAL_INDEX_CATCHUP | FLAG_STABLE_FIELD_IDS), + writer_feature_flags: self.writer_feature_flags + & (FLAG_MEM_WAL_INDEX_CATCHUP | FLAG_STABLE_FIELD_IDS), max_fragment_id: self.max_fragment_id, + max_allocated_field_id: self.max_allocated_field_id, transaction_file: Some(transaction_file), transaction_section: None, fragment_offsets: self.fragment_offsets.clone(), @@ -424,12 +434,17 @@ impl Manifest { } } - /// Get the max used field id + /// Get the highest field ID that may not be allocated again. /// /// This is different than [Schema::max_field_id] because it also considers /// the field ids in the data files that have been dropped from the schema, /// including overlay files referenced by fragments. pub fn max_field_id(&self) -> i32 { + self.max_allocated_field_id + .unwrap_or_else(|| self.max_referenced_field_id()) + } + + pub(crate) fn max_referenced_field_id(&self) -> i32 { let schema_max_id = self.schema.max_field_id().unwrap_or(-1); let fragment_max_id = self .fragments @@ -445,6 +460,26 @@ impl Manifest { schema_max_id.max(fragment_max_id) } + /// Whether the stable-field-ID allocation contract is active. + pub fn uses_stable_field_ids(&self) -> bool { + self.max_allocated_field_id.is_some() + } + + /// Activate stable field IDs at the maximum ID visible in this snapshot. + pub fn activate_stable_field_ids(&mut self) { + if self.max_allocated_field_id.is_none() { + self.max_allocated_field_id = Some(self.max_referenced_field_id()); + } + } + + /// Advance the persistent field-ID high-water mark to cover this manifest. + pub fn update_max_field_id(&mut self) { + let max_referenced_field_id = self.max_referenced_field_id(); + if let Some(max_allocated_field_id) = &mut self.max_allocated_field_id { + *max_allocated_field_id = (*max_allocated_field_id).max(max_referenced_field_id); + } + } + /// Return the fragments that are newer than the given manifest. /// Note this does not support recycling of fragment ids. pub fn fragments_since(&self, since: &Self) -> Result> { @@ -711,6 +746,9 @@ pub struct ManifestBuildConfig { /// It bypasses the "cannot enable stable row ids on existing dataset" guard and /// sets `manifest.next_row_id` to the provided value before activating the flag. pub migration_next_row_id: Option, + /// When `Some`, atomically activates stable field IDs. The boolean controls + /// whether the feature is also required of readers; writers are always gated. + pub stable_field_id_migration_requires_reader: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -988,6 +1026,7 @@ impl TryFrom for Manifest { reader_feature_flags: p.reader_feature_flags, writer_feature_flags: p.writer_feature_flags, max_fragment_id: p.max_fragment_id, + max_allocated_field_id: p.max_allocated_field_id, fragments, transaction_file: if p.transaction_file.is_empty() { None @@ -1050,6 +1089,7 @@ impl From<&Manifest> for pb::Manifest { reader_feature_flags: m.reader_feature_flags, writer_feature_flags: m.writer_feature_flags, max_fragment_id: m.max_fragment_id, + max_allocated_field_id: m.max_allocated_field_id, transaction_file: m.transaction_file.clone().unwrap_or_default(), next_row_id: m.next_row_id, data_format: Some(pb::manifest::DataStorageFormat { @@ -1136,7 +1176,7 @@ impl SelfDescribingFileReader for V1FileReader { #[cfg(test)] mod tests { - use crate::feature_flags::FLAG_USE_V2_FORMAT_DEPRECATED; + use crate::feature_flags::{FLAG_STABLE_FIELD_IDS, FLAG_USE_V2_FORMAT_DEPRECATED}; use crate::format::overlay::{DataOverlayFile, OverlayCoverage}; use crate::format::{DataFile, DeletionFile, DeletionFileType}; use std::num::NonZero; @@ -1600,6 +1640,66 @@ mod tests { assert_eq!(manifest.max_field_id(), 43); } + #[test] + fn stable_field_id_high_water_mark_survives_dropped_references_and_round_trip() { + let arrow_schema = ArrowSchema::new(vec![ArrowField::new( + "a", + arrow_schema::DataType::Int64, + false, + )]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let mut manifest = Manifest::new( + schema, + Arc::new(vec![]), + DataStorageFormat::default(), + HashMap::new(), + ); + manifest.activate_stable_field_ids(); + manifest.max_allocated_field_id = Some(43); + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + + assert_eq!(manifest.max_referenced_field_id(), 0); + assert_eq!(manifest.max_field_id(), 43); + + let recovered = Manifest::try_from(pb::Manifest::from(&manifest)).unwrap(); + assert_eq!(recovered.max_allocated_field_id, Some(43)); + assert_eq!(recovered.max_field_id(), 43); + assert_ne!(recovered.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(recovered.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } + + #[test] + fn shallow_clone_preserves_stable_field_id_allocation_state() { + let arrow_schema = ArrowSchema::new(vec![ArrowField::new( + "a", + arrow_schema::DataType::Int64, + false, + )]); + let schema = Schema::try_from(&arrow_schema).unwrap(); + let mut manifest = Manifest::new( + schema, + Arc::new(vec![]), + DataStorageFormat::default(), + HashMap::new(), + ); + manifest.max_allocated_field_id = Some(41); + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + + let cloned = manifest.shallow_clone( + Some("parent".to_string()), + "memory://parent".to_string(), + 7, + None, + String::new(), + ); + + assert_eq!(cloned.max_allocated_field_id, Some(41)); + assert_ne!(cloned.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(cloned.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } + #[test] fn test_config() { let arrow_schema = ArrowSchema::new(vec![ArrowField::new( diff --git a/rust/lance-table/src/transaction.rs b/rust/lance-table/src/transaction.rs index 74fba1b33ba..e55988d4605 100644 --- a/rust/lance-table/src/transaction.rs +++ b/rust/lance-table/src/transaction.rs @@ -47,7 +47,10 @@ pub use operation::{ pub use update_map::{ UpdateMap, UpdateMapEntry, translate_config_updates, translate_schema_metadata_updates, }; -pub use validate::validate_operation; +pub use validate::{ + validate_detached_stable_field_ids, validate_operation, validate_stable_field_id_manifest, + validate_stable_field_id_transition, +}; use crate::format::{IndexMetadata, Manifest}; use roaring::RoaringBitmap; diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index e47bc139e58..07142822484 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -11,8 +11,9 @@ //! metadata it stamps, the validation that runs before it. use crate::feature_flags::{ - FLAG_MEM_WAL_INDEX_CATCHUP, FLAG_STABLE_ROW_IDS, apply_feature_flags, - inherit_mem_wal_index_catchup, validate_mem_wal_index_catchup_flags, + FLAG_MEM_WAL_INDEX_CATCHUP, FLAG_STABLE_FIELD_IDS, FLAG_STABLE_ROW_IDS, apply_feature_flags, + can_read_dataset, can_write_dataset, inherit_mem_wal_index_catchup, + validate_mem_wal_index_catchup_flags, validate_stable_field_id_flags, }; use crate::format::overlay::TOMBSTONE_FIELD_ID; use crate::format::{ @@ -32,7 +33,9 @@ use crate::system_index::mem_wal::{ use crate::transaction::UpdateMode::{RewriteColumns, RewriteRows}; use crate::transaction::row_version::resolve_update_version_metadata; use crate::transaction::update_map::apply_update_map; -use crate::transaction::validate::merge_fragment_physically_rewritten; +use crate::transaction::validate::{ + merge_fragment_physically_rewritten, validate_stable_field_id_manifest, +}; use crate::transaction::{ CoverageIdentity, DataReplacementGroup, LogicalIndexSegments, Operation, ReadVersionState, RewriteGroup, Transaction, UpdatedFragmentOffsets, @@ -103,16 +106,42 @@ impl Transaction { .resolve_version_location(base_path, version, &object_store.inner) .await?; let mut manifest = read_manifest(object_store, &location.path, location.size).await?; + if !can_read_dataset(manifest.reader_feature_flags) { + return Err(Error::not_supported(format!( + "Restore target cannot be read by this version of Lance. Flags: {}", + manifest.reader_feature_flags + ))); + } + if !can_write_dataset(manifest.writer_feature_flags) { + return Err(Error::not_supported(format!( + "Restore target cannot be written by this version of Lance. Flags: {}", + manifest.writer_feature_flags + ))); + } // Read below the reader validation boundary, so nothing else refuses a // half-set manifest here: the flag reset would quietly drop the lone bit // and republish an undefined state as legacy. validate_mem_wal_index_catchup_flags(&manifest)?; + validate_stable_field_id_flags(&manifest)?; manifest.set_timestamp(config.timestamp_nanos); manifest.transaction_file = Some(tx_path.to_string()); let indices = read_manifest_indexes(object_store, &location, &manifest).await?; manifest.max_fragment_id = manifest .max_fragment_id .max(current_manifest.max_fragment_id); + if current_manifest.uses_stable_field_ids() { + let Some(restored_max_field_id) = manifest.max_allocated_field_id else { + return Err(Error::invalid_input(format!( + "Cannot restore version {version}: stable field IDs were activated after that version" + ))); + }; + manifest.max_allocated_field_id = + Some(restored_max_field_id.max(current_manifest.max_field_id())); + if current_manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0 { + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + } + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + } // A version from before catch-up was required carries MemWAL state this // protocol never validated -- catch-up values activation deliberately // cleared, or compaction progress it deliberately refused to trust. @@ -1352,6 +1381,28 @@ impl Transaction { ) }; + if current_manifest.is_none() { + manifest.activate_stable_field_ids(); + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + } + + if let Some(require_reader) = config.stable_field_id_migration_requires_reader { + if current_manifest + .map(|manifest| manifest.uses_stable_field_ids()) + .unwrap_or(false) + { + return Err(Error::invalid_input( + "Stable field IDs are already active for this dataset", + )); + } + manifest.activate_stable_field_ids(); + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + if require_reader { + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + } + } + manifest.tag.clone_from(&self.tag); if config.auto_set_feature_flags { @@ -1368,9 +1419,8 @@ impl Transaction { config.disable_transaction_file, )?; } - // Carried from the manifest this one is derived from. `new_from_previous` - // zeroes both feature words, so `apply_feature_flags` cannot see the - // previous state and every ordinary commit would otherwise drop the bit. + // Validate and explicitly carry the protocol across the derivation + // boundary. Constructors also preserve it as defense in depth. if let Some(current_manifest) = current_manifest { inherit_mem_wal_index_catchup(&mut manifest, current_manifest)?; } @@ -1416,6 +1466,7 @@ impl Transaction { manifest.set_timestamp(config.timestamp_nanos); manifest.update_max_fragment_id(); + manifest.update_max_field_id(); match &self.operation { Operation::Overwrite { @@ -1596,6 +1647,9 @@ impl Transaction { manifest.next_row_id = next_row_id; } + validate_stable_field_id_flags(&manifest)?; + validate_stable_field_id_manifest(&manifest)?; + Ok((manifest, final_indices)) } @@ -1651,6 +1705,57 @@ mod tests { ) } + #[test] + fn new_dataset_builds_with_stable_field_ids_by_default() { + let arrow_schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let mut schema = LanceSchema::try_from(&arrow_schema).unwrap(); + schema.try_set_field_id(None).unwrap(); + let transaction = Transaction::new( + 0, + Operation::Overwrite { + fragments: vec![], + schema, + config_upsert_values: None, + initial_bases: None, + }, + None, + ); + + let (manifest, _) = transaction + .build_manifest(None, vec![], "txn", &default_build_config()) + .unwrap(); + + assert_eq!(manifest.max_allocated_field_id, Some(0)); + assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } + + #[test] + fn activation_sets_writer_gate_when_auto_flags_are_disabled() { + let manifest = sample_manifest(); + let transaction = Transaction::new( + manifest.version, + Operation::UpdateConfig { + config_updates: None, + table_metadata_updates: None, + schema_metadata_updates: None, + field_metadata_updates: HashMap::new(), + }, + None, + ); + let mut config = default_build_config(); + config.auto_set_feature_flags = false; + config.stable_field_id_migration_requires_reader = Some(false); + + let (activated, _) = transaction + .build_manifest(Some(&manifest), vec![], "txn", &config) + .unwrap(); + + assert!(activated.uses_stable_field_ids()); + assert_eq!(activated.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(activated.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } + #[test] fn test_create_index_build_manifest_keeps_unremoved_same_name_indices() { let manifest = sample_manifest(); diff --git a/rust/lance-table/src/transaction/test_support.rs b/rust/lance-table/src/transaction/test_support.rs index 98806c2cbb2..eded723b462 100644 --- a/rust/lance-table/src/transaction/test_support.rs +++ b/rust/lance-table/src/transaction/test_support.rs @@ -30,6 +30,7 @@ pub fn default_build_config() -> ManifestBuildConfig { storage_format: None, disable_transaction_file: false, migration_next_row_id: None, + stable_field_id_migration_requires_reader: None, } } diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 505fbe89664..afd99e1a1d7 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -40,6 +40,8 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> } }; + validate_stable_field_id_operation(manifest, operation)?; + match operation { Operation::Append { fragments } => { // Fragments must contain all fields in the schema @@ -98,6 +100,165 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> } } +/// Validate stable-field-ID invariants that are independent of one operation. +pub fn validate_stable_field_id_manifest(manifest: &Manifest) -> Result<()> { + let Some(max_allocated_field_id) = manifest.max_allocated_field_id else { + return Ok(()); + }; + let max_referenced_field_id = manifest.max_referenced_field_id(); + if max_allocated_field_id < max_referenced_field_id { + return Err(Error::invalid_input(format!( + "Stable field-ID high-water mark {} is below referenced field ID {}", + max_allocated_field_id, max_referenced_field_id + ))); + } + Ok(()) +} + +/// Validate newly referenced field IDs before advancing the high-water mark. +/// +/// A data-only operation must not manufacture allocator state by putting an +/// otherwise unknown ID in a data-file or overlay mapping. IDs above the parent +/// high-water mark are legal only when the canonical successor schema contains +/// that newly allocated identity. Overwrite has no retained physical state, so +/// every non-negative reference must belong to its replacement schema. +pub fn validate_stable_field_id_transition( + parent: &Manifest, + successor: &Manifest, + operation: &Operation, +) -> Result<()> { + let Some(parent_max_field_id) = parent.max_allocated_field_id else { + return Ok(()); + }; + let Some(successor_max_field_id) = successor.max_allocated_field_id else { + return Err(Error::invalid_input( + "Stable field-ID activation marker is missing from the successor manifest", + )); + }; + if successor_max_field_id < parent_max_field_id { + return Err(Error::invalid_input(format!( + "Stable field-ID high-water mark decreases from {parent_max_field_id} to {successor_max_field_id}" + ))); + } + let successor_schema_ids = successor + .schema + .fields_pre_order() + .map(|field| field.id) + .collect::>(); + + if !matches!(operation, Operation::Restore { .. }) { + validate_dense_new_field_ids( + parent, + successor + .schema + .fields_pre_order() + .filter(|field| parent.schema.field_by_id(field.id).is_none()), + )?; + } + + for field_id in successor + .fragments + .iter() + .flat_map(|fragment| fragment.referenced_lance_files()) + .flat_map(|file| file.fields.iter()) + .copied() + .filter(|field_id| *field_id >= 0) + { + if matches!(operation, Operation::Overwrite { .. }) + && !successor_schema_ids.contains(&field_id) + { + return Err(Error::invalid_input(format!( + "Overwrite references field ID {field_id}, which is not in its replacement schema" + ))); + } + if field_id > parent_max_field_id && !successor_schema_ids.contains(&field_id) { + return Err(Error::invalid_input(format!( + "Data file or overlay references new field ID {field_id}, but the canonical successor schema does not contain that identity" + ))); + } + } + Ok(()) +} + +fn validate_dense_new_field_ids<'a>( + manifest: &Manifest, + new_fields: impl Iterator, +) -> Result<()> { + let expected_ids = i64::from(manifest.max_field_id()) + 1..; + for (expected, field) in expected_ids.zip(new_fields) { + if i64::from(field.id) != expected { + return Err(Error::invalid_input(format!( + "New field '{}' has ID {}, but stable field IDs must be densely allocated from {}", + field.name, field.id, expected + ))); + } + } + Ok(()) +} + +fn validate_stable_field_id_operation(manifest: &Manifest, operation: &Operation) -> Result<()> { + if !manifest.uses_stable_field_ids() { + return Ok(()); + } + validate_stable_field_id_manifest(manifest)?; + + let (schema, replaces_all_identities) = match operation { + Operation::Overwrite { schema, .. } => (schema, true), + Operation::Merge { schema, .. } | Operation::Project { schema, .. } => (schema, false), + _ => return Ok(()), + }; + schema.validate()?; + + if replaces_all_identities { + return validate_dense_new_field_ids(manifest, schema.fields_pre_order()); + } + + for field in schema.fields_pre_order() { + let Some(prior_field) = manifest.schema.field_by_id(field.id) else { + continue; + }; + if field.parent_id != prior_field.parent_id { + return Err(Error::invalid_input(format!( + "Field ID {} moves from parent {} to parent {}; stable field identity cannot move between parents", + field.id, prior_field.parent_id, field.parent_id + ))); + } + if field.logical_type != prior_field.logical_type { + return Err(Error::invalid_input(format!( + "Field ID {} changes logical type from {} to {}; type replacement must allocate a new field ID", + field.id, prior_field.logical_type, field.logical_type + ))); + } + } + + validate_dense_new_field_ids( + manifest, + schema + .fields_pre_order() + .filter(|field| manifest.schema.field_by_id(field.id).is_none()), + ) +} + +/// Reject detached schema changes once stable field identity is active. +pub fn validate_detached_stable_field_ids( + manifest: &Manifest, + operation: &Operation, +) -> Result<()> { + if !manifest.uses_stable_field_ids() { + return Ok(()); + } + match operation { + Operation::Merge { schema, .. } if schema == &manifest.schema => Ok(()), + Operation::Merge { .. } + | Operation::Project { .. } + | Operation::Overwrite { .. } + | Operation::Restore { .. } => Err(Error::invalid_input( + "Detached commits cannot change schema after stable field IDs are activated", + )), + _ => Ok(()), + } +} + // An overwrite's fragments are newly written, so they are given fresh ids at // commit time. A deletion file cannot come along for that ride: its path embeds // the fragment id, so renumbering the fragment would orphan the deletion vector @@ -530,6 +691,167 @@ mod tests { ) } + fn activated_manifest() -> Manifest { + let schema = one_field_schema(); + let mut manifest = manifest_with_file_fields(schema, vec![0]); + manifest.activate_stable_field_ids(); + manifest + } + + #[test] + fn stable_field_ids_require_dense_allocation_above_high_water_mark() { + let mut manifest = activated_manifest(); + manifest.max_allocated_field_id = Some(5); + let mut schema = manifest.schema.clone(); + let mut new_field = + LanceCoreField::try_from(&ArrowField::new("b", DataType::Int32, true)).unwrap(); + new_field.id = 6; + schema.fields.push(new_field); + let valid = Operation::Project { + schema: schema.clone(), + preserves_nullability: true, + }; + validate_operation(Some(&manifest), &valid).unwrap(); + + schema.fields.last_mut().unwrap().id = 7; + let skipped = Operation::Project { + schema, + preserves_nullability: true, + }; + let err = validate_operation(Some(&manifest), &skipped).unwrap_err(); + assert!( + err.to_string().contains("densely allocated from 6"), + "{err}" + ); + } + + #[test] + fn stable_field_ids_require_fresh_identity_for_type_replacement() { + let manifest = activated_manifest(); + let mut schema = manifest.schema.clone(); + schema.fields[0].logical_type = LogicalType::try_from(&DataType::Float32).unwrap(); + let operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + let err = validate_operation(Some(&manifest), &operation).unwrap_err(); + + assert!(err.to_string().contains("type replacement"), "{err}"); + } + + #[test] + fn stable_field_ids_require_overwrite_to_replace_every_identity() { + let manifest = activated_manifest(); + let schema = manifest.schema.clone(); + let operation = Operation::Overwrite { + fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0])], + schema, + config_upsert_values: None, + initial_bases: None, + }; + + let err = validate_operation(Some(&manifest), &operation).unwrap_err(); + + assert!( + err.to_string().contains("densely allocated from 1"), + "{err}" + ); + } + + #[test] + fn stable_field_id_manifest_rejects_high_water_mark_below_overlay_reference() { + let mut manifest = activated_manifest(); + Arc::make_mut(&mut manifest.fragments)[0] + .overlays + .push(DataOverlayFile { + data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![7], None), + coverage: OverlayCoverage::Shared(Arc::new(RoaringBitmap::from_iter([0_u32]))), + committed_version: 1, + }); + + let err = validate_stable_field_id_manifest(&manifest).unwrap_err(); + + assert!( + err.to_string().contains("below referenced field ID 7"), + "{err}" + ); + } + + #[test] + fn stable_field_id_transition_rejects_file_only_allocator_advance() { + let manifest = activated_manifest(); + let mut successor = Manifest::new_from_previous( + &manifest, + manifest.schema.clone(), + Arc::new(vec![fragment_with_file_fields(0, "new.lance", vec![0, 1])]), + ); + successor.max_allocated_field_id = manifest.max_allocated_field_id; + let operation = Operation::UpdateConfig { + config_updates: None, + table_metadata_updates: None, + schema_metadata_updates: None, + field_metadata_updates: HashMap::new(), + }; + + let err = + validate_stable_field_id_transition(&manifest, &successor, &operation).unwrap_err(); + + assert!( + err.to_string().contains("canonical successor schema"), + "{err}" + ); + } + + #[test] + fn stable_field_id_transition_rejects_decreasing_high_water_mark() { + let manifest = activated_manifest(); + let mut successor = Manifest::new_from_previous( + &manifest, + manifest.schema.clone(), + manifest.fragments.clone(), + ); + successor.max_allocated_field_id = Some(manifest.max_field_id() - 1); + let operation = Operation::UpdateConfig { + config_updates: None, + table_metadata_updates: None, + schema_metadata_updates: None, + field_metadata_updates: HashMap::new(), + }; + + let err = + validate_stable_field_id_transition(&manifest, &successor, &operation).unwrap_err(); + + assert!( + err.to_string().contains("high-water mark decreases"), + "{err}" + ); + } + + #[test] + fn detached_stable_field_ids_allow_data_only_merge_and_reject_schema_change() { + let manifest = activated_manifest(); + let data_only = Operation::Merge { + fragments: manifest.fragments.as_ref().clone(), + schema: manifest.schema.clone(), + preserves_nullability: true, + }; + validate_detached_stable_field_ids(&manifest, &data_only).unwrap(); + + let mut changed_schema = manifest.schema.clone(); + changed_schema.fields[0].name = "renamed".to_string(); + let schema_change = Operation::Project { + schema: changed_schema, + preserves_nullability: true, + }; + let err = validate_detached_stable_field_ids(&manifest, &schema_change).unwrap_err(); + assert!( + err.to_string() + .contains("Detached commits cannot change schema"), + "{err}" + ); + } + #[rstest::rstest] #[case::logical_type(DataType::Float32, true)] #[case::nullability(DataType::Int32, false)] diff --git a/rust/lance/src/blob.rs b/rust/lance/src/blob.rs index 6ad2a2ad26c..5c663de6241 100644 --- a/rust/lance/src/blob.rs +++ b/rust/lance/src/blob.rs @@ -1328,6 +1328,37 @@ mod tests { assert!(normalized.fields[1].children[1].id >= 0); } + #[test] + fn blob_runtime_and_descriptor_fields_do_not_enter_logical_field_id_space() { + let mut metadata = HashMap::new(); + metadata.insert(ARROW_EXT_NAME_KEY.to_string(), BLOB_V2_EXT_NAME.to_string()); + let prepared_field = prepared_blob_field_with_metadata("blob", true, metadata); + let prepared = LanceSchema::try_from(&ArrowSchema::new(vec![prepared_field])).unwrap(); + + // Prepared-only children such as `blob_id` and `blob_size` are writer + // representation details. Normalization retains IDs only for the + // persistent logical identities (`blob`, `data`, and `uri`). + let logical = prepared_to_logical_blob_schema(&prepared).unwrap(); + assert_eq!(logical.fields[0].id, 0); + assert_eq!(logical.fields[0].children[0].id, 1); + assert_eq!(logical.fields[0].children[1].id, 2); + assert_eq!(logical.max_field_id(), Some(2)); + + // Descriptor projection creates a file/read-local representation. Its + // children deliberately remain synthetic and cannot advance a manifest + // field-ID high-water mark. + let mut descriptor = logical; + descriptor.fields[0].unloaded_mut(); + assert_eq!(descriptor.fields[0].id, 0); + assert!( + descriptor.fields[0] + .children + .iter() + .all(|child| child.id == -1) + ); + assert_eq!(descriptor.max_field_id(), Some(0)); + } + #[tokio::test] async fn test_sidecar_writers_return_prepared_values() { let temp_dir = TempDir::default(); diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index c3746c7ff08..a8cabeb86ec 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -150,6 +150,7 @@ use lance_index::scalar::lance_format::LanceIndexStore; use lance_namespace::models::{DeclareTableRequest, DescribeTableRequest}; use lance_table::feature_flags::{ apply_feature_flags, can_read_dataset, validate_mem_wal_index_catchup_flags, + validate_stable_field_id_flags, }; use lance_table::io::deletion::{DELETIONS_DIR, relative_deletion_file_path}; use lance_table::rowids::{RowIdSequence, write_row_ids}; @@ -184,6 +185,20 @@ pub const DEFAULT_INDEX_CACHE_SIZE: usize = 6 * 1024 * 1024 * 1024; // smaller. pub const DEFAULT_METADATA_CACHE_SIZE: usize = 1024 * 1024 * 1024; +/// Mixed-version compatibility policy for stable-field-ID activation. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum StableFieldIdMigrationMode { + /// Require both new readers and new writers. This is the safe default when + /// pre-gate binaries may still access the dataset. + #[default] + ReadersAndWriters, + /// Require new writers while allowing legacy readers. + /// + /// Use only after every pre-gate writer has been retired. Older binaries do + /// not universally check writer feature flags on every commit path. + WritersOnly, +} + /// Lance Dataset #[derive(Clone)] pub struct Dataset { @@ -767,6 +782,7 @@ impl Dataset { }?; validate_mem_wal_index_catchup_flags(&manifest)?; + validate_stable_field_id_flags(&manifest)?; if !can_read_dataset(manifest.reader_feature_flags) { let message = format!( @@ -3227,6 +3243,50 @@ impl Dataset { Ok(()) } + /// Activate monotonic, non-reusable field IDs for a legacy dataset. + /// + /// The activation commit records the current maximum referenced field ID as + /// a persistent high-water mark. Later schema changes allocate above it even + /// after fields and their files are dropped. New datasets already use this + /// contract by default. Activation is one-way and idempotent. + /// + /// ``` + /// # use lance::{Dataset, Result}; + /// # use lance::dataset::StableFieldIdMigrationMode; + /// # async fn activate(dataset: &mut Dataset) -> Result<()> { + /// dataset + /// .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + /// .await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn migrate_to_stable_field_ids( + &mut self, + mode: StableFieldIdMigrationMode, + ) -> Result<()> { + if self.manifest.uses_stable_field_ids() { + return Ok(()); + } + + let transaction = Transaction::new( + self.manifest.version, + Operation::Merge { + fragments: self.manifest.fragments.as_ref().clone(), + schema: self.manifest.schema.clone(), + preserves_nullability: true, + }, + None, + ); + let require_reader = matches!(mode, StableFieldIdMigrationMode::ReadersAndWriters); + let new_ds = CommitBuilder::new(Arc::new(self.clone())) + .with_max_retries(0) + .with_stable_field_id_migration_activation(require_reader) + .execute(transaction) + .await?; + *self = new_ds; + Ok(()) + } + /// Shallow clone the target version into a new dataset at target_path. /// 'target_path': the uri string to clone the dataset into. /// 'version': the version cloned from, could be a version number or tag. @@ -3698,7 +3758,7 @@ impl Dataset { // Final schema is union of current schema, plus the RHS schema without // the right_on key. let mut new_schema: Schema = self.schema().merge(joiner.out_schema().as_ref())?; - new_schema.set_field_id(Some(self.manifest.max_field_id())); + new_schema.try_set_field_id(Some(self.manifest.max_field_id()))?; // Write new data file to each fragment. Parallelism is done over columns, // so no parallelism done at this level. @@ -4034,6 +4094,9 @@ pub(crate) struct ManifestWriteConfig { /// It bypasses the "cannot enable stable row ids on existing dataset" guard and /// sets `manifest.next_row_id` to the provided value before activating the flag. migration_next_row_id: Option, // default None + /// When `Some`, this commit activates stable field IDs. `true` also sets + /// the reader feature bit for fail-closed mixed-version safety. + stable_field_id_migration_requires_reader: Option, } impl Default for ManifestWriteConfig { @@ -4046,6 +4109,7 @@ impl Default for ManifestWriteConfig { use_legacy_format: None, storage_format: None, migration_next_row_id: None, + stable_field_id_migration_requires_reader: None, } } } @@ -4074,6 +4138,8 @@ impl ManifestWriteConfig { storage_format: self.storage_format.clone(), disable_transaction_file: self.disable_transaction_file, migration_next_row_id: self.migration_next_row_id, + stable_field_id_migration_requires_reader: self + .stable_field_id_migration_requires_reader, } } } @@ -4090,6 +4156,7 @@ pub(crate) async fn write_manifest_file( naming_scheme: ManifestNamingScheme, transaction: Option, ) -> std::result::Result { + manifest.update_max_field_id(); if config.auto_set_feature_flags { // build_manifest may have already set FLAG_STABLE_ROW_IDS on the manifest. // Preserve it here so this second apply_feature_flags call does not clear it diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 620d32e8694..e97703f3533 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -1918,7 +1918,7 @@ impl FileFragment { // Final schema is union of current schema, plus the RHS schema without // the right_on key. let mut new_schema: Schema = self.schema().merge(joiner.out_schema().as_ref())?; - new_schema.set_field_id(Some(max_field_id)); + new_schema.try_set_field_id(Some(max_field_id))?; let new_fragment = self .clone() diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 3643e8043ec..0c002318f4e 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -383,7 +383,7 @@ pub(super) async fn add_columns_to_fragments( return Err(e); } }; - schema.set_field_id(Some(dataset.manifest.max_field_id())); + schema.try_set_field_id(Some(dataset.manifest.max_field_id()))?; let preserves_nullability = !merge_introduces_required_field(dataset.schema(), &schema); @@ -733,10 +733,9 @@ pub(super) async fn alter_columns( let mut new_schema = dataset.schema().clone(); // Mapping of old to new fields that need to be casted. - let mut cast_fields: Vec<(Field, Field)> = Vec::new(); + let mut cast_sources: Vec = Vec::new(); let mut tightens_nullability = false; - let mut next_field_id = dataset.manifest.max_field_id() + 1; let version = dataset.manifest.data_storage_format.lance_file_format(); for alteration in alterations { @@ -783,12 +782,45 @@ pub(super) async fn alter_columns( field_dest.nullable, ); *field_dest = Field::try_from(&arrow_field)?; - field_dest.set_id(field_src.parent_id, &mut next_field_id); + // Keep the old id temporarily so the replacement can be located + // after every alteration has been applied. Fresh ids are assigned + // below in canonical schema order, independent of request order. + field_dest.id = field_src.id; + field_dest.parent_id = field_src.parent_id; - cast_fields.push((field_src.clone(), field_dest.clone())); + cast_sources.push(field_src.clone()); } } + let mut cast_fields = Vec::with_capacity(cast_sources.len()); + if !cast_sources.is_empty() { + let destination_paths = cast_sources + .iter() + .map(|source| new_schema.field_path(source.id)) + .collect::>>()?; + + for source in &cast_sources { + new_schema + .mut_field_by_id(source.id) + .expect("cast source must still identify its replacement") + .id = -1; + } + new_schema.try_set_field_id(Some(dataset.manifest.max_field_id()))?; + + cast_fields = cast_sources + .into_iter() + .zip(destination_paths) + .map(|(source, path)| { + let destination = new_schema.field(&path).ok_or_else(|| { + Error::internal(format!( + "cast replacement field '{path}' disappeared while assigning field ids" + )) + })?; + Ok((source, destination.clone())) + }) + .collect::>>()?; + } + new_schema.validate()?; // If any column being cast has an attached index, fail fast. Cast operations @@ -3613,6 +3645,7 @@ mod test { }), ) .await?; + assert!(dataset.manifest.uses_stable_field_ids()); assert_eq!(dataset.manifest.max_field_id(), 0); // Test we can add 1 column, drop it, then add another column. Validate @@ -3627,7 +3660,7 @@ mod test { assert_eq!(dataset.manifest.max_field_id(), 1); dataset.drop_columns(&["x"]).await?; - assert_eq!(dataset.manifest.max_field_id(), 0); + assert_eq!(dataset.manifest.max_field_id(), 1); dataset .add_columns( @@ -3636,7 +3669,7 @@ mod test { None, ) .await?; - assert_eq!(dataset.manifest.max_field_id(), 1); + assert_eq!(dataset.manifest.max_field_id(), 2); let data = dataset.scan().try_into_batch().await?; let expected_data = RecordBatch::try_new( @@ -3648,7 +3681,7 @@ mod test { )?; assert_eq!(data, expected_data); dataset.drop_columns(&["y"]).await?; - assert_eq!(dataset.manifest.max_field_id(), 0); + assert_eq!(dataset.manifest.max_field_id(), 2); // Test we can add 2 columns, drop 1, then add another column. Validate // the field ids are as expected. @@ -3662,12 +3695,12 @@ mod test { None, ) .await?; - assert_eq!(dataset.manifest.max_field_id(), 2); + assert_eq!(dataset.manifest.max_field_id(), 4); dataset.drop_columns(&["b"]).await?; // Even though we dropped a column, we still have the fragment with a and // b. So it should still act as if that field id is still in play. - assert_eq!(dataset.manifest.max_field_id(), 2); + assert_eq!(dataset.manifest.max_field_id(), 4); dataset .add_columns( @@ -3676,7 +3709,7 @@ mod test { None, ) .await?; - assert_eq!(dataset.manifest.max_field_id(), 3); + assert_eq!(dataset.manifest.max_field_id(), 5); let data = dataset.scan().try_into_batch().await?; let expected_schema = Arc::new(ArrowSchema::new(vec![ diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 1a080687f32..f10f3a742cb 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1250,7 +1250,7 @@ async fn test_write_manifest( let write_fut = require_send(write_fut); let mut dataset = write_fut.await.unwrap(); - // Check it has no flags + // New datasets enable stable field IDs by default. let manifest = read_manifest( dataset.object_store.as_ref(), &dataset @@ -1272,7 +1272,14 @@ async fn test_write_manifest( manifest.data_storage_format.version.to_manifest_string(), "stable" | "next" )); - assert_eq!(manifest.reader_feature_flags, 0); + assert_eq!( + manifest.reader_feature_flags, + feature_flags::FLAG_STABLE_FIELD_IDS + ); + assert_eq!( + manifest.writer_feature_flags, + feature_flags::FLAG_STABLE_FIELD_IDS + ); // Create one with deletions dataset.delete("i < 10").await.unwrap(); @@ -1293,11 +1300,11 @@ async fn test_write_manifest( .unwrap(); assert_eq!( manifest.writer_feature_flags, - feature_flags::FLAG_DELETION_FILES + feature_flags::FLAG_DELETION_FILES | feature_flags::FLAG_STABLE_FIELD_IDS ); assert_eq!( manifest.reader_feature_flags, - feature_flags::FLAG_DELETION_FILES + feature_flags::FLAG_DELETION_FILES | feature_flags::FLAG_STABLE_FIELD_IDS ); // Write with custom manifest @@ -1318,6 +1325,7 @@ async fn test_write_manifest( storage_format: None, disable_transaction_file: false, migration_next_row_id: None, + stable_field_id_migration_requires_reader: None, }, dataset.manifest_location.naming_scheme, None, @@ -1352,6 +1360,94 @@ async fn test_write_manifest( assert!(matches!(write_result, Err(Error::NotSupported { .. }))); } +#[tokio::test] +async fn test_clone_and_restore_reject_unknown_writer_requirements() { + let source_uri = TempStrDir::default(); + let shallow_clone_uri = TempStrDir::default(); + let deep_clone_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &source_uri, + None, + ) + .await + .unwrap(); + let location = dataset + .commit_handler + .resolve_latest_location(&dataset.base, dataset.object_store.as_ref()) + .await + .unwrap(); + let mut unknown_writer = + read_manifest(dataset.object_store.as_ref(), &location.path, location.size) + .await + .unwrap(); + unknown_writer.version += 1; + unknown_writer.writer_feature_flags |= feature_flags::FLAG_UNKNOWN; + let config = ManifestWriteConfig { + auto_set_feature_flags: false, + timestamp: None, + use_stable_row_ids: false, + use_legacy_format: None, + storage_format: None, + disable_transaction_file: false, + migration_next_row_id: None, + stable_field_id_migration_requires_reader: None, + }; + write_manifest_file( + dataset.object_store.as_ref(), + dataset.commit_handler.as_ref(), + &dataset.base, + &mut unknown_writer, + None, + &config, + dataset.manifest_location.naming_scheme, + None, + ) + .await + .unwrap(); + + let mut source = Dataset::open(&source_uri).await.unwrap(); + let err = source + .shallow_clone(shallow_clone_uri.as_str(), source.version().version, None) + .await + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "{err}"); + let err = source + .deep_clone(deep_clone_uri.as_str(), source.version().version, None) + .await + .unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "{err}"); + + let unknown_version = unknown_writer.version; + let mut supported_head = unknown_writer.clone(); + supported_head.version += 1; + supported_head.writer_feature_flags &= !feature_flags::FLAG_UNKNOWN; + write_manifest_file( + dataset.object_store.as_ref(), + dataset.commit_handler.as_ref(), + &dataset.base, + &mut supported_head, + None, + &config, + dataset.manifest_location.naming_scheme, + None, + ) + .await + .unwrap(); + + let latest = Dataset::open(&source_uri).await.unwrap(); + let mut target = latest.checkout_version(unknown_version).await.unwrap(); + let err = target.restore().await.unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "{err}"); +} + #[tokio::test] async fn test_rle_v2_v23_write_and_append() { let test_uri = TempStrDir::default(); diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index a9f58eb7c64..7a372ee7a80 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -4,19 +4,22 @@ use std::sync::Arc; use std::vec; -use crate::dataset::InsertBuilder; use crate::dataset::optimize::{CompactionOptions, compact_files}; +use crate::dataset::{ + ColumnAlteration, InsertBuilder, NewColumnTransform, StableFieldIdMigrationMode, +}; use crate::index::DatasetIndexExt; use crate::utils::test::copy_test_data_to_tmp; use crate::{Dataset, Result}; +use lance_core::utils::tempfile::TempStrDir; use lance_index::{IndexType, scalar::ScalarIndexParams}; -use lance_table::feature_flags::FLAG_STABLE_ROW_IDS; +use lance_table::feature_flags::{FLAG_STABLE_FIELD_IDS, FLAG_STABLE_ROW_IDS}; use lance_table::format::IndexMetadata; use crate::dataset::write::{WriteMode, WriteParams}; use arrow::compute::concat_batches; use arrow_array::RecordBatch; -use arrow_array::{Float32Array, Int64Array, RecordBatchIterator}; +use arrow_array::{Float32Array, Int32Array, Int64Array, RecordBatchIterator}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use lance_file::version::LanceFileVersion; @@ -578,6 +581,202 @@ async fn make_simple_dataset(uri: &str, n: i64) -> Dataset { .unwrap() } +#[tokio::test] +async fn test_new_datasets_use_stable_field_ids_and_migration_is_idempotent() { + let source_uri = TempStrDir::default(); + let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + assert!(dataset.manifest.uses_stable_field_ids()); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); + assert_ne!( + dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + let created_version = dataset.version().version; + + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .await + .unwrap(); + assert_eq!(dataset.version().version, created_version); + + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::WritersOnly) + .await + .unwrap(); + assert_eq!(dataset.version().version, created_version); +} + +#[tokio::test] +async fn test_stable_field_id_restore_boundary_and_high_water_mark() { + let source_uri = TempStrDir::default(); + let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + let activation_version = dataset.version().version; + + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(vec![ArrowField::new( + "new_field", + DataType::Int32, + true, + )]))), + None, + None, + ) + .await + .unwrap(); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(1)); + + let mut activation_snapshot = dataset.checkout_version(activation_version).await.unwrap(); + activation_snapshot.restore().await.unwrap(); + assert_eq!(activation_snapshot.manifest.max_allocated_field_id, Some(1)); + assert!(activation_snapshot.schema().field("new_field").is_none()); +} + +#[tokio::test] +async fn test_shallow_clone_preserves_stable_field_id_state() { + let source_uri = TempStrDir::default(); + let clone_uri = TempStrDir::default(); + let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .await + .unwrap(); + + let cloned = dataset + .shallow_clone(clone_uri.as_str(), dataset.version().version, None) + .await + .unwrap(); + + assert_eq!( + cloned.manifest.max_allocated_field_id, + dataset.manifest.max_allocated_field_id + ); + assert_ne!( + cloned.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + assert_ne!( + cloned.manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); +} + +#[tokio::test] +async fn test_overwrite_replaces_all_stable_field_identities() { + let source_uri = TempStrDir::default(); + let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .await + .unwrap(); + + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int64, false), + ArrowField::new("replacement", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from_iter_values(0..10)), + Arc::new(Int64Array::from_iter_values(10..20)), + ], + ) + .unwrap(); + let overwritten = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + source_uri.as_str(), + Some(WriteParams { + mode: WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!(overwritten.schema().field("id").unwrap().id, 1); + assert_eq!(overwritten.schema().field("replacement").unwrap().id, 2); + assert_eq!(overwritten.manifest.max_allocated_field_id, Some(2)); + assert!( + overwritten + .manifest + .fragments + .iter() + .flat_map(|fragment| fragment.files.iter()) + .all(|file| file.fields.iter().all(|field_id| *field_id >= 1)) + ); +} + +#[tokio::test] +async fn test_stable_field_id_rename_and_nullability_preserve_identity() { + let source_uri = TempStrDir::default(); + let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .await + .unwrap(); + + dataset + .alter_columns(&[ColumnAlteration::new("id".to_string()) + .rename("renamed".to_string()) + .set_nullable(true)]) + .await + .unwrap(); + + let renamed = dataset.schema().field("renamed").unwrap(); + assert_eq!(renamed.id, 0); + assert!(renamed.nullable); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); + + dataset + .alter_columns(&[ColumnAlteration::new("renamed".to_string()).cast_to(DataType::Int32)]) + .await + .unwrap(); + + assert_eq!(dataset.schema().field("renamed").unwrap().id, 1); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(1)); +} + +#[tokio::test] +async fn test_stable_field_id_multi_cast_uses_schema_order() { + let source_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, false), + ArrowField::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![3, 4])), + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + source_uri.as_str(), + None, + ) + .await + .unwrap(); + + dataset + .alter_columns(&[ + ColumnAlteration::new("b".to_string()).cast_to(DataType::Int64), + ColumnAlteration::new("a".to_string()).cast_to(DataType::Int64), + ]) + .await + .unwrap(); + + assert_eq!(dataset.schema().field("a").unwrap().id, 2); + assert_eq!(dataset.schema().field("b").unwrap().id, 3); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(3)); + dataset.validate().await.unwrap(); +} + #[tokio::test] async fn test_migrate_to_stable_row_ids_basic() { // Create a dataset without stable row IDs (the default). diff --git a/rust/lance/src/dataset/updater.rs b/rust/lance/src/dataset/updater.rs index c5cd8db9600..74c7a7445e2 100644 --- a/rust/lance/src/dataset/updater.rs +++ b/rust/lance/src/dataset/updater.rs @@ -195,7 +195,8 @@ impl Updater { // Need to infer the schema. let output_schema = batch.schema(); let mut final_schema = self.fragment.schema().merge(output_schema.as_ref())?; - final_schema.set_field_id(Some(self.fragment.dataset().manifest.max_field_id())); + final_schema + .try_set_field_id(Some(self.fragment.dataset().manifest.max_field_id()))?; self.final_schema = Some(final_schema); self.final_schema.as_ref().unwrap().validate()?; self.write_schema = Some(self.final_schema.as_ref().unwrap().project_by_schema( diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 81b0fa37109..58fe99f3095 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -1344,6 +1344,13 @@ pub(super) fn prepare_write_schema( OnMissing::Error, OnTypeMismatch::Error, )? + } else if let Some(dataset) = dataset + && matches!(params.mode, WriteMode::Overwrite) + && dataset.manifest.uses_stable_field_ids() + { + let mut schema = normalized_converted_schema; + schema.try_reassign_field_ids(Some(dataset.manifest.max_field_id()))?; + schema } else { normalized_converted_schema }; diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index 7ab5b17a9de..611cf560bd7 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -54,6 +54,8 @@ pub struct CommitBuilder<'a> { timeout: Option, /// When `Some`, this commit is the second step of `migrate_to_stable_row_ids`. migration_next_row_id: Option, + /// When `Some`, this commit atomically activates stable field IDs. + stable_field_id_migration_requires_reader: Option, } /// Default timeout applied to [`CommitBuilder::execute`] when none is set. @@ -78,6 +80,7 @@ impl<'a> CommitBuilder<'a> { transaction_properties: None, timeout: Some(DEFAULT_COMMIT_TIMEOUT), migration_next_row_id: None, + stable_field_id_migration_requires_reader: None, } } @@ -266,6 +269,14 @@ impl<'a> CommitBuilder<'a> { self } + pub(crate) fn with_stable_field_id_migration_activation( + mut self, + require_reader: bool, + ) -> Self { + self.stable_field_id_migration_requires_reader = Some(require_reader); + self + } + pub async fn execute(self, transaction: Transaction) -> Result { let timeout = self.timeout; if let Some(t) = timeout @@ -429,6 +440,8 @@ impl<'a> CommitBuilder<'a> { use_stable_row_ids, storage_format: self.storage_format.map(DataStorageFormat::new), migration_next_row_id: self.migration_next_row_id, + stable_field_id_migration_requires_reader: self + .stable_field_id_migration_requires_reader, ..Default::default() }; diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index f88217ae2c3..e6b2a04aa96 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -34,6 +34,7 @@ use lance_file::version::LanceFileVersion; use lance_index::metrics::NoOpMetricsCollector; use lance_io::utils::CachedFileSize; use lance_select::RowAddrTreeMap; +use lance_table::feature_flags::{can_write_dataset, validate_stable_field_id_flags}; use lance_table::format::{ DETACHED_VERSION_MASK, DeletionFile, Fragment, IndexMetadata, Manifest, WriterVersion, is_detached_version, list_index_files_with_sizes, pb, @@ -42,6 +43,10 @@ use lance_table::io::commit::{ CommitConfig, CommitError, CommitHandler, ManifestLocation, ManifestNamingScheme, }; use lance_table::io::manifest::read_manifest; +use lance_table::transaction::{ + validate_detached_stable_field_ids, validate_operation, validate_stable_field_id_manifest, + validate_stable_field_id_transition, +}; use rand::{Rng, rng}; use super::ObjectStore; @@ -389,6 +394,12 @@ async fn do_commit_new_dataset( &Session::default(), ) .await?; + if !can_write_dataset(source_manifest.writer_feature_flags) { + return Err(Error::not_supported(format!( + "Clone source cannot be written by this version of Lance. Flags: {}", + source_manifest.writer_feature_flags + ))); + } if *is_shallow { let new_base_id = source_manifest @@ -470,6 +481,12 @@ async fn do_commit_new_dataset( (manifest, indices) }; + if !manifest.uses_stable_field_ids() { + manifest.activate_stable_field_ids(); + manifest.reader_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; + manifest.writer_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; + } + let result = write_manifest_file( object_store, commit_handler, @@ -727,13 +744,24 @@ fn fix_schema(manifest: &mut Manifest) -> Result<()> { return Ok(()); } + if manifest.uses_stable_field_ids() { + return Err(Error::invalid_input( + "Cannot repair duplicate field IDs after stable field identity is activated; repair would change an existing identity", + )); + } + // Now, we need to remap the field ids to be unique. let mut old_field_id_mapping: HashMap = HashMap::new(); let mut fields_with_duplicate_ids = fields_with_duplicate_ids.into_iter().collect::>(); fields_with_duplicate_ids.sort_unstable(); - for (field_id_seed, field_id) in (manifest.max_field_id() + 1..).zip(fields_with_duplicate_ids) - { - old_field_id_mapping.insert(field_id, field_id_seed); + let next_field_ids = i64::from(manifest.max_field_id()) + 1..; + for (next_field_id, field_id) in next_field_ids.zip(fields_with_duplicate_ids) { + let assigned_field_id = i32::try_from(next_field_id).map_err(|_| { + Error::invalid_input( + "Cannot repair duplicate field IDs because the field-ID space is exhausted", + ) + })?; + old_field_id_mapping.insert(field_id, assigned_field_id); } let mut fragments = manifest.fragments.as_ref().clone(); @@ -1033,6 +1061,16 @@ pub(crate) async fn do_commit_detached_transaction( commit_config: &CommitConfig, retry_timeout: Duration, ) -> Result<(Manifest, ManifestLocation)> { + if !can_write_dataset(dataset.manifest.writer_feature_flags) { + return Err(Error::not_supported(format!( + "This dataset cannot be written by this version of Lance. Flags: {}", + dataset.manifest.writer_feature_flags + ))); + } + validate_stable_field_id_flags(&dataset.manifest)?; + validate_detached_stable_field_ids(&dataset.manifest, &transaction.operation)?; + validate_operation(Some(&dataset.manifest), &transaction.operation)?; + let pb_transaction = pb::Transaction::from(transaction); let inline_transaction = pb_transaction.encoded_len() <= MAX_INLINE_TRANSACTION_BYTES; @@ -1084,9 +1122,12 @@ pub(crate) async fn do_commit_detached_transaction( migrate_manifest(dataset, &mut manifest, /*recompute_stats=*/ false).await?; // fix_schema and check_storage_version are just for sanity-checking and consistency fix_schema(&mut manifest)?; + validate_stable_field_id_transition(&dataset.manifest, &manifest, &transaction.operation)?; + manifest.update_max_field_id(); check_storage_version(&mut manifest)?; check_column_indices(&manifest)?; check_fragment_ids(&manifest)?; + validate_stable_field_id_manifest(&manifest)?; // Runs after the coverage derivation and can replace a fragment bitmap // while keeping its UUID, so anything it narrowed loses its position. let recovered_coverage = migrate_indices(dataset, &mut indices).await?; @@ -1391,6 +1432,15 @@ pub(crate) async fn commit_transaction( transaction = rebase.finish(&dataset).await?; } + if !can_write_dataset(dataset.manifest.writer_feature_flags) { + return Err(Error::not_supported(format!( + "This dataset cannot be written by this version of Lance. Flags: {}", + dataset.manifest.writer_feature_flags + ))); + } + validate_stable_field_id_flags(&dataset.manifest)?; + validate_operation(Some(&dataset.manifest), &transaction.operation)?; + // Recomputed every attempt: the rebase above may have rewritten the // transaction. let pb_transaction = pb::Transaction::from(&transaction); @@ -1444,10 +1494,13 @@ pub(crate) async fn commit_transaction( migrate_manifest(&dataset, &mut manifest, recompute_stats).await?; fix_schema(&mut manifest)?; + validate_stable_field_id_transition(&dataset.manifest, &manifest, &transaction.operation)?; + manifest.update_max_field_id(); check_storage_version(&mut manifest)?; check_column_indices(&manifest)?; check_fragment_ids(&manifest)?; + validate_stable_field_id_manifest(&manifest)?; // Runs after the coverage derivation and can replace a fragment bitmap // while keeping its UUID, so anything it narrowed loses its position. @@ -2050,6 +2103,23 @@ mod tests { (test_dir, ds) } + #[tokio::test] + async fn every_commit_path_checks_required_writer_flags() { + let (_test_dir, mut dataset) = get_empty_dataset().await; + Arc::make_mut(&mut dataset.manifest).writer_feature_flags |= + lance_table::feature_flags::FLAG_UNKNOWN; + + let err = dataset + .update_config(HashMap::from([( + "test.required-writer-gate".to_string(), + Some("value".to_string()), + )])) + .await + .unwrap_err(); + + assert!(matches!(err, Error::NotSupported { .. }), "{err}"); + } + #[tokio::test] async fn test_good_concurrent_config_writes() { let (_tmpdir, dataset) = get_empty_dataset().await; @@ -2260,6 +2330,60 @@ mod tests { assert_eq!(manifest.fragments.as_ref(), &expected_fragments); } + #[test] + fn test_fix_schema_rejects_field_id_exhaustion() { + let mut field = + Field::try_from(ArrowField::new("a", arrow_schema::DataType::Int64, false)).unwrap(); + field.id = i32::MAX; + let schema = Schema { + fields: vec![field], + metadata: Default::default(), + }; + let mut fragment = Fragment::new(0); + fragment.files = vec![ + DataFile::new_legacy_from_fields("first", vec![i32::MAX], None), + DataFile::new_legacy_from_fields("second", vec![i32::MAX], None), + ]; + let mut manifest = Manifest::new( + schema, + Arc::new(vec![fragment]), + DataStorageFormat::default(), + HashMap::new(), + ); + + let err = fix_schema(&mut manifest).unwrap_err(); + + assert!(err.to_string().contains("space is exhausted"), "{err}"); + } + + #[test] + fn test_fix_schema_does_not_reassign_stable_field_identity() { + let mut field = + Field::try_from(ArrowField::new("a", arrow_schema::DataType::Int64, false)).unwrap(); + field.id = 0; + let schema = Schema { + fields: vec![field], + metadata: Default::default(), + }; + let mut fragment = Fragment::new(0); + fragment.files = vec![ + DataFile::new_legacy_from_fields("first", vec![0], None), + DataFile::new_legacy_from_fields("second", vec![0], None), + ]; + let mut manifest = Manifest::new( + schema, + Arc::new(vec![fragment]), + DataStorageFormat::default(), + HashMap::new(), + ); + manifest.activate_stable_field_ids(); + + let err = fix_schema(&mut manifest).unwrap_err(); + + assert!(err.to_string().contains("stable field identity"), "{err}"); + assert_eq!(manifest.schema.field("a").unwrap().id, 0); + } + /// A CommitHandler that always fails with OtherError, used to simulate /// a manifest write failure so we can verify orphaned transaction files /// are cleaned up. From 2964690545e1e7f644db9361a6b3a57ec63c8cff Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 20 Aug 2026 23:24:27 +0800 Subject: [PATCH 02/24] fix: canonicalize stable field ids at commit boundaries --- java/lance-jni/src/transaction.rs | 72 +++- .../org/lance/operation/OverwriteTest.java | 15 +- python/python/lance/dataset.py | 3 - .../python/tests/compat/compat_decorator.py | 12 + .../python/tests/compat/test_file_formats.py | 3 +- .../tests/compat/test_scalar_indices.py | 18 +- .../tests/compat/test_vector_indices.py | 10 +- python/python/tests/test_dataset.py | 13 +- python/src/transaction.rs | 32 +- rust/lance-table/src/transaction.rs | 1 + rust/lance-table/src/transaction/validate.rs | 406 +++++++++++++++++- rust/lance/src/dataset/fragment.rs | 26 +- rust/lance/src/dataset/schema_evolution.rs | 49 ++- .../src/dataset/tests/dataset_migrations.rs | 53 +++ rust/lance/src/dataset/write.rs | 25 +- rust/lance/src/dataset/write/commit.rs | 8 +- rust/lance/src/io/commit.rs | 9 +- 17 files changed, 676 insertions(+), 79 deletions(-) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index 53457fd2113..d56f37d4cbe 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -33,7 +33,7 @@ use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; use prost::Message; use prost_types::Any; use roaring::RoaringBitmap; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use uuid::Uuid; @@ -978,12 +978,46 @@ fn convert_schema_from_operation( Ok((schema, field_id_remap)) } -fn remap_fragment_field_ids(fragments: &mut [Fragment], field_id_remap: &HashMap) { +type DataFileIdentity = (Option, String); + +fn retained_file_identities( + dataset: Option<&mut BlockingDataset>, + read_version: u64, +) -> Result> { + let Some(dataset) = dataset else { + return Ok(HashSet::new()); + }; + let collect = |dataset: &BlockingDataset| { + dataset + .inner + .manifest() + .fragments + .iter() + .flat_map(|fragment| fragment.referenced_lance_files()) + .map(|file| (file.base_id, file.path.clone())) + .collect() + }; + if dataset.inner.version().version == read_version { + Ok(collect(dataset)) + } else { + let read_dataset = dataset.checkout_version(read_version)?; + Ok(collect(&read_dataset)) + } +} + +fn remap_fragment_field_ids( + fragments: &mut [Fragment], + field_id_remap: &HashMap, + retained_files: &HashSet, +) { if field_id_remap.is_empty() { return; } for fragment in fragments { for file in fragment.referenced_lance_files_mut() { + if retained_files.contains(&(file.base_id, file.path.clone())) { + continue; + } for field_id in Arc::make_mut(&mut file.fields) { if let Some(canonical_id) = field_id_remap.get(field_id) { *field_id = *canonical_id; @@ -1195,7 +1229,7 @@ fn convert_to_rust_operation( env: &mut JNIEnv<'_>, java_operation: &JObject<'_>, allocator: Option<&JObject<'_>>, - dataset: Option<&mut BlockingDataset>, + mut dataset: Option<&mut BlockingDataset>, read_version: u64, ) -> Result { let op_name = env.get_string_from_method(java_operation, "name")?; @@ -1346,11 +1380,11 @@ fn convert_to_rust_operation( "BufferAllocator is required for Overwrite operations".to_string(), ) })?, - dataset, + dataset.as_deref_mut(), read_version, true, )?; - remap_fragment_field_ids(&mut fragments, &field_id_remap); + remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); Operation::Overwrite { fragments, schema, @@ -1496,6 +1530,7 @@ fn convert_to_rust_operation( import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; + let retained_files = retained_file_identities(dataset.as_deref_mut(), read_version)?; let (schema, field_id_remap) = convert_schema_from_operation( env, java_operation, @@ -1508,7 +1543,7 @@ fn convert_to_rust_operation( read_version, false, )?; - remap_fragment_field_ids(&mut fragments, &field_id_remap); + remap_fragment_field_ids(&mut fragments, &field_id_remap, &retained_files); Operation::Merge { fragments, preserves_nullability: env @@ -1824,6 +1859,7 @@ mod tests { DataType as ArrowDataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, }; + use lance_table::format::DataFile; use std::{collections::HashMap, sync::Arc}; use super::*; @@ -1892,6 +1928,30 @@ mod tests { assert_eq!(schema.field("renamed").unwrap().id, 10); } + #[test] + fn merge_remap_skips_retained_files() { + let mut retained = Fragment::new(0); + retained.files.push(DataFile::new_legacy_from_fields( + "retained.lance", + vec![0, 1, 2], + None, + )); + let mut rewritten = Fragment::new(1); + rewritten.files.push(DataFile::new_legacy_from_fields( + "new.lance", + vec![1, 2], + None, + )); + let mut fragments = vec![retained, rewritten]; + let remap = HashMap::from([(1, 2), (2, 3)]); + let retained_files = HashSet::from([(None, "retained.lance".to_string())]); + + remap_fragment_field_ids(&mut fragments, &remap, &retained_files); + + assert_eq!(fragments[0].files[0].fields.as_ref(), &[0, 1, 2]); + assert_eq!(fragments[1].files[0].fields.as_ref(), &[2, 3]); + } + #[test] fn test_create_schema_from_arrow() { // base_schema has an existing field id diff --git a/java/src/test/java/org/lance/operation/OverwriteTest.java b/java/src/test/java/org/lance/operation/OverwriteTest.java index c1def711edc..26bfa1f911e 100644 --- a/java/src/test/java/org/lance/operation/OverwriteTest.java +++ b/java/src/test/java/org/lance/operation/OverwriteTest.java @@ -36,6 +36,7 @@ import java.util.Collections; import java.util.List; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -123,7 +124,19 @@ void testOverwrite(@TempDir Path tempDir) throws Exception { Schema schemaRes = scanner.schema(); assertEquals(testDataset.getSchema(), schemaRes); } - assertEquals(retryTxn, dataset.readTransaction().orElse(null)); + try (Transaction committedTxn = dataset.readTransaction().orElseThrow()) { + assertEquals(retryTxn.readVersion(), committedTxn.readVersion()); + assertEquals(retryTxn.uuid(), committedTxn.uuid()); + assertEquals(retryTxn.transactionProperties(), committedTxn.transactionProperties()); + Overwrite committedOverwrite = (Overwrite) committedTxn.operation(); + assertEquals(testDataset.getSchema(), committedOverwrite.schema()); + assertEquals( + Collections.singletonMap("config_key", "config_value"), + committedOverwrite.configUpsertValues().orElseThrow()); + assertArrayEquals( + new int[] {4, 5}, + committedOverwrite.fragments().get(0).getFiles().get(0).getFields()); + } } } } diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index e386e6cb854..4eb9348b0c8 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -5882,8 +5882,6 @@ class Overwrite(BaseOperation): initial_bases: Optional[List[DatasetBasePath]] = None def __post_init__(self): - if isinstance(self.new_schema, pa.Schema): - self.new_schema = LanceSchema.from_pyarrow(self.new_schema) LanceOperation._validate_fragments(self.fragments) @dataclass @@ -6125,7 +6123,6 @@ def __post_init__(self): "Please use a LanceSchema instead.", DeprecationWarning, ) - self.schema = LanceSchema.from_pyarrow(self.schema) LanceOperation._validate_fragments(self.fragments) @dataclass diff --git a/python/python/tests/compat/compat_decorator.py b/python/python/tests/compat/compat_decorator.py index fdfe09a6879..bb79d58c716 100644 --- a/python/python/tests/compat/compat_decorator.py +++ b/python/python/tests/compat/compat_decorator.py @@ -130,6 +130,8 @@ def last_beta_release(): if LAST_BETA_RELEASE is not None: VERSIONS.append(LAST_BETA_RELEASE) +STABLE_FIELD_IDS_MIN_READER_VERSION = Version("11.0.0b16") + class UpgradeDowngradeTest: """Base class for compatibility tests. @@ -170,6 +172,16 @@ def compat_env(self, version: str, method_name: str) -> Dict[str, str]: return {} +class DatasetUpgradeDowngradeTest(UpgradeDowngradeTest): + """Compatibility contract for datasets with stable field IDs enabled.""" + + def skip_downgrade(self, version: str) -> bool: + # New datasets require the stable-field-ID reader feature. Versions + # released before the feature landed must fail closed instead of + # interpreting the dataset with reusable field identities. + return Version(version) < STABLE_FIELD_IDS_MIN_READER_VERSION + + @contextmanager def _temporary_env(overrides: Optional[Dict[str, str]]): if not overrides: diff --git a/python/python/tests/compat/test_file_formats.py b/python/python/tests/compat/test_file_formats.py index 9ceafa81a49..be2b533411b 100644 --- a/python/python/tests/compat/test_file_formats.py +++ b/python/python/tests/compat/test_file_formats.py @@ -15,6 +15,7 @@ from lance.file import LanceFileReader, LanceFileWriter from .compat_decorator import ( + DatasetUpgradeDowngradeTest, UpgradeDowngradeTest, compat_test, ) @@ -117,7 +118,7 @@ def check_write(self): @compat_test(min_version="0.16.0") -class BasicTypesLegacy(UpgradeDowngradeTest): +class BasicTypesLegacy(DatasetUpgradeDowngradeTest): """Test legacy data storage version 0.1 compatibility.""" def __init__(self, path: Path): diff --git a/python/python/tests/compat/test_scalar_indices.py b/python/python/tests/compat/test_scalar_indices.py index 850c1212532..bb7468266d3 100644 --- a/python/python/tests/compat/test_scalar_indices.py +++ b/python/python/tests/compat/test_scalar_indices.py @@ -17,14 +17,14 @@ import pyarrow as pa from .compat_decorator import ( - UpgradeDowngradeTest, + DatasetUpgradeDowngradeTest, compat_test, ) from .util import safe_data_storage_version @compat_test(min_version="0.30.0") -class BTreeIndex(UpgradeDowngradeTest): +class BTreeIndex(DatasetUpgradeDowngradeTest): """Test BTREE scalar index compatibility (introduced in 0.20.0). Started fully working in 0.30.0 with various fixes. @@ -80,7 +80,7 @@ def check_write(self): @compat_test(min_version="0.22.0") -class BitmapLabelListIndex(UpgradeDowngradeTest): +class BitmapLabelListIndex(DatasetUpgradeDowngradeTest): """Test BITMAP and LABEL_LIST scalar index compatibility (introduced in 0.20.0). Started fully working in 0.22.0 with fixes to LABEL_LIST index. @@ -138,7 +138,7 @@ def check_write(self): @compat_test(min_version="0.36.0") -class NgramIndex(UpgradeDowngradeTest): +class NgramIndex(DatasetUpgradeDowngradeTest): """Test NGRAM index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -187,7 +187,7 @@ def check_write(self): @compat_test(min_version="0.36.0") -class ZonemapBloomfilterIndex(UpgradeDowngradeTest): +class ZonemapBloomfilterIndex(DatasetUpgradeDowngradeTest): """Test ZONEMAP and BLOOMFILTER index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -264,11 +264,11 @@ def check_write(self): def skip_downgrade(self, version: str) -> bool: # In 0.X the zonemap index did not properly handle NULL in filters - return version.startswith("0.") + return super().skip_downgrade(version) or version.startswith("0.") @compat_test(min_version="0.36.0") -class JsonIndex(UpgradeDowngradeTest): +class JsonIndex(DatasetUpgradeDowngradeTest): """Test JSON index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -324,7 +324,7 @@ def check_write(self): @compat_test(min_version="0.36.0") -class FtsIndex(UpgradeDowngradeTest): +class FtsIndex(DatasetUpgradeDowngradeTest): """Test FTS (full-text search) index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -379,7 +379,7 @@ def check_write(self): ds.optimize.compact_files() def skip_downgrade(self, version: str) -> bool: - return version.startswith("0.") + return super().skip_downgrade(version) or version.startswith("0.") def current_env(self, method_name: str) -> dict[str, str]: if method_name == "create": diff --git a/python/python/tests/compat/test_vector_indices.py b/python/python/tests/compat/test_vector_indices.py index e97d6be8bf6..04b9c675259 100644 --- a/python/python/tests/compat/test_vector_indices.py +++ b/python/python/tests/compat/test_vector_indices.py @@ -18,14 +18,14 @@ import pyarrow.compute as pc from .compat_decorator import ( - UpgradeDowngradeTest, + DatasetUpgradeDowngradeTest, compat_test, ) from .util import safe_data_storage_version @compat_test(min_version="0.29.1.beta2") -class PqVectorIndex(UpgradeDowngradeTest): +class PqVectorIndex(DatasetUpgradeDowngradeTest): """Test PQ (Product Quantization) vector index compatibility.""" def __init__(self, path: Path): @@ -101,7 +101,7 @@ def check_write(self): @compat_test(min_version="0.39.0") -class HnswPqVectorIndex(UpgradeDowngradeTest): +class HnswPqVectorIndex(DatasetUpgradeDowngradeTest): """Test IVF_HNSW_PQ vector index compatibility. Note: Only tests versions >= 0.39.0 because earlier versions don't support @@ -182,7 +182,7 @@ def check_write(self): @compat_test(min_version="0.39.0") -class HnswSqVectorIndex(UpgradeDowngradeTest): +class HnswSqVectorIndex(DatasetUpgradeDowngradeTest): """Test IVF_HNSW_SQ vector index compatibility. Note: Only tests versions >= 0.39.0 because earlier versions don't support @@ -263,7 +263,7 @@ def check_write(self): @compat_test(min_version="4.0.0-beta.8") -class IvfRqVectorIndex(UpgradeDowngradeTest): +class IvfRqVectorIndex(DatasetUpgradeDowngradeTest): """Test IVF_RQ vector index compatibility. V2 was introduced in v4.0.0-beta.8""" def __init__(self, path: Path): diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 9cddbf83dac..1ece8941919 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5333,7 +5333,8 @@ def test_data_overlay_dense(tmp_path: Path): assert result.column("id").to_pylist() == list(range(10)) -def test_data_overlay_newest_wins(tmp_path: Path): +def test_data_overlay_newest_wins(tmp_path: Path, monkeypatch): + monkeypatch.setenv("LANCE_ENABLE_UNSTABLE_DATA_OVERLAY_FILES", "1") base_dir = tmp_path / "test" table = pa.table( { @@ -5427,9 +5428,13 @@ def test_data_overlay_sparse_per_field(tmp_path: Path): assert result.column("val").to_pylist()[2] == 20 -def test_data_overlay_round_trips_through_fragment_metadata(tmp_path: Path): +def test_data_overlay_round_trips_through_fragment_metadata( + tmp_path: Path, monkeypatch +): import json + monkeypatch.setenv("LANCE_ENABLE_UNSTABLE_DATA_OVERLAY_FILES", "1") + base_dir = tmp_path / "test" table = pa.table( { @@ -5468,11 +5473,11 @@ def test_data_overlay_round_trips_through_fragment_metadata(tmp_path: Path): assert restored.overlays[0].offsets == [1, 4] assert restored.overlays[0].committed_version == overlay_version - # A commit that round-trips the fragment (here an Overwrite) must keep the + # A commit that round-trips the fragment (here a Merge) must keep the # overlays, so the overlay still resolves on read instead of being dropped. dataset = lance.LanceDataset.commit( dataset, - lance.LanceOperation.Overwrite(dataset.schema, [restored]), + lance.LanceOperation.Merge([restored], dataset.lance_schema, True), read_version=dataset.version, ) result = dataset.to_table() diff --git a/python/src/transaction.rs b/python/src/transaction.rs index ae48fbc7dfd..d2a60890c3b 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -10,9 +10,10 @@ use lance::dataset::transaction::{ DataOverlayGroup, DataReplacementGroup, Operation, RewriteGroup, RewrittenIndex, Transaction, UpdateMap, UpdateMapEntry, UpdateMode, UpdatedFragmentOffsets, }; -use lance::datatypes::Schema; +use lance::datatypes::{Field, Schema}; use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; use lance_table::format::{BasePath, DataFile, Fragment, IndexFile, IndexMetadata}; +use lance_table::transaction::TRANSACTION_SCHEMA_SOURCE_RAW_ARROW; use pyo3::exceptions::PyValueError; use pyo3::types::PySet; use pyo3::{Bound, FromPyObject, PyAny, PyResult, Python}; @@ -996,17 +997,20 @@ fn extract_schema(schema: &Bound<'_, PyAny>) -> PyResult { } fn convert_schema(arrow_schema: &ArrowSchema) -> PyResult { - let mut schema = Schema::try_from(arrow_schema).map_err(|e| { - PyValueError::new_err(format!( - "Failed to convert Arrow schema to Lance schema: {}", - e - )) - })?; - // A raw Arrow schema is not an allocation authority. Assign a canonical - // standalone sequence; dataset commits will either match existing - // identities or reject it in favor of a dataset-derived LanceSchema. - schema - .try_reassign_field_ids(None) - .map_err(|e| PyValueError::new_err(format!("Failed to assign field ids: {e}")))?; - Ok(schema) + let fields = arrow_schema + .fields + .iter() + .map(|field| Field::try_from(field.as_ref())) + .collect::>() + .map_err(|e| { + PyValueError::new_err(format!( + "Failed to convert Arrow schema to Lance schema: {e}" + )) + })?; + let mut metadata = arrow_schema.metadata.clone(); + metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + Ok(Schema { fields, metadata }) } diff --git a/rust/lance-table/src/transaction.rs b/rust/lance-table/src/transaction.rs index e55988d4605..8eb2135ea06 100644 --- a/rust/lance-table/src/transaction.rs +++ b/rust/lance-table/src/transaction.rs @@ -48,6 +48,7 @@ pub use update_map::{ UpdateMap, UpdateMapEntry, translate_config_updates, translate_schema_metadata_updates, }; pub use validate::{ + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW, canonicalize_stable_field_ids, validate_detached_stable_field_ids, validate_operation, validate_stable_field_id_manifest, validate_stable_field_id_transition, }; diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index afd99e1a1d7..7462e650222 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -15,6 +15,260 @@ use lance_core::{Error, Result}; use lance_file::version::ConcreteFileVersion; use std::collections::{HashMap, HashSet}; +type DataFileIdentity = (Option, String); + +/// Transient schema-metadata marker used by bindings for raw Arrow input. +pub const TRANSACTION_SCHEMA_SOURCE_RAW_ARROW: &str = "lance:transaction_schema_source_raw_arrow"; + +/// Canonicalize schema identities supplied by a transaction before validation. +/// +/// Arrow field-ID metadata is descriptive input, not allocation authority. New +/// datasets allocate from zero, while stable datasets preserve compatible +/// existing identities and allocate every new identity above the persisted +/// high-water mark. File mappings written against the incoming schema are +/// updated in the same step; files retained by a merge are never rewritten. +pub fn canonicalize_stable_field_ids( + manifest: Option<&Manifest>, + operation: &mut Operation, +) -> Result<()> { + let raw_arrow_schema = match operation { + Operation::Overwrite { schema, .. } + | Operation::Project { schema, .. } + | Operation::Merge { schema, .. } => schema + .metadata + .remove(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + .is_some(), + _ => false, + }; + if manifest.is_some_and(|manifest| !manifest.uses_stable_field_ids()) { + if raw_arrow_schema { + match operation { + Operation::Overwrite { schema, .. } + | Operation::Project { schema, .. } + | Operation::Merge { schema, .. } => { + // Legacy datasets retain the standalone Arrow conversion + // contract. Missing IDs still need to be assigned after the + // transient provenance marker has been removed. + schema.try_set_field_id(None)?; + schema.validate()?; + schema.verify_primary_key()?; + } + _ => {} + } + } + return Ok(()); + } + + match operation { + Operation::Overwrite { + schema, fragments, .. + } => { + let field_id_remap = + canonicalize_schema(manifest, schema, true, false, raw_arrow_schema)?; + remap_fragment_field_ids(fragments, &field_id_remap, &HashSet::new()); + } + Operation::Project { schema, .. } if raw_arrow_schema => { + let Some(manifest) = manifest else { + return Ok(()); + }; + canonicalize_schema(Some(manifest), schema, false, false, false)?; + } + Operation::Merge { + schema, fragments, .. + } => { + let Some(manifest) = manifest else { + return Ok(()); + }; + let retained_files = manifest + .fragments + .iter() + .flat_map(|fragment| fragment.referenced_lance_files()) + .map(|file| (file.base_id, file.path.clone())) + .collect(); + if raw_arrow_schema { + let field_id_remap = + canonicalize_schema(Some(manifest), schema, false, false, false)?; + remap_fragment_field_ids(fragments, &field_id_remap, &retained_files); + } + canonicalize_merge_replacements(manifest, schema, fragments, &retained_files)?; + } + _ => {} + } + Ok(()) +} + +fn canonicalize_merge_replacements( + manifest: &Manifest, + schema: &mut Schema, + fragments: &mut [Fragment], + retained_files: &HashSet, +) -> Result<()> { + let retained_field_ids = manifest + .fragments + .iter() + .flat_map(|fragment| fragment.referenced_lance_files()) + .flat_map(|file| file.fields.iter()) + .copied() + .collect::>(); + let replaced_field_ids = fragments + .iter() + .flat_map(|fragment| fragment.referenced_lance_files()) + .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone()))) + .flat_map(|file| file.fields.iter()) + .copied() + .filter(|field_id| retained_field_ids.contains(field_id)) + .collect::>(); + if replaced_field_ids.is_empty() { + return Ok(()); + } + + let original = schema.clone(); + let max_field_id = manifest.max_field_id(); + for field in &mut schema.fields { + clear_replaced_and_new_field_ids(field, max_field_id, &replaced_field_ids); + } + schema.try_set_field_id(Some(max_field_id))?; + schema.validate()?; + schema.verify_primary_key()?; + + let field_id_remap = original + .fields_pre_order() + .zip(schema.fields_pre_order()) + .filter_map(|(original, canonical)| { + (original.id != canonical.id).then_some((original.id, canonical.id)) + }) + .collect(); + remap_fragment_field_ids(fragments, &field_id_remap, retained_files); + Ok(()) +} + +fn clear_replaced_and_new_field_ids( + field: &mut Field, + max_field_id: i32, + replaced_field_ids: &HashSet, +) { + if field.id > max_field_id || replaced_field_ids.contains(&field.id) { + clear_field_ids(field); + return; + } + for child in &mut field.children { + clear_replaced_and_new_field_ids(child, max_field_id, replaced_field_ids); + } +} + +fn canonicalize_schema( + manifest: Option<&Manifest>, + schema: &mut Schema, + replaces_all_identities: bool, + allow_id_binding: bool, + remap_standalone_ids: bool, +) -> Result> { + let original = schema.clone(); + + let max_existing_id = manifest.map(Manifest::max_field_id); + if replaces_all_identities || manifest.is_none() { + schema.try_reassign_field_ids(max_existing_id)?; + } else if let Some(manifest) = manifest { + for field in &mut schema.fields { + canonicalize_field(field, -1, &manifest.schema, None, allow_id_binding); + } + schema.try_set_field_id(max_existing_id)?; + } + schema.validate()?; + schema.verify_primary_key()?; + + let mut field_id_remap = HashMap::new(); + for (ordinal, (original, canonical)) in original + .fields_pre_order() + .zip(schema.fields_pre_order()) + .enumerate() + { + // Standalone fragment writers assign zero-based IDs when raw Arrow + // fields omit metadata. Overwrite may later commit those files against + // an existing stable dataset, so preserve that source mapping too. + if remap_standalone_ids { + field_id_remap.insert(ordinal as i32, canonical.id); + } + if original.id >= 0 { + field_id_remap.insert(original.id, canonical.id); + } + } + field_id_remap.retain(|source, canonical| source != canonical); + Ok(field_id_remap) +} + +fn canonicalize_field( + field: &mut Field, + parent_id: i32, + base_schema: &Schema, + base_parent: Option<&Field>, + allow_id_binding: bool, +) { + let same_name = match base_parent { + Some(parent) => parent.children.iter().find(|base| base.name == field.name), + None => base_schema + .fields + .iter() + .find(|base| base.name == field.name), + }; + let by_id = (allow_id_binding && field.id >= 0) + .then(|| base_schema.field_by_id(field.id)) + .flatten() + .filter(|base| base.parent_id == parent_id); + let base_field = if allow_id_binding && field.id >= 0 { + by_id.filter(|base| base.logical_type == field.logical_type) + } else { + same_name.filter(|base| base.logical_type == field.logical_type) + }; + + let Some(base_field) = base_field else { + clear_field_ids(field); + return; + }; + + field.id = base_field.id; + field.parent_id = parent_id; + for child in &mut field.children { + canonicalize_field( + child, + field.id, + base_schema, + Some(base_field), + allow_id_binding, + ); + } +} + +fn clear_field_ids(field: &mut Field) { + field.id = -1; + field.parent_id = -1; + for child in &mut field.children { + clear_field_ids(child); + } +} + +fn remap_fragment_field_ids( + fragments: &mut [Fragment], + field_id_remap: &HashMap, + retained_files: &HashSet, +) { + if field_id_remap.is_empty() { + return; + } + for fragment in fragments { + for file in fragment.referenced_lance_files_mut() { + if retained_files.contains(&(file.base_id, file.path.clone())) { + continue; + } + for field_id in std::sync::Arc::make_mut(&mut file.fields) { + if let Some(canonical_id) = field_id_remap.get(field_id) { + *field_id = *canonical_id; + } + } + } + } +} + /// Validate the operation is valid for the given manifest. pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> Result<()> { let manifest = match (manifest, operation) { @@ -40,9 +294,7 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> } }; - validate_stable_field_id_operation(manifest, operation)?; - - match operation { + let result = match operation { Operation::Append { fragments } => { // Fragments must contain all fields in the schema schema_fragments_valid(Some(manifest), &manifest.schema, fragments) @@ -97,7 +349,9 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> Ok(()) } _ => Ok(()), - } + }; + result?; + validate_stable_field_id_operation(manifest, operation) } /// Validate stable-field-ID invariants that are independent of one operation. @@ -503,7 +757,7 @@ fn merge_schema_valid( fragment .files .iter() - .any(|file| file.fields.contains(&field.id)) + .any(|file| file_materializes_field(file.fields.as_ref(), field)) }); if !materialized { return Err(Error::invalid_input(format!( @@ -518,6 +772,14 @@ fn merge_schema_valid( Ok(()) } +fn file_materializes_field(file_field_ids: &[i32], field: &Field) -> bool { + file_field_ids.contains(&field.id) + || field + .children + .iter() + .any(|child| file_materializes_field(file_field_ids, child)) +} + fn is_field_binding_fully_rewritten( manifest: &Manifest, new_fragment_map: &HashMap, @@ -759,6 +1021,140 @@ mod tests { ); } + #[test] + fn canonicalize_overwrite_remaps_hostile_arrow_field_ids() { + let manifest = activated_manifest(); + let mut schema = one_field_schema(); + schema.fields[0].id = i32::MAX; + let mut operation = Operation::Overwrite { + fragments: vec![fragment_with_file_fields(0, "new.lance", vec![i32::MAX])], + schema, + config_upsert_values: None, + initial_bases: None, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Overwrite { + schema, fragments, .. + } = operation + else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, 1); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]); + } + + #[test] + fn canonicalize_overwrite_remaps_standalone_fragment_field_ids() { + let manifest = activated_manifest(); + let mut schema = one_field_schema(); + schema.fields[0].id = -1; + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut operation = Operation::Overwrite { + fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0])], + schema, + config_upsert_values: None, + initial_bases: None, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Overwrite { + schema, fragments, .. + } = operation + else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, 1); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]); + } + + #[test] + fn canonicalize_raw_arrow_project_does_not_bind_by_field_id() { + let manifest = activated_manifest(); + let mut schema = one_field_schema(); + schema.fields[0].name = "renamed".to_string(); + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Project { schema, .. } = operation else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, 1); + assert!( + !schema + .metadata + .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + ); + } + + #[test] + fn canonicalize_raw_arrow_schema_for_legacy_dataset() { + let manifest = manifest_with_file_fields(one_field_schema(), vec![0]); + let mut schema = one_field_schema(); + schema.fields[0].id = -1; + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Project { schema, .. } = operation else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, 0); + assert!( + !schema + .metadata + .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + ); + } + + #[test] + fn canonicalize_merge_allocates_fresh_identity_for_replaced_column() { + let manifest = activated_manifest(); + let mut merged_fragment = manifest.fragments[0].clone(); + merged_fragment.files.push(DataFile::new_legacy_from_fields( + "replacement.lance", + vec![0], + None, + )); + let mut operation = Operation::Merge { + fragments: vec![merged_fragment], + schema: manifest.schema.clone(), + preserves_nullability: true, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Merge { + schema, fragments, .. + } = operation + else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, 1); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]); + assert_eq!(fragments[0].files[1].fields.as_ref(), &[1]); + } + #[test] fn stable_field_id_manifest_rejects_high_water_mark_below_overlay_reference() { let mut manifest = activated_manifest(); diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index e97703f3533..6886515e352 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -5977,29 +5977,35 @@ mod tests { } let mut fragments: Vec = Vec::new(); + let rows_per_fragment = dataset_rows / paths.len(); for (idx, path) in paths.iter().enumerate() { - let f = FileFragment::create_from_file(path, &dataset, idx, None) + let f = FileFragment::create_from_file(path, &dataset, idx, Some(rows_per_fragment)) .await .unwrap(); fragments.push(f) } - let op = Operation::Overwrite { + let op = Operation::Merge { schema: schema.clone(), fragments, - config_upsert_values: None, - initial_bases: None, + preserves_nullability: true, }; - let new_dataset = - Dataset::commit(test_uri, op, None, None, None, Default::default(), false) - .await - .unwrap(); + let new_dataset = Dataset::commit( + test_uri, + op, + Some(dataset.version().version), + None, + None, + Default::default(), + false, + ) + .await + .unwrap(); assert_eq!(new_dataset.count_rows(None).await.unwrap(), dataset_rows); - // Fragments will have number of rows recorded in metadata, even though - // we passed `None` when constructing the `FileFragment`. + // Fragments retain the row count supplied when reconstructing them. let fragments = new_dataset.get_fragments(); assert_eq!(fragments.len(), 5); for f in fragments { diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 0c002318f4e..54a4255de17 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -2148,13 +2148,41 @@ mod test { ) .await?; dataset.validate().await?; + let checkpoint_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "double_id", + DataType::Int32, + false, + )])); + let checkpoint_schema_ref = checkpoint_schema.clone(); + let checkpoint_result = add_columns_impl( + &dataset.get_fragments(), + Some(vec!["id".to_string()]), + Box::new(move |batch: &RecordBatch| { + let id = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + Ok(RecordBatch::try_new( + checkpoint_schema_ref.clone(), + vec![Arc::new(Int32Array::from_iter_values( + id.values().iter().map(|i| i * 2), + ))], + )?) + }), + None, + None, + None, + ) + .await?; + let cached_fragment = checkpoint_result.fragments[0].clone(); - #[derive(Default)] struct RequestCounter { pub get_batch_requests: Mutex>, pub insert_batch_requests: Mutex>, pub get_fragment_requests: Mutex>, pub insert_fragment_requests: Mutex>, + pub cached_fragment: Fragment, } impl UDFCheckpointStore for RequestCounter { @@ -2183,16 +2211,7 @@ mod test { fn get_fragment(&self, fragment_id: u32) -> Result> { self.get_fragment_requests.lock().unwrap().push(fragment_id); if fragment_id == 0 { - Ok(Some(Fragment { - files: vec![], - id: 0, - overlays: vec![], - deletion_file: None, - row_id_meta: None, - physical_rows: Some(50), - last_updated_at_version_meta: None, - created_at_version_meta: None, - })) + Ok(Some(self.cached_fragment.clone())) } else { Ok(None) } @@ -2207,7 +2226,13 @@ mod test { } } - let request_counter = Arc::new(RequestCounter::default()); + let request_counter = Arc::new(RequestCounter { + get_batch_requests: Mutex::default(), + insert_batch_requests: Mutex::default(), + get_fragment_requests: Mutex::default(), + insert_fragment_requests: Mutex::default(), + cached_fragment, + }); let output_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( "double_id", diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index 7a372ee7a80..15cd0f6def6 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use std::collections::HashMap; use std::sync::Arc; use std::vec; @@ -323,6 +324,22 @@ async fn test_fix_v0_10_5_corrupt_schema() { ); } +#[tokio::test] +async fn test_deep_clone_repairs_legacy_schema_before_activation() { + let source_dir = copy_test_data_to_tmp("v0.10.5/corrupt_schema").unwrap(); + let clone_uri = TempStrDir::default(); + let mut source = Dataset::open(&source_dir.path_str()).await.unwrap(); + + let mut cloned = source + .deep_clone(clone_uri.as_str(), source.version().version, None) + .await + .unwrap(); + + cloned.delete("false").await.unwrap(); + cloned.validate().await.unwrap(); + assert!(cloned.manifest.uses_stable_field_ids()); +} + #[tokio::test] async fn test_fix_v0_21_0_corrupt_fragment_bitmap() { // In v0.21.0 and earlier, delta indices had a bug where the fragment bitmap @@ -777,6 +794,42 @@ async fn test_stable_field_id_multi_cast_uses_schema_order() { dataset.validate().await.unwrap(); } +#[tokio::test] +async fn test_new_dataset_ignores_hostile_arrow_field_id() { + let source_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, false).with_metadata(HashMap::from([( + "lance:field_id".to_string(), + i32::MAX.to_string(), + )])), + ])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))]).unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + source_uri.as_str(), + None, + ) + .await + .unwrap(); + + assert_eq!(dataset.schema().field("a").unwrap().id, 0); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(vec![ArrowField::new( + "b", + DataType::Int32, + true, + )]))), + None, + None, + ) + .await + .unwrap(); + assert_eq!(dataset.schema().field("b").unwrap().id, 1); +} + #[tokio::test] async fn test_migrate_to_stable_row_ids_basic() { // Create a dataset without stable row IDs (the default). diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 58fe99f3095..b4b0f95990d 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -1328,7 +1328,11 @@ pub(super) fn prepare_write_schema( params: &WriteParams, mut schema_compare_options: lance_core::datatypes::SchemaCompareOptions, ) -> Result { - let schema = if let Some(dataset) = dataset + let schema = if dataset.is_none() { + let mut schema = normalized_converted_schema; + schema.try_reassign_field_ids(None)?; + schema + } else if let Some(dataset) = dataset && matches!(params.mode, WriteMode::Append | WriteMode::Create) { schema_compare_options.compare_nullability = NullabilityComparison::Ignore; @@ -2337,7 +2341,7 @@ mod tests { let object_store = Arc::new(ObjectStore::memory()); let base_path = Path::from("test"); - let (fragments, _) = write_fragments_internal( + let (fragments, written_schema) = write_fragments_internal( ConcreteFileVersion::V1, None, object_store.clone(), @@ -2353,7 +2357,16 @@ mod tests { assert_eq!(fragments.len(), 1); let fragment = &fragments[0]; assert_eq!(fragment.files.len(), 1); - assert_eq!(fragment.files[0].fields.as_ref(), &[0, 1, 3]); + // New datasets canonicalize incoming field IDs before writing while + // preserving the schema's field order. + assert_eq!( + written_schema + .fields_pre_order() + .map(|field| field.id) + .collect::>(), + vec![0, 1, 2] + ); + assert_eq!(fragment.files[0].fields.as_ref(), &[0, 1, 2]); let path = base_path .clone() @@ -2364,16 +2377,16 @@ mod tests { &path, file_reader, None, - schema.clone(), + written_schema.clone(), 0, 0, - 3, + 2, None, ) .await .unwrap(); assert_eq!(reader.num_batches(), 1); - let batch = reader.read_batch(0, .., &schema).await.unwrap(); + let batch = reader.read_batch(0, .., &written_schema).await.unwrap(); assert_eq!(batch, data); } diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index 611cf560bd7..3bb67929902 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -30,6 +30,7 @@ use super::{WriteDestination, resolve_commit_handler}; use crate::dataset::branch_location::BranchLocation; use crate::dataset::transaction::validate_operation; use lance_core::utils::tracing::{DATASET_COMMITTED_EVENT, TRACE_DATASET_EVENTS}; +use lance_table::transaction::canonicalize_stable_field_ids; use tracing::info; /// Create a new commit from a [`Transaction`]. @@ -303,7 +304,7 @@ impl<'a> CommitBuilder<'a> { } } - async fn execute_inner(self, transaction: Transaction) -> Result { + async fn execute_inner(self, mut transaction: Transaction) -> Result { let session = self .session .or_else(|| self.dest.dataset().map(|ds| ds.session.clone())) @@ -387,6 +388,11 @@ impl<'a> CommitBuilder<'a> { )); } + canonicalize_stable_field_ids( + dest.dataset().map(|dataset| dataset.manifest.as_ref()), + &mut transaction.operation, + )?; + // Validate the operation before proceeding with the commit // This ensures that operations like Merge have proper validation for data integrity if let Some(dataset) = dest.dataset() { diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index e6b2a04aa96..cbe16812eb3 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -44,8 +44,8 @@ use lance_table::io::commit::{ }; use lance_table::io::manifest::read_manifest; use lance_table::transaction::{ - validate_detached_stable_field_ids, validate_operation, validate_stable_field_id_manifest, - validate_stable_field_id_transition, + canonicalize_stable_field_ids, validate_detached_stable_field_ids, validate_operation, + validate_stable_field_id_manifest, validate_stable_field_id_transition, }; use rand::{Rng, rng}; @@ -360,6 +360,9 @@ async fn do_commit_new_dataset( metadata_cache: &DSMetadataCache, store_registry: Arc, ) -> Result<(Manifest, ManifestLocation)> { + let mut transaction = transaction.clone(); + canonicalize_stable_field_ids(None, &mut transaction.operation)?; + let transaction = &transaction; let pb_transaction = pb::Transaction::from(transaction); let inline_transaction = pb_transaction.encoded_len() <= MAX_INLINE_TRANSACTION_BYTES; @@ -482,6 +485,7 @@ async fn do_commit_new_dataset( }; if !manifest.uses_stable_field_ids() { + fix_schema(&mut manifest)?; manifest.activate_stable_field_ids(); manifest.reader_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; @@ -1439,6 +1443,7 @@ pub(crate) async fn commit_transaction( ))); } validate_stable_field_id_flags(&dataset.manifest)?; + canonicalize_stable_field_ids(Some(&dataset.manifest), &mut transaction.operation)?; validate_operation(Some(&dataset.manifest), &transaction.operation)?; // Recomputed every attempt: the rebase above may have rewritten the From 7e46c86af6a7494bc6578b92b2a1c51a6a37693b Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 20 Aug 2026 23:46:11 +0800 Subject: [PATCH 03/24] fix: enforce stable field ids in namespace rewrites --- .../lance-namespace-impls/src/dir/manifest.rs | 38 ++++++++++++++++--- 1 file changed, 33 insertions(+), 5 deletions(-) diff --git a/rust/lance-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index 0a5ed29dc7f..3a4b4df0112 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -60,6 +60,9 @@ use lance_table::format::{Fragment, IndexMetadata, Manifest}; use lance_table::io::commit::{ CommitError, CommitHandler, commit_handler_from_url, write_manifest_file_to_path, }; +use lance_table::transaction::{ + canonicalize_stable_field_ids, validate_stable_field_id_transition, +}; use object_store::{Error as ObjectStoreError, path::Path}; use roaring::RoaringBitmap; use std::io::Cursor; @@ -1977,7 +1980,7 @@ impl ManifestNamespace { ..WriteParams::default() }; - let transaction = match InsertBuilder::new(dataset.clone()) + let mut transaction = match InsertBuilder::new(dataset.clone()) .with_params(&write_params) .execute_uncommitted_stream(output_stream) .await @@ -1993,10 +1996,7 @@ impl ManifestNamespace { let (mutation, index_data) = Self::take_manifest_rewrite_result(&shared)?; - let Operation::Overwrite { - fragments, schema, .. - } = &transaction.operation - else { + let Operation::Overwrite { fragments, .. } = &transaction.operation else { return Err(NamespaceError::Internal { message: "Manifest rewrite transaction is not an overwrite".to_string(), } @@ -2017,11 +2017,39 @@ impl ManifestNamespace { return Ok(mutation.result); } + if let Err(err) = + canonicalize_stable_field_ids(Some(dataset.manifest()), &mut transaction.operation) + { + self.cleanup_staged_manifest_files(&object_store, &staged_data_files, &[]) + .await; + return Err(err); + } + + let Operation::Overwrite { + fragments, schema, .. + } = &transaction.operation + else { + return Err(NamespaceError::Internal { + message: "Manifest rewrite transaction is not an overwrite".to_string(), + } + .into()); + }; + let mut manifest = Self::manifest_from_overwrite_transaction( dataset.manifest(), schema.clone(), fragments, ); + if let Err(err) = validate_stable_field_id_transition( + dataset.manifest(), + &manifest, + &transaction.operation, + ) { + self.cleanup_staged_manifest_files(&object_store, &staged_data_files, &[]) + .await; + return Err(err); + } + manifest.update_max_field_id(); let target_version = manifest.version; let index_uuids = [Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4()]; From 7d59c810c9ed99a73faf233dc5d4e7141872e106 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 00:45:56 +0800 Subject: [PATCH 04/24] fix: preserve stable identities across schema operations --- rust/lance-table/src/transaction/validate.rs | 225 ++++++++++++++---- rust/lance/src/dataset.rs | 8 +- .../src/dataset/tests/dataset_migrations.rs | 14 ++ rust/lance/src/io/commit.rs | 2 +- 4 files changed, 196 insertions(+), 53 deletions(-) diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 7462e650222..866c4485f46 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -17,6 +17,13 @@ use std::collections::{HashMap, HashSet}; type DataFileIdentity = (Option, String); +#[derive(Default)] +struct FieldIdRemap { + canonical_ids: HashSet, + explicit_ids: HashMap, + standalone_ids: HashMap, +} + /// Transient schema-metadata marker used by bindings for raw Arrow input. pub const TRANSACTION_SCHEMA_SOURCE_RAW_ARROW: &str = "lance:transaction_schema_source_raw_arrow"; @@ -71,7 +78,7 @@ pub fn canonicalize_stable_field_ids( let Some(manifest) = manifest else { return Ok(()); }; - canonicalize_schema(Some(manifest), schema, false, false, false)?; + canonicalize_raw_project_schema(manifest, schema)?; } Operation::Merge { schema, fragments, .. @@ -103,21 +110,21 @@ fn canonicalize_merge_replacements( fragments: &mut [Fragment], retained_files: &HashSet, ) -> Result<()> { - let retained_field_ids = manifest - .fragments - .iter() - .flat_map(|fragment| fragment.referenced_lance_files()) - .flat_map(|file| file.fields.iter()) - .copied() - .collect::>(); - let replaced_field_ids = fragments - .iter() - .flat_map(|fragment| fragment.referenced_lance_files()) - .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone()))) - .flat_map(|file| file.fields.iter()) - .copied() - .filter(|field_id| retained_field_ids.contains(field_id)) - .collect::>(); + let mut replaced_field_ids = HashSet::new(); + for fragment in fragments.iter() { + let retained_field_ids = fragment + .referenced_lance_files() + .filter(|file| retained_files.contains(&(file.base_id, file.path.clone()))) + .flat_map(|file| file.fields.iter().copied()) + .collect::>(); + replaced_field_ids.extend( + fragment + .referenced_lance_files() + .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone()))) + .flat_map(|file| file.fields.iter().copied()) + .filter(|field_id| retained_field_ids.contains(field_id)), + ); + } if replaced_field_ids.is_empty() { return Ok(()); } @@ -131,13 +138,15 @@ fn canonicalize_merge_replacements( schema.validate()?; schema.verify_primary_key()?; - let field_id_remap = original - .fields_pre_order() - .zip(schema.fields_pre_order()) - .filter_map(|(original, canonical)| { - (original.id != canonical.id).then_some((original.id, canonical.id)) - }) - .collect(); + let mut field_id_remap = FieldIdRemap::default(); + for (original, canonical) in original.fields_pre_order().zip(schema.fields_pre_order()) { + field_id_remap.canonical_ids.insert(canonical.id); + if original.id >= 0 { + field_id_remap + .explicit_ids + .insert(original.id, canonical.id); + } + } remap_fragment_field_ids(fragments, &field_id_remap, retained_files); Ok(()) } @@ -156,13 +165,31 @@ fn clear_replaced_and_new_field_ids( } } +fn canonicalize_raw_project_schema(manifest: &Manifest, schema: &mut Schema) -> Result<()> { + let mut unmatched_fields = Vec::new(); + for field in &mut schema.fields { + if !canonicalize_field(field, -1, &manifest.schema, None, false) { + unmatched_fields.push(field.name.clone()); + } + } + if !unmatched_fields.is_empty() { + return Err(Error::invalid_input(format!( + "Raw Arrow Project fields [{}] do not match existing field identities; Project cannot allocate new identities because it writes no data", + unmatched_fields.join(", ") + ))); + } + schema.validate()?; + schema.verify_primary_key()?; + Ok(()) +} + fn canonicalize_schema( manifest: Option<&Manifest>, schema: &mut Schema, replaces_all_identities: bool, allow_id_binding: bool, remap_standalone_ids: bool, -) -> Result> { +) -> Result { let original = schema.clone(); let max_existing_id = manifest.map(Manifest::max_field_id); @@ -177,23 +204,27 @@ fn canonicalize_schema( schema.validate()?; schema.verify_primary_key()?; - let mut field_id_remap = HashMap::new(); + let mut field_id_remap = FieldIdRemap::default(); for (ordinal, (original, canonical)) in original .fields_pre_order() .zip(schema.fields_pre_order()) .enumerate() { + field_id_remap.canonical_ids.insert(canonical.id); // Standalone fragment writers assign zero-based IDs when raw Arrow // fields omit metadata. Overwrite may later commit those files against // an existing stable dataset, so preserve that source mapping too. if remap_standalone_ids { - field_id_remap.insert(ordinal as i32, canonical.id); + field_id_remap + .standalone_ids + .insert(ordinal as i32, canonical.id); } if original.id >= 0 { - field_id_remap.insert(original.id, canonical.id); + field_id_remap + .explicit_ids + .insert(original.id, canonical.id); } } - field_id_remap.retain(|source, canonical| source != canonical); Ok(field_id_remap) } @@ -203,7 +234,7 @@ fn canonicalize_field( base_schema: &Schema, base_parent: Option<&Field>, allow_id_binding: bool, -) { +) -> bool { let same_name = match base_parent { Some(parent) => parent.children.iter().find(|base| base.name == field.name), None => base_schema @@ -223,20 +254,24 @@ fn canonicalize_field( let Some(base_field) = base_field else { clear_field_ids(field); - return; + return false; }; field.id = base_field.id; field.parent_id = parent_id; + let mut all_children_match = true; for child in &mut field.children { - canonicalize_field( + if !canonicalize_field( child, field.id, base_schema, Some(base_field), allow_id_binding, - ); + ) { + all_children_match = false; + } } + all_children_match } fn clear_field_ids(field: &mut Field) { @@ -249,19 +284,43 @@ fn clear_field_ids(field: &mut Field) { fn remap_fragment_field_ids( fragments: &mut [Fragment], - field_id_remap: &HashMap, + field_id_remap: &FieldIdRemap, retained_files: &HashSet, ) { - if field_id_remap.is_empty() { - return; - } for fragment in fragments { + // A dataset-aware staged overwrite already uses canonical IDs. Require + // the whole fragment to cover the canonical schema before skipping so + // an overlapping source ID in one split file cannot mask a remap. + let canonical_fragment = !field_id_remap.standalone_ids.is_empty() + && fragment + .referenced_lance_files() + .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone()))) + .flat_map(|file| file.fields.iter().copied()) + .filter(|field_id| *field_id >= 0) + .collect::>() + == field_id_remap.canonical_ids; for file in fragment.referenced_lance_files_mut() { if retained_files.contains(&(file.base_id, file.path.clone())) { continue; } + if canonical_fragment { + continue; + } + let non_negative_ids = file + .fields + .iter() + .copied() + .filter(|field_id| *field_id >= 0); + let source_ids = if non_negative_ids + .clone() + .all(|field_id| field_id_remap.explicit_ids.contains_key(&field_id)) + { + &field_id_remap.explicit_ids + } else { + &field_id_remap.standalone_ids + }; for field_id in std::sync::Arc::make_mut(&mut file.fields) { - if let Some(canonical_id) = field_id_remap.get(field_id) { + if let Some(canonical_id) = source_ids.get(field_id) { *field_id = *canonical_id; } } @@ -1074,30 +1133,61 @@ mod tests { } #[test] - fn canonicalize_raw_arrow_project_does_not_bind_by_field_id() { + fn canonicalize_overwrite_preserves_already_canonical_fragment_field_ids() { let manifest = activated_manifest(); - let mut schema = one_field_schema(); - schema.fields[0].name = "renamed".to_string(); + let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])) + .unwrap(); + schema.fields[0].id = 1; + schema.fields[1].id = 2; schema.metadata.insert( TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), String::new(), ); - let mut operation = Operation::Project { + let mut operation = Operation::Overwrite { + fragments: vec![fragment_with_file_fields(0, "new.lance", vec![1, 2])], schema, - preserves_nullability: true, + config_upsert_values: None, + initial_bases: None, }; canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); - let Operation::Project { schema, .. } = operation else { + let Operation::Overwrite { + schema, fragments, .. + } = operation + else { unreachable!(); }; - assert_eq!(schema.fields[0].id, 1); - assert!( - !schema - .metadata - .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + assert_eq!( + schema + .fields_pre_order() + .map(|field| field.id) + .collect::>(), + vec![1, 2] ); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[1, 2]); + } + + #[test] + fn canonicalize_raw_arrow_project_rejects_unmatched_field() { + let manifest = activated_manifest(); + let mut schema = one_field_schema(); + schema.fields[0].name = "renamed".to_string(); + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + let err = canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap_err(); + + assert!(err.to_string().contains("writes no data"), "{err}"); } #[test] @@ -1128,7 +1218,44 @@ mod tests { } #[test] - fn canonicalize_merge_allocates_fresh_identity_for_replaced_column() { + fn canonicalize_merge_preserves_identity_for_physical_rewrite() { + let mut manifest = activated_manifest(); + Arc::make_mut(&mut manifest.fragments).push(fragment_with_file_fields( + 1, + "retained.lance", + vec![0], + )); + let rewritten_fragment = fragment_with_file_fields(0, "replacement.lance", vec![0]); + let mut operation = Operation::Merge { + fragments: vec![rewritten_fragment, manifest.fragments[1].clone()], + schema: manifest.schema.clone(), + preserves_nullability: true, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Merge { + schema, fragments, .. + } = operation + else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, 0); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]); + assert_eq!(fragments[1].files[0].fields.as_ref(), &[0]); + validate_operation( + Some(&manifest), + &Operation::Merge { + fragments, + schema, + preserves_nullability: true, + }, + ) + .unwrap(); + } + + #[test] + fn canonicalize_merge_allocates_fresh_identity_for_overlaid_column() { let manifest = activated_manifest(); let mut merged_fragment = manifest.fragments[0].clone(); merged_fragment.files.push(DataFile::new_legacy_from_fields( diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index a8cabeb86ec..d1a48c2b470 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -134,7 +134,7 @@ use crate::datatypes::Schema; use crate::index::retain_supported_indices; use crate::io::commit::{ DEFAULT_COMMIT_RETRY_TIMEOUT, commit_detached_transaction, commit_new_dataset, - commit_transaction, detect_overlapping_fragments, + commit_transaction, detect_overlapping_fragments, fix_schema, }; use crate::session::Session; use crate::utils::temporal::{SystemTime, timestamp_to_nanos, utc_now}; @@ -3268,11 +3268,13 @@ impl Dataset { return Ok(()); } + let mut repaired_manifest = self.manifest.as_ref().clone(); + fix_schema(&mut repaired_manifest)?; let transaction = Transaction::new( self.manifest.version, Operation::Merge { - fragments: self.manifest.fragments.as_ref().clone(), - schema: self.manifest.schema.clone(), + fragments: repaired_manifest.fragments.as_ref().clone(), + schema: repaired_manifest.schema, preserves_nullability: true, }, None, diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index 15cd0f6def6..fa6e4de0ae4 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -340,6 +340,20 @@ async fn test_deep_clone_repairs_legacy_schema_before_activation() { assert!(cloned.manifest.uses_stable_field_ids()); } +#[tokio::test] +async fn test_stable_field_id_migration_repairs_legacy_schema_before_activation() { + let test_dir = copy_test_data_to_tmp("v0.10.5/corrupt_schema").unwrap(); + let mut dataset = Dataset::open(&test_dir.path_str()).await.unwrap(); + + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .await + .unwrap(); + + dataset.validate().await.unwrap(); + assert!(dataset.manifest.uses_stable_field_ids()); +} + #[tokio::test] async fn test_fix_v0_21_0_corrupt_fragment_bitmap() { // In v0.21.0 and earlier, delta indices had a bug where the fragment bitmap diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index cbe16812eb3..c9e8e363ae0 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -725,7 +725,7 @@ fn check_column_indices(manifest: &Manifest) -> Result<()> { /// Fix schema in case of duplicate field ids. /// /// See test dataset v0.10.5/corrupt_schema -fn fix_schema(manifest: &mut Manifest) -> Result<()> { +pub(crate) fn fix_schema(manifest: &mut Manifest) -> Result<()> { // We can short-circuit if there is only one file per fragment or no fragments. if manifest.fragments.iter().all(|f| f.files.len() <= 1) { return Ok(()); From 2b45dbe9a6879c31c2efdf7abf8ed0e55bf2fd5e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 01:50:28 +0800 Subject: [PATCH 05/24] fix: reject ambiguous stable field mappings --- python/python/tests/test_dataset.py | 24 +++ rust/lance-table/src/transaction/validate.rs | 174 ++++++++++++++----- 2 files changed, 156 insertions(+), 42 deletions(-) diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 1ece8941919..75edec7859c 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5596,6 +5596,30 @@ def test_schema_project_raw_arrow_cannot_allocate_field_id(tmp_path: Path): assert dataset.to_table() == pa.table({"a": pa.array([1, 2], pa.int32())}) +def test_overwrite_rejects_ambiguous_raw_arrow_field_ids(tmp_path: Path): + dataset = lance.write_dataset( + pa.table({"old": pa.array([1, 2], pa.int32())}), tmp_path + ) + schema = pa.schema( + [ + pa.field("b", pa.int32(), metadata={b"lance:field_id": b"2"}), + pa.field("c", pa.int32(), metadata={b"lance:field_id": b"1"}), + ] + ) + table = pa.Table.from_arrays( + [pa.array([10, 11], pa.int32()), pa.array([20, 21], pa.int32())], + schema=schema, + ) + fragment = lance.fragment.LanceFragment.create(tmp_path, table, mode="overwrite") + + with pytest.raises(OSError, match="ambiguous raw Arrow field IDs"): + lance.LanceDataset.commit( + dataset, + lance.LanceOperation.Overwrite(schema, [fragment]), + read_version=dataset.version, + ) + + def test_schema_project_rename_column(tmp_path: Path): table = pa.Table.from_pydict({"a": range(100, 200), "b": range(300, 400)}) base_dir = tmp_path / "test" diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 866c4485f46..b364e2b3be2 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -21,7 +21,7 @@ type DataFileIdentity = (Option, String); struct FieldIdRemap { canonical_ids: HashSet, explicit_ids: HashMap, - standalone_ids: HashMap, + raw_source_ids: HashMap, } /// Transient schema-metadata marker used by bindings for raw Arrow input. @@ -72,7 +72,7 @@ pub fn canonicalize_stable_field_ids( } => { let field_id_remap = canonicalize_schema(manifest, schema, true, false, raw_arrow_schema)?; - remap_fragment_field_ids(fragments, &field_id_remap, &HashSet::new()); + remap_fragment_field_ids(fragments, &field_id_remap, &HashSet::new())?; } Operation::Project { schema, .. } if raw_arrow_schema => { let Some(manifest) = manifest else { @@ -95,7 +95,7 @@ pub fn canonicalize_stable_field_ids( if raw_arrow_schema { let field_id_remap = canonicalize_schema(Some(manifest), schema, false, false, false)?; - remap_fragment_field_ids(fragments, &field_id_remap, &retained_files); + remap_fragment_field_ids(fragments, &field_id_remap, &retained_files)?; } canonicalize_merge_replacements(manifest, schema, fragments, &retained_files)?; } @@ -147,7 +147,7 @@ fn canonicalize_merge_replacements( .insert(original.id, canonical.id); } } - remap_fragment_field_ids(fragments, &field_id_remap, retained_files); + remap_fragment_field_ids(fragments, &field_id_remap, retained_files)?; Ok(()) } @@ -188,9 +188,16 @@ fn canonicalize_schema( schema: &mut Schema, replaces_all_identities: bool, allow_id_binding: bool, - remap_standalone_ids: bool, + remap_raw_source_ids: bool, ) -> Result { let original = schema.clone(); + let raw_source = if remap_raw_source_ids { + let mut raw_source = original.clone(); + raw_source.try_set_field_id(None)?; + Some(raw_source) + } else { + None + }; let max_existing_id = manifest.map(Manifest::max_field_id); if replaces_all_identities || manifest.is_none() { @@ -205,26 +212,22 @@ fn canonicalize_schema( schema.verify_primary_key()?; let mut field_id_remap = FieldIdRemap::default(); - for (ordinal, (original, canonical)) in original - .fields_pre_order() - .zip(schema.fields_pre_order()) - .enumerate() - { + for (original, canonical) in original.fields_pre_order().zip(schema.fields_pre_order()) { field_id_remap.canonical_ids.insert(canonical.id); - // Standalone fragment writers assign zero-based IDs when raw Arrow - // fields omit metadata. Overwrite may later commit those files against - // an existing stable dataset, so preserve that source mapping too. - if remap_standalone_ids { - field_id_remap - .standalone_ids - .insert(ordinal as i32, canonical.id); - } if original.id >= 0 { field_id_remap .explicit_ids .insert(original.id, canonical.id); } } + if let Some(raw_source) = raw_source { + field_id_remap.raw_source_ids.extend( + raw_source + .fields_pre_order() + .zip(schema.fields_pre_order()) + .map(|(source, canonical)| (source.id, canonical.id)), + ); + } Ok(field_id_remap) } @@ -286,39 +289,54 @@ fn remap_fragment_field_ids( fragments: &mut [Fragment], field_id_remap: &FieldIdRemap, retained_files: &HashSet, -) { +) -> Result<()> { for fragment in fragments { - // A dataset-aware staged overwrite already uses canonical IDs. Require - // the whole fragment to cover the canonical schema before skipping so - // an overlapping source ID in one split file cannot mask a remap. - let canonical_fragment = !field_id_remap.standalone_ids.is_empty() - && fragment + let source_ids = if field_id_remap.raw_source_ids.is_empty() { + &field_id_remap.explicit_ids + } else { + // Raw Arrow overwrite fragments may have been written either by a + // standalone writer using the source schema IDs or by a + // dataset-aware writer using canonical IDs. Resolve that namespace + // once for the whole fragment so split files cannot disagree. If + // both interpretations are possible and produce different + // identities then there is no safe mapping without provenance. + let fragment_field_ids = fragment .referenced_lance_files() .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone()))) .flat_map(|file| file.fields.iter().copied()) .filter(|field_id| *field_id >= 0) - .collect::>() - == field_id_remap.canonical_ids; + .collect::>(); + let canonical_source = fragment_field_ids + .iter() + .all(|field_id| field_id_remap.canonical_ids.contains(field_id)); + let raw_source = fragment_field_ids + .iter() + .all(|field_id| field_id_remap.raw_source_ids.contains_key(field_id)); + let raw_changes_identity = fragment_field_ids + .iter() + .any(|field_id| field_id_remap.raw_source_ids.get(field_id) != Some(field_id)); + + match (canonical_source, raw_source, raw_changes_identity) { + (true, true, true) => { + return Err(Error::invalid_input(format!( + "Fragment {} has ambiguous raw Arrow field IDs; its file mappings can be interpreted as either source or canonical identities", + fragment.id + ))); + } + (true, _, _) => continue, + (false, true, _) => &field_id_remap.raw_source_ids, + (false, false, _) => { + return Err(Error::invalid_input(format!( + "Fragment {} field IDs do not match either the raw Arrow source schema or the canonical replacement schema", + fragment.id + ))); + } + } + }; for file in fragment.referenced_lance_files_mut() { if retained_files.contains(&(file.base_id, file.path.clone())) { continue; } - if canonical_fragment { - continue; - } - let non_negative_ids = file - .fields - .iter() - .copied() - .filter(|field_id| *field_id >= 0); - let source_ids = if non_negative_ids - .clone() - .all(|field_id| field_id_remap.explicit_ids.contains_key(&field_id)) - { - &field_id_remap.explicit_ids - } else { - &field_id_remap.standalone_ids - }; for field_id in std::sync::Arc::make_mut(&mut file.fields) { if let Some(canonical_id) = source_ids.get(field_id) { *field_id = *canonical_id; @@ -326,6 +344,7 @@ fn remap_fragment_field_ids( } } } + Ok(()) } /// Validate the operation is valid for the given manifest. @@ -1171,6 +1190,77 @@ mod tests { assert_eq!(fragments[0].files[0].fields.as_ref(), &[1, 2]); } + #[test] + fn canonicalize_overwrite_uses_one_raw_namespace_for_split_files() { + let manifest = activated_manifest(); + let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])) + .unwrap(); + for field in &mut schema.fields { + field.id = -1; + } + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut fragment = fragment_with_file_fields(0, "b.lance", vec![0]); + fragment + .files + .push(DataFile::new_legacy_from_fields("c.lance", vec![1], None)); + let mut operation = Operation::Overwrite { + fragments: vec![fragment], + schema, + config_upsert_values: None, + initial_bases: None, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Overwrite { + schema, fragments, .. + } = operation + else { + unreachable!(); + }; + assert_eq!( + schema + .fields_pre_order() + .map(|field| field.id) + .collect::>(), + vec![1, 2] + ); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]); + assert_eq!(fragments[0].files[1].fields.as_ref(), &[2]); + } + + #[test] + fn canonicalize_overwrite_rejects_ambiguous_raw_field_ids() { + let manifest = activated_manifest(); + let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])) + .unwrap(); + schema.fields[0].id = 2; + schema.fields[1].id = 1; + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut operation = Operation::Overwrite { + fragments: vec![fragment_with_file_fields(0, "new.lance", vec![2, 1])], + schema, + config_upsert_values: None, + initial_bases: None, + }; + + let err = canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap_err(); + + assert!(err.to_string().contains("ambiguous raw Arrow field IDs")); + } + #[test] fn canonicalize_raw_arrow_project_rejects_unmatched_field() { let manifest = activated_manifest(); From 885148443f9cb6db0ad2da226ba6ded1a999f90a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 02:12:37 +0800 Subject: [PATCH 06/24] fix: centralize binding field id remapping --- java/lance-jni/src/transaction.rs | 115 ++++-------------- .../java/org/lance/operation/MergeTest.java | 68 +++++++++++ python/python/tests/test_dataset.py | 50 ++++++++ rust/lance-table/src/transaction/validate.rs | 35 +++++- 4 files changed, 174 insertions(+), 94 deletions(-) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index d56f37d4cbe..a30e1179711 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -30,10 +30,11 @@ use lance_file::version::{LanceFileVersion, V2_FORMAT_2_0, V2_FORMAT_2_1, V2_FOR use lance_io::object_store::{LanceNamespaceStorageOptionsProvider, StorageOptionsProvider}; use lance_table::io::commit::CommitHandler; use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; +use lance_table::transaction::TRANSACTION_SCHEMA_SOURCE_RAW_ARROW; use prost::Message; use prost_types::Any; use roaring::RoaringBitmap; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::sync::Arc; use uuid::Uuid; @@ -922,7 +923,8 @@ fn convert_schema_from_operation( dataset: Option<&mut BlockingDataset>, read_version: u64, replaces_all_identities: bool, -) -> Result<(LanceSchema, HashMap)> { + commit_allocator_is_authoritative: bool, +) -> Result { let schema_ptr = env .call_method( java_operation, @@ -935,13 +937,21 @@ fn convert_schema_from_operation( let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) }; let arrow_schema = Schema::try_from(&c_schema)?; - let original_schema = LanceSchema::try_from(&arrow_schema).map_err(|e| { + let mut original_schema = LanceSchema::try_from(&arrow_schema).map_err(|e| { Error::input_error(format!( "Failed to convert Arrow schema to Lance schema: {}", e )) })?; + if commit_allocator_is_authoritative { + original_schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + return Ok(original_schema); + } + let schema = if let Some(dataset) = dataset { // Derive field ids based on the transaction read dataset schema. let read_schema = { @@ -968,63 +978,7 @@ fn convert_schema_from_operation( LanceSchema::from_arrow_schema(&arrow_schema, None, None, true, true)? }; - let field_id_remap = original_schema - .fields_pre_order() - .zip(schema.fields_pre_order()) - .filter_map(|(original, canonical)| { - (original.id >= 0 && original.id != canonical.id).then_some((original.id, canonical.id)) - }) - .collect(); - Ok((schema, field_id_remap)) -} - -type DataFileIdentity = (Option, String); - -fn retained_file_identities( - dataset: Option<&mut BlockingDataset>, - read_version: u64, -) -> Result> { - let Some(dataset) = dataset else { - return Ok(HashSet::new()); - }; - let collect = |dataset: &BlockingDataset| { - dataset - .inner - .manifest() - .fragments - .iter() - .flat_map(|fragment| fragment.referenced_lance_files()) - .map(|file| (file.base_id, file.path.clone())) - .collect() - }; - if dataset.inner.version().version == read_version { - Ok(collect(dataset)) - } else { - let read_dataset = dataset.checkout_version(read_version)?; - Ok(collect(&read_dataset)) - } -} - -fn remap_fragment_field_ids( - fragments: &mut [Fragment], - field_id_remap: &HashMap, - retained_files: &HashSet, -) { - if field_id_remap.is_empty() { - return; - } - for fragment in fragments { - for file in fragment.referenced_lance_files_mut() { - if retained_files.contains(&(file.base_id, file.path.clone())) { - continue; - } - for field_id in Arc::make_mut(&mut file.fields) { - if let Some(canonical_id) = field_id_remap.get(field_id) { - *field_id = *canonical_id; - } - } - } - } + Ok(schema) } trait SchemaExt { @@ -1235,7 +1189,7 @@ fn convert_to_rust_operation( let op_name = env.get_string_from_method(java_operation, "name")?; let op = match op_name.as_str() { "Project" => { - let (schema, _) = convert_schema_from_operation( + let schema = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1246,6 +1200,7 @@ fn convert_to_rust_operation( dataset, read_version, false, + false, )?; Operation::Project { preserves_nullability: env @@ -1359,7 +1314,7 @@ fn convert_to_rust_operation( } } "Overwrite" => { - let mut fragments: Vec = + let fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; @@ -1372,7 +1327,7 @@ fn convert_to_rust_operation( to_rust_map(env, &config_upsert_values) }, )?; - let (schema, field_id_remap) = convert_schema_from_operation( + let schema = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1383,8 +1338,8 @@ fn convert_to_rust_operation( dataset.as_deref_mut(), read_version, true, + true, )?; - remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); Operation::Overwrite { fragments, schema, @@ -1526,12 +1481,11 @@ fn convert_to_rust_operation( Operation::DataReplacement { replacements } } "Merge" => { - let mut fragments: Vec = + let fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; - let retained_files = retained_file_identities(dataset.as_deref_mut(), read_version)?; - let (schema, field_id_remap) = convert_schema_from_operation( + let schema = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1542,8 +1496,8 @@ fn convert_to_rust_operation( dataset, read_version, false, + true, )?; - remap_fragment_field_ids(&mut fragments, &field_id_remap, &retained_files); Operation::Merge { fragments, preserves_nullability: env @@ -1859,7 +1813,6 @@ mod tests { DataType as ArrowDataType, Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema, }; - use lance_table::format::DataFile; use std::{collections::HashMap, sync::Arc}; use super::*; @@ -1928,30 +1881,6 @@ mod tests { assert_eq!(schema.field("renamed").unwrap().id, 10); } - #[test] - fn merge_remap_skips_retained_files() { - let mut retained = Fragment::new(0); - retained.files.push(DataFile::new_legacy_from_fields( - "retained.lance", - vec![0, 1, 2], - None, - )); - let mut rewritten = Fragment::new(1); - rewritten.files.push(DataFile::new_legacy_from_fields( - "new.lance", - vec![1, 2], - None, - )); - let mut fragments = vec![retained, rewritten]; - let remap = HashMap::from([(1, 2), (2, 3)]); - let retained_files = HashSet::from([(None, "retained.lance".to_string())]); - - remap_fragment_field_ids(&mut fragments, &remap, &retained_files); - - assert_eq!(fragments[0].files[0].fields.as_ref(), &[0, 1, 2]); - assert_eq!(fragments[1].files[0].fields.as_ref(), &[2, 3]); - } - #[test] fn test_create_schema_from_arrow() { // base_schema has an existing field id diff --git a/java/src/test/java/org/lance/operation/MergeTest.java b/java/src/test/java/org/lance/operation/MergeTest.java index 9f69eeaa738..943a2f0c2ec 100644 --- a/java/src/test/java/org/lance/operation/MergeTest.java +++ b/java/src/test/java/org/lance/operation/MergeTest.java @@ -363,6 +363,74 @@ void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Ex } } + @Test + void testMergeRejectsAmbiguousRawArrowFieldIds(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("testMergeRejectsAmbiguousRawArrowFieldIds").toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + + try (Dataset initialDataset = createAndAppendRows(testDataset, 2)) { + Field bField = + new Field( + "b", new FieldType(true, new ArrowType.Int(32, true), null, fieldMeta(3)), null); + Field cField = + new Field( + "c", new FieldType(true, new ArrowType.Int(32, true), null, fieldMeta(2)), null); + Schema newColumnsSchema = new Schema(Arrays.asList(bField, cField), null); + Schema evolvedSchema = + new Schema( + Arrays.asList( + testDataset.getSchema().getFields().get(0), + testDataset.getSchema().getFields().get(1), + bField, + cField), + null); + + try (VectorSchemaRoot root = VectorSchemaRoot.create(newColumnsSchema, allocator)) { + root.allocateNew(); + IntVector bVector = (IntVector) root.getVector("b"); + IntVector cVector = (IntVector) root.getVector("c"); + bVector.setSafe(0, 10); + bVector.setSafe(1, 11); + cVector.setSafe(0, 20); + cVector.setSafe(1, 21); + root.setRowCount(2); + + DataFile newDataFile = + writeLanceDataFile( + dataset.allocator(), datasetPath, root, new int[] {2, 3}, new int[] {0, 1}); + FragmentMetadata fragmentMeta = initialDataset.getFragment(0).metadata(); + List dataFiles = fragmentMeta.getFiles(); + dataFiles.add(newDataFile); + FragmentMetadata evolvedFragment = + new FragmentMetadata( + fragmentMeta.getId(), + dataFiles, + fragmentMeta.getPhysicalRows(), + fragmentMeta.getDeletionFile(), + fragmentMeta.getRowIdMeta()); + + try (Transaction transaction = + new Transaction.Builder() + .readVersion(initialDataset.version()) + .operation( + Merge.builder() + .fragments(Collections.singletonList(evolvedFragment)) + .schema(evolvedSchema) + .build()) + .build()) { + IllegalArgumentException error = + Assertions.assertThrows( + IllegalArgumentException.class, + () -> new CommitBuilder(initialDataset).execute(transaction).close()); + Assertions.assertTrue(error.getMessage().contains("ambiguous raw Arrow field IDs")); + } + } + } + } + } + private Map fieldMeta(int fieldId) { Map idMeta = new HashMap<>(); idMeta.put("lance:field_id", String.valueOf(fieldId)); diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 75edec7859c..247e5b14a91 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5620,6 +5620,56 @@ def test_overwrite_rejects_ambiguous_raw_arrow_field_ids(tmp_path: Path): ) +def test_merge_rejects_ambiguous_raw_arrow_field_ids(tmp_path: Path): + dataset = lance.write_dataset( + pa.table({"old": pa.array([1, 2], pa.int32())}), tmp_path + ) + staged = lance.fragment.LanceFragment.create( + tmp_path, + pa.table( + { + "b": pa.array([10, 11], pa.int32()), + "c": pa.array([20, 21], pa.int32()), + } + ), + mode="overwrite", + ) + staged_file = staged.files[0] + canonical_file = lance.fragment.DataFile( + path=staged_file.path, + fields=[1, 2], + column_indices=staged_file.column_indices, + file_major_version=staged_file.file_major_version, + file_minor_version=staged_file.file_minor_version, + file_size_bytes=staged_file.file_size_bytes, + base_id=staged_file.base_id, + ) + current = dataset.get_fragments()[0].metadata + merged = lance.fragment.FragmentMetadata( + id=current.id, + files=[*current.files, canonical_file], + physical_rows=current.physical_rows, + deletion_file=current.deletion_file, + row_id_meta=current.row_id_meta, + ) + schema = pa.schema( + [ + pa.field("old", pa.int32()), + pa.field("b", pa.int32(), metadata={b"lance:field_id": b"2"}), + pa.field("c", pa.int32(), metadata={b"lance:field_id": b"1"}), + ] + ) + + with pytest.deprecated_call(): + operation = lance.LanceOperation.Merge([merged], schema) + with pytest.raises(OSError, match="ambiguous raw Arrow field IDs"): + lance.LanceDataset.commit( + dataset, + operation, + read_version=dataset.version, + ) + + def test_schema_project_rename_column(tmp_path: Path): table = pa.Table.from_pydict({"a": range(100, 200), "b": range(300, 400)}) base_dir = tmp_path / "test" diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index b364e2b3be2..3535bd06279 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -94,7 +94,7 @@ pub fn canonicalize_stable_field_ids( .collect(); if raw_arrow_schema { let field_id_remap = - canonicalize_schema(Some(manifest), schema, false, false, false)?; + canonicalize_schema(Some(manifest), schema, false, false, true)?; remap_fragment_field_ids(fragments, &field_id_remap, &retained_files)?; } canonicalize_merge_replacements(manifest, schema, fragments, &retained_files)?; @@ -1372,6 +1372,39 @@ mod tests { assert_eq!(fragments[0].files[1].fields.as_ref(), &[1]); } + #[test] + fn canonicalize_merge_rejects_ambiguous_raw_field_ids() { + let manifest = activated_manifest(); + let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("c", DataType::Int32, true), + ])) + .unwrap(); + schema.fields[0].id = 0; + schema.fields[1].id = 2; + schema.fields[2].id = 1; + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut merged_fragment = manifest.fragments[0].clone(); + merged_fragment.files.push(DataFile::new_legacy_from_fields( + "new.lance", + vec![1, 2], + None, + )); + let mut operation = Operation::Merge { + fragments: vec![merged_fragment], + schema, + preserves_nullability: true, + }; + + let err = canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap_err(); + + assert!(err.to_string().contains("ambiguous raw Arrow field IDs")); + } + #[test] fn stable_field_id_manifest_rejects_high_water_mark_below_overlay_reference() { let mut manifest = activated_manifest(); From bb427807eaa45c078d64c95192d232105bc08a1b Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 02:22:02 +0800 Subject: [PATCH 07/24] fix(java): preserve legacy field id remapping --- java/lance-jni/src/transaction.rs | 122 ++++++++++++++---- .../java/org/lance/operation/MergeTest.java | 59 +++++++++ 2 files changed, 156 insertions(+), 25 deletions(-) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index a30e1179711..8c237ea194f 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -34,7 +34,7 @@ use lance_table::transaction::TRANSACTION_SCHEMA_SOURCE_RAW_ARROW; use prost::Message; use prost_types::Any; use roaring::RoaringBitmap; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use uuid::Uuid; @@ -924,7 +924,7 @@ fn convert_schema_from_operation( read_version: u64, replaces_all_identities: bool, commit_allocator_is_authoritative: bool, -) -> Result { +) -> Result<(LanceSchema, HashMap)> { let schema_ptr = env .call_method( java_operation, @@ -944,27 +944,36 @@ fn convert_schema_from_operation( )) })?; - if commit_allocator_is_authoritative { + let read_context = match dataset { + Some(dataset) if dataset.inner.version().version == read_version => Some(( + dataset.inner.schema().clone(), + dataset.inner.manifest().max_field_id(), + dataset.inner.manifest().uses_stable_field_ids(), + )), + Some(dataset) => { + let read_dataset = dataset.checkout_version(read_version)?; + Some(( + read_dataset.inner.schema().clone(), + read_dataset.inner.manifest().max_field_id(), + read_dataset.inner.manifest().uses_stable_field_ids(), + )) + } + None => None, + }; + + if commit_allocator_is_authoritative + && read_context + .as_ref() + .is_none_or(|(_, _, stable_field_ids)| *stable_field_ids) + { original_schema.metadata.insert( TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), String::new(), ); - return Ok(original_schema); + return Ok((original_schema, HashMap::new())); } - let schema = if let Some(dataset) = dataset { - // Derive field ids based on the transaction read dataset schema. - let read_schema = { - if dataset.inner.version().version == read_version { - dataset.inner.schema().clone() - } else { - let read_dataset = dataset.checkout_version(read_version)?; - read_dataset.inner.schema().clone() - } - }; - - let max_field_id = dataset.inner.manifest().max_field_id(); - let stable_field_ids = dataset.inner.manifest().uses_stable_field_ids(); + let schema = if let Some((read_schema, max_field_id, stable_field_ids)) = read_context { LanceSchema::from_arrow_schema( &arrow_schema, Some(read_schema), @@ -978,7 +987,63 @@ fn convert_schema_from_operation( LanceSchema::from_arrow_schema(&arrow_schema, None, None, true, true)? }; - Ok(schema) + let field_id_remap = original_schema + .fields_pre_order() + .zip(schema.fields_pre_order()) + .filter_map(|(original, canonical)| { + (original.id >= 0 && original.id != canonical.id).then_some((original.id, canonical.id)) + }) + .collect(); + Ok((schema, field_id_remap)) +} + +type DataFileIdentity = (Option, String); + +fn retained_file_identities( + dataset: Option<&mut BlockingDataset>, + read_version: u64, +) -> Result> { + let Some(dataset) = dataset else { + return Ok(HashSet::new()); + }; + let collect = |dataset: &BlockingDataset| { + dataset + .inner + .manifest() + .fragments + .iter() + .flat_map(|fragment| fragment.referenced_lance_files()) + .map(|file| (file.base_id, file.path.clone())) + .collect() + }; + if dataset.inner.version().version == read_version { + Ok(collect(dataset)) + } else { + let read_dataset = dataset.checkout_version(read_version)?; + Ok(collect(&read_dataset)) + } +} + +fn remap_fragment_field_ids( + fragments: &mut [Fragment], + field_id_remap: &HashMap, + retained_files: &HashSet, +) { + if field_id_remap.is_empty() { + return; + } + for fragment in fragments { + for file in fragment.referenced_lance_files_mut() { + if retained_files.contains(&(file.base_id, file.path.clone())) { + continue; + } + for field_id in Arc::make_mut(&mut file.fields) { + if let Some(canonical_id) = field_id_remap.get(field_id) { + *field_id = *canonical_id; + } + } + } + } } trait SchemaExt { @@ -1189,7 +1254,7 @@ fn convert_to_rust_operation( let op_name = env.get_string_from_method(java_operation, "name")?; let op = match op_name.as_str() { "Project" => { - let schema = convert_schema_from_operation( + let (schema, _) = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1197,7 +1262,7 @@ fn convert_to_rust_operation( "BufferAllocator is required for Project operations".to_string(), ) })?, - dataset, + dataset.as_deref_mut(), read_version, false, false, @@ -1314,7 +1379,7 @@ fn convert_to_rust_operation( } } "Overwrite" => { - let fragments: Vec = + let mut fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; @@ -1327,7 +1392,7 @@ fn convert_to_rust_operation( to_rust_map(env, &config_upsert_values) }, )?; - let schema = convert_schema_from_operation( + let (schema, field_id_remap) = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1340,6 +1405,7 @@ fn convert_to_rust_operation( true, true, )?; + remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); Operation::Overwrite { fragments, schema, @@ -1481,11 +1547,11 @@ fn convert_to_rust_operation( Operation::DataReplacement { replacements } } "Merge" => { - let fragments: Vec = + let mut fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; - let schema = convert_schema_from_operation( + let (schema, field_id_remap) = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1493,11 +1559,17 @@ fn convert_to_rust_operation( "BufferAllocator is required for Merge operations".to_string(), ) })?, - dataset, + dataset.as_deref_mut(), read_version, false, true, )?; + let retained_files = if field_id_remap.is_empty() { + HashSet::new() + } else { + retained_file_identities(dataset, read_version)? + }; + remap_fragment_field_ids(&mut fragments, &field_id_remap, &retained_files); Operation::Merge { fragments, preserves_nullability: env diff --git a/java/src/test/java/org/lance/operation/MergeTest.java b/java/src/test/java/org/lance/operation/MergeTest.java index 943a2f0c2ec..11aef49a353 100644 --- a/java/src/test/java/org/lance/operation/MergeTest.java +++ b/java/src/test/java/org/lance/operation/MergeTest.java @@ -22,6 +22,7 @@ import org.lance.ipc.LanceScanner; import org.lance.schema.LanceField; import org.lance.schema.LanceSchema; +import org.lance.schema.SqlExpressions; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.IntVector; @@ -37,13 +38,18 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import java.io.IOException; import java.nio.charset.StandardCharsets; +import java.nio.file.Files; import java.nio.file.Path; +import java.nio.file.StandardCopyOption; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.stream.Stream; public class MergeTest extends OperationTestBase { @@ -431,6 +437,59 @@ void testMergeRejectsAmbiguousRawArrowFieldIds(@TempDir Path tempDir) throws Exc } } + @Test + void testLegacyMergeInheritsNonContiguousFieldIds(@TempDir Path tempDir) throws Exception { + Path source = + Path.of("..", "test_data", "v0.10.5", "corrupt_schema").toAbsolutePath().normalize(); + Path datasetPath = tempDir.resolve("legacy"); + copyDirectory(source, datasetPath); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset legacy = Dataset.open(datasetPath.toString(), allocator)) { + legacy.dropColumns(Collections.singletonList("y")); + long baseVersion = legacy.version(); + legacy.addColumns( + new SqlExpressions.Builder().withExpression("z", "x + 1").build(), Optional.empty()); + + Merge generated; + try (Transaction transaction = legacy.readTransaction().orElseThrow()) { + generated = (Merge) transaction.operation(); + } + + try (Dataset restored = legacy.checkoutVersion(baseVersion)) { + restored.restore(); + try (Transaction transaction = + new Transaction.Builder() + .readVersion(restored.version()) + .operation( + Merge.builder() + .fragments(generated.fragments()) + .schema(generated.schema()) + .build()) + .build(); + Dataset merged = new CommitBuilder(restored).execute(transaction)) { + Assertions.assertEquals(0, findField(merged.getLanceSchema().fields(), "x").getId()); + Assertions.assertEquals(4, findField(merged.getLanceSchema().fields(), "b").getId()); + Assertions.assertEquals(5, findField(merged.getLanceSchema().fields(), "c").getId()); + Assertions.assertEquals(6, findField(merged.getLanceSchema().fields(), "z").getId()); + } + } + } + } + + private void copyDirectory(Path source, Path target) throws IOException { + try (Stream paths = Files.walk(source)) { + for (Path path : (Iterable) paths::iterator) { + Path destination = target.resolve(source.relativize(path)); + if (Files.isDirectory(path)) { + Files.createDirectories(destination); + } else { + Files.copy(path, destination, StandardCopyOption.REPLACE_EXISTING); + } + } + } + } + private Map fieldMeta(int fieldId) { Map idMeta = new HashMap<>(); idMeta.put("lance:field_id", String.valueOf(fieldId)); From d199a9a20b5f123bf7ec846f065bc222f0ed144c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 03:48:18 +0800 Subject: [PATCH 08/24] fix(java): validate project field identities --- java/lance-jni/src/transaction.rs | 65 ++++++--- .../java/org/lance/operation/ProjectTest.java | 125 ++++++++++++++++++ rust/lance-table/src/transaction/validate.rs | 30 ++++- 3 files changed, 202 insertions(+), 18 deletions(-) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index 8c237ea194f..5a631c17d1d 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -916,14 +916,26 @@ fn convert_to_rust_transaction( .build()) } +#[derive(Clone, Copy)] +enum RawArrowFieldIdMode { + AssignMissing, + // Preserve the distinction between explicit IDs and fields without identity metadata. + ExplicitOnly, +} + +struct SchemaConversionOptions { + replaces_all_identities: bool, + commit_allocator_is_authoritative: bool, + raw_field_id_mode: RawArrowFieldIdMode, +} + fn convert_schema_from_operation( env: &mut JNIEnv, java_operation: &JObject, java_allocator: &JObject, dataset: Option<&mut BlockingDataset>, read_version: u64, - replaces_all_identities: bool, - commit_allocator_is_authoritative: bool, + options: SchemaConversionOptions, ) -> Result<(LanceSchema, HashMap)> { let schema_ptr = env .call_method( @@ -937,12 +949,24 @@ fn convert_schema_from_operation( let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) }; let arrow_schema = Schema::try_from(&c_schema)?; - let mut original_schema = LanceSchema::try_from(&arrow_schema).map_err(|e| { - Error::input_error(format!( - "Failed to convert Arrow schema to Lance schema: {}", - e - )) - })?; + let mut original_schema = + if matches!(options.raw_field_id_mode, RawArrowFieldIdMode::ExplicitOnly) { + LanceSchema { + fields: arrow_schema + .fields + .iter() + .map(|field| Field::try_from(field.as_ref())) + .collect::>()?, + metadata: arrow_schema.metadata.clone(), + } + } else { + LanceSchema::try_from(&arrow_schema).map_err(|e| { + Error::input_error(format!( + "Failed to convert Arrow schema to Lance schema: {}", + e + )) + })? + }; let read_context = match dataset { Some(dataset) if dataset.inner.version().version == read_version => Some(( @@ -961,7 +985,7 @@ fn convert_schema_from_operation( None => None, }; - if commit_allocator_is_authoritative + if options.commit_allocator_is_authoritative && read_context .as_ref() .is_none_or(|(_, _, stable_field_ids)| *stable_field_ids) @@ -979,7 +1003,7 @@ fn convert_schema_from_operation( Some(read_schema), Some(max_field_id), stable_field_ids, - replaces_all_identities, + options.replaces_all_identities, )? } else { // New datasets use stable field IDs by default, so Arrow metadata is @@ -1264,8 +1288,11 @@ fn convert_to_rust_operation( })?, dataset.as_deref_mut(), read_version, - false, - false, + SchemaConversionOptions { + replaces_all_identities: false, + commit_allocator_is_authoritative: true, + raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, + }, )?; Operation::Project { preserves_nullability: env @@ -1402,8 +1429,11 @@ fn convert_to_rust_operation( })?, dataset.as_deref_mut(), read_version, - true, - true, + SchemaConversionOptions { + replaces_all_identities: true, + commit_allocator_is_authoritative: true, + raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, + }, )?; remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); Operation::Overwrite { @@ -1561,8 +1591,11 @@ fn convert_to_rust_operation( })?, dataset.as_deref_mut(), read_version, - false, - true, + SchemaConversionOptions { + replaces_all_identities: false, + commit_allocator_is_authoritative: true, + raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, + }, )?; let retained_files = if field_id_remap.is_empty() { HashSet::new() diff --git a/java/src/test/java/org/lance/operation/ProjectTest.java b/java/src/test/java/org/lance/operation/ProjectTest.java index bd3dd7d2960..71e2729a2db 100644 --- a/java/src/test/java/org/lance/operation/ProjectTest.java +++ b/java/src/test/java/org/lance/operation/ProjectTest.java @@ -17,20 +17,29 @@ import org.lance.Dataset; import org.lance.TestUtils; import org.lance.Transaction; +import org.lance.schema.LanceField; import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import java.io.IOException; +import java.nio.file.Files; import java.nio.file.Path; +import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; +import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; public class ProjectTest extends OperationTestBase { @@ -103,4 +112,120 @@ void testPreservesNullabilityEqualityAndRoundTrip(@TempDir Path tempDir) { } } } + + @Test + void testProjectRejectsUnmaterializedStableField(@TempDir Path tempDir) { + String datasetPath = tempDir.resolve("testProjectRejectsUnmaterializedStableField").toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + dataset = testDataset.createEmptyDataset(); + + Field existing = dataset.getSchema().getFields().get(0); + Field unmaterialized = + new Field("unmaterialized", existing.getFieldType(), existing.getChildren()); + try (Transaction transaction = + new Transaction.Builder() + .readVersion(dataset.version()) + .operation( + Project.builder() + .schema(new Schema(Collections.singletonList(unmaterialized))) + .build()) + .build()) { + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> new CommitBuilder(dataset).execute(transaction)); + assertTrue(error.getMessage().contains("writes no data")); + } + } + } + + @Test + void testProjectPreservesExplicitRenameIdentity(@TempDir Path tempDir) { + String datasetPath = tempDir.resolve("testProjectPreservesExplicitRenameIdentity").toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + dataset = testDataset.createEmptyDataset(); + + Field existing = dataset.getSchema().getFields().get(0); + int existingId = + dataset.getLanceSchema().fields().stream() + .filter(field -> field.getName().equals(existing.getName())) + .findFirst() + .map(LanceField::getId) + .orElseThrow(); + Map metadata = new HashMap<>(existing.getMetadata()); + metadata.put("lance:field_id", String.valueOf(existingId)); + Field renamed = + new Field( + "renamed", + new FieldType( + existing.isNullable(), existing.getType(), existing.getDictionary(), metadata), + existing.getChildren()); + + try (Transaction transaction = + new Transaction.Builder() + .readVersion(dataset.version()) + .operation( + Project.builder() + .schema(new Schema(Collections.singletonList(renamed))) + .build()) + .build(); + Dataset committed = new CommitBuilder(dataset).execute(transaction)) { + LanceField committedField = committed.getLanceSchema().fields().get(0); + assertEquals("renamed", committedField.getName()); + assertEquals(existingId, committedField.getId()); + } + } + } + + @Test + void testLegacyProjectPreservesNonContiguousFieldIds(@TempDir Path tempDir) throws Exception { + Path source = + Path.of("..", "test_data", "v0.10.5", "corrupt_schema").toAbsolutePath().normalize(); + Path datasetPath = tempDir.resolve("legacy"); + copyDirectory(source, datasetPath); + + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset legacy = Dataset.open(datasetPath.toString(), allocator)) { + legacy.dropColumns(Collections.singletonList("y")); + List projectedFields = new ArrayList<>(legacy.getSchema().getFields()); + Collections.reverse(projectedFields); + + try (Transaction transaction = + new Transaction.Builder() + .readVersion(legacy.version()) + .operation(Project.builder().schema(new Schema(projectedFields)).build()) + .build(); + Dataset projected = new CommitBuilder(legacy).execute(transaction)) { + List fields = projected.getLanceSchema().fields(); + assertEquals(5, findField(fields, "c").getId()); + assertEquals(4, findField(fields, "b").getId()); + assertEquals(0, findField(fields, "x").getId()); + } + } + } + + private void copyDirectory(Path source, Path target) throws IOException { + try (Stream paths = Files.walk(source)) { + for (Path path : (Iterable) paths::iterator) { + Path destination = target.resolve(source.relativize(path)); + if (Files.isDirectory(path)) { + Files.createDirectories(destination); + } else { + Files.copy(path, destination, StandardCopyOption.REPLACE_EXISTING); + } + } + } + } + + private LanceField findField(List fields, String fieldName) { + return fields.stream() + .filter(field -> field.getName().equals(fieldName)) + .findFirst() + .orElseThrow( + () -> new IllegalStateException(String.format("field '%s' not found", fieldName))); + } } diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 3535bd06279..6e1f483c0b4 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -168,7 +168,7 @@ fn clear_replaced_and_new_field_ids( fn canonicalize_raw_project_schema(manifest: &Manifest, schema: &mut Schema) -> Result<()> { let mut unmatched_fields = Vec::new(); for field in &mut schema.fields { - if !canonicalize_field(field, -1, &manifest.schema, None, false) { + if !canonicalize_field(field, -1, &manifest.schema, None, true) { unmatched_fields.push(field.name.clone()); } } @@ -250,7 +250,9 @@ fn canonicalize_field( .flatten() .filter(|base| base.parent_id == parent_id); let base_field = if allow_id_binding && field.id >= 0 { - by_id.filter(|base| base.logical_type == field.logical_type) + by_id + .filter(|base| base.logical_type == field.logical_type) + .or_else(|| same_name.filter(|base| base.logical_type == field.logical_type)) } else { same_name.filter(|base| base.logical_type == field.logical_type) }; @@ -1266,6 +1268,7 @@ mod tests { let manifest = activated_manifest(); let mut schema = one_field_schema(); schema.fields[0].name = "renamed".to_string(); + schema.fields[0].id = -1; schema.metadata.insert( TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), String::new(), @@ -1280,6 +1283,29 @@ mod tests { assert!(err.to_string().contains("writes no data"), "{err}"); } + #[test] + fn canonicalize_raw_arrow_project_preserves_explicit_existing_identity() { + let manifest = activated_manifest(); + let mut schema = one_field_schema(); + schema.fields[0].name = "renamed".to_string(); + schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Project { schema, .. } = operation else { + unreachable!(); + }; + assert_eq!(schema.fields[0].name, "renamed"); + assert_eq!(schema.fields[0].id, 0); + } + #[test] fn canonicalize_raw_arrow_schema_for_legacy_dataset() { let manifest = manifest_with_file_fields(one_field_schema(), vec![0]); From 3b041d1668f35ad942475ea5df10c271cec06fee Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 21 Aug 2026 04:23:08 +0800 Subject: [PATCH 09/24] test(java): verify legacy project values --- .../java/org/lance/operation/ProjectTest.java | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/java/src/test/java/org/lance/operation/ProjectTest.java b/java/src/test/java/org/lance/operation/ProjectTest.java index 71e2729a2db..a62db31eb9f 100644 --- a/java/src/test/java/org/lance/operation/ProjectTest.java +++ b/java/src/test/java/org/lance/operation/ProjectTest.java @@ -17,9 +17,12 @@ import org.lance.Dataset; import org.lance.TestUtils; import org.lance.Transaction; +import org.lance.ipc.LanceScanner; import org.lance.schema.LanceField; import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.FieldType; import org.apache.arrow.vector.types.pojo.Schema; @@ -204,6 +207,21 @@ void testLegacyProjectPreservesNonContiguousFieldIds(@TempDir Path tempDir) thro assertEquals(5, findField(fields, "c").getId()); assertEquals(4, findField(fields, "b").getId()); assertEquals(0, findField(fields, "x").getId()); + + try (LanceScanner scanner = projected.newScan(); + ArrowReader reader = scanner.scanBatches()) { + assertTrue(reader.loadNextBatch()); + VectorSchemaRoot root = reader.getVectorSchemaRoot(); + assertEquals("c", root.getSchema().getFields().get(0).getName()); + assertEquals("b", root.getSchema().getFields().get(1).getName()); + assertEquals("x", root.getSchema().getFields().get(2).getName()); + assertEquals(0L, root.getVector("c").getObject(0)); + assertEquals(0L, root.getVector("b").getObject(0)); + assertEquals(0L, root.getVector("x").getObject(0)); + assertEquals(5L, root.getVector("c").getObject(1)); + assertEquals(4L, root.getVector("b").getObject(1)); + assertEquals(1L, root.getVector("x").getObject(1)); + } } } } From 981ff8e563961c79cfe26db914c949cd23da7d6e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 31 Aug 2026 17:39:40 +0800 Subject: [PATCH 10/24] fix: address stable field ID review feedback --- docs/src/format/table/schema.md | 16 +- java/lance-jni/src/transaction.rs | 218 +++------------- .../org/lance/operation/OverwriteTest.java | 2 +- protos/table.proto | 5 +- .../python/tests/compat/compat_decorator.py | 55 +++- .../tests/compat/test_scalar_indices.py | 9 +- .../tests/compat/test_vector_indices.py | 5 +- python/python/tests/test_dataset.py | 26 +- .../src/transaction/manifest_build.rs | 23 +- rust/lance-table/src/transaction/validate.rs | 39 +-- rust/lance/src/dataset.rs | 11 +- rust/lance/src/dataset/fragment.rs | 26 +- .../src/dataset/tests/dataset_migrations.rs | 47 +++- rust/lance/src/io/commit.rs | 240 +++++++++++++++++- 14 files changed, 433 insertions(+), 289 deletions(-) diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index 2ad0d25f5fb..1bb050ebcfa 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -238,9 +238,15 @@ files. Activation provides a forward guarantee only; it cannot reconstruct ident dropped or reused in older snapshots. New datasets activate this contract in their initial manifest and require both readers and writers -to understand it. Existing legacy datasets remain unchanged until an explicit migration commit; -controlled deployments may use a writer-only migration only after every pre-gate writer has been -retired. +to understand it. The reader gate is required because older Lance versions do not enforce unknown +writer feature flags on every mutation path; allowing them to open the dataset could let a legacy +commit clear the writer gate and resume reusable allocation. Existing legacy datasets remain +unchanged until an explicit migration commit. Controlled deployments may use writer-only migration +only after every pre-gate writer has been retired. + +Activation is one-way within a branch ancestry. After migration, restore cannot target a version +from before activation because that version does not carry the high-water mark needed to preserve +retired identities. ### Field ID Properties @@ -329,8 +335,8 @@ Field IDs enable efficient schema evolution: - **Reorder Columns**: Change field order in schema; IDs remain the same - **Metadata or Nullability Change**: Preserve the field ID - **Type Replacement**: Allocate a new field ID and retire the old identity -- **Overwrite**: On an activated dataset, replace all field identities with fresh IDs above the - current high-water mark +- **Overwrite**: Preserve compatible logical identities; allocate new IDs for added fields and type + replacements The use of field IDs ensures that data files can be correctly interpreted even as the schema changes over time. diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index 5a631c17d1d..d7782269dec 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -924,8 +924,6 @@ enum RawArrowFieldIdMode { } struct SchemaConversionOptions { - replaces_all_identities: bool, - commit_allocator_is_authoritative: bool, raw_field_id_mode: RawArrowFieldIdMode, } @@ -985,10 +983,9 @@ fn convert_schema_from_operation( None => None, }; - if options.commit_allocator_is_authoritative - && read_context - .as_ref() - .is_none_or(|(_, _, stable_field_ids)| *stable_field_ids) + if read_context + .as_ref() + .is_none_or(|(_, _, stable_field_ids)| *stable_field_ids) { original_schema.metadata.insert( TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), @@ -997,19 +994,9 @@ fn convert_schema_from_operation( return Ok((original_schema, HashMap::new())); } - let schema = if let Some((read_schema, max_field_id, stable_field_ids)) = read_context { - LanceSchema::from_arrow_schema( - &arrow_schema, - Some(read_schema), - Some(max_field_id), - stable_field_ids, - options.replaces_all_identities, - )? - } else { - // New datasets use stable field IDs by default, so Arrow metadata is - // not an allocation authority even before a manifest exists. - LanceSchema::from_arrow_schema(&arrow_schema, None, None, true, true)? - }; + let (read_schema, max_field_id, _) = read_context.expect("legacy dataset context"); + let schema = + LanceSchema::from_arrow_schema(&arrow_schema, Some(read_schema), Some(max_field_id))?; let field_id_remap = original_schema .fields_pre_order() @@ -1091,17 +1078,11 @@ trait SchemaExt { max_existing_id: Option, ) -> Result<()>; - /// Create a schema from `arrow_schema`. - /// - /// For an existing dataset, Arrow field-ID metadata is not an allocation - /// authority: matching identities inherit IDs from `base_schema` and every - /// new identity is allocated above `max_existing_id`. + /// Create a schema from `arrow_schema` using the legacy Java conversion rules. fn from_arrow_schema( arrow_schema: &Schema, base_schema: Option, max_existing_id: Option, - dataset_allocator_is_authoritative: bool, - replaces_all_identities: bool, ) -> Result; } @@ -1132,8 +1113,6 @@ impl SchemaExt for LanceSchema { arrow_schema: &Schema, base_schema: Option, max_existing_id: Option, - dataset_allocator_is_authoritative: bool, - replaces_all_identities: bool, ) -> Result { let mut schema = Self { fields: arrow_schema @@ -1143,25 +1122,7 @@ impl SchemaExt for LanceSchema { .collect::>()?, metadata: arrow_schema.metadata.clone(), }; - if dataset_allocator_is_authoritative { - if replaces_all_identities || base_schema.is_none() { - schema.try_reassign_field_ids(max_existing_id)?; - } else if let Some(base_schema) = &base_schema { - for field in &mut schema.fields { - let base_field = if field.id >= 0 { - base_schema - .field_by_id(field.id) - .or_else(|| base_schema.field(&field.name)) - } else { - base_schema.field(&field.name) - }; - field.set_stable_field_id_from_field(-1, base_schema, base_field)?; - } - schema.try_set_field_id(max_existing_id)?; - } - } else { - schema.set_field_id_from_schema(base_schema, max_existing_id)?; - } + schema.set_field_id_from_schema(base_schema, max_existing_id)?; schema.validate()?; schema.verify_primary_key()?; @@ -1170,15 +1131,6 @@ impl SchemaExt for LanceSchema { } trait FieldExt { - fn clear_field_ids(&mut self); - - fn set_stable_field_id_from_field( - &mut self, - parent_id: i32, - base_schema: &LanceSchema, - base_field: Option<&Field>, - ) -> lance_core::Result<()>; - /// Recursively set field ID and parent ID for this field and all its children. fn set_field_id_from_field( &mut self, @@ -1188,51 +1140,6 @@ trait FieldExt { } impl FieldExt for Field { - fn clear_field_ids(&mut self) { - self.id = -1; - self.parent_id = -1; - for child in &mut self.children { - child.clear_field_ids(); - } - } - - fn set_stable_field_id_from_field( - &mut self, - parent_id: i32, - base_schema: &LanceSchema, - base_field: Option<&Field>, - ) -> lance_core::Result<()> { - let Some(base_field) = base_field.filter(|base| base.logical_type == self.logical_type) - else { - self.clear_field_ids(); - self.parent_id = parent_id; - return Ok(()); - }; - - self.id = base_field.id; - self.parent_id = parent_id; - for child in &mut self.children { - let base_child = if child.id >= 0 { - base_schema - .field_by_id(child.id) - .filter(|base| base.parent_id == base_field.id) - .or_else(|| { - base_field - .children - .iter() - .find(|base| base.name == child.name) - }) - } else { - base_field - .children - .iter() - .find(|base| base.name == child.name) - }; - child.set_stable_field_id_from_field(self.id, base_schema, base_child)?; - } - Ok(()) - } - fn set_field_id_from_field( &mut self, parent_id: i32, @@ -1289,8 +1196,6 @@ fn convert_to_rust_operation( dataset.as_deref_mut(), read_version, SchemaConversionOptions { - replaces_all_identities: false, - commit_allocator_is_authoritative: true, raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, }, )?; @@ -1430,8 +1335,6 @@ fn convert_to_rust_operation( dataset.as_deref_mut(), read_version, SchemaConversionOptions { - replaces_all_identities: true, - commit_allocator_is_authoritative: true, raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, }, )?; @@ -1592,8 +1495,6 @@ fn convert_to_rust_operation( dataset.as_deref_mut(), read_version, SchemaConversionOptions { - replaces_all_identities: false, - commit_allocator_is_authoritative: true, raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, }, )?; @@ -1924,20 +1825,6 @@ mod tests { pub const LANCE_FIELD_ID_KEY: &str = "lance:field_id"; - #[test] - fn java_schema_conversion_rejects_field_id_exhaustion() { - let arrow_schema = ArrowSchema::new(vec![ArrowField::new( - "new_field", - ArrowDataType::Int32, - false, - )]); - - let err = LanceSchema::from_arrow_schema(&arrow_schema, None, Some(i32::MAX), true, true) - .unwrap_err(); - - assert!(err.to_string().contains("IDs are exhausted"), "{err}"); - } - #[test] fn legacy_java_schema_conversion_preserves_arrow_field_ids() { let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); @@ -1958,34 +1845,12 @@ mod tests { ]); let schema = - LanceSchema::from_arrow_schema(&arrow_schema, Some(base_schema), Some(0), false, false) - .unwrap(); + LanceSchema::from_arrow_schema(&arrow_schema, Some(base_schema), Some(0)).unwrap(); assert_eq!(schema.field("a").unwrap().id, 5); assert_eq!(schema.field("b").unwrap().id, 9); } - #[test] - fn stable_java_schema_conversion_preserves_rename_identity() { - let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); - base.id = 10; - let base_schema = LanceSchema { - fields: vec![base], - metadata: HashMap::new(), - }; - let arrow_schema = ArrowSchema::new(vec![ - ArrowField::new("renamed", ArrowDataType::Int32, false).with_metadata(HashMap::from([ - (LANCE_FIELD_ID_KEY.to_string(), "10".to_string()), - ])), - ]); - - let schema = - LanceSchema::from_arrow_schema(&arrow_schema, Some(base_schema), Some(10), true, false) - .unwrap(); - - assert_eq!(schema.field("renamed").unwrap().id, 10); - } - #[test] fn test_create_schema_from_arrow() { // base_schema has an existing field id @@ -2072,19 +1937,22 @@ mod tests { metadata: HashMap::from([("base_schema_k".to_string(), "base_schema_v".to_string())]), }; - // new_schema specifies field IDs in Arrow metadata to prove they are - // ignored for an existing dataset: - // - field a: inherits from base_schema + // new_schema specifies: + // - field a: manual field id // - field b: no id -> should inherit from base_schema // - field c: new field -> should be assigned based on max_field_id - // - existing nested fields inherit; new nested fields use the allocator + // - struct s: parent+child(x) manual, child(y) inherit, child(z) max_field_id + // - list l: parent manual, child(item) inherit + // - list l2: parent manual, child(item) max_field_id + // - map m: parent manual, child(entries/key/value) inherit + // - map m2: parent manual, child(entries/key/value) max_field_id let mut a_meta = HashMap::new(); a_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "5".to_string()); let arrow_a = ArrowField::new("a", ArrowDataType::Int32, false).with_metadata(a_meta); let arrow_b = ArrowField::new("b", ArrowDataType::Int32, false); let arrow_c = ArrowField::new("c", ArrowDataType::Int32, false); - // struct s: metadata IDs on the parent and child x are ignored. + // struct s: manual parent + manual child x let mut s_meta = HashMap::new(); s_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "50".to_string()); let mut x_meta = HashMap::new(); @@ -2100,7 +1968,7 @@ mod tests { ) .with_metadata(s_meta); - // list l: the metadata parent ID is ignored; both identities inherit. + // list l: parent manual, item inherit let mut l_meta = HashMap::new(); l_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "60".to_string()); let arrow_l = ArrowField::new( @@ -2114,7 +1982,7 @@ mod tests { ) .with_metadata(l_meta); - // list l2: the metadata parent ID is ignored; both identities are new. + // list l2: parent manual, item max_field_id (no base match) let mut l2_meta = HashMap::new(); l2_meta.insert(LANCE_FIELD_ID_KEY.to_string(), "61".to_string()); let arrow_l2 = ArrowField::new( @@ -2128,7 +1996,7 @@ mod tests { ) .with_metadata(l2_meta); - // map m: the metadata parent ID is ignored; all identities inherit. + // map m: parent manual, entries/key/value inherit let map_entries = ArrowField::new( "entries", ArrowDataType::Struct(ArrowFields::from(vec![ @@ -2142,7 +2010,7 @@ mod tests { let arrow_m = ArrowField::new("m", ArrowDataType::Map(Arc::new(map_entries), false), true) .with_metadata(m_meta); - // map m2: metadata IDs are ignored; the full new subtree is allocated. + // map m2: parent manual, entries/key/value max_field_id let map_entries = ArrowField::new( "entries", ArrowDataType::Struct(ArrowFields::from(vec![ @@ -2164,18 +2032,12 @@ mod tests { HashMap::from([("new_schema_k".to_string(), "new_schema_v".to_string())]), ); - let schema = LanceSchema::from_arrow_schema( - &arrow_schema, - Some(base_schema), - Some(100), - true, - false, - ) - .unwrap(); + let schema = + LanceSchema::from_arrow_schema(&arrow_schema, Some(base_schema), Some(100)).unwrap(); - // 1. Arrow metadata cannot override an existing identity. + // 1. Manually specified field id let got_a = schema.field("a").unwrap(); - assert_eq!(got_a.id, 10); + assert_eq!(got_a.id, 5); assert!(!got_a.metadata.contains_key(LANCE_FIELD_ID_KEY)); // 2. Inherit field id + metadata from base_schema (field b) @@ -2186,31 +2048,31 @@ mod tests { let got_c = schema.field("c").unwrap(); assert_eq!(got_c.id, 101); - // 4. struct: existing parent and children inherit; z is allocated. + // 4. struct: parent+child(x) manual, child(y) inherit, child(z) max_field_id let got_s = schema.field("s").unwrap(); - assert_eq!(got_s.id, 20); + assert_eq!(got_s.id, 50); let got_sx = schema.field("s.x").unwrap(); - assert_eq!(got_sx.id, 21); + assert_eq!(got_sx.id, 51); let got_sy = schema.field("s.y").unwrap(); assert_eq!(got_sy.id, 22); let got_sz = schema.field("s.z").unwrap(); assert_eq!(got_sz.id, 102); - // 5. list l: parent and item inherit. + // 5. list l: parent manual, item inherit let got_l = schema.field("l").unwrap(); - assert_eq!(got_l.id, 30); + assert_eq!(got_l.id, 60); let got_li = schema.field("l.item").unwrap(); assert_eq!(got_li.id, 31); - // 6. list l2: parent and item are newly allocated. + // 6. list l2: parent manual, item max_field_id let got_l2 = schema.field("l2").unwrap(); - assert_eq!(got_l2.id, 103); + assert_eq!(got_l2.id, 61); let got_l2i = schema.field("l2.item").unwrap(); - assert_eq!(got_l2i.id, 104); + assert_eq!(got_l2i.id, 103); - // 7. map m: parent, entries, key, and value inherit. + // 7. map m: parent manual, entries/key/value inherit let got_m = schema.field("m").unwrap(); - assert_eq!(got_m.id, 40); + assert_eq!(got_m.id, 70); let got_me = schema.field("m.entries").unwrap(); assert_eq!(got_me.id, 41); let got_mk = schema.field("m.entries.key").unwrap(); @@ -2218,15 +2080,15 @@ mod tests { let got_mv = schema.field("m.entries.value").unwrap(); assert_eq!(got_mv.id, 43); - // 8. map m2: the full new subtree is allocated in pre-order. + // 8. map m2: parent manual, entries/key/value max_field_id let got_m2 = schema.field("m2").unwrap(); - assert_eq!(got_m2.id, 105); + assert_eq!(got_m2.id, 71); let got_m2e = schema.field("m2.entries").unwrap(); - assert_eq!(got_m2e.id, 106); + assert_eq!(got_m2e.id, 104); let got_m2k = schema.field("m2.entries.key").unwrap(); - assert_eq!(got_m2k.id, 107); + assert_eq!(got_m2k.id, 105); let got_m2v = schema.field("m2.entries.value").unwrap(); - assert_eq!(got_m2v.id, 108); + assert_eq!(got_m2v.id, 106); // 9. Schema metadata: when new_schema.metadata is non-empty, use new_schema metadata assert_eq!( diff --git a/java/src/test/java/org/lance/operation/OverwriteTest.java b/java/src/test/java/org/lance/operation/OverwriteTest.java index 26bfa1f911e..24317271a93 100644 --- a/java/src/test/java/org/lance/operation/OverwriteTest.java +++ b/java/src/test/java/org/lance/operation/OverwriteTest.java @@ -134,7 +134,7 @@ void testOverwrite(@TempDir Path tempDir) throws Exception { Collections.singletonMap("config_key", "config_value"), committedOverwrite.configUpsertValues().orElseThrow()); assertArrayEquals( - new int[] {4, 5}, + new int[] {0, 1}, committedOverwrite.fragments().get(0).getFiles().get(0).getFields()); } } diff --git a/protos/table.proto b/protos/table.proto index 9df1e6f0850..02894066c2b 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -125,9 +125,8 @@ message Manifest { // MemWAL SSTables still hold. Writers must refuse it too: one that does not // maintain index_catchup can change an index without withdrawing the // position recorded for it. Setting it is one-way. - // * 1 << 8: field IDs use a persistent allocation high-water mark and are - // never reused. This bit is set for readers only by fail-closed activation; - // writers always require the corresponding writer bit. + // * 1 << 8: readers must understand the stable field-ID contract. This is an + // optional fail-closed policy; stable field IDs always require the writer bit. uint64 reader_feature_flags = 9; // Feature flags for writers. diff --git a/python/python/tests/compat/compat_decorator.py b/python/python/tests/compat/compat_decorator.py index bb79d58c716..d9c62b7bacc 100644 --- a/python/python/tests/compat/compat_decorator.py +++ b/python/python/tests/compat/compat_decorator.py @@ -130,8 +130,6 @@ def last_beta_release(): if LAST_BETA_RELEASE is not None: VERSIONS.append(LAST_BETA_RELEASE) -STABLE_FIELD_IDS_MIN_READER_VERSION = Version("11.0.0b16") - class UpgradeDowngradeTest: """Base class for compatibility tests. @@ -155,10 +153,18 @@ def skip_read_after_current_write(self, version: str) -> bool: """Return True to skip the old-version read after current-version writes.""" return False + def expect_read_after_current_write_failure(self, version: str) -> bool: + """Return True when the old reader must reject current-version data.""" + return False + def skip_write_after_current_write(self, version: str) -> bool: """Return True to skip the old-version write after current-version writes.""" return False + def expect_write_after_current_write_failure(self, version: str) -> bool: + """Return True when the old writer must reject current-version data.""" + return False + def skip_downgrade(self, version: str) -> bool: """Return True to skip the current-write -> old-read downgrade test.""" return False @@ -173,13 +179,18 @@ def compat_env(self, version: str, method_name: str) -> Dict[str, str]: class DatasetUpgradeDowngradeTest(UpgradeDowngradeTest): - """Compatibility contract for datasets with stable field IDs enabled.""" + """Compatibility contract for new datasets with stable field IDs.""" - def skip_downgrade(self, version: str) -> bool: - # New datasets require the stable-field-ID reader feature. Versions - # released before the feature landed must fail closed instead of - # interpreting the dataset with reusable field identities. - return Version(version) < STABLE_FIELD_IDS_MIN_READER_VERSION + def expect_read_after_current_write_failure(self, version: str) -> bool: + # Every released version selected by this suite predates the reader + # feature bit. Once a supporting release enters the matrix this fails + # visibly, instead of silently skipping a valid reader combination. + return True + + def expect_write_after_current_write_failure(self, version: str) -> bool: + # Opening the dataset must fail before an old writer can enter a commit + # path that does not enforce unknown writer feature flags. + return True @contextmanager @@ -336,7 +347,7 @@ def _make_test_function(cls, param_names, test_type): if test_type == "downgrade": func_body = f''' def test_func({sig_params}): - """Test that old Lance version can read data written by current version.""" + """Test the old-version contract for data written by the current version.""" from pathlib import Path obj = cls(tmp_path / "data.lance", {init_params}) obj.compat_version = version @@ -350,9 +361,31 @@ def test_func({sig_params}): # Old version: verify can read venv = venv_factory.get_venv(version) if not obj.skip_read_after_current_write(version): - venv.execute_method(obj, "check_read", obj.compat_env(version, "check_read")) + if obj.expect_read_after_current_write_failure(version): + with pytest.raises( + RuntimeError, + match="cannot be read by this version|Flags", + ): + venv.execute_method( + obj, "check_read", obj.compat_env(version, "check_read") + ) + else: + venv.execute_method( + obj, "check_read", obj.compat_env(version, "check_read") + ) if not obj.skip_write_after_current_write(version): - venv.execute_method(obj, "check_write", obj.compat_env(version, "check_write")) + if obj.expect_write_after_current_write_failure(version): + with pytest.raises( + RuntimeError, + match="cannot be (?:read|written) by this version|Flags", + ): + venv.execute_method( + obj, "check_write", obj.compat_env(version, "check_write") + ) + else: + venv.execute_method( + obj, "check_write", obj.compat_env(version, "check_write") + ) ''' else: # upgrade_downgrade func_body = f''' diff --git a/python/python/tests/compat/test_scalar_indices.py b/python/python/tests/compat/test_scalar_indices.py index bb7468266d3..cebbdd72b9d 100644 --- a/python/python/tests/compat/test_scalar_indices.py +++ b/python/python/tests/compat/test_scalar_indices.py @@ -16,10 +16,7 @@ import lance import pyarrow as pa -from .compat_decorator import ( - DatasetUpgradeDowngradeTest, - compat_test, -) +from .compat_decorator import DatasetUpgradeDowngradeTest, compat_test from .util import safe_data_storage_version @@ -264,7 +261,7 @@ def check_write(self): def skip_downgrade(self, version: str) -> bool: # In 0.X the zonemap index did not properly handle NULL in filters - return super().skip_downgrade(version) or version.startswith("0.") + return version.startswith("0.") @compat_test(min_version="0.36.0") @@ -379,7 +376,7 @@ def check_write(self): ds.optimize.compact_files() def skip_downgrade(self, version: str) -> bool: - return super().skip_downgrade(version) or version.startswith("0.") + return version.startswith("0.") def current_env(self, method_name: str) -> dict[str, str]: if method_name == "create": diff --git a/python/python/tests/compat/test_vector_indices.py b/python/python/tests/compat/test_vector_indices.py index 04b9c675259..0b7239ec5fb 100644 --- a/python/python/tests/compat/test_vector_indices.py +++ b/python/python/tests/compat/test_vector_indices.py @@ -17,10 +17,7 @@ import pyarrow as pa import pyarrow.compute as pc -from .compat_decorator import ( - DatasetUpgradeDowngradeTest, - compat_test, -) +from .compat_decorator import DatasetUpgradeDowngradeTest, compat_test from .util import safe_data_storage_version diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 247e5b14a91..c43c0dff441 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5232,6 +5232,25 @@ def test_detached_commits(tmp_path: Path): assert detached2.to_table() == pa.table({"x": [0, 1, 3]}) +def test_detached_raw_arrow_merge_strips_transaction_metadata(tmp_path: Path): + dataset = lance.write_dataset(pa.table({"x": [0, 1]}), tmp_path) + fragment = dataset.get_fragments()[0].metadata + with pytest.deprecated_call(): + operation = lance.LanceOperation.Merge([fragment], dataset.schema, True) + + detached = lance.LanceDataset.commit( + dataset, + operation, + read_version=dataset.version, + detached=True, + ) + + assert detached.to_table() == dataset.to_table() + assert b"lance:transaction_schema_source_raw_arrow" not in ( + detached.schema.metadata or {} + ) + + def test_dataset_drop(tmp_path: Path): table = pa.table({"x": [0]}) lance.write_dataset(table, tmp_path) @@ -5473,11 +5492,11 @@ def test_data_overlay_round_trips_through_fragment_metadata( assert restored.overlays[0].offsets == [1, 4] assert restored.overlays[0].committed_version == overlay_version - # A commit that round-trips the fragment (here a Merge) must keep the + # A commit that round-trips the fragment (here an Overwrite) must keep the # overlays, so the overlay still resolves on read instead of being dropped. dataset = lance.LanceDataset.commit( dataset, - lance.LanceOperation.Merge([restored], dataset.lance_schema, True), + lance.LanceOperation.Overwrite(dataset.schema, [restored]), read_version=dataset.version, ) result = dataset.to_table() @@ -5594,6 +5613,9 @@ def test_schema_project_raw_arrow_cannot_allocate_field_id(tmp_path: Path): assert dataset.lance_schema.fields()[0].id() == 0 assert dataset.to_table() == pa.table({"a": pa.array([1, 2], pa.int32())}) + assert b"lance:transaction_schema_source_raw_arrow" not in ( + dataset.schema.metadata or {} + ) def test_overwrite_rejects_ambiguous_raw_arrow_field_ids(tmp_path: Path): diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 07142822484..ffe188336b6 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1388,18 +1388,19 @@ impl Transaction { } if let Some(require_reader) = config.stable_field_id_migration_requires_reader { - if current_manifest + let already_active = current_manifest .map(|manifest| manifest.uses_stable_field_ids()) - .unwrap_or(false) - { - return Err(Error::invalid_input( - "Stable field IDs are already active for this dataset", - )); - } - manifest.activate_stable_field_ids(); - manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - if require_reader { - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + .unwrap_or(false); + if already_active { + if require_reader { + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + } + } else { + manifest.activate_stable_field_ids(); + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + if require_reader { + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + } } } diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 6e1f483c0b4..a57b842c620 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -70,8 +70,7 @@ pub fn canonicalize_stable_field_ids( Operation::Overwrite { schema, fragments, .. } => { - let field_id_remap = - canonicalize_schema(manifest, schema, true, false, raw_arrow_schema)?; + let field_id_remap = canonicalize_schema(manifest, schema, true, raw_arrow_schema)?; remap_fragment_field_ids(fragments, &field_id_remap, &HashSet::new())?; } Operation::Project { schema, .. } if raw_arrow_schema => { @@ -93,8 +92,7 @@ pub fn canonicalize_stable_field_ids( .map(|file| (file.base_id, file.path.clone())) .collect(); if raw_arrow_schema { - let field_id_remap = - canonicalize_schema(Some(manifest), schema, false, false, true)?; + let field_id_remap = canonicalize_schema(Some(manifest), schema, false, true)?; remap_fragment_field_ids(fragments, &field_id_remap, &retained_files)?; } canonicalize_merge_replacements(manifest, schema, fragments, &retained_files)?; @@ -186,7 +184,6 @@ fn canonicalize_raw_project_schema(manifest: &Manifest, schema: &mut Schema) -> fn canonicalize_schema( manifest: Option<&Manifest>, schema: &mut Schema, - replaces_all_identities: bool, allow_id_binding: bool, remap_raw_source_ids: bool, ) -> Result { @@ -200,7 +197,7 @@ fn canonicalize_schema( }; let max_existing_id = manifest.map(Manifest::max_field_id); - if replaces_all_identities || manifest.is_none() { + if manifest.is_none() { schema.try_reassign_field_ids(max_existing_id)?; } else if let Some(manifest) = manifest { for field in &mut schema.fields { @@ -536,17 +533,14 @@ fn validate_stable_field_id_operation(manifest: &Manifest, operation: &Operation } validate_stable_field_id_manifest(manifest)?; - let (schema, replaces_all_identities) = match operation { - Operation::Overwrite { schema, .. } => (schema, true), - Operation::Merge { schema, .. } | Operation::Project { schema, .. } => (schema, false), - _ => return Ok(()), + let (Operation::Overwrite { schema, .. } + | Operation::Merge { schema, .. } + | Operation::Project { schema, .. }) = operation + else { + return Ok(()); }; schema.validate()?; - if replaces_all_identities { - return validate_dense_new_field_ids(manifest, schema.fields_pre_order()); - } - for field in schema.fields_pre_order() { let Some(prior_field) = manifest.schema.field_by_id(field.id) else { continue; @@ -1083,7 +1077,7 @@ mod tests { } #[test] - fn stable_field_ids_require_overwrite_to_replace_every_identity() { + fn stable_field_ids_allow_overwrite_to_preserve_compatible_identity() { let manifest = activated_manifest(); let schema = manifest.schema.clone(); let operation = Operation::Overwrite { @@ -1093,12 +1087,7 @@ mod tests { initial_bases: None, }; - let err = validate_operation(Some(&manifest), &operation).unwrap_err(); - - assert!( - err.to_string().contains("densely allocated from 1"), - "{err}" - ); + validate_operation(Some(&manifest), &operation).unwrap(); } #[test] @@ -1121,8 +1110,8 @@ mod tests { else { unreachable!(); }; - assert_eq!(schema.fields[0].id, 1); - assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]); + assert_eq!(schema.fields[0].id, 0); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]); } #[test] @@ -1149,8 +1138,8 @@ mod tests { else { unreachable!(); }; - assert_eq!(schema.fields[0].id, 1); - assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]); + assert_eq!(schema.fields[0].id, 0); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]); } #[test] diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index d1a48c2b470..87d6b328d86 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -3248,7 +3248,9 @@ impl Dataset { /// The activation commit records the current maximum referenced field ID as /// a persistent high-water mark. Later schema changes allocate above it even /// after fields and their files are dropped. New datasets already use this - /// contract by default. Activation is one-way and idempotent. + /// contract by default with reader and writer gates. `ReadersAndWriters` + /// can add a fail-closed reader gate to a writer-only activated dataset. + /// Activation and reader gating are one-way and idempotent. /// /// ``` /// # use lance::{Dataset, Result}; @@ -3264,7 +3266,11 @@ impl Dataset { &mut self, mode: StableFieldIdMigrationMode, ) -> Result<()> { - if self.manifest.uses_stable_field_ids() { + let require_reader = matches!(mode, StableFieldIdMigrationMode::ReadersAndWriters); + let reader_already_required = self.manifest.reader_feature_flags + & lance_table::feature_flags::FLAG_STABLE_FIELD_IDS + != 0; + if self.manifest.uses_stable_field_ids() && (!require_reader || reader_already_required) { return Ok(()); } @@ -3279,7 +3285,6 @@ impl Dataset { }, None, ); - let require_reader = matches!(mode, StableFieldIdMigrationMode::ReadersAndWriters); let new_ds = CommitBuilder::new(Arc::new(self.clone())) .with_max_retries(0) .with_stable_field_id_migration_activation(require_reader) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 6886515e352..e97703f3533 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -5977,35 +5977,29 @@ mod tests { } let mut fragments: Vec = Vec::new(); - let rows_per_fragment = dataset_rows / paths.len(); for (idx, path) in paths.iter().enumerate() { - let f = FileFragment::create_from_file(path, &dataset, idx, Some(rows_per_fragment)) + let f = FileFragment::create_from_file(path, &dataset, idx, None) .await .unwrap(); fragments.push(f) } - let op = Operation::Merge { + let op = Operation::Overwrite { schema: schema.clone(), fragments, - preserves_nullability: true, + config_upsert_values: None, + initial_bases: None, }; - let new_dataset = Dataset::commit( - test_uri, - op, - Some(dataset.version().version), - None, - None, - Default::default(), - false, - ) - .await - .unwrap(); + let new_dataset = + Dataset::commit(test_uri, op, None, None, None, Default::default(), false) + .await + .unwrap(); assert_eq!(new_dataset.count_rows(None).await.unwrap(), dataset_rows); - // Fragments retain the row count supplied when reconstructing them. + // Fragments will have number of rows recorded in metadata, even though + // we passed `None` when constructing the `FileFragment`. let fragments = new_dataset.get_fragments(); assert_eq!(fragments.len(), 5); for f in fragments { diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index fa6e4de0ae4..2721ebb468f 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -346,12 +346,31 @@ async fn test_stable_field_id_migration_repairs_legacy_schema_before_activation( let mut dataset = Dataset::open(&test_dir.path_str()).await.unwrap(); dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::WritersOnly) .await .unwrap(); dataset.validate().await.unwrap(); assert!(dataset.manifest.uses_stable_field_ids()); + assert_eq!( + dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + + let writer_only_version = dataset.version().version; + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .await + .unwrap(); + assert_eq!(dataset.version().version, writer_only_version + 1); + assert_ne!( + dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); } #[tokio::test] @@ -628,6 +647,16 @@ async fn test_new_datasets_use_stable_field_ids_and_migration_is_idempotent() { ); let created_version = dataset.version().version; + dataset + .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) + .await + .unwrap(); + assert_eq!(dataset.version().version, created_version); + assert_ne!( + dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + dataset .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) .await @@ -697,7 +726,7 @@ async fn test_shallow_clone_preserves_stable_field_id_state() { } #[tokio::test] -async fn test_overwrite_replaces_all_stable_field_identities() { +async fn test_overwrite_preserves_compatible_stable_field_identities() { let source_uri = TempStrDir::default(); let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; dataset @@ -728,17 +757,9 @@ async fn test_overwrite_replaces_all_stable_field_identities() { .await .unwrap(); - assert_eq!(overwritten.schema().field("id").unwrap().id, 1); - assert_eq!(overwritten.schema().field("replacement").unwrap().id, 2); - assert_eq!(overwritten.manifest.max_allocated_field_id, Some(2)); - assert!( - overwritten - .manifest - .fragments - .iter() - .flat_map(|fragment| fragment.files.iter()) - .all(|file| file.fields.iter().all(|field_id| *field_id >= 1)) - ); + assert_eq!(overwritten.schema().field("id").unwrap().id, 0); + assert_eq!(overwritten.schema().field("replacement").unwrap().id, 1); + assert_eq!(overwritten.manifest.max_allocated_field_id, Some(1)); } #[tokio::test] diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index c9e8e363ae0..742a59b1dd9 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -1072,6 +1072,9 @@ pub(crate) async fn do_commit_detached_transaction( ))); } validate_stable_field_id_flags(&dataset.manifest)?; + let mut transaction = transaction.clone(); + canonicalize_stable_field_ids(Some(&dataset.manifest), &mut transaction.operation)?; + let transaction = &transaction; validate_detached_stable_field_ids(&dataset.manifest, &transaction.operation)?; validate_operation(Some(&dataset.manifest), &transaction.operation)?; @@ -1443,12 +1446,17 @@ pub(crate) async fn commit_transaction( ))); } validate_stable_field_id_flags(&dataset.manifest)?; - canonicalize_stable_field_ids(Some(&dataset.manifest), &mut transaction.operation)?; - validate_operation(Some(&dataset.manifest), &transaction.operation)?; + + // Keep the rebased transaction as the source operation for the next + // attempt. Canonicalization consumes transient binding provenance, so + // mutating that source would make a retry depend on the first attempt. + let mut attempt_transaction = transaction.clone(); + canonicalize_stable_field_ids(Some(&dataset.manifest), &mut attempt_transaction.operation)?; + validate_operation(Some(&dataset.manifest), &attempt_transaction.operation)?; // Recomputed every attempt: the rebase above may have rewritten the // transaction. - let pb_transaction = pb::Transaction::from(&transaction); + let pb_transaction = pb::Transaction::from(&attempt_transaction); let inline_transaction = pb_transaction.encoded_len() <= MAX_INLINE_TRANSACTION_BYTES; current_transaction_file = if !write_config.disable_transaction_file() { @@ -1465,7 +1473,7 @@ pub(crate) async fn commit_transaction( )); } // Build an up-to-date manifest from the transaction and current manifest - let (mut manifest, mut indices) = match transaction.operation { + let (mut manifest, mut indices) = match attempt_transaction.operation { Operation::Restore { version } => { Transaction::restore_old_manifest( object_store, @@ -1478,7 +1486,7 @@ pub(crate) async fn commit_transaction( ) .await? } - _ => transaction.build_manifest_with_read_version( + _ => attempt_transaction.build_manifest_with_read_version( Some(dataset.manifest.as_ref()), dataset.load_indices().await?.as_ref().clone(), transaction_file, @@ -1499,7 +1507,11 @@ pub(crate) async fn commit_transaction( migrate_manifest(&dataset, &mut manifest, recompute_stats).await?; fix_schema(&mut manifest)?; - validate_stable_field_id_transition(&dataset.manifest, &manifest, &transaction.operation)?; + validate_stable_field_id_transition( + &dataset.manifest, + &manifest, + &attempt_transaction.operation, + )?; manifest.update_max_field_id(); check_storage_version(&mut manifest)?; @@ -1537,7 +1549,7 @@ pub(crate) async fn commit_transaction( Ok(manifest_location) => { record_successful_commit( &dataset, - &transaction, + &attempt_transaction, &manifest, &manifest_location, indices, @@ -1558,7 +1570,7 @@ pub(crate) async fn commit_transaction( commit_handler, &dataset.base, target_version, - &transaction, + &attempt_transaction, ) .await { @@ -1569,7 +1581,7 @@ pub(crate) async fn commit_transaction( let committed_manifest = *committed_manifest; record_successful_commit( &dataset, - &transaction, + &attempt_transaction, &committed_manifest, &location, indices, @@ -1628,7 +1640,7 @@ pub(crate) async fn commit_transaction( commit_handler, &dataset.base, target_version, - &transaction, + &attempt_transaction, ) .await { @@ -1639,7 +1651,7 @@ pub(crate) async fn commit_transaction( let committed_manifest = *committed_manifest; record_successful_commit( &dataset, - &transaction, + &attempt_transaction, &committed_manifest, &location, indices, @@ -1704,6 +1716,7 @@ mod tests { CommitLease, CommitLock, ManifestWriter, RenameCommitHandler, UnsafeCommitHandler, commit_handler_from_url, }; + use lance_table::transaction::TRANSACTION_SCHEMA_SOURCE_RAW_ARROW; use lance_testing::datagen::generate_random_array; use super::*; @@ -1834,6 +1847,211 @@ mod tests { test_commit_handler(handler, true).await; } + #[derive(Debug)] + struct InjectForeignCommitHandler { + foreign: Mutex>, + } + + #[async_trait::async_trait] + impl CommitHandler for InjectForeignCommitHandler { + fn is_version_not_found_definitive(&self) -> bool { + true + } + + async fn commit( + &self, + manifest: &mut Manifest, + indices: Option>, + base_path: &Path, + object_store: &ObjectStore, + manifest_writer: ManifestWriter, + naming_scheme: ManifestNamingScheme, + transaction: Option, + ) -> std::result::Result { + let foreign = self.foreign.lock().unwrap().take(); + if let Some((mut foreign_manifest, foreign_transaction)) = foreign { + foreign_manifest.version = manifest.version; + RenameCommitHandler + .commit( + &mut foreign_manifest, + None, + base_path, + object_store, + manifest_writer, + naming_scheme, + Some(foreign_transaction), + ) + .await?; + return Err(CommitError::CommitConflict); + } + + RenameCommitHandler + .commit( + manifest, + indices, + base_path, + object_store, + manifest_writer, + naming_scheme, + transaction, + ) + .await + } + } + + fn inject_foreign_commit_handler( + manifest: Manifest, + transaction: &Transaction, + ) -> Arc { + let transaction = pb::Transaction::from(transaction).into(); + Arc::new(InjectForeignCommitHandler { + foreign: Mutex::new(Some((manifest, transaction))), + }) + } + + #[tokio::test] + async fn raw_arrow_project_retry_matches_single_attempt_ids() { + let tmp = TempStrDir::default(); + let uri = tmp.as_str(); + let dataset = Dataset::write( + RecordBatchIterator::new( + vec![Ok(simple_batch(&simple_schema(), vec![1, 2, 3]))], + simple_schema(), + ), + uri, + None, + ) + .await + .unwrap(); + + let mut foreign_manifest = dataset.manifest.as_ref().clone(); + foreign_manifest.max_fragment_id = Some(foreign_manifest.max_fragment_id.unwrap_or(0) + 1); + let foreign_transaction = Transaction::new( + dataset.version().version, + Operation::ReserveFragments { num_fragments: 1 }, + None, + ); + let handler = inject_foreign_commit_handler(foreign_manifest.clone(), &foreign_transaction); + + let raw_schema = Schema { + fields: vec![Field::new_arrow("x", DataType::Int32, false).unwrap()], + metadata: HashMap::from([( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + )]), + }; + let mut expected = Operation::Project { + schema: raw_schema.clone(), + preserves_nullability: true, + }; + canonicalize_stable_field_ids(Some(&foreign_manifest), &mut expected).unwrap(); + let Operation::Project { + schema: expected_schema, + .. + } = expected + else { + unreachable!(); + }; + + let committed = Dataset::commit( + uri, + Operation::Project { + schema: raw_schema, + preserves_nullability: true, + }, + Some(dataset.version().version), + None, + Some(handler), + Default::default(), + false, + ) + .await + .unwrap(); + + assert_eq!(committed.schema(), &expected_schema); + assert!( + !committed + .schema() + .metadata + .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + ); + } + + #[tokio::test] + async fn raw_arrow_merge_retry_matches_single_attempt_ids() { + let tmp = TempStrDir::default(); + let uri = tmp.as_str(); + let dataset = Dataset::write( + RecordBatchIterator::new( + vec![Ok(simple_batch(&simple_schema(), vec![1, 2, 3]))], + simple_schema(), + ), + uri, + None, + ) + .await + .unwrap(); + + let mut foreign_manifest = dataset.manifest.as_ref().clone(); + foreign_manifest.max_fragment_id = Some(foreign_manifest.max_fragment_id.unwrap_or(0) + 1); + let foreign_transaction = Transaction::new( + dataset.version().version, + Operation::ReserveFragments { num_fragments: 1 }, + None, + ); + let handler = inject_foreign_commit_handler(foreign_manifest.clone(), &foreign_transaction); + + let mut merged_fragment = dataset.manifest.fragments[0].clone(); + let mut new_file = merged_fragment.files[0].clone(); + new_file.path = "retry-new-column.lance".to_string(); + new_file.fields = Arc::from([1]); + merged_fragment.files.push(new_file); + let raw_schema = Schema { + fields: vec![ + Field::new_arrow("x", DataType::Int32, false).unwrap(), + Field::new_arrow("new_column", DataType::Int32, true).unwrap(), + ], + metadata: HashMap::from([( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + )]), + }; + let operation = Operation::Merge { + fragments: vec![merged_fragment], + schema: raw_schema, + preserves_nullability: true, + }; + let mut expected = operation.clone(); + canonicalize_stable_field_ids(Some(&foreign_manifest), &mut expected).unwrap(); + + let committed = Dataset::commit( + uri, + operation, + Some(dataset.version().version), + None, + Some(handler), + Default::default(), + false, + ) + .await + .unwrap(); + + let Operation::Merge { + schema: expected_schema, + fragments: expected_fragments, + .. + } = expected + else { + unreachable!(); + }; + assert_eq!(committed.schema(), &expected_schema); + assert_eq!(committed.schema().field("new_column").unwrap().id, 1); + assert_eq!( + committed.manifest.fragments[0].files[1].fields, + expected_fragments[0].files[1].fields + ); + } + #[tokio::test] async fn test_unsafe_commit_handler() { let handler = Arc::new(UnsafeCommitHandler); From 4a17fc287c4843525ded9f6af0602443cef26f1d Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 31 Aug 2026 18:36:27 +0800 Subject: [PATCH 11/24] fix: complete stable field ID main integration --- rust/lance-table/src/feature_flags.rs | 2 -- rust/lance-table/src/transaction/manifest_build.rs | 1 + rust/lance/src/dataset/tests/dataset_io.rs | 1 + 3 files changed, 2 insertions(+), 2 deletions(-) diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index cb5f3c8cd32..218cdb123e7 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -637,7 +637,5 @@ mod tests { assert!(can_write_dataset(FLAG_COVERED_INDEX_METADATA)); assert!(!can_read_dataset(FLAG_MIXED_DATA_FILE_VERSIONS)); assert!(!can_write_dataset(FLAG_MIXED_DATA_FILE_VERSIONS)); - assert!(FLAG_MIXED_DATA_FILE_VERSIONS < FLAG_STABLE_FIELD_IDS); - assert!(FLAG_STABLE_FIELD_IDS < FLAG_UNKNOWN); } } diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 26d97b91b1f..44ff9ef4a98 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -13,6 +13,7 @@ use crate::feature_flags::{ FLAG_COVERED_INDEX_METADATA, FLAG_STABLE_FIELD_IDS, FLAG_STABLE_ROW_IDS, apply_feature_flags, ensure_can_read_manifest, ensure_can_write_manifest, inherit_sticky_feature_flags, + validate_stable_field_id_flags, }; use crate::format::overlay::TOMBSTONE_FIELD_ID; use crate::format::{ diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 508a802dcde..badb9e6d629 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -3541,6 +3541,7 @@ async fn write_manifest_file_rejects_a_nullable_primary_key() { storage_format: None, disable_transaction_file: false, migration_next_row_id: None, + stable_field_id_migration_requires_reader: None, }, dataset.manifest_location.naming_scheme, None, From aff26e74265c97e7240df2c740b12deb42cfba6c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 31 Aug 2026 20:27:36 +0800 Subject: [PATCH 12/24] fix: make stable field IDs writer-only --- docs/src/format/table/schema.md | 11 ++-- docs/src/format/table/versioning.md | 2 +- protos/table.proto | 4 +- .../python/tests/compat/compat_decorator.py | 47 ++------------- rust/lance-table/src/feature_flags.rs | 33 +++++------ rust/lance-table/src/format/manifest.rs | 16 ++--- .../src/transaction/manifest_build.rs | 19 ++---- .../src/transaction/test_support.rs | 2 +- rust/lance/src/dataset.rs | 48 ++++----------- rust/lance/src/dataset/tests/dataset_io.rs | 11 ++-- .../src/dataset/tests/dataset_migrations.rs | 59 ++++--------------- rust/lance/src/dataset/write/commit.rs | 16 ++--- rust/lance/src/io/commit.rs | 1 - 13 files changed, 73 insertions(+), 196 deletions(-) diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index 1bb050ebcfa..6d7ed3d633c 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -237,12 +237,11 @@ maximum non-negative ID referenced by the canonical manifest schema, base data f files. Activation provides a forward guarantee only; it cannot reconstruct identities that were dropped or reused in older snapshots. -New datasets activate this contract in their initial manifest and require both readers and writers -to understand it. The reader gate is required because older Lance versions do not enforce unknown -writer feature flags on every mutation path; allowing them to open the dataset could let a legacy -commit clear the writer gate and resume reusable allocation. Existing legacy datasets remain -unchanged until an explicit migration commit. Controlled deployments may use writer-only migration -only after every pre-gate writer has been retired. +New datasets activate this contract in their initial manifest. The corresponding feature flag is a +writer requirement because stable field IDs change allocation but not how readers interpret schema +or data files. Existing legacy datasets remain unchanged until an explicit migration commit. Before +activation, operators must ensure that every writer accessing the dataset understands and preserves +the stable field-ID contract. Activation is one-way within a branch ancestry. After migration, restore cannot target a version from before activation because that version does not carry the high-water mark needed to preserve diff --git a/docs/src/format/table/versioning.md b/docs/src/format/table/versioning.md index 12af7db1cff..8e1093ea9dc 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -32,7 +32,7 @@ they should return an "unsupported" error on any read or write operation. | 64 | `FLAG_UNSTABLE_DATA_OVERLAY_FILES` | Yes | Yes | Fragments may carry data overlay files. Unstable: release builds reject it unless explicitly opted in. | | 128 | `FLAG_COVERED_INDEX_METADATA` | Yes | Yes | Some index declares covering columns (`IndexMetadata.covering_fields`), so `fields` means keyed columns followed by carried ones. An implementation without this flag selects an index by membership of `fields` and would answer a query on a merely-carried column with an index keyed on a different one. | | 256 | `FLAG_MIXED_DATA_FILE_VERSIONS` | Yes | Yes | Reserved for datasets that may reference recognized V2 data files with different exact versions. | -| 512 | `FLAG_STABLE_FIELD_IDS` | Conditional | Yes | The manifest carries a persistent field-ID high-water mark. The safe activation mode also requires readers; a controlled rollout may require writers only. See [Field IDs](schema.md#field-ids). | +| 512 | `FLAG_STABLE_FIELD_IDS` | No | Yes | The manifest carries a persistent field-ID high-water mark that writers must use when allocating field IDs. See [Field IDs](schema.md#field-ids). | diff --git a/protos/table.proto b/protos/table.proto index f040cf969e6..6afea29afbe 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -128,8 +128,6 @@ message Manifest { // * 1 << 8: reserved for datasets that may reference recognized V2 data files // with different exact versions. Implementations that do not support the // per-file exact-version contract must treat this bit as unknown. - // * 1 << 9: readers must understand the stable field-ID contract. This is an - // optional fail-closed policy; stable field IDs always require the writer bit. uint64 reader_feature_flags = 9; // Feature flags for writers. @@ -140,6 +138,8 @@ message Manifest { // // The flag identities are the same as for reader_feature_flags, but the values of // reader_feature_flags and writer_feature_flags are not required to be identical. + // * 1 << 9: newly allocated field IDs must be greater than + // max_allocated_field_id, and the high-water mark must advance atomically. uint64 writer_feature_flags = 10; // The highest fragment ID that has been used so far. diff --git a/python/python/tests/compat/compat_decorator.py b/python/python/tests/compat/compat_decorator.py index f7ff52a6f02..9cb22013fd4 100644 --- a/python/python/tests/compat/compat_decorator.py +++ b/python/python/tests/compat/compat_decorator.py @@ -184,18 +184,10 @@ def skip_read_after_current_write(self, version: str) -> bool: """Return True to skip the old-version read after current-version writes.""" return False - def expect_read_after_current_write_failure(self, version: str) -> bool: - """Return True when the old reader must reject current-version data.""" - return False - def skip_write_after_current_write(self, version: str) -> bool: """Return True to skip the old-version write after current-version writes.""" return False - def expect_write_after_current_write_failure(self, version: str) -> bool: - """Return True when the old writer must reject current-version data.""" - return False - def skip_downgrade(self, version: str) -> bool: """Return True to skip the current-write -> old-read downgrade test.""" return False @@ -212,15 +204,10 @@ def compat_env(self, version: str, method_name: str) -> Dict[str, str]: class DatasetUpgradeDowngradeTest(UpgradeDowngradeTest): """Compatibility contract for new datasets with stable field IDs.""" - def expect_read_after_current_write_failure(self, version: str) -> bool: - # Every released version selected by this suite predates the reader - # feature bit. Once a supporting release enters the matrix this fails - # visibly, instead of silently skipping a valid reader combination. - return True - - def expect_write_after_current_write_failure(self, version: str) -> bool: - # Opening the dataset must fail before an old writer can enter a commit - # path that does not enforce unknown writer feature flags. + def skip_write_after_current_write(self, version: str) -> bool: + # Older readers remain compatible. Older writers are outside the + # contract because released commit paths do not all enforce unknown + # writer flags. return True @@ -402,31 +389,9 @@ def test_func({sig_params}): # Old version: verify can read venv = venv_factory.get_venv(version) if not obj.skip_read_after_current_write(version): - if obj.expect_read_after_current_write_failure(version): - with pytest.raises( - RuntimeError, - match="cannot be read by this version|Flags", - ): - venv.execute_method( - obj, "check_read", obj.compat_env(version, "check_read") - ) - else: - venv.execute_method( - obj, "check_read", obj.compat_env(version, "check_read") - ) + venv.execute_method(obj, "check_read", obj.compat_env(version, "check_read")) if not obj.skip_write_after_current_write(version): - if obj.expect_write_after_current_write_failure(version): - with pytest.raises( - RuntimeError, - match="cannot be (?:read|written) by this version|Flags", - ): - venv.execute_method( - obj, "check_write", obj.compat_env(version, "check_write") - ) - else: - venv.execute_method( - obj, "check_write", obj.compat_env(version, "check_write") - ) + venv.execute_method(obj, "check_write", obj.compat_env(version, "check_write")) ''' else: # upgrade_downgrade func_body = f''' diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 4e83f5ea9ad..81d5fac978c 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -54,9 +54,8 @@ pub const FLAG_COVERED_INDEX_METADATA: u64 = 1 << 7; /// different exact versions. pub const FLAG_MIXED_DATA_FILE_VERSIONS: u64 = 1 << 8; /// Field IDs are allocated from a persistent high-water mark and are never reused. -/// -/// This is always a writer requirement. It is also a reader requirement when an -/// activation chooses fail-closed compatibility with pre-gate binaries. +/// Writers must understand this allocation contract. It does not change how +/// readers interpret the schema or data files. pub const FLAG_STABLE_FIELD_IDS: u64 = 1 << 9; /// The first bit that is unknown as a feature flag pub const FLAG_UNKNOWN: u64 = 1 << 10; @@ -91,14 +90,6 @@ pub fn apply_feature_flags( let covered_index_metadata = (manifest.reader_feature_flags | manifest.writer_feature_flags) & FLAG_COVERED_INDEX_METADATA; let sticky_paired_flags = validated_sticky_paired_flags(manifest)?; - let stable_field_ids_reader = if manifest.max_allocated_field_id.is_some() - && manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0 - { - FLAG_STABLE_FIELD_IDS - } else { - 0 - }; - // Reset flags manifest.reader_feature_flags = 0; manifest.writer_feature_flags = 0; @@ -159,7 +150,6 @@ pub fn apply_feature_flags( if manifest.max_allocated_field_id.is_some() { manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - manifest.reader_feature_flags |= stable_field_ids_reader; } manifest.reader_feature_flags |= covered_index_metadata; @@ -290,16 +280,21 @@ pub fn validate_paired_feature_flags(manifest: &Manifest) -> Result<()> { /// Refuse a manifest whose stable-field-ID marker and required flags disagree. /// /// The high-water mark is the activation marker and the writer bit keeps -/// pre-feature writers away. The reader bit is optional deployment policy, but -/// it cannot be set on a legacy manifest. +/// pre-feature writers away. Stable field IDs do not require reader support. pub fn validate_stable_field_id_flags(manifest: &Manifest) -> Result<()> { let activated = manifest.max_allocated_field_id.is_some(); let reader = manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0; let writer = manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS != 0; - if activated != writer || (reader && !activated) { + if activated != writer { + return Err(Error::corrupt_file_named( + "manifest", + "Manifest stable-field-ID high-water mark and writer feature flag disagree", + )); + } + if reader { return Err(Error::corrupt_file_named( "manifest", - "Manifest stable-field-ID high-water mark and reader/writer feature flags disagree", + "Stable field IDs are a writer-only feature and must not set the reader feature flag", )); } Ok(()) @@ -581,7 +576,7 @@ mod tests { } #[test] - fn apply_feature_flags_preserves_fail_closed_stable_field_id_gate() { + fn apply_feature_flags_clears_stable_field_id_reader_gate() { let mut manifest = empty_manifest(); manifest.activate_stable_field_ids(); manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; @@ -590,7 +585,7 @@ mod tests { apply_feature_flags(&mut manifest, false, false).unwrap(); apply_feature_flags(&mut manifest, false, false).unwrap(); - assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_eq!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } @@ -610,7 +605,7 @@ mod tests { validate_stable_field_id_flags(&writer_only).unwrap(); writer_only.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - validate_stable_field_id_flags(&writer_only).unwrap(); + assert!(validate_stable_field_id_flags(&writer_only).is_err()); } fn empty_manifest() -> Manifest { diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index 1319f8a17f2..8c414a809db 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -226,8 +226,7 @@ impl Manifest { index_section: None, // Caller should update index if they want to keep them. timestamp_nanos: 0, // This will be set on commit tag: None, - reader_feature_flags: previous.reader_feature_flags - & (STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS), + reader_feature_flags: previous.reader_feature_flags & STICKY_PAIRED_FLAGS, writer_feature_flags: previous.writer_feature_flags & (STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS), max_fragment_id: previous.max_fragment_id, @@ -296,7 +295,7 @@ impl Manifest { // Sticky capabilities are also retained because the clone keeps the // source file identities that require them. reader_feature_flags: self.reader_feature_flags - & (FLAG_COVERED_INDEX_METADATA | STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS), + & (FLAG_COVERED_INDEX_METADATA | STICKY_PAIRED_FLAGS), writer_feature_flags: self.writer_feature_flags & (FLAG_COVERED_INDEX_METADATA | STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS), max_fragment_id: self.max_fragment_id, @@ -755,9 +754,8 @@ pub struct ManifestBuildConfig { /// It bypasses the "cannot enable stable row ids on existing dataset" guard and /// sets `manifest.next_row_id` to the provided value before activating the flag. pub migration_next_row_id: Option, - /// When `Some`, atomically activates stable field IDs. The boolean controls - /// whether the feature is also required of readers; writers are always gated. - pub stable_field_id_migration_requires_reader: Option, + /// Whether this commit atomically activates stable field IDs. + pub activate_stable_field_ids: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -1665,7 +1663,6 @@ mod tests { ); manifest.activate_stable_field_ids(); manifest.max_allocated_field_id = Some(43); - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; assert_eq!(manifest.max_referenced_field_id(), 0); @@ -1674,7 +1671,7 @@ mod tests { let recovered = Manifest::try_from(pb::Manifest::from(&manifest)).unwrap(); assert_eq!(recovered.max_allocated_field_id, Some(43)); assert_eq!(recovered.max_field_id(), 43); - assert_ne!(recovered.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_eq!(recovered.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); assert_ne!(recovered.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } @@ -1693,7 +1690,6 @@ mod tests { HashMap::new(), ); manifest.max_allocated_field_id = Some(41); - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; let cloned = manifest.shallow_clone( @@ -1705,7 +1701,7 @@ mod tests { ); assert_eq!(cloned.max_allocated_field_id, Some(41)); - assert_ne!(cloned.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_eq!(cloned.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); assert_ne!(cloned.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 44ff9ef4a98..a8377180279 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -126,9 +126,6 @@ impl Transaction { }; manifest.max_allocated_field_id = Some(restored_max_field_id.max(current_manifest.max_field_id())); - if current_manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0 { - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - } manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; } // Row ids are a high-water mark like fragment ids: rewinding hands old ids to new rows. @@ -1308,24 +1305,16 @@ impl Transaction { if current_manifest.is_none() { manifest.activate_stable_field_ids(); - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; } - if let Some(require_reader) = config.stable_field_id_migration_requires_reader { + if config.activate_stable_field_ids { let already_active = current_manifest .map(|manifest| manifest.uses_stable_field_ids()) .unwrap_or(false); - if already_active { - if require_reader { - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - } - } else { + if !already_active { manifest.activate_stable_field_ids(); manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - if require_reader { - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - } } } @@ -1637,7 +1626,7 @@ mod tests { .unwrap(); assert_eq!(manifest.max_allocated_field_id, Some(0)); - assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_eq!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } @@ -1656,7 +1645,7 @@ mod tests { ); let mut config = default_build_config(); config.auto_set_feature_flags = false; - config.stable_field_id_migration_requires_reader = Some(false); + config.activate_stable_field_ids = true; let (activated, _) = transaction .build_manifest(Some(&manifest), vec![], "txn", &config) diff --git a/rust/lance-table/src/transaction/test_support.rs b/rust/lance-table/src/transaction/test_support.rs index 0304f567e81..63284d65fad 100644 --- a/rust/lance-table/src/transaction/test_support.rs +++ b/rust/lance-table/src/transaction/test_support.rs @@ -30,7 +30,7 @@ pub fn default_build_config() -> ManifestBuildConfig { storage_format: None, disable_transaction_file: false, migration_next_row_id: None, - stable_field_id_migration_requires_reader: None, + activate_stable_field_ids: false, } } diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index d42cd9b109f..b4957e76278 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -208,20 +208,6 @@ pub const DEFAULT_INDEX_CACHE_SIZE: usize = 6 * 1024 * 1024 * 1024; // smaller. pub const DEFAULT_METADATA_CACHE_SIZE: usize = 1024 * 1024 * 1024; -/// Mixed-version compatibility policy for stable-field-ID activation. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum StableFieldIdMigrationMode { - /// Require both new readers and new writers. This is the safe default when - /// pre-gate binaries may still access the dataset. - #[default] - ReadersAndWriters, - /// Require new writers while allowing legacy readers. - /// - /// Use only after every pre-gate writer has been retired. Older binaries do - /// not universally check writer feature flags on every commit path. - WritersOnly, -} - /// Lance Dataset #[derive(Clone)] pub struct Dataset { @@ -3267,29 +3253,19 @@ impl Dataset { /// The activation commit records the current maximum referenced field ID as /// a persistent high-water mark. Later schema changes allocate above it even /// after fields and their files are dropped. New datasets already use this - /// contract by default with reader and writer gates. `ReadersAndWriters` - /// can add a fail-closed reader gate to a writer-only activated dataset. - /// Activation and reader gating are one-way and idempotent. + /// contract by default. Activation is one-way and idempotent. Before + /// migrating, ensure that every writer accessing the dataset understands + /// the stable field-ID writer feature flag. /// /// ``` /// # use lance::{Dataset, Result}; - /// # use lance::dataset::StableFieldIdMigrationMode; /// # async fn activate(dataset: &mut Dataset) -> Result<()> { - /// dataset - /// .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) - /// .await?; + /// dataset.migrate_to_stable_field_ids().await?; /// # Ok(()) /// # } /// ``` - pub async fn migrate_to_stable_field_ids( - &mut self, - mode: StableFieldIdMigrationMode, - ) -> Result<()> { - let require_reader = matches!(mode, StableFieldIdMigrationMode::ReadersAndWriters); - let reader_already_required = self.manifest.reader_feature_flags - & lance_table::feature_flags::FLAG_STABLE_FIELD_IDS - != 0; - if self.manifest.uses_stable_field_ids() && (!require_reader || reader_already_required) { + pub async fn migrate_to_stable_field_ids(&mut self) -> Result<()> { + if self.manifest.uses_stable_field_ids() { return Ok(()); } @@ -3306,7 +3282,7 @@ impl Dataset { ); let new_ds = CommitBuilder::new(Arc::new(self.clone())) .with_max_retries(0) - .with_stable_field_id_migration_activation(require_reader) + .with_stable_field_id_migration_activation() .execute(transaction) .await?; *self = new_ds; @@ -4130,9 +4106,8 @@ pub(crate) struct ManifestWriteConfig { /// It bypasses the "cannot enable stable row ids on existing dataset" guard and /// sets `manifest.next_row_id` to the provided value before activating the flag. migration_next_row_id: Option, // default None - /// When `Some`, this commit activates stable field IDs. `true` also sets - /// the reader feature bit for fail-closed mixed-version safety. - stable_field_id_migration_requires_reader: Option, + /// Whether this commit activates stable field IDs. + activate_stable_field_ids: bool, } impl Default for ManifestWriteConfig { @@ -4145,7 +4120,7 @@ impl Default for ManifestWriteConfig { use_legacy_format: None, storage_format: None, migration_next_row_id: None, - stable_field_id_migration_requires_reader: None, + activate_stable_field_ids: false, } } } @@ -4174,8 +4149,7 @@ impl ManifestWriteConfig { storage_format: self.storage_format.clone(), disable_transaction_file: self.disable_transaction_file, migration_next_row_id: self.migration_next_row_id, - stable_field_id_migration_requires_reader: self - .stable_field_id_migration_requires_reader, + activate_stable_field_ids: self.activate_stable_field_ids, } } } diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index badb9e6d629..b984bb645b4 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1323,10 +1323,7 @@ async fn test_write_manifest( manifest.data_storage_format.version.to_manifest_string(), "stable" | "next" )); - assert_eq!( - manifest.reader_feature_flags, - feature_flags::FLAG_STABLE_FIELD_IDS - ); + assert_eq!(manifest.reader_feature_flags, 0); assert_eq!( manifest.writer_feature_flags, feature_flags::FLAG_STABLE_FIELD_IDS @@ -1355,7 +1352,7 @@ async fn test_write_manifest( ); assert_eq!( manifest.reader_feature_flags, - feature_flags::FLAG_DELETION_FILES | feature_flags::FLAG_STABLE_FIELD_IDS + feature_flags::FLAG_DELETION_FILES ); // Write with custom manifest @@ -1376,7 +1373,7 @@ async fn test_write_manifest( storage_format: None, disable_transaction_file: false, migration_next_row_id: None, - stable_field_id_migration_requires_reader: None, + activate_stable_field_ids: false, }, dataset.manifest_location.naming_scheme, None, @@ -3541,7 +3538,7 @@ async fn write_manifest_file_rejects_a_nullable_primary_key() { storage_format: None, disable_transaction_file: false, migration_next_row_id: None, - stable_field_id_migration_requires_reader: None, + activate_stable_field_ids: false, }, dataset.manifest_location.naming_scheme, None, diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index f80c33631a5..a44ec5d1869 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -6,9 +6,7 @@ use std::sync::Arc; use std::vec; use crate::dataset::optimize::{CompactionOptions, compact_files}; -use crate::dataset::{ - ColumnAlteration, InsertBuilder, NewColumnTransform, StableFieldIdMigrationMode, -}; +use crate::dataset::{ColumnAlteration, InsertBuilder, NewColumnTransform}; use crate::index::DatasetIndexExt; use crate::utils::test::copy_test_data_to_tmp; use crate::{Dataset, Result}; @@ -348,10 +346,7 @@ async fn test_stable_field_id_migration_repairs_legacy_schema_before_activation( let test_dir = copy_test_data_to_tmp("v0.10.5/corrupt_schema").unwrap(); let mut dataset = Dataset::open(&test_dir.path_str()).await.unwrap(); - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::WritersOnly) - .await - .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); dataset.validate().await.unwrap(); assert!(dataset.manifest.uses_stable_field_ids()); @@ -364,16 +359,9 @@ async fn test_stable_field_id_migration_repairs_legacy_schema_before_activation( 0 ); - let writer_only_version = dataset.version().version; - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) - .await - .unwrap(); - assert_eq!(dataset.version().version, writer_only_version + 1); - assert_ne!( - dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, - 0 - ); + let activation_version = dataset.version().version; + dataset.migrate_to_stable_field_ids().await.unwrap(); + assert_eq!(dataset.version().version, activation_version); } #[tokio::test] @@ -696,7 +684,7 @@ async fn test_new_datasets_use_stable_field_ids_and_migration_is_idempotent() { let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; assert!(dataset.manifest.uses_stable_field_ids()); assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); - assert_ne!( + assert_eq!( dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); @@ -706,26 +694,14 @@ async fn test_new_datasets_use_stable_field_ids_and_migration_is_idempotent() { ); let created_version = dataset.version().version; - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) - .await - .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); assert_eq!(dataset.version().version, created_version); - assert_ne!( + assert_eq!( dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) - .await - .unwrap(); - assert_eq!(dataset.version().version, created_version); - - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::WritersOnly) - .await - .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); assert_eq!(dataset.version().version, created_version); } @@ -760,10 +736,7 @@ async fn test_shallow_clone_preserves_stable_field_id_state() { let source_uri = TempStrDir::default(); let clone_uri = TempStrDir::default(); let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) - .await - .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); let cloned = dataset .shallow_clone(clone_uri.as_str(), dataset.version().version, None) @@ -774,7 +747,7 @@ async fn test_shallow_clone_preserves_stable_field_id_state() { cloned.manifest.max_allocated_field_id, dataset.manifest.max_allocated_field_id ); - assert_ne!( + assert_eq!( cloned.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); @@ -788,10 +761,7 @@ async fn test_shallow_clone_preserves_stable_field_id_state() { async fn test_overwrite_preserves_compatible_stable_field_identities() { let source_uri = TempStrDir::default(); let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) - .await - .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); let schema = Arc::new(ArrowSchema::new(vec![ ArrowField::new("id", DataType::Int64, false), @@ -825,10 +795,7 @@ async fn test_overwrite_preserves_compatible_stable_field_identities() { async fn test_stable_field_id_rename_and_nullability_preserve_identity() { let source_uri = TempStrDir::default(); let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; - dataset - .migrate_to_stable_field_ids(StableFieldIdMigrationMode::ReadersAndWriters) - .await - .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); dataset .alter_columns(&[ColumnAlteration::new("id".to_string()) diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index 45c77d5c236..21cfcb46ebb 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -55,8 +55,8 @@ pub struct CommitBuilder<'a> { timeout: Option, /// When `Some`, this commit is the second step of `migrate_to_stable_row_ids`. migration_next_row_id: Option, - /// When `Some`, this commit atomically activates stable field IDs. - stable_field_id_migration_requires_reader: Option, + /// Whether this commit atomically activates stable field IDs. + activate_stable_field_ids: bool, } /// Default timeout applied to [`CommitBuilder::execute`] when none is set. @@ -81,7 +81,7 @@ impl<'a> CommitBuilder<'a> { transaction_properties: None, timeout: Some(DEFAULT_COMMIT_TIMEOUT), migration_next_row_id: None, - stable_field_id_migration_requires_reader: None, + activate_stable_field_ids: false, } } @@ -270,11 +270,8 @@ impl<'a> CommitBuilder<'a> { self } - pub(crate) fn with_stable_field_id_migration_activation( - mut self, - require_reader: bool, - ) -> Self { - self.stable_field_id_migration_requires_reader = Some(require_reader); + pub(crate) fn with_stable_field_id_migration_activation(mut self) -> Self { + self.activate_stable_field_ids = true; self } @@ -446,8 +443,7 @@ impl<'a> CommitBuilder<'a> { use_stable_row_ids, storage_format: self.storage_format.map(DataStorageFormat::new), migration_next_row_id: self.migration_next_row_id, - stable_field_id_migration_requires_reader: self - .stable_field_id_migration_requires_reader, + activate_stable_field_ids: self.activate_stable_field_ids, ..Default::default() }; diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 324403c0285..98b088a0d2a 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -497,7 +497,6 @@ async fn do_commit_new_dataset( if !manifest.uses_stable_field_ids() { fix_schema(&mut manifest)?; manifest.activate_stable_field_ids(); - manifest.reader_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; } From 3a560ff6f300a3a557d071c7b456a110931a0dc8 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 31 Aug 2026 20:54:46 +0800 Subject: [PATCH 13/24] fix: preserve raw field bindings across commit retries --- rust/lance/src/dataset/write/commit.rs | 17 +---------------- rust/lance/src/io/commit.rs | 6 ++++-- 2 files changed, 5 insertions(+), 18 deletions(-) diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index 21cfcb46ebb..2cf1a05bb0b 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -28,9 +28,7 @@ use crate::{ use super::{WriteDestination, resolve_commit_handler}; use crate::dataset::branch_location::BranchLocation; -use crate::dataset::transaction::validate_operation; use lance_core::utils::tracing::{DATASET_COMMITTED_EVENT, TRACE_DATASET_EVENTS}; -use lance_table::transaction::canonicalize_stable_field_ids; use tracing::info; /// Create a new commit from a [`Transaction`]. @@ -301,7 +299,7 @@ impl<'a> CommitBuilder<'a> { } } - async fn execute_inner(self, mut transaction: Transaction) -> Result { + async fn execute_inner(self, transaction: Transaction) -> Result { let session = self .session .or_else(|| self.dest.dataset().map(|ds| ds.session.clone())) @@ -385,19 +383,6 @@ impl<'a> CommitBuilder<'a> { )); } - canonicalize_stable_field_ids( - dest.dataset().map(|dataset| dataset.manifest.as_ref()), - &mut transaction.operation, - )?; - - // Validate the operation before proceeding with the commit - // This ensures that operations like Merge have proper validation for data integrity - if let Some(dataset) = dest.dataset() { - validate_operation(Some(&dataset.manifest), &transaction.operation)?; - } else { - validate_operation(None, &transaction.operation)?; - } - let (metadata_cache, index_cache) = match &dest { WriteDestination::Dataset(ds) => (ds.metadata_cache.clone(), ds.index_cache.clone()), WriteDestination::Uri(uri) => ( diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 98b088a0d2a..ba33a5f25d3 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -363,6 +363,7 @@ async fn do_commit_new_dataset( let mut transaction = transaction.clone(); canonicalize_stable_field_ids(None, &mut transaction.operation)?; let transaction = &transaction; + validate_operation(None, &transaction.operation)?; let pb_transaction = pb::Transaction::from(transaction); let inline_transaction = pb_transaction.encoded_len() <= MAX_INLINE_TRANSACTION_BYTES; // Classified from the operation itself. Reading it back off the inline @@ -1992,7 +1993,7 @@ mod tests { } #[tokio::test] - async fn raw_arrow_merge_retry_matches_single_attempt_ids() { + async fn raw_arrow_merge_retry_rebinds_after_allocator_advance() { let tmp = TempStrDir::default(); let uri = tmp.as_str(); let dataset = Dataset::write( @@ -2008,6 +2009,7 @@ mod tests { let mut foreign_manifest = dataset.manifest.as_ref().clone(); foreign_manifest.max_fragment_id = Some(foreign_manifest.max_fragment_id.unwrap_or(0) + 1); + foreign_manifest.max_allocated_field_id = Some(foreign_manifest.max_field_id() + 1); let foreign_transaction = Transaction::new( dataset.version().version, Operation::ReserveFragments { num_fragments: 1 }, @@ -2059,7 +2061,7 @@ mod tests { unreachable!(); }; assert_eq!(committed.schema(), &expected_schema); - assert_eq!(committed.schema().field("new_column").unwrap().id, 1); + assert_eq!(committed.schema().field("new_column").unwrap().id, 2); assert_eq!( committed.manifest.fragments[0].files[1].fields, expected_fragments[0].files[1].fields From d952779ec290809295bf7d800eb69dcd7a585b69 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 31 Aug 2026 23:51:30 +0800 Subject: [PATCH 14/24] fix: fence automatic stable field ID activation --- docs/src/format/table/schema.md | 12 ++-- docs/src/format/table/versioning.md | 2 +- protos/table.proto | 5 ++ .../python/tests/compat/compat_decorator.py | 7 +- .../python/tests/compat/test_file_formats.py | 30 +++++++++ rust/lance-table/src/feature_flags.rs | 66 ++++++++++++++----- rust/lance-table/src/format/manifest.rs | 61 +++++++++-------- .../src/transaction/manifest_build.rs | 3 +- rust/lance-table/src/transaction/validate.rs | 34 ++++++++++ rust/lance/src/dataset/tests/dataset_io.rs | 7 +- .../src/dataset/tests/dataset_migrations.rs | 18 ++++- rust/lance/src/io/commit.rs | 1 + 12 files changed, 187 insertions(+), 59 deletions(-) diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index 6d7ed3d633c..9ca3a3c36f0 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -237,11 +237,13 @@ maximum non-negative ID referenced by the canonical manifest schema, base data f files. Activation provides a forward guarantee only; it cannot reconstruct identities that were dropped or reused in older snapshots. -New datasets activate this contract in their initial manifest. The corresponding feature flag is a -writer requirement because stable field IDs change allocation but not how readers interpret schema -or data files. Existing legacy datasets remain unchanged until an explicit migration commit. Before -activation, operators must ensure that every writer accessing the dataset understands and preserves -the stable field-ID contract. +New datasets activate this contract in their initial manifest and set the corresponding reader and +writer feature bits. The reader bit is a compatibility fence for released runtimes whose generic +commit path did not enforce unknown writer bits; stable field IDs do not otherwise change how fields +are read. Existing legacy datasets remain unchanged until an explicit migration commit. Explicit +migration sets only the writer bit and requires operators to retire older writers before activation. +The reader fence, once set by automatic activation, is retained by every later commit, clone, and +restore in that branch ancestry. Activation is one-way within a branch ancestry. After migration, restore cannot target a version from before activation because that version does not carry the high-water mark needed to preserve diff --git a/docs/src/format/table/versioning.md b/docs/src/format/table/versioning.md index 8e1093ea9dc..0eaf44484d0 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -32,7 +32,7 @@ they should return an "unsupported" error on any read or write operation. | 64 | `FLAG_UNSTABLE_DATA_OVERLAY_FILES` | Yes | Yes | Fragments may carry data overlay files. Unstable: release builds reject it unless explicitly opted in. | | 128 | `FLAG_COVERED_INDEX_METADATA` | Yes | Yes | Some index declares covering columns (`IndexMetadata.covering_fields`), so `fields` means keyed columns followed by carried ones. An implementation without this flag selects an index by membership of `fields` and would answer a query on a merely-carried column with an index keyed on a different one. | | 256 | `FLAG_MIXED_DATA_FILE_VERSIONS` | Yes | Yes | Reserved for datasets that may reference recognized V2 data files with different exact versions. | -| 512 | `FLAG_STABLE_FIELD_IDS` | No | Yes | The manifest carries a persistent field-ID high-water mark that writers must use when allocating field IDs. See [Field IDs](schema.md#field-ids). | +| 512 | `FLAG_STABLE_FIELD_IDS` | Automatic activation | Yes | The manifest carries a persistent field-ID high-water mark that writers must use when allocating field IDs. Automatic activation also fences released runtimes whose write paths did not consistently enforce writer flags; explicit migration is writer-only. See [Field IDs](schema.md#field-ids). | diff --git a/protos/table.proto b/protos/table.proto index 6afea29afbe..3f2d9dde5ab 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -128,6 +128,9 @@ message Manifest { // * 1 << 8: reserved for datasets that may reference recognized V2 data files // with different exact versions. Implementations that do not support the // per-file exact-version contract must treat this bit as unknown. + // * 1 << 9: the stable field-ID contract was activated automatically. The + // reader fence excludes released runtimes whose generic commit path did not + // enforce unknown writer flags; it does not change how fields are read. uint64 reader_feature_flags = 9; // Feature flags for writers. @@ -140,6 +143,8 @@ message Manifest { // reader_feature_flags and writer_feature_flags are not required to be identical. // * 1 << 9: newly allocated field IDs must be greater than // max_allocated_field_id, and the high-water mark must advance atomically. + // Automatic activation sets this bit in both flag words. Explicit migration + // may set only the writer bit after older writers have been retired. uint64 writer_feature_flags = 10; // The highest fragment ID that has been used so far. diff --git a/python/python/tests/compat/compat_decorator.py b/python/python/tests/compat/compat_decorator.py index 9cb22013fd4..263eac32bb6 100644 --- a/python/python/tests/compat/compat_decorator.py +++ b/python/python/tests/compat/compat_decorator.py @@ -204,10 +204,9 @@ def compat_env(self, version: str, method_name: str) -> Dict[str, str]: class DatasetUpgradeDowngradeTest(UpgradeDowngradeTest): """Compatibility contract for new datasets with stable field IDs.""" - def skip_write_after_current_write(self, version: str) -> bool: - # Older readers remain compatible. Older writers are outside the - # contract because released commit paths do not all enforce unknown - # writer flags. + def skip_downgrade(self, version: str) -> bool: + # Automatically activated datasets fence released runtimes from opening + # them. A dedicated v11 regression verifies this refusal. return True diff --git a/python/python/tests/compat/test_file_formats.py b/python/python/tests/compat/test_file_formats.py index be2b533411b..93d6b2d989a 100644 --- a/python/python/tests/compat/test_file_formats.py +++ b/python/python/tests/compat/test_file_formats.py @@ -11,6 +11,7 @@ from pathlib import Path import lance +import pyarrow as pa import pytest from lance.file import LanceFileReader, LanceFileWriter @@ -139,3 +140,32 @@ def check_write(self): lance.write_dataset( build_basic_types(), self.path, data_storage_version="0.1", mode="append" ) + + +class StableFieldIdWriterFence: + def __init__(self, path: Path): + self.path = path + + def create(self): + dataset = lance.write_dataset(pa.table({"a": [1]}), self.path) + dataset.add_columns({"retired": "a"}) + dataset.drop_columns(["retired"]) + + def assert_write_rejected(self): + with pytest.raises(ValueError, match="cannot be read by this version of Lance"): + lance.dataset(self.path) + + def add_replacement_and_check(self): + dataset = lance.dataset(self.path) + dataset.add_columns({"replacement": "a"}) + assert dataset.lance_schema.field("replacement").id() == 2 + + +@pytest.mark.compat +def test_released_writer_cannot_erase_field_id_high_water(venv_factory, tmp_path): + case = StableFieldIdWriterFence(tmp_path / "data.lance") + case.create() + + venv_factory.get_venv("11.0.0").execute_method(case, "assert_write_rejected") + + case.add_replacement_and_check() diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 81d5fac978c..1207b38f61f 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -70,6 +70,8 @@ const _: () = assert!(FLAG_MIXED_DATA_FILE_VERSIONS < FLAG_UNKNOWN); const _: () = assert!(FLAG_STABLE_FIELD_IDS < FLAG_UNKNOWN); pub(crate) const STICKY_PAIRED_FLAGS: u64 = FLAG_MIXED_DATA_FILE_VERSIONS; +pub(crate) const STICKY_READER_FLAGS: u64 = STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS; +pub(crate) const STICKY_WRITER_FLAGS: u64 = STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS; /// Environment variable that opts a release build into reading and writing data /// overlay files before the feature is generally released. @@ -90,6 +92,7 @@ pub fn apply_feature_flags( let covered_index_metadata = (manifest.reader_feature_flags | manifest.writer_feature_flags) & FLAG_COVERED_INDEX_METADATA; let sticky_paired_flags = validated_sticky_paired_flags(manifest)?; + let stable_field_id_reader_fence = manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS; // Reset flags manifest.reader_feature_flags = 0; manifest.writer_feature_flags = 0; @@ -149,6 +152,7 @@ pub fn apply_feature_flags( } if manifest.max_allocated_field_id.is_some() { + manifest.reader_feature_flags |= stable_field_id_reader_fence; manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; } @@ -160,20 +164,23 @@ pub fn apply_feature_flags( Ok(()) } -/// Carry sticky paired capabilities from the manifest a new one is derived -/// from. +/// Carry sticky capabilities from the manifest a new one is derived from. /// /// [`apply_feature_flags`] carries these bits across its own reset, but it only /// ever sees one manifest. Constructors preserve these flags, and this helper /// also validates that the source is not half-set before a derived manifest is /// committed. /// -/// A half-set state is refused rather than normalized: one bit set means a -/// legacy reader or a legacy writer is still permitted, which is neither mode. +/// Stable field IDs permit writer-only fencing after an explicit migration, +/// while automatic activation also carries a reader fence to exclude released +/// writers that did not enforce unknown writer flags on every commit path. pub fn inherit_sticky_feature_flags(destination: &mut Manifest, source: &Manifest) -> Result<()> { let sticky_flags = validated_sticky_paired_flags(source)?; - destination.reader_feature_flags |= sticky_flags; - destination.writer_feature_flags |= sticky_flags; + validate_stable_field_id_flags(source)?; + destination.reader_feature_flags |= + sticky_flags | (source.reader_feature_flags & FLAG_STABLE_FIELD_IDS); + destination.writer_feature_flags |= + sticky_flags | (source.writer_feature_flags & FLAG_STABLE_FIELD_IDS); Ok(()) } @@ -279,8 +286,11 @@ pub fn validate_paired_feature_flags(manifest: &Manifest) -> Result<()> { /// Refuse a manifest whose stable-field-ID marker and required flags disagree. /// -/// The high-water mark is the activation marker and the writer bit keeps -/// pre-feature writers away. Stable field IDs do not require reader support. +/// The high-water mark is the activation marker and always requires the writer +/// bit. Automatic activation also sets the reader bit as a compatibility fence +/// against released writers that did not enforce unknown writer flags on every +/// commit path. Explicit migration may remain writer-only once older writers +/// have been retired. pub fn validate_stable_field_id_flags(manifest: &Manifest) -> Result<()> { let activated = manifest.max_allocated_field_id.is_some(); let reader = manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0; @@ -291,10 +301,10 @@ pub fn validate_stable_field_id_flags(manifest: &Manifest) -> Result<()> { "Manifest stable-field-ID high-water mark and writer feature flag disagree", )); } - if reader { + if reader && !writer { return Err(Error::corrupt_file_named( "manifest", - "Stable field IDs are a writer-only feature and must not set the reader feature flag", + "Manifest has the stable-field-ID reader fence without the activation marker and writer feature flag", )); } Ok(()) @@ -507,6 +517,27 @@ mod tests { ); } + #[test] + fn inheriting_preserves_stable_field_id_fence_mode() { + for reader_fenced in [false, true] { + let mut source = empty_manifest(); + source.activate_stable_field_ids(); + source.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + if reader_fenced { + source.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + } + let mut destination = empty_manifest(); + + inherit_sticky_feature_flags(&mut destination, &source).unwrap(); + + assert_eq!( + destination.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0, + reader_fenced + ); + assert_ne!(destination.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } + } + #[test] fn inheriting_refuses_a_half_set_source() { for (reader, writer) in [ @@ -565,7 +596,7 @@ mod tests { } #[test] - fn apply_feature_flags_sets_writer_gate_for_stable_field_ids() { + fn apply_feature_flags_sets_writer_gate_for_explicit_stable_field_id_activation() { let mut manifest = empty_manifest(); manifest.activate_stable_field_ids(); @@ -576,7 +607,7 @@ mod tests { } #[test] - fn apply_feature_flags_clears_stable_field_id_reader_gate() { + fn apply_feature_flags_preserves_automatic_stable_field_id_reader_fence() { let mut manifest = empty_manifest(); manifest.activate_stable_field_ids(); manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; @@ -585,7 +616,7 @@ mod tests { apply_feature_flags(&mut manifest, false, false).unwrap(); apply_feature_flags(&mut manifest, false, false).unwrap(); - assert_eq!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } @@ -604,8 +635,13 @@ mod tests { writer_only.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; validate_stable_field_id_flags(&writer_only).unwrap(); - writer_only.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - assert!(validate_stable_field_id_flags(&writer_only).is_err()); + let mut paired = writer_only.clone(); + paired.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + validate_stable_field_id_flags(&paired).unwrap(); + + let mut reader_without_activation = empty_manifest(); + reader_without_activation.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + assert!(validate_stable_field_id_flags(&reader_without_activation).is_err()); } fn empty_manifest() -> Manifest { diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index 8c414a809db..e6bb69f872b 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -18,9 +18,7 @@ use std::ops::Range; use std::sync::Arc; use super::Fragment; -use crate::feature_flags::{ - FLAG_COVERED_INDEX_METADATA, FLAG_STABLE_FIELD_IDS, STICKY_PAIRED_FLAGS, -}; +use crate::feature_flags::{FLAG_COVERED_INDEX_METADATA, STICKY_READER_FLAGS, STICKY_WRITER_FLAGS}; use crate::feature_flags::{FLAG_STABLE_ROW_IDS, has_deprecated_v2_feature_flag}; use crate::format::fragment::DataFileFieldInterner; use crate::format::pb; @@ -226,9 +224,8 @@ impl Manifest { index_section: None, // Caller should update index if they want to keep them. timestamp_nanos: 0, // This will be set on commit tag: None, - reader_feature_flags: previous.reader_feature_flags & STICKY_PAIRED_FLAGS, - writer_feature_flags: previous.writer_feature_flags - & (STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS), + reader_feature_flags: previous.reader_feature_flags & STICKY_READER_FLAGS, + writer_feature_flags: previous.writer_feature_flags & STICKY_WRITER_FLAGS, max_fragment_id: previous.max_fragment_id, max_allocated_field_id: previous.max_allocated_field_id, transaction_file: None, @@ -295,9 +292,9 @@ impl Manifest { // Sticky capabilities are also retained because the clone keeps the // source file identities that require them. reader_feature_flags: self.reader_feature_flags - & (FLAG_COVERED_INDEX_METADATA | STICKY_PAIRED_FLAGS), + & (FLAG_COVERED_INDEX_METADATA | STICKY_READER_FLAGS), writer_feature_flags: self.writer_feature_flags - & (FLAG_COVERED_INDEX_METADATA | STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS), + & (FLAG_COVERED_INDEX_METADATA | STICKY_WRITER_FLAGS), max_fragment_id: self.max_fragment_id, max_allocated_field_id: self.max_allocated_field_id, transaction_file: Some(transaction_file), @@ -1676,33 +1673,41 @@ mod tests { } #[test] - fn shallow_clone_preserves_stable_field_id_allocation_state() { + fn shallow_clone_preserves_stable_field_id_allocation_state_and_fence_mode() { let arrow_schema = ArrowSchema::new(vec![ArrowField::new( "a", arrow_schema::DataType::Int64, false, )]); let schema = Schema::try_from(&arrow_schema).unwrap(); - let mut manifest = Manifest::new( - schema, - Arc::new(vec![]), - DataStorageFormat::default(), - HashMap::new(), - ); - manifest.max_allocated_field_id = Some(41); - manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - - let cloned = manifest.shallow_clone( - Some("parent".to_string()), - "memory://parent".to_string(), - 7, - None, - String::new(), - ); + for reader_fenced in [false, true] { + let mut manifest = Manifest::new( + schema.clone(), + Arc::new(vec![]), + DataStorageFormat::default(), + HashMap::new(), + ); + manifest.max_allocated_field_id = Some(41); + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + if reader_fenced { + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + } - assert_eq!(cloned.max_allocated_field_id, Some(41)); - assert_eq!(cloned.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); - assert_ne!(cloned.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + let cloned = manifest.shallow_clone( + Some("parent".to_string()), + "memory://parent".to_string(), + 7, + None, + String::new(), + ); + + assert_eq!(cloned.max_allocated_field_id, Some(41)); + assert_eq!( + cloned.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0, + reader_fenced + ); + assert_ne!(cloned.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + } } #[test] diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index a8377180279..abeace843f7 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1305,6 +1305,7 @@ impl Transaction { if current_manifest.is_none() { manifest.activate_stable_field_ids(); + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; } @@ -1626,7 +1627,7 @@ mod tests { .unwrap(); assert_eq!(manifest.max_allocated_field_id, Some(0)); - assert_eq!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index a57b842c620..28519a8c67b 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -7,6 +7,7 @@ //! a fragment list that disagrees with the schema, a merge that silently dropped //! or rewrote data files — before any manifest is written. +use crate::feature_flags::FLAG_STABLE_FIELD_IDS; use crate::format::{Fragment, Manifest}; use crate::io::deletion::relative_deletion_file_path; use crate::transaction::{Operation, UpdateMode, UpdatedFragmentOffsets}; @@ -471,6 +472,13 @@ pub fn validate_stable_field_id_transition( "Stable field-ID high-water mark decreases from {parent_max_field_id} to {successor_max_field_id}" ))); } + if parent.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0 + && successor.reader_feature_flags & FLAG_STABLE_FIELD_IDS == 0 + { + return Err(Error::invalid_input( + "Stable field-ID reader fence cannot be removed after automatic activation", + )); + } let successor_schema_ids = successor .schema .fields_pre_order() @@ -1489,6 +1497,32 @@ mod tests { ); } + #[test] + fn stable_field_id_transition_rejects_removing_reader_fence() { + let mut manifest = activated_manifest(); + manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; + let mut successor = Manifest::new_from_previous( + &manifest, + manifest.schema.clone(), + manifest.fragments.clone(), + ); + successor.reader_feature_flags &= !FLAG_STABLE_FIELD_IDS; + let operation = Operation::UpdateConfig { + config_updates: None, + table_metadata_updates: None, + schema_metadata_updates: None, + field_metadata_updates: HashMap::new(), + }; + + let err = + validate_stable_field_id_transition(&manifest, &successor, &operation).unwrap_err(); + + assert!( + err.to_string().contains("reader fence cannot be removed"), + "{err}" + ); + } + #[test] fn detached_stable_field_ids_allow_data_only_merge_and_reject_schema_change() { let manifest = activated_manifest(); diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index b984bb645b4..06a05172611 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1323,7 +1323,10 @@ async fn test_write_manifest( manifest.data_storage_format.version.to_manifest_string(), "stable" | "next" )); - assert_eq!(manifest.reader_feature_flags, 0); + assert_eq!( + manifest.reader_feature_flags, + feature_flags::FLAG_STABLE_FIELD_IDS + ); assert_eq!( manifest.writer_feature_flags, feature_flags::FLAG_STABLE_FIELD_IDS @@ -1352,7 +1355,7 @@ async fn test_write_manifest( ); assert_eq!( manifest.reader_feature_flags, - feature_flags::FLAG_DELETION_FILES + feature_flags::FLAG_DELETION_FILES | feature_flags::FLAG_STABLE_FIELD_IDS ); // Write with custom manifest diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index a44ec5d1869..d615425e158 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -339,6 +339,14 @@ async fn test_deep_clone_repairs_legacy_schema_before_activation() { cloned.delete("false").await.unwrap(); cloned.validate().await.unwrap(); assert!(cloned.manifest.uses_stable_field_ids()); + assert_ne!( + cloned.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + assert_ne!( + cloned.manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); } #[tokio::test] @@ -684,7 +692,7 @@ async fn test_new_datasets_use_stable_field_ids_and_migration_is_idempotent() { let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; assert!(dataset.manifest.uses_stable_field_ids()); assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); - assert_eq!( + assert_ne!( dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); @@ -696,7 +704,7 @@ async fn test_new_datasets_use_stable_field_ids_and_migration_is_idempotent() { dataset.migrate_to_stable_field_ids().await.unwrap(); assert_eq!(dataset.version().version, created_version); - assert_eq!( + assert_ne!( dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); @@ -728,6 +736,10 @@ async fn test_stable_field_id_restore_boundary_and_high_water_mark() { let mut activation_snapshot = dataset.checkout_version(activation_version).await.unwrap(); activation_snapshot.restore().await.unwrap(); assert_eq!(activation_snapshot.manifest.max_allocated_field_id, Some(1)); + assert_ne!( + activation_snapshot.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); assert!(activation_snapshot.schema().field("new_field").is_none()); } @@ -747,7 +759,7 @@ async fn test_shallow_clone_preserves_stable_field_id_state() { cloned.manifest.max_allocated_field_id, dataset.manifest.max_allocated_field_id ); - assert_eq!( + assert_ne!( cloned.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index ba33a5f25d3..2ca8cb54fad 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -498,6 +498,7 @@ async fn do_commit_new_dataset( if !manifest.uses_stable_field_ids() { fix_schema(&mut manifest)?; manifest.activate_stable_field_ids(); + manifest.reader_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; } From 399260c94c594f6d30728e90832924c7f726cf29 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 1 Sep 2026 01:31:39 +0800 Subject: [PATCH 15/24] fix: align stable field IDs in overwrite fragments --- .../tests/compat/test_scalar_indices.py | 4 ++-- python/python/tests/test_multi_base.py | 5 +++++ rust/lance-core/src/datatypes/schema.rs | 4 ++-- rust/lance/src/dataset/write.rs | 21 ++++++++++++++++--- 4 files changed, 27 insertions(+), 7 deletions(-) diff --git a/python/python/tests/compat/test_scalar_indices.py b/python/python/tests/compat/test_scalar_indices.py index cebbdd72b9d..b9fcb9def27 100644 --- a/python/python/tests/compat/test_scalar_indices.py +++ b/python/python/tests/compat/test_scalar_indices.py @@ -261,7 +261,7 @@ def check_write(self): def skip_downgrade(self, version: str) -> bool: # In 0.X the zonemap index did not properly handle NULL in filters - return version.startswith("0.") + return super().skip_downgrade(version) or version.startswith("0.") @compat_test(min_version="0.36.0") @@ -376,7 +376,7 @@ def check_write(self): ds.optimize.compact_files() def skip_downgrade(self, version: str) -> bool: - return version.startswith("0.") + return super().skip_downgrade(version) or version.startswith("0.") def current_env(self, method_name: str) -> dict[str, str]: if method_name == "create": diff --git a/python/python/tests/test_multi_base.py b/python/python/tests/test_multi_base.py index 128a8905842..d905bb92a75 100644 --- a/python/python/tests/test_multi_base.py +++ b/python/python/tests/test_multi_base.py @@ -1225,6 +1225,11 @@ def test_write_fragments_overwrite_mode_with_target_bases(self): ) assert len(fragments) > 0 + assert all( + data_file.fields == [0, 1] + for fragment in fragments + for data_file in fragment.files + ) # Commit with Overwrite operation operation = lance.LanceOperation.Overwrite( diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index 883a2ce5f07..c6d7d96f0cd 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -715,8 +715,8 @@ impl Schema { /// Replace every field ID with a fresh checked allocation. /// - /// The first assigned ID is one greater than `max_existing_id`. This is used - /// when an operation replaces the complete schema identity, such as overwrite. + /// The first assigned ID is one greater than `max_existing_id`. Use this when + /// every input field must receive a new identity. pub fn try_reassign_field_ids(&mut self, max_existing_id: Option) -> Result<()> { let field_count = self.fields_pre_order().count() as i64; let first_id = i64::from(max_existing_id.unwrap_or(-1)) + 1; diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 1650de454d2..6eb9592c2ea 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -29,6 +29,7 @@ use lance_io::traits::Writer; use lance_table::format::{BasePath, DataFile, Fragment, IndexMetadata}; use lance_table::io::commit::{CommitHandler, commit_handler_from_url}; use lance_table::io::manifest::ManifestDescribing; +use lance_table::transaction::{Operation, canonicalize_stable_field_ids}; use object_store::path::Path; use std::borrow::Cow; use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; @@ -1527,9 +1528,23 @@ pub(super) fn prepare_write_schema( && matches!(params.mode, WriteMode::Overwrite) && dataset.manifest.uses_stable_field_ids() { - let mut schema = normalized_converted_schema; - schema.try_reassign_field_ids(Some(dataset.manifest.max_field_id()))?; - schema + // Uncommitted fragment APIs return files without the schema used to + // write them, so their mappings must already use commit-time IDs. + let mut operation = Operation::Overwrite { + fragments: Vec::new(), + schema: normalized_converted_schema, + config_upsert_values: None, + initial_bases: None, + }; + canonicalize_stable_field_ids(Some(&dataset.manifest), &mut operation)?; + match operation { + Operation::Overwrite { schema, .. } => schema, + _ => { + return Err(Error::internal( + "Stable field-ID canonicalization changed an Overwrite operation", + )); + } + } } else { normalized_converted_schema }; From 6fde30fa09b85dc544e77835de727362017c2f40 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 1 Sep 2026 14:39:38 +0800 Subject: [PATCH 16/24] fix: enforce stable field ID commit boundaries --- docs/src/format/table/schema.md | 11 +-- docs/src/format/table/versioning.md | 2 +- .../org/lance/operation/SchemaOperation.java | 7 +- protos/table.proto | 6 +- .../python/tests/compat/compat_decorator.py | 7 +- .../python/tests/compat/test_file_formats.py | 30 ------ python/python/tests/test_dataset.py | 90 ----------------- rust/lance-table/src/feature_flags.rs | 65 ++++++------ rust/lance-table/src/format/manifest.rs | 48 ++++----- .../src/transaction/manifest_build.rs | 14 +-- rust/lance-table/src/transaction/validate.rs | 71 +++++++------- rust/lance/src/dataset.rs | 8 +- rust/lance/src/dataset/tests/dataset_io.rs | 16 +-- .../src/dataset/tests/dataset_migrations.rs | 98 ++++++++++++++++--- rust/lance/src/dataset/write.rs | 13 ++- rust/lance/src/io/commit.rs | 3 - 16 files changed, 201 insertions(+), 288 deletions(-) diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index 9ca3a3c36f0..0a8b2c9a1ab 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -237,13 +237,10 @@ maximum non-negative ID referenced by the canonical manifest schema, base data f files. Activation provides a forward guarantee only; it cannot reconstruct identities that were dropped or reused in older snapshots. -New datasets activate this contract in their initial manifest and set the corresponding reader and -writer feature bits. The reader bit is a compatibility fence for released runtimes whose generic -commit path did not enforce unknown writer bits; stable field IDs do not otherwise change how fields -are read. Existing legacy datasets remain unchanged until an explicit migration commit. Explicit -migration sets only the writer bit and requires operators to retire older writers before activation. -The reader fence, once set by automatic activation, is retained by every later commit, clone, and -restore in that branch ancestry. +Datasets retain legacy allocation semantics until an explicit migration commit. Migration sets the +writer feature bit and requires operators to retire every older writer before activation, including +writers that commit from an already serialized manifest. Stable field IDs do not change how fields +are read and therefore do not set a reader feature bit. Activation is one-way within a branch ancestry. After migration, restore cannot target a version from before activation because that version does not carry the high-water mark needed to preserve diff --git a/docs/src/format/table/versioning.md b/docs/src/format/table/versioning.md index 0eaf44484d0..ee1aa312c1a 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -32,7 +32,7 @@ they should return an "unsupported" error on any read or write operation. | 64 | `FLAG_UNSTABLE_DATA_OVERLAY_FILES` | Yes | Yes | Fragments may carry data overlay files. Unstable: release builds reject it unless explicitly opted in. | | 128 | `FLAG_COVERED_INDEX_METADATA` | Yes | Yes | Some index declares covering columns (`IndexMetadata.covering_fields`), so `fields` means keyed columns followed by carried ones. An implementation without this flag selects an index by membership of `fields` and would answer a query on a merely-carried column with an index keyed on a different one. | | 256 | `FLAG_MIXED_DATA_FILE_VERSIONS` | Yes | Yes | Reserved for datasets that may reference recognized V2 data files with different exact versions. | -| 512 | `FLAG_STABLE_FIELD_IDS` | Automatic activation | Yes | The manifest carries a persistent field-ID high-water mark that writers must use when allocating field IDs. Automatic activation also fences released runtimes whose write paths did not consistently enforce writer flags; explicit migration is writer-only. See [Field IDs](schema.md#field-ids). | +| 512 | `FLAG_STABLE_FIELD_IDS` | No | Yes | After explicit migration, the manifest carries a persistent field-ID high-water mark that writers must use when allocating field IDs. Operators must retire older writers before activation. See [Field IDs](schema.md#field-ids). | diff --git a/java/src/main/java/org/lance/operation/SchemaOperation.java b/java/src/main/java/org/lance/operation/SchemaOperation.java index 2fbcc289001..9e27ed1976e 100644 --- a/java/src/main/java/org/lance/operation/SchemaOperation.java +++ b/java/src/main/java/org/lance/operation/SchemaOperation.java @@ -31,10 +31,9 @@ *

  • Otherwise, allocate based on the max field id of the dataset. * * - *

    New datasets use stable field IDs. On those datasets, metadata may identify an existing field - * for rename or reorder, but it cannot allocate a new identity. New IDs are assigned by the dataset - * allocator, and field mappings in fragments committed by the same operation are remapped to the - * canonical IDs. + *

    Datasets use stable field IDs only after explicit migration. On those datasets, compatible + * existing fields keep their identities, while new IDs are assigned by the dataset allocator and + * field mappings in fragments committed by the same operation are remapped to the canonical IDs. */ public abstract class SchemaOperation implements Operation { private final Schema schema; diff --git a/protos/table.proto b/protos/table.proto index 3f2d9dde5ab..50786545fa0 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -128,9 +128,6 @@ message Manifest { // * 1 << 8: reserved for datasets that may reference recognized V2 data files // with different exact versions. Implementations that do not support the // per-file exact-version contract must treat this bit as unknown. - // * 1 << 9: the stable field-ID contract was activated automatically. The - // reader fence excludes released runtimes whose generic commit path did not - // enforce unknown writer flags; it does not change how fields are read. uint64 reader_feature_flags = 9; // Feature flags for writers. @@ -143,8 +140,7 @@ message Manifest { // reader_feature_flags and writer_feature_flags are not required to be identical. // * 1 << 9: newly allocated field IDs must be greater than // max_allocated_field_id, and the high-water mark must advance atomically. - // Automatic activation sets this bit in both flag words. Explicit migration - // may set only the writer bit after older writers have been retired. + // Explicit migration sets this bit after older writers have been retired. uint64 writer_feature_flags = 10; // The highest fragment ID that has been used so far. diff --git a/python/python/tests/compat/compat_decorator.py b/python/python/tests/compat/compat_decorator.py index 263eac32bb6..1ebdba59a6a 100644 --- a/python/python/tests/compat/compat_decorator.py +++ b/python/python/tests/compat/compat_decorator.py @@ -202,12 +202,7 @@ def compat_env(self, version: str, method_name: str) -> Dict[str, str]: class DatasetUpgradeDowngradeTest(UpgradeDowngradeTest): - """Compatibility contract for new datasets with stable field IDs.""" - - def skip_downgrade(self, version: str) -> bool: - # Automatically activated datasets fence released runtimes from opening - # them. A dedicated v11 regression verifies this refusal. - return True + """Compatibility contract for datasets.""" @contextmanager diff --git a/python/python/tests/compat/test_file_formats.py b/python/python/tests/compat/test_file_formats.py index 93d6b2d989a..be2b533411b 100644 --- a/python/python/tests/compat/test_file_formats.py +++ b/python/python/tests/compat/test_file_formats.py @@ -11,7 +11,6 @@ from pathlib import Path import lance -import pyarrow as pa import pytest from lance.file import LanceFileReader, LanceFileWriter @@ -140,32 +139,3 @@ def check_write(self): lance.write_dataset( build_basic_types(), self.path, data_storage_version="0.1", mode="append" ) - - -class StableFieldIdWriterFence: - def __init__(self, path: Path): - self.path = path - - def create(self): - dataset = lance.write_dataset(pa.table({"a": [1]}), self.path) - dataset.add_columns({"retired": "a"}) - dataset.drop_columns(["retired"]) - - def assert_write_rejected(self): - with pytest.raises(ValueError, match="cannot be read by this version of Lance"): - lance.dataset(self.path) - - def add_replacement_and_check(self): - dataset = lance.dataset(self.path) - dataset.add_columns({"replacement": "a"}) - assert dataset.lance_schema.field("replacement").id() == 2 - - -@pytest.mark.compat -def test_released_writer_cannot_erase_field_id_high_water(venv_factory, tmp_path): - case = StableFieldIdWriterFence(tmp_path / "data.lance") - case.create() - - venv_factory.get_venv("11.0.0").execute_method(case, "assert_write_rejected") - - case.add_replacement_and_check() diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 8dd857bef5d..a6fa811f558 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -5781,96 +5781,6 @@ def test_schema_project_drop_column(tmp_path: Path): assert tbl == expected -def test_schema_project_raw_arrow_cannot_allocate_field_id(tmp_path: Path): - dataset = lance.write_dataset( - pa.table({"a": pa.array([1, 2], pa.int32())}), tmp_path - ) - schema = pa.schema([pa.field("a", pa.int32(), metadata={b"lance:field_id": b"42"})]) - - project = lance.LanceOperation.Project(schema) - dataset = lance.LanceDataset.commit(dataset, project, read_version=dataset.version) - - assert dataset.lance_schema.fields()[0].id() == 0 - assert dataset.to_table() == pa.table({"a": pa.array([1, 2], pa.int32())}) - assert b"lance:transaction_schema_source_raw_arrow" not in ( - dataset.schema.metadata or {} - ) - - -def test_overwrite_rejects_ambiguous_raw_arrow_field_ids(tmp_path: Path): - dataset = lance.write_dataset( - pa.table({"old": pa.array([1, 2], pa.int32())}), tmp_path - ) - schema = pa.schema( - [ - pa.field("b", pa.int32(), metadata={b"lance:field_id": b"2"}), - pa.field("c", pa.int32(), metadata={b"lance:field_id": b"1"}), - ] - ) - table = pa.Table.from_arrays( - [pa.array([10, 11], pa.int32()), pa.array([20, 21], pa.int32())], - schema=schema, - ) - fragment = lance.fragment.LanceFragment.create(tmp_path, table, mode="overwrite") - - with pytest.raises(OSError, match="ambiguous raw Arrow field IDs"): - lance.LanceDataset.commit( - dataset, - lance.LanceOperation.Overwrite(schema, [fragment]), - read_version=dataset.version, - ) - - -def test_merge_rejects_ambiguous_raw_arrow_field_ids(tmp_path: Path): - dataset = lance.write_dataset( - pa.table({"old": pa.array([1, 2], pa.int32())}), tmp_path - ) - staged = lance.fragment.LanceFragment.create( - tmp_path, - pa.table( - { - "b": pa.array([10, 11], pa.int32()), - "c": pa.array([20, 21], pa.int32()), - } - ), - mode="overwrite", - ) - staged_file = staged.files[0] - canonical_file = lance.fragment.DataFile( - path=staged_file.path, - fields=[1, 2], - column_indices=staged_file.column_indices, - file_major_version=staged_file.file_major_version, - file_minor_version=staged_file.file_minor_version, - file_size_bytes=staged_file.file_size_bytes, - base_id=staged_file.base_id, - ) - current = dataset.get_fragments()[0].metadata - merged = lance.fragment.FragmentMetadata( - id=current.id, - files=[*current.files, canonical_file], - physical_rows=current.physical_rows, - deletion_file=current.deletion_file, - row_id_meta=current.row_id_meta, - ) - schema = pa.schema( - [ - pa.field("old", pa.int32()), - pa.field("b", pa.int32(), metadata={b"lance:field_id": b"2"}), - pa.field("c", pa.int32(), metadata={b"lance:field_id": b"1"}), - ] - ) - - with pytest.deprecated_call(): - operation = lance.LanceOperation.Merge([merged], schema) - with pytest.raises(OSError, match="ambiguous raw Arrow field IDs"): - lance.LanceDataset.commit( - dataset, - operation, - read_version=dataset.version, - ) - - def test_schema_project_rename_column(tmp_path: Path): table = pa.Table.from_pydict({"a": range(100, 200), "b": range(300, 400)}) base_dir = tmp_path / "test" diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index 1207b38f61f..a0761f6008a 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -70,7 +70,7 @@ const _: () = assert!(FLAG_MIXED_DATA_FILE_VERSIONS < FLAG_UNKNOWN); const _: () = assert!(FLAG_STABLE_FIELD_IDS < FLAG_UNKNOWN); pub(crate) const STICKY_PAIRED_FLAGS: u64 = FLAG_MIXED_DATA_FILE_VERSIONS; -pub(crate) const STICKY_READER_FLAGS: u64 = STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS; +pub(crate) const STICKY_READER_FLAGS: u64 = STICKY_PAIRED_FLAGS; pub(crate) const STICKY_WRITER_FLAGS: u64 = STICKY_PAIRED_FLAGS | FLAG_STABLE_FIELD_IDS; /// Environment variable that opts a release build into reading and writing data @@ -92,7 +92,11 @@ pub fn apply_feature_flags( let covered_index_metadata = (manifest.reader_feature_flags | manifest.writer_feature_flags) & FLAG_COVERED_INDEX_METADATA; let sticky_paired_flags = validated_sticky_paired_flags(manifest)?; - let stable_field_id_reader_fence = manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS; + let stable_field_ids = manifest.max_allocated_field_id.is_some(); + if stable_field_ids { + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + } + validate_stable_field_id_flags(manifest)?; // Reset flags manifest.reader_feature_flags = 0; manifest.writer_feature_flags = 0; @@ -151,8 +155,7 @@ pub fn apply_feature_flags( manifest.writer_feature_flags |= FLAG_DISABLE_TRANSACTION_FILE; } - if manifest.max_allocated_field_id.is_some() { - manifest.reader_feature_flags |= stable_field_id_reader_fence; + if stable_field_ids { manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; } @@ -171,14 +174,11 @@ pub fn apply_feature_flags( /// also validates that the source is not half-set before a derived manifest is /// committed. /// -/// Stable field IDs permit writer-only fencing after an explicit migration, -/// while automatic activation also carries a reader fence to exclude released -/// writers that did not enforce unknown writer flags on every commit path. +/// Stable field IDs are activated explicitly and only require writer support. pub fn inherit_sticky_feature_flags(destination: &mut Manifest, source: &Manifest) -> Result<()> { let sticky_flags = validated_sticky_paired_flags(source)?; validate_stable_field_id_flags(source)?; - destination.reader_feature_flags |= - sticky_flags | (source.reader_feature_flags & FLAG_STABLE_FIELD_IDS); + destination.reader_feature_flags |= sticky_flags; destination.writer_feature_flags |= sticky_flags | (source.writer_feature_flags & FLAG_STABLE_FIELD_IDS); Ok(()) @@ -287,10 +287,8 @@ pub fn validate_paired_feature_flags(manifest: &Manifest) -> Result<()> { /// Refuse a manifest whose stable-field-ID marker and required flags disagree. /// /// The high-water mark is the activation marker and always requires the writer -/// bit. Automatic activation also sets the reader bit as a compatibility fence -/// against released writers that did not enforce unknown writer flags on every -/// commit path. Explicit migration may remain writer-only once older writers -/// have been retired. +/// bit. Stable field IDs do not change read semantics, so the reader bit is not +/// a valid activation mode. pub fn validate_stable_field_id_flags(manifest: &Manifest) -> Result<()> { let activated = manifest.max_allocated_field_id.is_some(); let reader = manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0; @@ -301,10 +299,10 @@ pub fn validate_stable_field_id_flags(manifest: &Manifest) -> Result<()> { "Manifest stable-field-ID high-water mark and writer feature flag disagree", )); } - if reader && !writer { + if reader { return Err(Error::corrupt_file_named( "manifest", - "Manifest has the stable-field-ID reader fence without the activation marker and writer feature flag", + "Manifest has a stable-field-ID reader feature flag, but stable field IDs only require writer support", )); } Ok(()) @@ -518,24 +516,16 @@ mod tests { } #[test] - fn inheriting_preserves_stable_field_id_fence_mode() { - for reader_fenced in [false, true] { - let mut source = empty_manifest(); - source.activate_stable_field_ids(); - source.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - if reader_fenced { - source.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - } - let mut destination = empty_manifest(); + fn inheriting_preserves_stable_field_id_writer_gate() { + let mut source = empty_manifest(); + source.activate_stable_field_ids(); + source.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + let mut destination = empty_manifest(); - inherit_sticky_feature_flags(&mut destination, &source).unwrap(); + inherit_sticky_feature_flags(&mut destination, &source).unwrap(); - assert_eq!( - destination.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0, - reader_fenced - ); - assert_ne!(destination.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); - } + assert_eq!(destination.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(destination.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } #[test] @@ -607,17 +597,18 @@ mod tests { } #[test] - fn apply_feature_flags_preserves_automatic_stable_field_id_reader_fence() { + fn apply_feature_flags_rejects_stable_field_id_reader_flag() { let mut manifest = empty_manifest(); manifest.activate_stable_field_ids(); manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - apply_feature_flags(&mut manifest, false, false).unwrap(); - apply_feature_flags(&mut manifest, false, false).unwrap(); + let err = apply_feature_flags(&mut manifest, false, false).unwrap_err(); - assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); - assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert!( + err.to_string().contains("only require writer support"), + "{err}" + ); } #[test] @@ -637,7 +628,7 @@ mod tests { let mut paired = writer_only.clone(); paired.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - validate_stable_field_id_flags(&paired).unwrap(); + assert!(validate_stable_field_id_flags(&paired).is_err()); let mut reader_without_activation = empty_manifest(); reader_without_activation.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index e6bb69f872b..a134f626ef5 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -1673,41 +1673,33 @@ mod tests { } #[test] - fn shallow_clone_preserves_stable_field_id_allocation_state_and_fence_mode() { + fn shallow_clone_preserves_stable_field_id_allocation_state() { let arrow_schema = ArrowSchema::new(vec![ArrowField::new( "a", arrow_schema::DataType::Int64, false, )]); let schema = Schema::try_from(&arrow_schema).unwrap(); - for reader_fenced in [false, true] { - let mut manifest = Manifest::new( - schema.clone(), - Arc::new(vec![]), - DataStorageFormat::default(), - HashMap::new(), - ); - manifest.max_allocated_field_id = Some(41); - manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - if reader_fenced { - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - } + let mut manifest = Manifest::new( + schema, + Arc::new(vec![]), + DataStorageFormat::default(), + HashMap::new(), + ); + manifest.max_allocated_field_id = Some(41); + manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - let cloned = manifest.shallow_clone( - Some("parent".to_string()), - "memory://parent".to_string(), - 7, - None, - String::new(), - ); - - assert_eq!(cloned.max_allocated_field_id, Some(41)); - assert_eq!( - cloned.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0, - reader_fenced - ); - assert_ne!(cloned.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); - } + let cloned = manifest.shallow_clone( + Some("parent".to_string()), + "memory://parent".to_string(), + 7, + None, + String::new(), + ); + + assert_eq!(cloned.max_allocated_field_id, Some(41)); + assert_eq!(cloned.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_ne!(cloned.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } #[test] diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index abeace843f7..b0b3811f8f2 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1303,12 +1303,6 @@ impl Transaction { ) }; - if current_manifest.is_none() { - manifest.activate_stable_field_ids(); - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; - } - if config.activate_stable_field_ids { let already_active = current_manifest .map(|manifest| manifest.uses_stable_field_ids()) @@ -1607,7 +1601,7 @@ mod tests { } #[test] - fn new_dataset_builds_with_stable_field_ids_by_default() { + fn new_dataset_preserves_legacy_field_id_allocation_by_default() { let arrow_schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); let mut schema = LanceSchema::try_from(&arrow_schema).unwrap(); schema.try_set_field_id(None).unwrap(); @@ -1626,9 +1620,9 @@ mod tests { .build_manifest(None, vec![], "txn", &default_build_config()) .unwrap(); - assert_eq!(manifest.max_allocated_field_id, Some(0)); - assert_ne!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); - assert_ne!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_eq!(manifest.max_allocated_field_id, None); + assert_eq!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + assert_eq!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); } #[test] diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 28519a8c67b..74b992723e8 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -7,7 +7,6 @@ //! a fragment list that disagrees with the schema, a merge that silently dropped //! or rewrote data files — before any manifest is written. -use crate::feature_flags::FLAG_STABLE_FIELD_IDS; use crate::format::{Fragment, Manifest}; use crate::io::deletion::relative_deletion_file_path; use crate::transaction::{Operation, UpdateMode, UpdatedFragmentOffsets}; @@ -71,7 +70,8 @@ pub fn canonicalize_stable_field_ids( Operation::Overwrite { schema, fragments, .. } => { - let field_id_remap = canonicalize_schema(manifest, schema, true, raw_arrow_schema)?; + let field_id_remap = + canonicalize_schema(manifest, schema, !raw_arrow_schema, raw_arrow_schema)?; remap_fragment_field_ids(fragments, &field_id_remap, &HashSet::new())?; } Operation::Project { schema, .. } if raw_arrow_schema => { @@ -472,13 +472,6 @@ pub fn validate_stable_field_id_transition( "Stable field-ID high-water mark decreases from {parent_max_field_id} to {successor_max_field_id}" ))); } - if parent.reader_feature_flags & FLAG_STABLE_FIELD_IDS != 0 - && successor.reader_feature_flags & FLAG_STABLE_FIELD_IDS == 0 - { - return Err(Error::invalid_input( - "Stable field-ID reader fence cannot be removed after automatic activation", - )); - } let successor_schema_ids = successor .schema .fields_pre_order() @@ -1150,6 +1143,40 @@ mod tests { assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]); } + #[test] + fn canonicalize_raw_arrow_overwrite_matches_reordered_fields_by_name() { + let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("b", DataType::Int32, true), + ])) + .unwrap(); + let mut manifest = manifest_with_file_fields(schema, vec![0, 1]); + manifest.activate_stable_field_ids(); + let mut raw_schema = LanceSchema::try_from(&ArrowSchema::new(vec![ + ArrowField::new("b", DataType::Int32, true), + ArrowField::new("a", DataType::Int32, true), + ])) + .unwrap(); + raw_schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); + let mut operation = Operation::Overwrite { + fragments: vec![], + schema: raw_schema, + config_upsert_values: None, + initial_bases: None, + }; + + canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Overwrite { schema, .. } = operation else { + unreachable!(); + }; + assert_eq!(schema.field("b").unwrap().id, 1); + assert_eq!(schema.field("a").unwrap().id, 0); + } + #[test] fn canonicalize_overwrite_preserves_already_canonical_fragment_field_ids() { let manifest = activated_manifest(); @@ -1497,32 +1524,6 @@ mod tests { ); } - #[test] - fn stable_field_id_transition_rejects_removing_reader_fence() { - let mut manifest = activated_manifest(); - manifest.reader_feature_flags |= FLAG_STABLE_FIELD_IDS; - let mut successor = Manifest::new_from_previous( - &manifest, - manifest.schema.clone(), - manifest.fragments.clone(), - ); - successor.reader_feature_flags &= !FLAG_STABLE_FIELD_IDS; - let operation = Operation::UpdateConfig { - config_updates: None, - table_metadata_updates: None, - schema_metadata_updates: None, - field_metadata_updates: HashMap::new(), - }; - - let err = - validate_stable_field_id_transition(&manifest, &successor, &operation).unwrap_err(); - - assert!( - err.to_string().contains("reader fence cannot be removed"), - "{err}" - ); - } - #[test] fn detached_stable_field_ids_allow_data_only_merge_and_reject_schema_change() { let manifest = activated_manifest(); diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index b4957e76278..0064aa190fa 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -3252,10 +3252,10 @@ impl Dataset { /// /// The activation commit records the current maximum referenced field ID as /// a persistent high-water mark. Later schema changes allocate above it even - /// after fields and their files are dropped. New datasets already use this - /// contract by default. Activation is one-way and idempotent. Before - /// migrating, ensure that every writer accessing the dataset understands - /// the stable field-ID writer feature flag. + /// after fields and their files are dropped. Activation is one-way and + /// idempotent. Before migrating, retire every writer that does not understand + /// the stable field-ID writer feature flag, including writers that commit + /// from an already serialized manifest. /// /// ``` /// # use lance::{Dataset, Result}; diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index 06a05172611..ff32459025a 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1301,7 +1301,7 @@ async fn test_write_manifest( let write_fut = require_send(write_fut); let mut dataset = write_fut.await.unwrap(); - // New datasets enable stable field IDs by default. + // New datasets retain legacy field-ID allocation until explicitly migrated. let manifest = read_manifest( dataset.object_store.as_ref(), &dataset @@ -1323,14 +1323,8 @@ async fn test_write_manifest( manifest.data_storage_format.version.to_manifest_string(), "stable" | "next" )); - assert_eq!( - manifest.reader_feature_flags, - feature_flags::FLAG_STABLE_FIELD_IDS - ); - assert_eq!( - manifest.writer_feature_flags, - feature_flags::FLAG_STABLE_FIELD_IDS - ); + assert_eq!(manifest.reader_feature_flags, 0); + assert_eq!(manifest.writer_feature_flags, 0); // Create one with deletions dataset.delete("i < 10").await.unwrap(); @@ -1351,11 +1345,11 @@ async fn test_write_manifest( .unwrap(); assert_eq!( manifest.writer_feature_flags, - feature_flags::FLAG_DELETION_FILES | feature_flags::FLAG_STABLE_FIELD_IDS + feature_flags::FLAG_DELETION_FILES ); assert_eq!( manifest.reader_feature_flags, - feature_flags::FLAG_DELETION_FILES | feature_flags::FLAG_STABLE_FIELD_IDS + feature_flags::FLAG_DELETION_FILES ); // Write with custom manifest diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index d615425e158..406bf57d766 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -326,7 +326,7 @@ async fn test_fix_v0_10_5_corrupt_schema() { } #[tokio::test] -async fn test_deep_clone_repairs_legacy_schema_before_activation() { +async fn test_deep_clone_repairs_legacy_schema_without_activation() { let source_dir = copy_test_data_to_tmp("v0.10.5/corrupt_schema").unwrap(); let clone_uri = TempStrDir::default(); let mut source = Dataset::open(&source_dir.path_str()).await.unwrap(); @@ -338,12 +338,12 @@ async fn test_deep_clone_repairs_legacy_schema_before_activation() { cloned.delete("false").await.unwrap(); cloned.validate().await.unwrap(); - assert!(cloned.manifest.uses_stable_field_ids()); - assert_ne!( + assert!(!cloned.manifest.uses_stable_field_ids()); + assert_eq!( cloned.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); - assert_ne!( + assert_eq!( cloned.manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); @@ -687,36 +687,44 @@ async fn make_simple_dataset(uri: &str, n: i64) -> Dataset { } #[tokio::test] -async fn test_new_datasets_use_stable_field_ids_and_migration_is_idempotent() { +async fn test_new_datasets_use_legacy_field_ids_until_explicit_migration() { let source_uri = TempStrDir::default(); let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; - assert!(dataset.manifest.uses_stable_field_ids()); - assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); - assert_ne!( + assert!(!dataset.manifest.uses_stable_field_ids()); + assert_eq!(dataset.manifest.max_allocated_field_id, None); + assert_eq!( dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); - assert_ne!( + assert_eq!( dataset.manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); let created_version = dataset.version().version; dataset.migrate_to_stable_field_ids().await.unwrap(); - assert_eq!(dataset.version().version, created_version); - assert_ne!( + assert_eq!(dataset.version().version, created_version + 1); + assert!(dataset.manifest.uses_stable_field_ids()); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); + assert_eq!( dataset.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, + 0 + ); + let activation_version = dataset.version().version; dataset.migrate_to_stable_field_ids().await.unwrap(); - assert_eq!(dataset.version().version, created_version); + assert_eq!(dataset.version().version, activation_version); } #[tokio::test] async fn test_stable_field_id_restore_boundary_and_high_water_mark() { let source_uri = TempStrDir::default(); let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + dataset.migrate_to_stable_field_ids().await.unwrap(); let activation_version = dataset.version().version; dataset @@ -736,7 +744,7 @@ async fn test_stable_field_id_restore_boundary_and_high_water_mark() { let mut activation_snapshot = dataset.checkout_version(activation_version).await.unwrap(); activation_snapshot.restore().await.unwrap(); assert_eq!(activation_snapshot.manifest.max_allocated_field_id, Some(1)); - assert_ne!( + assert_eq!( activation_snapshot.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); @@ -759,7 +767,7 @@ async fn test_shallow_clone_preserves_stable_field_id_state() { cloned.manifest.max_allocated_field_id, dataset.manifest.max_allocated_field_id ); - assert_ne!( + assert_eq!( cloned.manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0 ); @@ -803,6 +811,65 @@ async fn test_overwrite_preserves_compatible_stable_field_identities() { assert_eq!(overwritten.manifest.max_allocated_field_id, Some(1)); } +#[tokio::test] +async fn test_raw_arrow_overwrite_preserves_reordered_stable_field_identities() { + let source_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("a", DataType::Int64, false), + ArrowField::new("b", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(Int64Array::from(vec![3, 4])), + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)], schema), + source_uri.as_str(), + None, + ) + .await + .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); + + let reordered_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("b", DataType::Int64, false), + ArrowField::new("a", DataType::Int64, false), + ])); + let reordered_batch = RecordBatch::try_new( + reordered_schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![30, 40])), + Arc::new(Int64Array::from(vec![10, 20])), + ], + ) + .unwrap(); + let overwritten = Dataset::write( + RecordBatchIterator::new(vec![Ok(reordered_batch)], reordered_schema), + source_uri.as_str(), + Some(WriteParams { + mode: WriteMode::Overwrite, + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!(overwritten.schema().field("b").unwrap().id, 1); + assert_eq!(overwritten.schema().field("a").unwrap().id, 0); + assert!( + overwritten + .manifest + .fragments + .iter() + .flat_map(|fragment| &fragment.files) + .all(|file| file.fields.as_ref() == [1, 0]) + ); +} + #[tokio::test] async fn test_stable_field_id_rename_and_nullability_preserve_identity() { let source_uri = TempStrDir::default(); @@ -852,6 +919,7 @@ async fn test_stable_field_id_multi_cast_uses_schema_order() { ) .await .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); dataset .alter_columns(&[ @@ -887,7 +955,7 @@ async fn test_new_dataset_ignores_hostile_arrow_field_id() { .unwrap(); assert_eq!(dataset.schema().field("a").unwrap().id, 0); - assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); + assert_eq!(dataset.manifest.max_allocated_field_id, None); dataset .add_columns( NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(vec![ArrowField::new( diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 6eb9592c2ea..56ebb873d6b 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -29,7 +29,9 @@ use lance_io::traits::Writer; use lance_table::format::{BasePath, DataFile, Fragment, IndexMetadata}; use lance_table::io::commit::{CommitHandler, commit_handler_from_url}; use lance_table::io::manifest::ManifestDescribing; -use lance_table::transaction::{Operation, canonicalize_stable_field_ids}; +use lance_table::transaction::{ + Operation, TRANSACTION_SCHEMA_SOURCE_RAW_ARROW, canonicalize_stable_field_ids, +}; use object_store::path::Path; use std::borrow::Cow; use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; @@ -1530,9 +1532,16 @@ pub(super) fn prepare_write_schema( { // Uncommitted fragment APIs return files without the schema used to // write them, so their mappings must already use commit-time IDs. + // The converted Arrow schema carries positional IDs, not trusted Lance + // identities, so preserve its provenance through canonicalization. + let mut raw_schema = normalized_converted_schema; + raw_schema.metadata.insert( + TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), + String::new(), + ); let mut operation = Operation::Overwrite { fragments: Vec::new(), - schema: normalized_converted_schema, + schema: raw_schema, config_upsert_values: None, initial_bases: None, }; diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 9dea373ff77..8a4eee4d348 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -497,9 +497,6 @@ async fn do_commit_new_dataset( if !manifest.uses_stable_field_ids() { fix_schema(&mut manifest)?; - manifest.activate_stable_field_ids(); - manifest.reader_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; - manifest.writer_feature_flags |= lance_table::feature_flags::FLAG_STABLE_FIELD_IDS; } let result = write_manifest_file( From 678e85590b316966386dbfdd6d53b97c8e33b80d Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 2 Sep 2026 18:19:26 +0800 Subject: [PATCH 17/24] fix: align stable field ID activation contract --- docs/src/format/table/index.md | 17 +- docs/src/format/table/schema.md | 139 ++++++++++------ docs/src/format/table/versioning.md | 2 +- java/lance-jni/src/transaction.rs | 157 +++++++++++++----- .../java/org/lance/operation/MergeTest.java | 80 +-------- .../java/org/lance/operation/ProjectTest.java | 29 ---- protos/table.proto | 18 +- rust/lance-table/src/transaction/validate.rs | 32 ++-- rust/lance/src/dataset.rs | 4 +- rust/lance/src/dataset/schema_evolution.rs | 1 + rust/lance/src/io/commit.rs | 6 +- 11 files changed, 265 insertions(+), 220 deletions(-) diff --git a/docs/src/format/table/index.md b/docs/src/format/table/index.md index a0037dd3f50..7f4a0c5fdd4 100644 --- a/docs/src/format/table/index.md +++ b/docs/src/format/table/index.md @@ -28,11 +28,18 @@ A manifest describes a single version of the dataset. It contains the complete schema definition including nested fields, the list of data fragments comprising this version, a monotonically increasing version number, and an optional reference to the index section that describes a list of index metadata. -`max_allocated_field_id` is optional allocator state. Its presence activates the stable field-ID -contract for this branch ancestry. It is a high-water mark: every non-negative field ID in the -canonical manifest schema, base data-file mappings, and overlay mappings must be less than or equal -to it, and a successor manifest must never lower it. Schema-preserving rewrites retain the value; -new identities allocate above it with checked arithmetic. See [Field IDs](schema.md#field-ids). +`max_allocated_field_id` is optional. If a manifest sets it, the dataset uses stable field IDs from +that version onward. See [Field IDs](schema.md#field-ids). + +The field is a high-water mark. It starts at the largest field ID that the activation manifest +references. It then records the largest ID assigned after activation. When a manifest sets it: + +- Every field ID of 0 or greater in the manifest schema, a `DataFile.fields` mapping, or an overlay + mapping must be less than or equal to `max_allocated_field_id`. +- A writer must not lower `max_allocated_field_id` from the value in the previous manifest. +- A writer that adds fields must assign IDs greater than the previous `max_allocated_field_id` and + set the new high-water mark to at least the largest ID it assigned. The writer must fail if an + assigned ID does not fit in an `int32`.

    Manifest protobuf message diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index 0a8b2c9a1ab..dfa1f0f4e04 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -221,47 +221,54 @@ Assigned IDs with parent relationships: Note: A `parent_id` of -1 indicates a top-level field. For nested fields, `parent_id` references the ID of the parent field. Child fields reference their parent via `parent_id` rather than being stored as separate "children" arrays in the protobuf message (though the Rust in-memory representation maintains a children vector for convenience). **New field assignment (incremental):** -When fields are added later (e.g., through schema evolution), they receive incrementally assigned -IDs. On a dataset with the stable field-ID contract activated, allocation starts immediately after -the manifest's persistent `max_allocated_field_id`; IDs retired by drop or replacement are never -reused. - -Field-ID metadata supplied on an incoming Arrow schema is not an allocation authority. When that -schema is merged into an activated dataset, IDs for newly introduced logical fields are cleared and -assigned by the dataset allocator; callers cannot select or reserve IDs through Arrow metadata. - -The presence of `Manifest.max_allocated_field_id` is the activation marker. If it is absent, the -dataset has legacy allocation semantics and an implementation may derive the next ID from fields -still referenced by the current snapshot. Activation initializes the high-water mark from the -maximum non-negative ID referenced by the canonical manifest schema, base data files, and overlay -files. Activation provides a forward guarantee only; it cannot reconstruct identities that were -dropped or reused in older snapshots. - -Datasets retain legacy allocation semantics until an explicit migration commit. Migration sets the -writer feature bit and requires operators to retire every older writer before activation, including -writers that commit from an already serialized manifest. Stable field IDs do not change how fields -are read and therefore do not set a reader feature bit. - -Activation is one-way within a branch ancestry. After migration, restore cannot target a version -from before activation because that version does not carry the high-water mark needed to preserve -retired identities. + +`Manifest.max_allocated_field_id` selects between two behaviors: + +- If the manifest does not set the field, the dataset uses the legacy behavior. A writer may choose + the next ID from fields the current version still references. It may therefore reuse the ID of a + dropped field. +- If the manifest sets the field, the dataset uses stable field IDs. A writer assigns each new field + an ID greater than `max_allocated_field_id`. It does not reuse an ID dropped or replaced after + activation. + +For stable field IDs, a caller cannot choose the ID of a new field. An Arrow schema may carry +field-ID metadata, but the writer discards that metadata for new fields and assigns the IDs. The IDs +do not have to be consecutive, which leaves room for a future reservation mechanism. + +The first manifest that sets `max_allocated_field_id` initializes it to the largest field ID of 0 +or greater in the manifest schema, base data files, and overlay files. Earlier versions keep the +legacy behavior. Activation cannot recover an ID that an earlier version dropped or reused. + +`max_allocated_field_id` stores the allocator state. `FLAG_STABLE_FIELD_IDS` tells writers that they +must honor that state. A legacy manifest sets neither value. A stable manifest sets both. A manifest +that sets only one is invalid. The reader flag for stable field IDs must remain unset because the +feature does not change read behavior. + +A dataset changes to stable field IDs only through an explicit migration commit. + +A dataset cannot return to the legacy behavior. After activation, a restore must fail if it targets +a version that does not set `max_allocated_field_id`. That version does not record retired field +IDs, so a later commit could reuse one. ### Field ID Properties -- **Immutable after activation**: An identity keeps its ID throughout the activated branch ancestry -- **Unique**: Each field within a table has a unique ID -- **Never reused after activation**: Dropped and replaced identities permanently retire their IDs -- **Monotonic**: New identities are allocated densely above the persistent high-water mark -- **Sparse**: Field IDs may not form a contiguous sequence after schema evolution +- **Stable**: A field keeps the same ID for as long as the field exists. +- **Unique**: No two fields in one dataset version have the same ID. +- **Sparse**: The field IDs in one version do not have to be consecutive. + +When `max_allocated_field_id` is set, two more properties apply: -The guarantee is scoped to one dataset branch ancestry. Field IDs are not globally unique across -datasets or independently evolving branches. A persistent cross-dataset or cross-branch reference -must carry the corresponding dataset and ancestry identity. +- **Not reused**: After activation, no later version uses the ID of a dropped or replaced field. +- **Increasing**: Every new ID is greater than the activation high-water mark and every ID assigned + after activation. -Two branches may allocate the same integer after their common ancestor. A future branch-merge -operation must reject the merge when both sides created post-ancestor identities whose bindings -differ; it must not silently choose one binding, renumber already-persisted fields, or merge them by -name. Branch-local collision avoidance requires a separate allocator design. +A field ID is unique within one branch of one dataset. It is not unique across datasets or across +branches that changed independently. A reference stored outside the dataset must name the dataset +and branch as well as the field ID. + +Two branches can assign the same field ID after they diverge. Lance does not yet merge branches. A +future merge operation must fail if the branches assigned the same ID to different fields. It must +not pick one field, change an ID stored by an existing version, or match the fields by name. ### Using Field IDs @@ -328,11 +335,13 @@ The complete schema is represented as a collection of top-level fields plus meta Field IDs enable efficient schema evolution: - **Add Column**: Assign a new field ID and add to schema -- **Drop Column**: Remove the field and permanently retire its ID after activation +- **Drop Column**: Remove the field from the schema; when `max_allocated_field_id` is set, later + versions must not reuse its ID - **Rename Column**: Change field name; ID remains the same - **Reorder Columns**: Change field order in schema; IDs remain the same - **Metadata or Nullability Change**: Preserve the field ID -- **Type Replacement**: Allocate a new field ID and retire the old identity +- **Type Replacement**: A cast creates a replacement field with a new ID and retires the old + identity. This keeps one logical type bound to an ID in every version that references it - **Overwrite**: Preserve compatible logical identities; allocate new IDs for added fields and type replacements @@ -340,16 +349,52 @@ The use of field IDs ensures that data files can be correctly interpreted even a ### Blob Identity Namespace -A Blob column's canonical logical fields in the manifest schema participate in stable field-ID -allocation. The top-level Blob field is the public binding identity; logical children already -present in the manifest schema are also allocated and retired normally. +The rules above apply to a Blob field in the manifest schema and to its logical children. For +example, assume `image` has field ID 0, `data` has ID 1, and `uri` has ID 2. Writer input and the +manifest schema have this logical shape: + +```python +pa.schema([ + pa.field( + "image", + pa.struct([ + pa.field("data", pa.large_binary()), + pa.field("uri", pa.string()), + ]), + metadata={b"ARROW:extension:name": b"lance.blob.v2"}, + ), +]) +``` + +The writer may temporarily add `kind`, `blob_id`, `blob_size`, and `position`. A Lance data file +stores this descriptor shape: + +```python +pa.schema([ + pa.field( + "image", + pa.struct([ + pa.field("kind", pa.uint8(), nullable=False), + pa.field("position", pa.uint64(), nullable=False), + pa.field("size", pa.uint64(), nullable=False), + pa.field("blob_id", pa.uint32(), nullable=False), + pa.field("blob_uri", pa.string(), nullable=False), + ]), + ), +]) +``` + +The data file maps this column with `DataFile.fields = [0]`. The descriptor children are file +details. Their IDs may be `-1` or file-local, and they do not change `max_allocated_field_id`. + +A descriptor scan returns the stored struct. A materialized scan returns this public shape: + +```python +pa.schema([pa.field("image", pa.large_binary())]) +``` -Blob writer-prepared fields and stored descriptor fields are representation details. Synthetic -children such as `kind`, `blob_id`, `blob_size`, `position`, `size`, and `blob_uri` do not enter the -dataset field-ID namespace unless they are part of the canonical manifest schema or a -`DataFile.fields` mapping. Their IDs may remain `-1` or use a file-local namespace, and they do not -advance `max_allocated_field_id`. The `blob_id` value identifies a sidecar object and is unrelated -to schema field IDs. +Both scans refer to the top-level field ID 0. The synthetic descriptor children do not become +dataset fields. The `blob_id` value identifies a stored Blob object; it is not a field ID. ## Example Schemas diff --git a/docs/src/format/table/versioning.md b/docs/src/format/table/versioning.md index ee1aa312c1a..38d7380a904 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -32,7 +32,7 @@ they should return an "unsupported" error on any read or write operation. | 64 | `FLAG_UNSTABLE_DATA_OVERLAY_FILES` | Yes | Yes | Fragments may carry data overlay files. Unstable: release builds reject it unless explicitly opted in. | | 128 | `FLAG_COVERED_INDEX_METADATA` | Yes | Yes | Some index declares covering columns (`IndexMetadata.covering_fields`), so `fields` means keyed columns followed by carried ones. An implementation without this flag selects an index by membership of `fields` and would answer a query on a merely-carried column with an index keyed on a different one. | | 256 | `FLAG_MIXED_DATA_FILE_VERSIONS` | Yes | Yes | Reserved for datasets that may reference recognized V2 data files with different exact versions. | -| 512 | `FLAG_STABLE_FIELD_IDS` | No | Yes | After explicit migration, the manifest carries a persistent field-ID high-water mark that writers must use when allocating field IDs. Operators must retire older writers before activation. See [Field IDs](schema.md#field-ids). | +| 512 | `FLAG_STABLE_FIELD_IDS` | No | Yes | The manifest sets `max_allocated_field_id`, and a writer must assign new field IDs above it. See [Field IDs](schema.md#field-ids). | diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index d55f89103e1..b0ee408c1a6 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -928,30 +928,24 @@ enum RawArrowFieldIdMode { ExplicitOnly, } +#[derive(Clone, Copy)] +enum LegacyFieldIdMode { + Inherit, + Standalone, +} + struct SchemaConversionOptions { raw_field_id_mode: RawArrowFieldIdMode, + legacy_field_id_mode: LegacyFieldIdMode, } -fn convert_schema_from_operation( - env: &mut JNIEnv, - java_operation: &JObject, - java_allocator: &JObject, - dataset: Option<&mut BlockingDataset>, - read_version: u64, +type SchemaReadContext = (LanceSchema, i32, bool); + +fn convert_arrow_schema( + arrow_schema: &Schema, + read_context: Option, options: SchemaConversionOptions, ) -> Result<(LanceSchema, HashMap)> { - let schema_ptr = env - .call_method( - java_operation, - "exportSchema", - "(Lorg/apache/arrow/memory/BufferAllocator;)J", - &[JValue::Object(java_allocator)], - )? - .j()?; - let c_schema_ptr = schema_ptr as *mut FFI_ArrowSchema; - let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) }; - - let arrow_schema = Schema::try_from(&c_schema)?; let mut original_schema = if matches!(options.raw_field_id_mode, RawArrowFieldIdMode::ExplicitOnly) { LanceSchema { @@ -963,7 +957,7 @@ fn convert_schema_from_operation( metadata: arrow_schema.metadata.clone(), } } else { - LanceSchema::try_from(&arrow_schema).map_err(|e| { + LanceSchema::try_from(arrow_schema).map_err(|e| { Error::input_error(format!( "Failed to convert Arrow schema to Lance schema: {}", e @@ -971,23 +965,6 @@ fn convert_schema_from_operation( })? }; - let read_context = match dataset { - Some(dataset) if dataset.inner.version().version == read_version => Some(( - dataset.inner.schema().clone(), - dataset.inner.manifest().max_field_id(), - dataset.inner.manifest().uses_stable_field_ids(), - )), - Some(dataset) => { - let read_dataset = dataset.checkout_version(read_version)?; - Some(( - read_dataset.inner.schema().clone(), - read_dataset.inner.manifest().max_field_id(), - read_dataset.inner.manifest().uses_stable_field_ids(), - )) - } - None => None, - }; - if read_context .as_ref() .is_none_or(|(_, _, stable_field_ids)| *stable_field_ids) @@ -999,9 +976,13 @@ fn convert_schema_from_operation( return Ok((original_schema, HashMap::new())); } + if matches!(options.legacy_field_id_mode, LegacyFieldIdMode::Standalone) { + return Ok((original_schema, HashMap::new())); + } + let (read_schema, max_field_id, _) = read_context.expect("legacy dataset context"); let schema = - LanceSchema::from_arrow_schema(&arrow_schema, Some(read_schema), Some(max_field_id))?; + LanceSchema::from_arrow_schema(arrow_schema, Some(read_schema), Some(max_field_id))?; let field_id_remap = original_schema .fields_pre_order() @@ -1013,6 +994,46 @@ fn convert_schema_from_operation( Ok((schema, field_id_remap)) } +fn convert_schema_from_operation( + env: &mut JNIEnv, + java_operation: &JObject, + java_allocator: &JObject, + dataset: Option<&mut BlockingDataset>, + read_version: u64, + options: SchemaConversionOptions, +) -> Result<(LanceSchema, HashMap)> { + let schema_ptr = env + .call_method( + java_operation, + "exportSchema", + "(Lorg/apache/arrow/memory/BufferAllocator;)J", + &[JValue::Object(java_allocator)], + )? + .j()?; + let c_schema_ptr = schema_ptr as *mut FFI_ArrowSchema; + let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) }; + let arrow_schema = Schema::try_from(&c_schema)?; + + let read_context = match dataset { + Some(dataset) if dataset.inner.version().version == read_version => Some(( + dataset.inner.schema().clone(), + dataset.inner.manifest().max_field_id(), + dataset.inner.manifest().uses_stable_field_ids(), + )), + Some(dataset) => { + let read_dataset = dataset.checkout_version(read_version)?; + Some(( + read_dataset.inner.schema().clone(), + read_dataset.inner.manifest().max_field_id(), + read_dataset.inner.manifest().uses_stable_field_ids(), + )) + } + None => None, + }; + + convert_arrow_schema(&arrow_schema, read_context, options) +} + type DataFileIdentity = (Option, String); fn retained_file_identities( @@ -1202,6 +1223,7 @@ fn convert_to_rust_operation( read_version, SchemaConversionOptions { raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, + legacy_field_id_mode: LegacyFieldIdMode::Inherit, }, )?; Operation::Project { @@ -1341,6 +1363,7 @@ fn convert_to_rust_operation( read_version, SchemaConversionOptions { raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, + legacy_field_id_mode: LegacyFieldIdMode::Standalone, }, )?; remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); @@ -1501,6 +1524,7 @@ fn convert_to_rust_operation( read_version, SchemaConversionOptions { raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, + legacy_field_id_mode: LegacyFieldIdMode::Inherit, }, )?; let retained_files = if field_id_remap.is_empty() { @@ -1856,6 +1880,65 @@ mod tests { assert_eq!(schema.field("b").unwrap().id, 9); } + #[test] + fn stable_java_schema_conversion_marks_raw_arrow_input() { + let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); + base.id = 0; + let base_schema = LanceSchema { + fields: vec![base], + metadata: HashMap::new(), + }; + let arrow_schema = + ArrowSchema::new(vec![ArrowField::new("a", ArrowDataType::Int32, false)]); + + let (schema, field_id_remap) = convert_arrow_schema( + &arrow_schema, + Some((base_schema, 0, true)), + SchemaConversionOptions { + raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, + legacy_field_id_mode: LegacyFieldIdMode::Inherit, + }, + ) + .unwrap(); + + assert!( + schema + .metadata + .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + ); + assert_eq!(schema.field("a").unwrap().id, -1); + assert!(field_id_remap.is_empty()); + } + + #[test] + fn legacy_java_overwrite_allows_type_replacement() { + let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); + base.id = 0; + let base_schema = LanceSchema { + fields: vec![base], + metadata: HashMap::new(), + }; + let arrow_schema = ArrowSchema::new(vec![ArrowField::new("a", ArrowDataType::Utf8, false)]); + + let (schema, field_id_remap) = convert_arrow_schema( + &arrow_schema, + Some((base_schema, 0, false)), + SchemaConversionOptions { + raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, + legacy_field_id_mode: LegacyFieldIdMode::Standalone, + }, + ) + .unwrap(); + + assert_eq!(schema.field("a").unwrap().data_type(), ArrowDataType::Utf8); + assert!( + !schema + .metadata + .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + ); + assert!(field_id_remap.is_empty()); + } + #[test] fn test_create_schema_from_arrow() { // base_schema has an existing field id diff --git a/java/src/test/java/org/lance/operation/MergeTest.java b/java/src/test/java/org/lance/operation/MergeTest.java index 11aef49a353..1fb210b987e 100644 --- a/java/src/test/java/org/lance/operation/MergeTest.java +++ b/java/src/test/java/org/lance/operation/MergeTest.java @@ -198,10 +198,6 @@ void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Ex int addressFieldId = maxFieldId + 20; int cityFieldId = maxFieldId + 30; int countryFieldId = maxFieldId + 40; - int canonicalAgeFieldId = maxFieldId + 1; - int canonicalAddressFieldId = maxFieldId + 2; - int canonicalCityFieldId = maxFieldId + 3; - int canonicalCountryFieldId = maxFieldId + 4; Field idField = new Field( @@ -311,18 +307,18 @@ void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Ex // Verify field id. LanceField evolvedAgeField = findField(evolvedDataset.getLanceSchema().fields(), "age"); - Assertions.assertEquals(canonicalAgeFieldId, evolvedAgeField.getId()); + Assertions.assertEquals(ageFieldId, evolvedAgeField.getId()); LanceField evolvedAddressField = findField(evolvedDataset.getLanceSchema().fields(), "address"); - Assertions.assertEquals(canonicalAddressFieldId, evolvedAddressField.getId()); + Assertions.assertEquals(addressFieldId, evolvedAddressField.getId()); LanceField evolvedCityField = findField(evolvedAddressField.getChildren(), "city"); - Assertions.assertEquals(canonicalCityFieldId, evolvedCityField.getId()); + Assertions.assertEquals(cityFieldId, evolvedCityField.getId()); LanceField evolvedCountryField = findField(evolvedAddressField.getChildren(), "country"); - Assertions.assertEquals(canonicalCountryFieldId, evolvedCountryField.getId()); + Assertions.assertEquals(countryFieldId, evolvedCountryField.getId()); // Verify merged data try (LanceScanner scanner = evolvedDataset.newScan()) { @@ -369,74 +365,6 @@ void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Ex } } - @Test - void testMergeRejectsAmbiguousRawArrowFieldIds(@TempDir Path tempDir) throws Exception { - String datasetPath = tempDir.resolve("testMergeRejectsAmbiguousRawArrowFieldIds").toString(); - try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { - TestUtils.SimpleTestDataset testDataset = - new TestUtils.SimpleTestDataset(allocator, datasetPath); - - try (Dataset initialDataset = createAndAppendRows(testDataset, 2)) { - Field bField = - new Field( - "b", new FieldType(true, new ArrowType.Int(32, true), null, fieldMeta(3)), null); - Field cField = - new Field( - "c", new FieldType(true, new ArrowType.Int(32, true), null, fieldMeta(2)), null); - Schema newColumnsSchema = new Schema(Arrays.asList(bField, cField), null); - Schema evolvedSchema = - new Schema( - Arrays.asList( - testDataset.getSchema().getFields().get(0), - testDataset.getSchema().getFields().get(1), - bField, - cField), - null); - - try (VectorSchemaRoot root = VectorSchemaRoot.create(newColumnsSchema, allocator)) { - root.allocateNew(); - IntVector bVector = (IntVector) root.getVector("b"); - IntVector cVector = (IntVector) root.getVector("c"); - bVector.setSafe(0, 10); - bVector.setSafe(1, 11); - cVector.setSafe(0, 20); - cVector.setSafe(1, 21); - root.setRowCount(2); - - DataFile newDataFile = - writeLanceDataFile( - dataset.allocator(), datasetPath, root, new int[] {2, 3}, new int[] {0, 1}); - FragmentMetadata fragmentMeta = initialDataset.getFragment(0).metadata(); - List dataFiles = fragmentMeta.getFiles(); - dataFiles.add(newDataFile); - FragmentMetadata evolvedFragment = - new FragmentMetadata( - fragmentMeta.getId(), - dataFiles, - fragmentMeta.getPhysicalRows(), - fragmentMeta.getDeletionFile(), - fragmentMeta.getRowIdMeta()); - - try (Transaction transaction = - new Transaction.Builder() - .readVersion(initialDataset.version()) - .operation( - Merge.builder() - .fragments(Collections.singletonList(evolvedFragment)) - .schema(evolvedSchema) - .build()) - .build()) { - IllegalArgumentException error = - Assertions.assertThrows( - IllegalArgumentException.class, - () -> new CommitBuilder(initialDataset).execute(transaction).close()); - Assertions.assertTrue(error.getMessage().contains("ambiguous raw Arrow field IDs")); - } - } - } - } - } - @Test void testLegacyMergeInheritsNonContiguousFieldIds(@TempDir Path tempDir) throws Exception { Path source = diff --git a/java/src/test/java/org/lance/operation/ProjectTest.java b/java/src/test/java/org/lance/operation/ProjectTest.java index a62db31eb9f..aaf59fe88a7 100644 --- a/java/src/test/java/org/lance/operation/ProjectTest.java +++ b/java/src/test/java/org/lance/operation/ProjectTest.java @@ -42,7 +42,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; public class ProjectTest extends OperationTestBase { @@ -116,34 +115,6 @@ void testPreservesNullabilityEqualityAndRoundTrip(@TempDir Path tempDir) { } } - @Test - void testProjectRejectsUnmaterializedStableField(@TempDir Path tempDir) { - String datasetPath = tempDir.resolve("testProjectRejectsUnmaterializedStableField").toString(); - try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { - TestUtils.SimpleTestDataset testDataset = - new TestUtils.SimpleTestDataset(allocator, datasetPath); - dataset = testDataset.createEmptyDataset(); - - Field existing = dataset.getSchema().getFields().get(0); - Field unmaterialized = - new Field("unmaterialized", existing.getFieldType(), existing.getChildren()); - try (Transaction transaction = - new Transaction.Builder() - .readVersion(dataset.version()) - .operation( - Project.builder() - .schema(new Schema(Collections.singletonList(unmaterialized))) - .build()) - .build()) { - IllegalArgumentException error = - assertThrows( - IllegalArgumentException.class, - () -> new CommitBuilder(dataset).execute(transaction)); - assertTrue(error.getMessage().contains("writes no data")); - } - } - } - @Test void testProjectPreservesExplicitRenameIdentity(@TempDir Path tempDir) { String datasetPath = tempDir.resolve("testProjectPreservesExplicitRenameIdentity").toString(); diff --git a/protos/table.proto b/protos/table.proto index 50786545fa0..e83bc6156ad 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -138,9 +138,9 @@ message Manifest { // // The flag identities are the same as for reader_feature_flags, but the values of // reader_feature_flags and writer_feature_flags are not required to be identical. - // * 1 << 9: newly allocated field IDs must be greater than - // max_allocated_field_id, and the high-water mark must advance atomically. - // Explicit migration sets this bit after older writers have been retired. + // * 1 << 9: the manifest also sets max_allocated_field_id. A writer must + // assign new field IDs above that value and advance it in the same commit. + // The field and flag must either both be set or both be absent. uint64 writer_feature_flags = 10; // The highest fragment ID that has been used so far. @@ -224,12 +224,14 @@ message Manifest { // The branch of the dataset. None means main branch. optional string branch = 20; - // The highest field ID allocated since stable field identity was activated. + // The stable field-ID allocator high-water mark. // - // Presence activates the stable-field-ID contract. The value may name a field - // that is absent from the current schema and files because dropped IDs are - // permanently retired. Absence means this is a legacy dataset whose field-ID - // allocator is still derived from currently referenced fields. + // The first manifest that sets this field initializes it to the largest + // field ID the manifest references. Later writers assign new IDs above it. + // A dropped ID remains covered by this value and must not be reused. + // + // This field and FLAG_STABLE_FIELD_IDS in writer_feature_flags must either + // both be set or both be absent. A mismatch is invalid. optional int32 max_allocated_field_id = 22; } // Manifest diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 74b992723e8..ed6eb45dc86 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -479,7 +479,7 @@ pub fn validate_stable_field_id_transition( .collect::>(); if !matches!(operation, Operation::Restore { .. }) { - validate_dense_new_field_ids( + validate_new_field_ids( parent, successor .schema @@ -512,16 +512,16 @@ pub fn validate_stable_field_id_transition( Ok(()) } -fn validate_dense_new_field_ids<'a>( +fn validate_new_field_ids<'a>( manifest: &Manifest, new_fields: impl Iterator, ) -> Result<()> { - let expected_ids = i64::from(manifest.max_field_id()) + 1..; - for (expected, field) in expected_ids.zip(new_fields) { - if i64::from(field.id) != expected { + let max_allocated_field_id = manifest.max_field_id(); + for field in new_fields { + if field.id <= max_allocated_field_id { return Err(Error::invalid_input(format!( - "New field '{}' has ID {}, but stable field IDs must be densely allocated from {}", - field.name, field.id, expected + "New field '{}' has ID {}, but stable field IDs must be greater than the high-water mark {}", + field.name, field.id, max_allocated_field_id ))); } } @@ -560,7 +560,7 @@ fn validate_stable_field_id_operation(manifest: &Manifest, operation: &Operation } } - validate_dense_new_field_ids( + validate_new_field_ids( manifest, schema .fields_pre_order() @@ -1036,7 +1036,7 @@ mod tests { } #[test] - fn stable_field_ids_require_dense_allocation_above_high_water_mark() { + fn stable_field_ids_allow_reserved_ids_above_high_water_mark() { let mut manifest = activated_manifest(); manifest.max_allocated_field_id = Some(5); let mut schema = manifest.schema.clone(); @@ -1051,13 +1051,21 @@ mod tests { validate_operation(Some(&manifest), &valid).unwrap(); schema.fields.last_mut().unwrap().id = 7; - let skipped = Operation::Project { + let reserved = Operation::Project { + schema: schema.clone(), + preserves_nullability: true, + }; + validate_operation(Some(&manifest), &reserved).unwrap(); + + schema.fields.last_mut().unwrap().id = 5; + let reused = Operation::Project { schema, preserves_nullability: true, }; - let err = validate_operation(Some(&manifest), &skipped).unwrap_err(); + let err = validate_operation(Some(&manifest), &reused).unwrap_err(); assert!( - err.to_string().contains("densely allocated from 6"), + err.to_string() + .contains("greater than the high-water mark 5"), "{err}" ); } diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 0064aa190fa..a081c1dbdc7 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -3253,9 +3253,7 @@ impl Dataset { /// The activation commit records the current maximum referenced field ID as /// a persistent high-water mark. Later schema changes allocate above it even /// after fields and their files are dropped. Activation is one-way and - /// idempotent. Before migrating, retire every writer that does not understand - /// the stable field-ID writer feature flag, including writers that commit - /// from an already serialized manifest. + /// idempotent. /// /// ``` /// # use lance::{Dataset, Result}; diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 168858d0fdc..f0ce4493de4 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -3678,6 +3678,7 @@ mod test { }), ) .await?; + dataset.migrate_to_stable_field_ids().await?; assert!(dataset.manifest.uses_stable_field_ids()); assert_eq!(dataset.manifest.max_field_id(), 0); diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 8a4eee4d348..a507a131ffa 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -1926,7 +1926,7 @@ mod tests { async fn raw_arrow_project_retry_matches_single_attempt_ids() { let tmp = TempStrDir::default(); let uri = tmp.as_str(); - let dataset = Dataset::write( + let mut dataset = Dataset::write( RecordBatchIterator::new( vec![Ok(simple_batch(&simple_schema(), vec![1, 2, 3]))], simple_schema(), @@ -1936,6 +1936,7 @@ mod tests { ) .await .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); let mut foreign_manifest = dataset.manifest.as_ref().clone(); foreign_manifest.max_fragment_id = Some(foreign_manifest.max_fragment_id.unwrap_or(0) + 1); @@ -1994,7 +1995,7 @@ mod tests { async fn raw_arrow_merge_retry_rebinds_after_allocator_advance() { let tmp = TempStrDir::default(); let uri = tmp.as_str(); - let dataset = Dataset::write( + let mut dataset = Dataset::write( RecordBatchIterator::new( vec![Ok(simple_batch(&simple_schema(), vec![1, 2, 3]))], simple_schema(), @@ -2004,6 +2005,7 @@ mod tests { ) .await .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); let mut foreign_manifest = dataset.manifest.as_ref().clone(); foreign_manifest.max_fragment_id = Some(foreign_manifest.max_fragment_id.unwrap_or(0) + 1); From 79273432cc684cda3b5b3ac7792fbca187395463 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 21 Sep 2026 20:34:16 +0800 Subject: [PATCH 18/24] fix: clarify stable field IDs and make schema input explicit --- docs/src/format/table/schema.md | 14 +- java/lance-jni/src/blocking_dataset.rs | 3 + java/lance-jni/src/transaction.rs | 184 ++++++++++++------ rust/lance-core/src/datatypes/schema.rs | 83 ++++---- .../lance-namespace-impls/src/dir/manifest.rs | 10 +- rust/lance-table/src/transaction.rs | 5 +- .../src/transaction/manifest_build.rs | 2 + rust/lance-table/src/transaction/validate.rs | 131 ++++++------- rust/lance/src/dataset.rs | 9 +- rust/lance/src/dataset/tests/dataset_io.rs | 2 + .../src/dataset/tests/dataset_migrations.rs | 107 ++++++---- rust/lance/src/dataset/write.rs | 19 +- rust/lance/src/dataset/write/commit.rs | 102 ++++++++++ rust/lance/src/io/commit.rs | 101 +++++----- 14 files changed, 483 insertions(+), 289 deletions(-) diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index dfa1f0f4e04..c0e0ee5d2eb 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -247,8 +247,11 @@ feature does not change read behavior. A dataset changes to stable field IDs only through an explicit migration commit. A dataset cannot return to the legacy behavior. After activation, a restore must fail if it targets -a version that does not set `max_allocated_field_id`. That version does not record retired field -IDs, so a later commit could reuse one. +a version that does not set `max_allocated_field_id`. Before activation, different fields may have +used the same ID in different versions. For example, an old version may assign ID 1 to an integer +field `x`, while the activation version assigns it to a string field `y`. Restoring the old version +would make ID 1 refer to `x` again. Keeping the current high-water mark prevents future allocation +from reusing IDs, but does not resolve this existing conflict. Reading old versions remains supported. ### Field ID Properties @@ -384,8 +387,11 @@ pa.schema([ ]) ``` -The data file maps this column with `DataFile.fields = [0]`. The descriptor children are file -details. Their IDs may be `-1` or file-local, and they do not change `max_allocated_field_id`. +The entire descriptor is encoded in one physical column. The data file maps it with +`DataFile.fields = [0]` and `DataFile.column_indices = [k]`, where `k` is that column's index in the +Lance file. The reader locates column `k` and decodes the descriptor using its Blob page layout; +the descriptor children do not have separate entries in either mapping. Their IDs may be `-1` or +file-local, and they do not change `max_allocated_field_id`. A descriptor scan returns the stored struct. A materialized scan returns this public shape: diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index ba7a8c6aabb..f6a2166cfb6 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -59,6 +59,7 @@ use lance_namespace::LanceNamespace; use lance_table::io::commit::CommitHandler; use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; use lance_table::io::commit::{ManifestLocation, ManifestNamingScheme}; +use lance_table::transaction::SchemaInputKind; use std::collections::HashMap; use std::future::IntoFuture; use std::iter::empty; @@ -396,6 +397,7 @@ impl BlockingDataset { pub fn commit_transaction( &mut self, transaction: Transaction, + schema_input_kind: SchemaInputKind, store_params: ObjectStoreParams, detached: bool, enable_v2_manifest_paths: bool, @@ -407,6 +409,7 @@ impl BlockingDataset { commit_timeout: Option, ) -> Result { let mut builder = CommitBuilder::new(Arc::new(self.clone().inner)) + .with_schema_input_kind(schema_input_kind) .with_store_params(store_params) .with_detached(detached) .enable_v2_manifest_paths(enable_v2_manifest_paths) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index b0ee408c1a6..5938f362c6b 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -30,7 +30,7 @@ use lance_file::version::{LanceFileVersion, V2_FORMAT_2_0, V2_FORMAT_2_1, V2_FOR use lance_io::object_store::{LanceNamespaceStorageOptionsProvider, StorageOptionsProvider}; use lance_table::io::commit::CommitHandler; use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; -use lance_table::transaction::TRANSACTION_SCHEMA_SOURCE_RAW_ARROW; +use lance_table::transaction::SchemaInputKind; use prost::Message; use prost_types::Any; use roaring::RoaringBitmap; @@ -838,7 +838,7 @@ fn inner_commit_to_dataset<'local>( unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; BlockingDataset::new(dataset_guard.inner.clone()) }; - let transaction = convert_to_rust_transaction( + let (transaction, schema_input_kind) = convert_to_rust_transaction( env, java_transaction, Some(&java_allocator), @@ -869,6 +869,7 @@ fn inner_commit_to_dataset<'local>( unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; dataset_guard.commit_transaction( transaction, + schema_input_kind, store_params, detached, enable_v2_manifest_paths, @@ -888,7 +889,7 @@ fn convert_to_rust_transaction( java_transaction: JObject, allocator: Option<&JObject>, dataset: Option<&mut BlockingDataset>, -) -> Result { +) -> Result<(Transaction, SchemaInputKind)> { let read_ver = env.get_u64_from_method(&java_transaction, "readVersion")?; let uuid = env.get_string_from_method(&java_transaction, "uuid")?; let op = env @@ -899,7 +900,8 @@ fn convert_to_rust_transaction( &[], )? .l()?; - let op = convert_to_rust_operation(env, &op, allocator, dataset, read_ver)?; + let (op, schema_input_kind) = + convert_to_rust_operation(env, &op, allocator, dataset, read_ver)?; let tag = env.get_optional_from_method(&java_transaction, "tag", |env, tag_obj| { let tag_str = JString::from(tag_obj); @@ -914,11 +916,14 @@ fn convert_to_rust_transaction( to_rust_map(env, &transaction_properties) }, )?; - Ok(TransactionBuilder::new(read_ver, op) - .uuid(uuid) - .tag(tag) - .transaction_properties(transaction_properties.map(Arc::new)) - .build()) + Ok(( + TransactionBuilder::new(read_ver, op) + .uuid(uuid) + .tag(tag) + .transaction_properties(transaction_properties.map(Arc::new)) + .build(), + schema_input_kind, + )) } #[derive(Clone, Copy)] @@ -941,43 +946,53 @@ struct SchemaConversionOptions { type SchemaReadContext = (LanceSchema, i32, bool); +struct ConvertedSchema { + schema: LanceSchema, + field_id_remap: HashMap, + input_kind: SchemaInputKind, +} + fn convert_arrow_schema( arrow_schema: &Schema, read_context: Option, options: SchemaConversionOptions, -) -> Result<(LanceSchema, HashMap)> { - let mut original_schema = - if matches!(options.raw_field_id_mode, RawArrowFieldIdMode::ExplicitOnly) { - LanceSchema { - fields: arrow_schema - .fields - .iter() - .map(|field| Field::try_from(field.as_ref())) - .collect::>()?, - metadata: arrow_schema.metadata.clone(), - } - } else { - LanceSchema::try_from(arrow_schema).map_err(|e| { - Error::input_error(format!( - "Failed to convert Arrow schema to Lance schema: {}", - e - )) - })? - }; +) -> Result { + let original_schema = if matches!(options.raw_field_id_mode, RawArrowFieldIdMode::ExplicitOnly) + { + LanceSchema { + fields: arrow_schema + .fields + .iter() + .map(|field| Field::try_from(field.as_ref())) + .collect::>()?, + metadata: arrow_schema.metadata.clone(), + } + } else { + LanceSchema::try_from(arrow_schema).map_err(|e| { + Error::input_error(format!( + "Failed to convert Arrow schema to Lance schema: {}", + e + )) + })? + }; if read_context .as_ref() .is_none_or(|(_, _, stable_field_ids)| *stable_field_ids) { - original_schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); - return Ok((original_schema, HashMap::new())); + return Ok(ConvertedSchema { + schema: original_schema, + field_id_remap: HashMap::new(), + input_kind: SchemaInputKind::Arrow, + }); } if matches!(options.legacy_field_id_mode, LegacyFieldIdMode::Standalone) { - return Ok((original_schema, HashMap::new())); + return Ok(ConvertedSchema { + schema: original_schema, + field_id_remap: HashMap::new(), + input_kind: SchemaInputKind::Lance, + }); } let (read_schema, max_field_id, _) = read_context.expect("legacy dataset context"); @@ -991,7 +1006,11 @@ fn convert_arrow_schema( (original.id >= 0 && original.id != canonical.id).then_some((original.id, canonical.id)) }) .collect(); - Ok((schema, field_id_remap)) + Ok(ConvertedSchema { + schema, + field_id_remap, + input_kind: SchemaInputKind::Lance, + }) } fn convert_schema_from_operation( @@ -1001,7 +1020,7 @@ fn convert_schema_from_operation( dataset: Option<&mut BlockingDataset>, read_version: u64, options: SchemaConversionOptions, -) -> Result<(LanceSchema, HashMap)> { +) -> Result { let schema_ptr = env .call_method( java_operation, @@ -1207,11 +1226,14 @@ fn convert_to_rust_operation( allocator: Option<&JObject<'_>>, mut dataset: Option<&mut BlockingDataset>, read_version: u64, -) -> Result { +) -> Result<(Operation, SchemaInputKind)> { + let mut schema_input_kind = SchemaInputKind::Lance; let op_name = env.get_string_from_method(java_operation, "name")?; let op = match op_name.as_str() { "Project" => { - let (schema, _) = convert_schema_from_operation( + let ConvertedSchema { + schema, input_kind, .. + } = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1226,6 +1248,7 @@ fn convert_to_rust_operation( legacy_field_id_mode: LegacyFieldIdMode::Inherit, }, )?; + schema_input_kind = input_kind; Operation::Project { preserves_nullability: env .get_boolean_from_method(java_operation, "preservesNullability")?, @@ -1351,7 +1374,11 @@ fn convert_to_rust_operation( to_rust_map(env, &config_upsert_values) }, )?; - let (schema, field_id_remap) = convert_schema_from_operation( + let ConvertedSchema { + schema, + field_id_remap, + input_kind, + } = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1366,6 +1393,7 @@ fn convert_to_rust_operation( legacy_field_id_mode: LegacyFieldIdMode::Standalone, }, )?; + schema_input_kind = input_kind; remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); Operation::Overwrite { fragments, @@ -1512,7 +1540,11 @@ fn convert_to_rust_operation( import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; - let (schema, field_id_remap) = convert_schema_from_operation( + let ConvertedSchema { + schema, + field_id_remap, + input_kind, + } = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1527,6 +1559,7 @@ fn convert_to_rust_operation( legacy_field_id_mode: LegacyFieldIdMode::Inherit, }, )?; + schema_input_kind = input_kind; let retained_files = if field_id_remap.is_empty() { HashSet::new() } else { @@ -1544,13 +1577,13 @@ fn convert_to_rust_operation( let version: u64 = env .call_method(java_operation, "version", "()J", &[])? .j()? as u64; - return Ok(Operation::Restore { version }); + Operation::Restore { version } } "ReserveFragments" => { let num_fragments = env .call_method(java_operation, "numFragments", "()I", &[])? .i()? as u32; - return Ok(Operation::ReserveFragments { num_fragments }); + Operation::ReserveFragments { num_fragments } } "CreateIndex" => { let new_indices = @@ -1561,14 +1594,14 @@ fn convert_to_rust_operation( import_vec_from_method(env, java_operation, "getRemovedIndices", |env, index| { index.extract_object(env) })?; - return Ok(Operation::CreateIndex { + Operation::CreateIndex { new_indices, removed_indices, - }); + } } _ => unimplemented!(), }; - Ok(op) + Ok((op, schema_input_kind)) } fn extract_update_map(env: &mut JNIEnv, update_map_obj: &JObject) -> Result> { @@ -1804,11 +1837,12 @@ fn inner_commit_to_uri<'local>( } else { Some(allocator_obj) }; - let transaction = + let (transaction, schema_input_kind) = convert_to_rust_transaction(env, java_transaction, allocator_ref.as_ref(), ds.as_mut())?; // Build CommitBuilder with URI let mut builder = CommitBuilder::new(&*uri_str) + .with_schema_input_kind(schema_input_kind) .with_store_params(store_params) .with_detached(detached) .enable_v2_manifest_paths(enable_v2_manifest_paths) @@ -1881,7 +1915,7 @@ mod tests { } #[test] - fn stable_java_schema_conversion_marks_raw_arrow_input() { + fn stable_java_schema_conversion_preserves_unassigned_ids() { let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); base.id = 0; let base_schema = LanceSchema { @@ -1891,7 +1925,11 @@ mod tests { let arrow_schema = ArrowSchema::new(vec![ArrowField::new("a", ArrowDataType::Int32, false)]); - let (schema, field_id_remap) = convert_arrow_schema( + let ConvertedSchema { + schema, + field_id_remap, + input_kind, + } = convert_arrow_schema( &arrow_schema, Some((base_schema, 0, true)), SchemaConversionOptions { @@ -1901,15 +1939,42 @@ mod tests { ) .unwrap(); - assert!( - schema - .metadata - .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) - ); + assert_eq!(input_kind, SchemaInputKind::Arrow); + assert!(schema.metadata.is_empty()); assert_eq!(schema.field("a").unwrap().id, -1); assert!(field_id_remap.is_empty()); } + #[test] + fn stable_java_project_preserves_explicit_ids_and_metadata() { + let base_schema = LanceSchema::try_from(&ArrowSchema::new(vec![ArrowField::new( + "a", + ArrowDataType::Int32, + false, + )])) + .unwrap(); + let metadata = HashMap::from([("source".to_string(), "user metadata".to_string())]); + let arrow_schema = ArrowSchema::new(vec![ + ArrowField::new("renamed", ArrowDataType::Int32, false).with_metadata(HashMap::from([ + (LANCE_FIELD_ID_KEY.to_string(), "0".to_string()), + ])), + ]) + .with_metadata(metadata.clone()); + let converted = convert_arrow_schema( + &arrow_schema, + Some((base_schema, 0, true)), + SchemaConversionOptions { + raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, + legacy_field_id_mode: LegacyFieldIdMode::Inherit, + }, + ) + .unwrap(); + assert_eq!(converted.input_kind, SchemaInputKind::Arrow); + assert_eq!(converted.schema.field("renamed").unwrap().id, 0); + assert_eq!(converted.schema.metadata, metadata); + assert!(converted.field_id_remap.is_empty()); + } + #[test] fn legacy_java_overwrite_allows_type_replacement() { let mut base = Field::new_arrow("a", ArrowDataType::Int32, false).unwrap(); @@ -1920,7 +1985,11 @@ mod tests { }; let arrow_schema = ArrowSchema::new(vec![ArrowField::new("a", ArrowDataType::Utf8, false)]); - let (schema, field_id_remap) = convert_arrow_schema( + let ConvertedSchema { + schema, + field_id_remap, + input_kind, + } = convert_arrow_schema( &arrow_schema, Some((base_schema, 0, false)), SchemaConversionOptions { @@ -1931,11 +2000,8 @@ mod tests { .unwrap(); assert_eq!(schema.field("a").unwrap().data_type(), ArrowDataType::Utf8); - assert!( - !schema - .metadata - .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) - ); + assert_eq!(input_kind, SchemaInputKind::Lance); + assert!(schema.metadata.is_empty()); assert!(field_id_remap.is_empty()); } diff --git a/rust/lance-core/src/datatypes/schema.rs b/rust/lance-core/src/datatypes/schema.rs index c6d7d96f0cd..a49800151e5 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -697,16 +697,11 @@ impl Schema { /// /// Existing IDs are preserved. New IDs start after both this schema's /// maximum ID and `max_existing_id`. + /// If allocation fails, discard the partially updated schema. pub fn try_set_field_id(&mut self, max_existing_id: Option) -> Result<()> { let schema_max_id = self.max_field_id().unwrap_or(-1); let max_existing_id = max_existing_id.unwrap_or(-1); let mut current_id = i64::from(schema_max_id.max(max_existing_id)) + 1; - let unassigned_count = self.fields_pre_order().filter(|field| field.id < 0).count() as i64; - if unassigned_count > 0 && current_id + unassigned_count - 1 > i64::from(i32::MAX) { - return Err(Error::invalid_input( - "No further field ID can be allocated because IDs are exhausted", - )); - } for field in &mut self.fields { field.try_set_id(-1, &mut current_id)?; } @@ -717,14 +712,8 @@ impl Schema { /// /// The first assigned ID is one greater than `max_existing_id`. Use this when /// every input field must receive a new identity. + /// If allocation fails, discard the partially updated schema. pub fn try_reassign_field_ids(&mut self, max_existing_id: Option) -> Result<()> { - let field_count = self.fields_pre_order().count() as i64; - let first_id = i64::from(max_existing_id.unwrap_or(-1)) + 1; - if field_count > 0 && first_id + field_count - 1 > i64::from(i32::MAX) { - return Err(Error::invalid_input( - "No further field ID can be allocated because IDs are exhausted", - )); - } self.reset_id(); self.try_set_field_id(max_existing_id) } @@ -1790,37 +1779,45 @@ mod tests { use super::*; - #[test] - fn checked_field_id_allocation_is_atomic_on_exhaustion() { - let mut schema = Schema { - fields: vec![ - Field::try_from(&ArrowField::new("a", ArrowDataType::Int32, false)).unwrap(), - Field::try_from(&ArrowField::new("b", ArrowDataType::Int32, false)).unwrap(), - ], - metadata: HashMap::new(), - }; - - let err = schema.try_set_field_id(Some(i32::MAX - 1)).unwrap_err(); - - assert!(err.to_string().contains("IDs are exhausted"), "{err}"); - assert!(schema.fields.iter().all(|field| field.id == -1)); - } - - #[test] - fn checked_field_id_reassignment_is_atomic_on_exhaustion() { - let arrow_schema = ArrowSchema::new(vec![ - ArrowField::new("a", ArrowDataType::Int32, false), - ArrowField::new("b", ArrowDataType::Int32, false), - ]); + #[rstest::rstest] + #[case::last_id(i32::MAX - 1, 1, true)] + #[case::last_two_ids(i32::MAX - 2, 2, true)] + #[case::exhausted_mid_allocation(i32::MAX - 1, 2, false)] + #[case::exhausted_before_allocation(i32::MAX, 1, false)] + #[case::no_allocation_needed(i32::MAX, 0, true)] + fn checked_field_id_allocation_bounds( + #[case] max_existing_id: i32, + #[case] field_count: usize, + #[case] succeeds: bool, + #[values(false, true)] reassign: bool, + ) { + let arrow_schema = ArrowSchema::new( + (0..field_count) + .map(|i| ArrowField::new(format!("field_{i}"), ArrowDataType::Int32, false)) + .collect::>(), + ); let mut schema = Schema::try_from(&arrow_schema).unwrap(); - let original = schema.clone(); - - let err = schema - .try_reassign_field_ids(Some(i32::MAX - 1)) - .unwrap_err(); - - assert!(err.to_string().contains("IDs are exhausted"), "{err}"); - assert_eq!(schema, original); + let result = if reassign { + schema.try_reassign_field_ids(Some(max_existing_id)) + } else { + schema.reset_id(); + schema.try_set_field_id(Some(max_existing_id)) + }; + if succeeds { + result.unwrap(); + for (i, field) in schema.fields.iter().enumerate() { + assert_eq!( + i64::from(field.id), + i64::from(max_existing_id) + 1 + i as i64 + ); + } + // An exhausted ID space must still allow schemas with all IDs assigned. + schema.try_set_field_id(Some(i32::MAX)).unwrap(); + } else { + let err = result.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("IDs are exhausted"), "{err}"); + } } #[test] diff --git a/rust/lance-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index 2741010438d..df32fa321b0 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -61,7 +61,7 @@ use lance_table::io::commit::{ CommitError, CommitHandler, commit_handler_from_url, write_manifest_file_to_path, }; use lance_table::transaction::{ - canonicalize_stable_field_ids, validate_stable_field_id_transition, + SchemaInputKind, canonicalize_stable_field_ids, validate_stable_field_id_transition, }; use object_store::{Error as ObjectStoreError, path::Path}; use roaring::RoaringBitmap; @@ -2021,9 +2021,11 @@ impl ManifestNamespace { return Ok(mutation.result); } - if let Err(err) = - canonicalize_stable_field_ids(Some(dataset.manifest()), &mut transaction.operation) - { + if let Err(err) = canonicalize_stable_field_ids( + Some(dataset.manifest()), + &mut transaction.operation, + SchemaInputKind::Lance, + ) { self.cleanup_staged_manifest_files(&object_store, &staged_data_files, &[]) .await; return Err(err); diff --git a/rust/lance-table/src/transaction.rs b/rust/lance-table/src/transaction.rs index 8eb2135ea06..ad97cc205f5 100644 --- a/rust/lance-table/src/transaction.rs +++ b/rust/lance-table/src/transaction.rs @@ -48,9 +48,8 @@ pub use update_map::{ UpdateMap, UpdateMapEntry, translate_config_updates, translate_schema_metadata_updates, }; pub use validate::{ - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW, canonicalize_stable_field_ids, - validate_detached_stable_field_ids, validate_operation, validate_stable_field_id_manifest, - validate_stable_field_id_transition, + SchemaInputKind, canonicalize_stable_field_ids, validate_detached_stable_field_ids, + validate_operation, validate_stable_field_id_manifest, validate_stable_field_id_transition, }; use crate::format::{IndexMetadata, Manifest}; diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index b0b3811f8f2..d3601fcc463 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -119,6 +119,8 @@ impl Transaction { .max_fragment_id .max(current_manifest.max_fragment_id); if current_manifest.uses_stable_field_ids() { + // Before activation, different fields could share an ID across versions. + // Keeping today's high-water mark cannot prevent restoring such a collision. let Some(restored_max_field_id) = manifest.max_allocated_field_id else { return Err(Error::invalid_input(format!( "Cannot restore version {version}: stable field IDs were activated after that version" diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index ed6eb45dc86..064236c670f 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -24,29 +24,37 @@ struct FieldIdRemap { raw_source_ids: HashMap, } -/// Transient schema-metadata marker used by bindings for raw Arrow input. -pub const TRANSACTION_SCHEMA_SOURCE_RAW_ARROW: &str = "lance:transaction_schema_source_raw_arrow"; +/// How to interpret field IDs in the schema supplied to a commit. +/// +/// This is input context, not part of the persisted schema or transaction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum SchemaInputKind { + /// IDs describe Lance fields, including fields preserved across renames. + #[default] + Lance, + /// IDs come from Arrow conversion or caller metadata. Overwrite and Merge + /// match existing fields by name and type, not by these input IDs. Project + /// may use explicit IDs to rename existing fields; missing IDs must remain + /// unassigned rather than receiving positional IDs during conversion. + Arrow, +} -/// Canonicalize schema identities supplied by a transaction before validation. +/// Assign field IDs for an operation before validation. /// -/// Arrow field-ID metadata is descriptive input, not allocation authority. New -/// datasets allocate from zero, while stable datasets preserve compatible -/// existing identities and allocate every new identity above the persisted -/// high-water mark. File mappings written against the incoming schema are -/// updated in the same step; files retained by a merge are never rewritten. +/// - New datasets receive field IDs starting at zero. +/// - Stable datasets preserve IDs for compatible existing fields and assign new +/// fields IDs above the persisted high-water mark. +/// - Arrow field-ID metadata cannot choose the ID of a new field. +/// - Raw Arrow Project operations must match existing fields: they write no data +/// and cannot add fields. +/// - New files' field mappings follow the assigned IDs; files retained by a merge +/// are never remapped. pub fn canonicalize_stable_field_ids( manifest: Option<&Manifest>, operation: &mut Operation, + input_kind: SchemaInputKind, ) -> Result<()> { - let raw_arrow_schema = match operation { - Operation::Overwrite { schema, .. } - | Operation::Project { schema, .. } - | Operation::Merge { schema, .. } => schema - .metadata - .remove(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) - .is_some(), - _ => false, - }; + let raw_arrow_schema = input_kind == SchemaInputKind::Arrow; if manifest.is_some_and(|manifest| !manifest.uses_stable_field_ids()) { if raw_arrow_schema { match operation { @@ -54,8 +62,7 @@ pub fn canonicalize_stable_field_ids( | Operation::Project { schema, .. } | Operation::Merge { schema, .. } => { // Legacy datasets retain the standalone Arrow conversion - // contract. Missing IDs still need to be assigned after the - // transient provenance marker has been removed. + // contract. Missing IDs still need to be assigned. schema.try_set_field_id(None)?; schema.validate()?; schema.verify_primary_key()?; @@ -1063,6 +1070,7 @@ mod tests { preserves_nullability: true, }; let err = validate_operation(Some(&manifest), &reused).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); assert!( err.to_string() .contains("greater than the high-water mark 5"), @@ -1111,7 +1119,8 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Lance) + .unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1128,10 +1137,6 @@ mod tests { let manifest = activated_manifest(); let mut schema = one_field_schema(); schema.fields[0].id = -1; - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut operation = Operation::Overwrite { fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0])], schema, @@ -1139,7 +1144,8 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1165,10 +1171,9 @@ mod tests { ArrowField::new("a", DataType::Int32, true), ])) .unwrap(); - raw_schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); + raw_schema + .metadata + .insert("source".to_string(), "user metadata".to_string()); let mut operation = Operation::Overwrite { fragments: vec![], schema: raw_schema, @@ -1176,13 +1181,15 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap(); let Operation::Overwrite { schema, .. } = operation else { unreachable!(); }; assert_eq!(schema.field("b").unwrap().id, 1); assert_eq!(schema.field("a").unwrap().id, 0); + assert_eq!(schema.metadata.get("source").unwrap(), "user metadata"); } #[test] @@ -1195,10 +1202,6 @@ mod tests { .unwrap(); schema.fields[0].id = 1; schema.fields[1].id = 2; - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut operation = Operation::Overwrite { fragments: vec![fragment_with_file_fields(0, "new.lance", vec![1, 2])], schema, @@ -1206,7 +1209,8 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1235,10 +1239,6 @@ mod tests { for field in &mut schema.fields { field.id = -1; } - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut fragment = fragment_with_file_fields(0, "b.lance", vec![0]); fragment .files @@ -1250,7 +1250,8 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1279,10 +1280,6 @@ mod tests { .unwrap(); schema.fields[0].id = 2; schema.fields[1].id = 1; - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut operation = Operation::Overwrite { fragments: vec![fragment_with_file_fields(0, "new.lance", vec![2, 1])], schema, @@ -1290,7 +1287,9 @@ mod tests { initial_bases: None, }; - let err = canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap_err(); + let err = + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap_err(); assert!(err.to_string().contains("ambiguous raw Arrow field IDs")); } @@ -1301,16 +1300,14 @@ mod tests { let mut schema = one_field_schema(); schema.fields[0].name = "renamed".to_string(); schema.fields[0].id = -1; - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut operation = Operation::Project { schema, preserves_nullability: true, }; - let err = canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap_err(); + let err = + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap_err(); assert!(err.to_string().contains("writes no data"), "{err}"); } @@ -1320,16 +1317,13 @@ mod tests { let manifest = activated_manifest(); let mut schema = one_field_schema(); schema.fields[0].name = "renamed".to_string(); - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut operation = Operation::Project { schema, preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap(); let Operation::Project { schema, .. } = operation else { unreachable!(); @@ -1343,26 +1337,19 @@ mod tests { let manifest = manifest_with_file_fields(one_field_schema(), vec![0]); let mut schema = one_field_schema(); schema.fields[0].id = -1; - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut operation = Operation::Project { schema, preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap(); let Operation::Project { schema, .. } = operation else { unreachable!(); }; assert_eq!(schema.fields[0].id, 0); - assert!( - !schema - .metadata - .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) - ); + assert!(schema.metadata.is_empty()); } #[test] @@ -1380,7 +1367,8 @@ mod tests { preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Lance) + .unwrap(); let Operation::Merge { schema, fragments, .. @@ -1417,7 +1405,8 @@ mod tests { preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Lance) + .unwrap(); let Operation::Merge { schema, fragments, .. @@ -1442,10 +1431,6 @@ mod tests { schema.fields[0].id = 0; schema.fields[1].id = 2; schema.fields[2].id = 1; - schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); let mut merged_fragment = manifest.fragments[0].clone(); merged_fragment.files.push(DataFile::new_legacy_from_fields( "new.lance", @@ -1458,7 +1443,9 @@ mod tests { preserves_nullability: true, }; - let err = canonicalize_stable_field_ids(Some(&manifest), &mut operation).unwrap_err(); + let err = + canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) + .unwrap_err(); assert!(err.to_string().contains("ambiguous raw Arrow field IDs")); } diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index a081c1dbdc7..26fc2970cbc 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -103,9 +103,9 @@ mod take; pub mod transaction { pub use lance_table::transaction::{ DataOverlayGroup, DataReplacementGroup, Operation, ReadVersionState, RewriteGroup, - RewrittenIndex, Transaction, TransactionBuilder, UpdateMap, UpdateMapEntry, UpdateMode, - UpdatedFragmentOffsets, translate_config_updates, translate_schema_metadata_updates, - validate_operation, + RewrittenIndex, SchemaInputKind, Transaction, TransactionBuilder, UpdateMap, + UpdateMapEntry, UpdateMode, UpdatedFragmentOffsets, translate_config_updates, + translate_schema_metadata_updates, validate_operation, }; } pub mod udtf; @@ -4106,6 +4106,8 @@ pub(crate) struct ManifestWriteConfig { migration_next_row_id: Option, // default None /// Whether this commit activates stable field IDs. activate_stable_field_ids: bool, + /// How to interpret the schema supplied to this commit; never persisted. + pub(crate) schema_input_kind: lance_table::transaction::SchemaInputKind, } impl Default for ManifestWriteConfig { @@ -4119,6 +4121,7 @@ impl Default for ManifestWriteConfig { storage_format: None, migration_next_row_id: None, activate_stable_field_ids: false, + schema_input_kind: Default::default(), } } } diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index ff32459025a..db06f3f76e7 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1371,6 +1371,7 @@ async fn test_write_manifest( disable_transaction_file: false, migration_next_row_id: None, activate_stable_field_ids: false, + schema_input_kind: Default::default(), }, dataset.manifest_location.naming_scheme, None, @@ -3536,6 +3537,7 @@ async fn write_manifest_file_rejects_a_nullable_primary_key() { disable_transaction_file: false, migration_next_row_id: None, activate_stable_field_ids: false, + schema_input_kind: Default::default(), }, dataset.manifest_location.naming_scheme, None, diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index 406bf57d766..7b519b3c83d 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -18,10 +18,10 @@ use lance_table::rowids::read_row_ids; use crate::dataset::write::{WriteMode, WriteParams}; use arrow::compute::concat_batches; -use arrow_array::RecordBatch; use arrow_array::{ Array, Float32Array, Int32Array, Int64Array, ListArray, RecordBatchIterator, UInt32Array, }; +use arrow_array::{RecordBatch, record_batch}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use lance_file::version::LanceFileVersion; @@ -722,8 +722,22 @@ async fn test_new_datasets_use_legacy_field_ids_until_explicit_migration() { #[tokio::test] async fn test_stable_field_id_restore_boundary_and_high_water_mark() { - let source_uri = TempStrDir::default(); - let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + let source_uri = "memory://"; + let dataset = make_simple_dataset(source_uri, 2).await; + let legacy_version = dataset.version().version; + let legacy_id = dataset.schema().field("id").unwrap().id; + let batch = record_batch!(("replacement", Utf8, ["three", "four"])).unwrap(); + let expected = batch.clone(); + let mut dataset = InsertBuilder::new(Arc::new(dataset)) + .with_params(&WriteParams { + mode: WriteMode::Overwrite, + max_rows_per_file: 1, + ..Default::default() + }) + .execute(vec![batch]) + .await + .unwrap(); + assert_eq!(dataset.schema().field("replacement").unwrap().id, legacy_id); dataset.migrate_to_stable_field_ids().await.unwrap(); let activation_version = dataset.version().version; @@ -749,6 +763,37 @@ async fn test_stable_field_id_restore_boundary_and_high_water_mark() { 0 ); assert!(activation_snapshot.schema().field("new_field").is_none()); + + // The old snapshot is readable, but publishing it would bind the same ID to + // a different field even if we retained the current high-water mark. + let current_version = activation_snapshot.version().version; + let mut legacy_snapshot = dataset.checkout_version(legacy_version).await.unwrap(); + assert_eq!(legacy_snapshot.schema().field("id").unwrap().id, legacy_id); + assert_eq!( + legacy_snapshot + .scan() + .try_into_batch() + .await + .unwrap() + .num_rows(), + 2 + ); + let err = legacy_snapshot.restore().await.unwrap_err(); + assert!( + matches!(err, lance_core::Error::InvalidInput { .. }), + "{err}" + ); + assert!( + err.to_string() + .contains("stable field IDs were activated after that version"), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.version().version, current_version); + assert_eq!(dataset.schema().field("replacement").unwrap().id, legacy_id); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(1)); + assert_eq!(dataset.manifest.fragments.len(), 2); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); } #[tokio::test] @@ -813,53 +858,29 @@ async fn test_overwrite_preserves_compatible_stable_field_identities() { #[tokio::test] async fn test_raw_arrow_overwrite_preserves_reordered_stable_field_identities() { - let source_uri = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("a", DataType::Int64, false), - ArrowField::new("b", DataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int64Array::from(vec![1, 2])), - Arc::new(Int64Array::from(vec![3, 4])), - ], - ) - .unwrap(); - let mut dataset = Dataset::write( - RecordBatchIterator::new(vec![Ok(batch)], schema), - source_uri.as_str(), - None, - ) - .await - .unwrap(); + let batch = record_batch!(("a", Int64, [1, 2]), ("b", Int64, [3, 4])).unwrap(); + let mut dataset = InsertBuilder::new("memory://") + .execute(vec![batch]) + .await + .unwrap(); dataset.migrate_to_stable_field_ids().await.unwrap(); - let reordered_schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("b", DataType::Int64, false), - ArrowField::new("a", DataType::Int64, false), - ])); - let reordered_batch = RecordBatch::try_new( - reordered_schema.clone(), - vec![ - Arc::new(Int64Array::from(vec![30, 40])), - Arc::new(Int64Array::from(vec![10, 20])), - ], - ) - .unwrap(); - let overwritten = Dataset::write( - RecordBatchIterator::new(vec![Ok(reordered_batch)], reordered_schema), - source_uri.as_str(), - Some(WriteParams { + let reordered_batch = record_batch!(("b", Int64, [30, 40]), ("a", Int64, [10, 20])).unwrap(); + let expected = reordered_batch.clone(); + let overwritten = InsertBuilder::new(Arc::new(dataset)) + .with_params(&WriteParams { mode: WriteMode::Overwrite, + max_rows_per_file: 1, ..Default::default() - }), - ) - .await - .unwrap(); + }) + .execute(vec![reordered_batch]) + .await + .unwrap(); assert_eq!(overwritten.schema().field("b").unwrap().id, 1); assert_eq!(overwritten.schema().field("a").unwrap().id, 0); + assert_eq!(overwritten.manifest.fragments.len(), 2); + assert_eq!(overwritten.scan().try_into_batch().await.unwrap(), expected); assert!( overwritten .manifest diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 56ebb873d6b..97085e54ada 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -29,9 +29,7 @@ use lance_io::traits::Writer; use lance_table::format::{BasePath, DataFile, Fragment, IndexMetadata}; use lance_table::io::commit::{CommitHandler, commit_handler_from_url}; use lance_table::io::manifest::ManifestDescribing; -use lance_table::transaction::{ - Operation, TRANSACTION_SCHEMA_SOURCE_RAW_ARROW, canonicalize_stable_field_ids, -}; +use lance_table::transaction::{Operation, SchemaInputKind, canonicalize_stable_field_ids}; use object_store::path::Path; use std::borrow::Cow; use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; @@ -1533,19 +1531,18 @@ pub(super) fn prepare_write_schema( // Uncommitted fragment APIs return files without the schema used to // write them, so their mappings must already use commit-time IDs. // The converted Arrow schema carries positional IDs, not trusted Lance - // identities, so preserve its provenance through canonicalization. - let mut raw_schema = normalized_converted_schema; - raw_schema.metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); + // field IDs, so explicitly select Arrow input semantics. let mut operation = Operation::Overwrite { fragments: Vec::new(), - schema: raw_schema, + schema: normalized_converted_schema, config_upsert_values: None, initial_bases: None, }; - canonicalize_stable_field_ids(Some(&dataset.manifest), &mut operation)?; + canonicalize_stable_field_ids( + Some(&dataset.manifest), + &mut operation, + SchemaInputKind::Arrow, + )?; match operation { Operation::Overwrite { schema, .. } => schema, _ => { diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index 2cf1a05bb0b..cf0d6399aad 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -11,6 +11,7 @@ use lance_select::RowAddrTreeMap; use lance_table::{ format::{DataStorageFormat, is_detached_version}, io::commit::{CommitConfig, CommitHandler, ManifestNamingScheme}, + transaction::SchemaInputKind, }; use crate::io::commit::DEFAULT_COMMIT_RETRY_TIMEOUT; @@ -55,6 +56,7 @@ pub struct CommitBuilder<'a> { migration_next_row_id: Option, /// Whether this commit atomically activates stable field IDs. activate_stable_field_ids: bool, + schema_input_kind: SchemaInputKind, } /// Default timeout applied to [`CommitBuilder::execute`] when none is set. @@ -80,9 +82,27 @@ impl<'a> CommitBuilder<'a> { timeout: Some(DEFAULT_COMMIT_TIMEOUT), migration_next_row_id: None, activate_stable_field_ids: false, + schema_input_kind: SchemaInputKind::Lance, } } + /// Interpret transaction schema IDs according to their input source. + /// + /// Defaults to [`SchemaInputKind::Lance`]. Arrow Project inputs must retain + /// explicit IDs and leave missing IDs unassigned. This setting applies to + /// every commit attempt and is not persisted in the transaction. + /// + /// ```no_run + /// # use lance::dataset::CommitBuilder; + /// use lance::dataset::transaction::SchemaInputKind; + /// let builder = CommitBuilder::new("memory://") + /// .with_schema_input_kind(SchemaInputKind::Arrow); + /// ``` + pub fn with_schema_input_kind(mut self, input_kind: SchemaInputKind) -> Self { + self.schema_input_kind = input_kind; + self + } + /// Whether to use stable row ids. This makes the `_rowid` column stable /// after compaction, but not updates. /// @@ -429,6 +449,7 @@ impl<'a> CommitBuilder<'a> { storage_format: self.storage_format.map(DataStorageFormat::new), migration_next_row_id: self.migration_next_row_id, activate_stable_field_ids: self.activate_stable_field_ids, + schema_input_kind: self.schema_input_kind, ..Default::default() }; @@ -598,6 +619,7 @@ pub struct BatchCommitResult { #[cfg(test)] mod tests { use arrow::array::{Int32Array, RecordBatch}; + use arrow_array::record_batch; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use lance_io::utils::CachedFileSize; @@ -616,6 +638,86 @@ mod tests { use super::*; + #[tokio::test] + async fn raw_arrow_new_dataset_preserves_user_metadata() { + let mut field = + lance_core::datatypes::Field::new_arrow("a", DataType::Int32, false).unwrap(); + field.id = 42; + let metadata = HashMap::from([("source".to_string(), "user metadata".to_string())]); + let schema = lance_core::datatypes::Schema { + fields: vec![field], + metadata: metadata.clone(), + }; + let transaction = Transaction::new( + 0, + Operation::Overwrite { + schema, + fragments: vec![], + config_upsert_values: None, + initial_bases: None, + }, + None, + ); + let dataset = CommitBuilder::new("memory://") + .with_schema_input_kind(SchemaInputKind::Arrow) + .execute(transaction) + .await + .unwrap(); + assert_eq!(dataset.schema().field("a").unwrap().id, 0); + assert_eq!(dataset.schema().metadata, metadata); + let committed = dataset.read_transaction().await.unwrap().unwrap(); + let Operation::Overwrite { schema, .. } = committed.operation else { + panic!("expected Overwrite"); + }; + assert_eq!(schema.field("a").unwrap().id, 0); + assert_eq!(schema.metadata, metadata); + } + + #[rstest::rstest] + #[tokio::test] + async fn raw_arrow_merge_uses_commit_input_kind(#[values(false, true)] detached: bool) { + let batch = record_batch!(("a", Int32, [1, 2]), ("b", Int32, [3, 4])).unwrap(); + let mut dataset = InsertBuilder::new("memory://") + .with_params(&WriteParams { + max_rows_per_file: 1, + ..Default::default() + }) + .execute(vec![batch.clone()]) + .await + .unwrap(); + dataset.migrate_to_stable_field_ids().await.unwrap(); + let mut raw_schema = dataset.schema().clone(); + for field in &mut raw_schema.fields { + field.id += 10; + } + let transaction = Transaction::new( + dataset.version().version, + Operation::Merge { + schema: raw_schema, + fragments: dataset.manifest.fragments.as_ref().clone(), + preserves_nullability: true, + }, + None, + ); + let committed = CommitBuilder::new(Arc::new(dataset.clone())) + .with_schema_input_kind(SchemaInputKind::Arrow) + .with_detached(detached) + .execute(transaction) + .await + .unwrap(); + assert_eq!(is_detached_version(committed.version().version), detached); + assert_eq!(committed.schema(), dataset.schema()); + assert_eq!(committed.manifest.fragments.len(), 2); + assert_eq!(committed.manifest.fragments, dataset.manifest.fragments); + assert_eq!(committed.scan().try_into_batch().await.unwrap(), batch); + let persisted = committed.read_transaction().await.unwrap().unwrap(); + let Operation::Merge { schema, .. } = persisted.operation else { + panic!("expected Merge"); + }; + assert_eq!(&schema, dataset.schema()); + assert_eq!(schema.metadata, dataset.schema().metadata); + } + fn sample_fragment() -> Fragment { let (major_version, minor_version) = LanceFileVersion::Stable.resolve().to_data_file_numbers(); diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index a507a131ffa..81a049e2acd 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -361,7 +361,11 @@ async fn do_commit_new_dataset( store_registry: Arc, ) -> Result<(Manifest, ManifestLocation)> { let mut transaction = transaction.clone(); - canonicalize_stable_field_ids(None, &mut transaction.operation)?; + canonicalize_stable_field_ids( + None, + &mut transaction.operation, + write_config.schema_input_kind, + )?; let transaction = &transaction; validate_operation(None, &transaction.operation)?; let pb_transaction = pb::Transaction::from(transaction); @@ -1093,7 +1097,11 @@ pub(crate) async fn do_commit_detached_transaction( ) -> Result<(Manifest, ManifestLocation)> { ensure_can_write_manifest(&dataset.manifest)?; let mut transaction = transaction.clone(); - canonicalize_stable_field_ids(Some(&dataset.manifest), &mut transaction.operation)?; + canonicalize_stable_field_ids( + Some(&dataset.manifest), + &mut transaction.operation, + write_config.schema_input_kind, + )?; let transaction = &transaction; validate_detached_stable_field_ids(&dataset.manifest, &transaction.operation)?; validate_operation(Some(&dataset.manifest), &transaction.operation)?; @@ -1455,11 +1463,14 @@ pub(crate) async fn commit_transaction( ensure_can_write_manifest(&dataset.manifest)?; } - // Keep the rebased transaction as the source operation for the next - // attempt. Canonicalization consumes transient binding provenance, so - // mutating that source would make a retry depend on the first attempt. + // Assign IDs on an attempt-local copy: a retry must bind the original + // input IDs and file mappings against the new manifest's high-water mark. let mut attempt_transaction = transaction.clone(); - canonicalize_stable_field_ids(Some(&dataset.manifest), &mut attempt_transaction.operation)?; + canonicalize_stable_field_ids( + Some(&dataset.manifest), + &mut attempt_transaction.operation, + write_config.schema_input_kind, + )?; validate_operation(Some(&dataset.manifest), &attempt_transaction.operation)?; // Recomputed every attempt: the rebase above may have rewritten the @@ -1728,13 +1739,13 @@ mod tests { CommitLease, CommitLock, ManifestWriter, RenameCommitHandler, UnsafeCommitHandler, commit_handler_from_url, }; - use lance_table::transaction::TRANSACTION_SCHEMA_SOURCE_RAW_ARROW; + use lance_table::transaction::SchemaInputKind; use lance_testing::datagen::generate_random_array; use super::*; use crate::Dataset; - use crate::dataset::{WriteMode, WriteParams}; + use crate::dataset::{CommitBuilder, WriteMode, WriteParams}; use crate::index::DatasetIndexExt; use crate::index::vector::VectorIndexParams; use crate::utils::test::{DatagenExt, FragmentCount, FragmentRowCount}; @@ -1949,16 +1960,18 @@ mod tests { let raw_schema = Schema { fields: vec![Field::new_arrow("x", DataType::Int32, false).unwrap()], - metadata: HashMap::from([( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - )]), + metadata: HashMap::from([("input-note".to_string(), "project".to_string())]), }; let mut expected = Operation::Project { schema: raw_schema.clone(), preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&foreign_manifest), &mut expected).unwrap(); + canonicalize_stable_field_ids( + Some(&foreign_manifest), + &mut expected, + SchemaInputKind::Arrow, + ) + .unwrap(); let Operation::Project { schema: expected_schema, .. @@ -1967,27 +1980,24 @@ mod tests { unreachable!(); }; - let committed = Dataset::commit( - uri, - Operation::Project { - schema: raw_schema, - preserves_nullability: true, - }, - Some(dataset.version().version), - None, - Some(handler), - Default::default(), - false, - ) - .await - .unwrap(); + let committed = CommitBuilder::new(Arc::new(dataset.clone())) + .with_commit_handler(handler) + .with_schema_input_kind(SchemaInputKind::Arrow) + .execute(Transaction::new( + dataset.version().version, + Operation::Project { + schema: raw_schema, + preserves_nullability: true, + }, + None, + )) + .await + .unwrap(); assert_eq!(committed.schema(), &expected_schema); - assert!( - !committed - .schema() - .metadata - .contains_key(TRANSACTION_SCHEMA_SOURCE_RAW_ARROW) + assert_eq!( + committed.schema().metadata.get("input-note").unwrap(), + "project" ); } @@ -2027,10 +2037,7 @@ mod tests { Field::new_arrow("x", DataType::Int32, false).unwrap(), Field::new_arrow("new_column", DataType::Int32, true).unwrap(), ], - metadata: HashMap::from([( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - )]), + metadata: HashMap::new(), }; let operation = Operation::Merge { fragments: vec![merged_fragment], @@ -2038,20 +2045,20 @@ mod tests { preserves_nullability: true, }; let mut expected = operation.clone(); - canonicalize_stable_field_ids(Some(&foreign_manifest), &mut expected).unwrap(); - - let committed = Dataset::commit( - uri, - operation, - Some(dataset.version().version), - None, - Some(handler), - Default::default(), - false, + canonicalize_stable_field_ids( + Some(&foreign_manifest), + &mut expected, + SchemaInputKind::Arrow, ) - .await .unwrap(); + let committed = CommitBuilder::new(Arc::new(dataset.clone())) + .with_commit_handler(handler) + .with_schema_input_kind(SchemaInputKind::Arrow) + .execute(Transaction::new(dataset.version().version, operation, None)) + .await + .unwrap(); + let Operation::Merge { schema: expected_schema, fragments: expected_fragments, From eed6d732425cf8d589fbb72c6deba6f2cf70ae17 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 21 Sep 2026 21:24:01 +0800 Subject: [PATCH 19/24] refactor: resolve Arrow field IDs at conversion boundaries --- java/lance-jni/src/blocking_dataset.rs | 3 - java/lance-jni/src/transaction.rs | 136 ++++------ .../lance-namespace-impls/src/dir/manifest.rs | 4 +- rust/lance-table/src/transaction.rs | 2 +- rust/lance-table/src/transaction/validate.rs | 241 +++++++++++------- rust/lance/src/dataset.rs | 9 +- rust/lance/src/dataset/tests/dataset_io.rs | 2 - rust/lance/src/dataset/write.rs | 11 +- rust/lance/src/dataset/write/commit.rs | 32 +-- rust/lance/src/io/commit.rs | 75 +++--- 10 files changed, 254 insertions(+), 261 deletions(-) diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index f6a2166cfb6..ba7a8c6aabb 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -59,7 +59,6 @@ use lance_namespace::LanceNamespace; use lance_table::io::commit::CommitHandler; use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; use lance_table::io::commit::{ManifestLocation, ManifestNamingScheme}; -use lance_table::transaction::SchemaInputKind; use std::collections::HashMap; use std::future::IntoFuture; use std::iter::empty; @@ -397,7 +396,6 @@ impl BlockingDataset { pub fn commit_transaction( &mut self, transaction: Transaction, - schema_input_kind: SchemaInputKind, store_params: ObjectStoreParams, detached: bool, enable_v2_manifest_paths: bool, @@ -409,7 +407,6 @@ impl BlockingDataset { commit_timeout: Option, ) -> Result { let mut builder = CommitBuilder::new(Arc::new(self.clone().inner)) - .with_schema_input_kind(schema_input_kind) .with_store_params(store_params) .with_detached(detached) .enable_v2_manifest_paths(enable_v2_manifest_paths) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index 5938f362c6b..485435a4462 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -28,9 +28,10 @@ use lance_core::datatypes::Field; use lance_core::datatypes::Schema as LanceSchema; use lance_file::version::{LanceFileVersion, V2_FORMAT_2_0, V2_FORMAT_2_1, V2_FORMAT_2_2}; use lance_io::object_store::{LanceNamespaceStorageOptionsProvider, StorageOptionsProvider}; +use lance_table::format::Manifest; use lance_table::io::commit::CommitHandler; use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; -use lance_table::transaction::SchemaInputKind; +use lance_table::transaction::resolve_arrow_field_ids; use prost::Message; use prost_types::Any; use roaring::RoaringBitmap; @@ -838,7 +839,7 @@ fn inner_commit_to_dataset<'local>( unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; BlockingDataset::new(dataset_guard.inner.clone()) }; - let (transaction, schema_input_kind) = convert_to_rust_transaction( + let transaction = convert_to_rust_transaction( env, java_transaction, Some(&java_allocator), @@ -869,7 +870,6 @@ fn inner_commit_to_dataset<'local>( unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; dataset_guard.commit_transaction( transaction, - schema_input_kind, store_params, detached, enable_v2_manifest_paths, @@ -889,7 +889,7 @@ fn convert_to_rust_transaction( java_transaction: JObject, allocator: Option<&JObject>, dataset: Option<&mut BlockingDataset>, -) -> Result<(Transaction, SchemaInputKind)> { +) -> Result { let read_ver = env.get_u64_from_method(&java_transaction, "readVersion")?; let uuid = env.get_string_from_method(&java_transaction, "uuid")?; let op = env @@ -900,8 +900,7 @@ fn convert_to_rust_transaction( &[], )? .l()?; - let (op, schema_input_kind) = - convert_to_rust_operation(env, &op, allocator, dataset, read_ver)?; + let op = convert_to_rust_operation(env, &op, allocator, dataset, read_ver)?; let tag = env.get_optional_from_method(&java_transaction, "tag", |env, tag_obj| { let tag_str = JString::from(tag_obj); @@ -916,14 +915,11 @@ fn convert_to_rust_transaction( to_rust_map(env, &transaction_properties) }, )?; - Ok(( - TransactionBuilder::new(read_ver, op) - .uuid(uuid) - .tag(tag) - .transaction_properties(transaction_properties.map(Arc::new)) - .build(), - schema_input_kind, - )) + Ok(TransactionBuilder::new(read_ver, op) + .uuid(uuid) + .tag(tag) + .transaction_properties(transaction_properties.map(Arc::new)) + .build()) } #[derive(Clone, Copy)] @@ -949,7 +945,6 @@ type SchemaReadContext = (LanceSchema, i32, bool); struct ConvertedSchema { schema: LanceSchema, field_id_remap: HashMap, - input_kind: SchemaInputKind, } fn convert_arrow_schema( @@ -983,7 +978,6 @@ fn convert_arrow_schema( return Ok(ConvertedSchema { schema: original_schema, field_id_remap: HashMap::new(), - input_kind: SchemaInputKind::Arrow, }); } @@ -991,7 +985,6 @@ fn convert_arrow_schema( return Ok(ConvertedSchema { schema: original_schema, field_id_remap: HashMap::new(), - input_kind: SchemaInputKind::Lance, }); } @@ -1009,7 +1002,6 @@ fn convert_arrow_schema( Ok(ConvertedSchema { schema, field_id_remap, - input_kind: SchemaInputKind::Lance, }) } @@ -1017,8 +1009,7 @@ fn convert_schema_from_operation( env: &mut JNIEnv, java_operation: &JObject, java_allocator: &JObject, - dataset: Option<&mut BlockingDataset>, - read_version: u64, + manifest: Option<&Manifest>, options: SchemaConversionOptions, ) -> Result { let schema_ptr = env @@ -1033,53 +1024,19 @@ fn convert_schema_from_operation( let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) }; let arrow_schema = Schema::try_from(&c_schema)?; - let read_context = match dataset { - Some(dataset) if dataset.inner.version().version == read_version => Some(( - dataset.inner.schema().clone(), - dataset.inner.manifest().max_field_id(), - dataset.inner.manifest().uses_stable_field_ids(), - )), - Some(dataset) => { - let read_dataset = dataset.checkout_version(read_version)?; - Some(( - read_dataset.inner.schema().clone(), - read_dataset.inner.manifest().max_field_id(), - read_dataset.inner.manifest().uses_stable_field_ids(), - )) - } - None => None, - }; + let read_context = manifest.map(|manifest| { + ( + manifest.schema.clone(), + manifest.max_field_id(), + manifest.uses_stable_field_ids(), + ) + }); convert_arrow_schema(&arrow_schema, read_context, options) } type DataFileIdentity = (Option, String); -fn retained_file_identities( - dataset: Option<&mut BlockingDataset>, - read_version: u64, -) -> Result> { - let Some(dataset) = dataset else { - return Ok(HashSet::new()); - }; - let collect = |dataset: &BlockingDataset| { - dataset - .inner - .manifest() - .fragments - .iter() - .flat_map(|fragment| fragment.referenced_lance_files()) - .map(|file| (file.base_id, file.path.clone())) - .collect() - }; - if dataset.inner.version().version == read_version { - Ok(collect(dataset)) - } else { - let read_dataset = dataset.checkout_version(read_version)?; - Ok(collect(&read_dataset)) - } -} - fn remap_fragment_field_ids( fragments: &mut [Fragment], field_id_remap: &HashMap, @@ -1224,16 +1181,27 @@ fn convert_to_rust_operation( env: &mut JNIEnv<'_>, java_operation: &JObject<'_>, allocator: Option<&JObject<'_>>, - mut dataset: Option<&mut BlockingDataset>, + dataset: Option<&mut BlockingDataset>, read_version: u64, -) -> Result<(Operation, SchemaInputKind)> { - let mut schema_input_kind = SchemaInputKind::Lance; +) -> Result { let op_name = env.get_string_from_method(java_operation, "name")?; - let op = match op_name.as_str() { + let read_dataset = if matches!(op_name.as_str(), "Project" | "Overwrite" | "Merge") { + match dataset { + Some(dataset) if dataset.inner.version().version != read_version => { + Some(dataset.checkout_version(read_version)?) + } + Some(dataset) => Some(dataset.clone()), + None => None, + } + } else { + None + }; + let manifest = read_dataset + .as_ref() + .map(|dataset| dataset.inner.manifest()); + let mut op = match op_name.as_str() { "Project" => { - let ConvertedSchema { - schema, input_kind, .. - } = convert_schema_from_operation( + let ConvertedSchema { schema, .. } = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1241,14 +1209,12 @@ fn convert_to_rust_operation( "BufferAllocator is required for Project operations".to_string(), ) })?, - dataset.as_deref_mut(), - read_version, + manifest, SchemaConversionOptions { raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, legacy_field_id_mode: LegacyFieldIdMode::Inherit, }, )?; - schema_input_kind = input_kind; Operation::Project { preserves_nullability: env .get_boolean_from_method(java_operation, "preservesNullability")?, @@ -1377,7 +1343,6 @@ fn convert_to_rust_operation( let ConvertedSchema { schema, field_id_remap, - input_kind, } = convert_schema_from_operation( env, java_operation, @@ -1386,14 +1351,12 @@ fn convert_to_rust_operation( "BufferAllocator is required for Overwrite operations".to_string(), ) })?, - dataset.as_deref_mut(), - read_version, + manifest, SchemaConversionOptions { raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, legacy_field_id_mode: LegacyFieldIdMode::Standalone, }, )?; - schema_input_kind = input_kind; remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); Operation::Overwrite { fragments, @@ -1543,7 +1506,6 @@ fn convert_to_rust_operation( let ConvertedSchema { schema, field_id_remap, - input_kind, } = convert_schema_from_operation( env, java_operation, @@ -1552,18 +1514,21 @@ fn convert_to_rust_operation( "BufferAllocator is required for Merge operations".to_string(), ) })?, - dataset.as_deref_mut(), - read_version, + manifest, SchemaConversionOptions { raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, legacy_field_id_mode: LegacyFieldIdMode::Inherit, }, )?; - schema_input_kind = input_kind; let retained_files = if field_id_remap.is_empty() { HashSet::new() } else { - retained_file_identities(dataset, read_version)? + manifest + .into_iter() + .flat_map(|manifest| manifest.fragments.iter()) + .flat_map(|fragment| fragment.referenced_lance_files()) + .map(|file| (file.base_id, file.path.clone())) + .collect() }; remap_fragment_field_ids(&mut fragments, &field_id_remap, &retained_files); Operation::Merge { @@ -1601,7 +1566,8 @@ fn convert_to_rust_operation( } _ => unimplemented!(), }; - Ok((op, schema_input_kind)) + resolve_arrow_field_ids(manifest, &mut op)?; + Ok(op) } fn extract_update_map(env: &mut JNIEnv, update_map_obj: &JObject) -> Result> { @@ -1837,12 +1803,11 @@ fn inner_commit_to_uri<'local>( } else { Some(allocator_obj) }; - let (transaction, schema_input_kind) = + let transaction = convert_to_rust_transaction(env, java_transaction, allocator_ref.as_ref(), ds.as_mut())?; // Build CommitBuilder with URI let mut builder = CommitBuilder::new(&*uri_str) - .with_schema_input_kind(schema_input_kind) .with_store_params(store_params) .with_detached(detached) .enable_v2_manifest_paths(enable_v2_manifest_paths) @@ -1928,7 +1893,6 @@ mod tests { let ConvertedSchema { schema, field_id_remap, - input_kind, } = convert_arrow_schema( &arrow_schema, Some((base_schema, 0, true)), @@ -1939,7 +1903,6 @@ mod tests { ) .unwrap(); - assert_eq!(input_kind, SchemaInputKind::Arrow); assert!(schema.metadata.is_empty()); assert_eq!(schema.field("a").unwrap().id, -1); assert!(field_id_remap.is_empty()); @@ -1969,7 +1932,6 @@ mod tests { }, ) .unwrap(); - assert_eq!(converted.input_kind, SchemaInputKind::Arrow); assert_eq!(converted.schema.field("renamed").unwrap().id, 0); assert_eq!(converted.schema.metadata, metadata); assert!(converted.field_id_remap.is_empty()); @@ -1988,7 +1950,6 @@ mod tests { let ConvertedSchema { schema, field_id_remap, - input_kind, } = convert_arrow_schema( &arrow_schema, Some((base_schema, 0, false)), @@ -2000,7 +1961,6 @@ mod tests { .unwrap(); assert_eq!(schema.field("a").unwrap().data_type(), ArrowDataType::Utf8); - assert_eq!(input_kind, SchemaInputKind::Lance); assert!(schema.metadata.is_empty()); assert!(field_id_remap.is_empty()); } diff --git a/rust/lance-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index df32fa321b0..ad03219ab5b 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -61,7 +61,7 @@ use lance_table::io::commit::{ CommitError, CommitHandler, commit_handler_from_url, write_manifest_file_to_path, }; use lance_table::transaction::{ - SchemaInputKind, canonicalize_stable_field_ids, validate_stable_field_id_transition, + canonicalize_stable_field_ids, validate_stable_field_id_transition, }; use object_store::{Error as ObjectStoreError, path::Path}; use roaring::RoaringBitmap; @@ -2024,7 +2024,7 @@ impl ManifestNamespace { if let Err(err) = canonicalize_stable_field_ids( Some(dataset.manifest()), &mut transaction.operation, - SchemaInputKind::Lance, + None, ) { self.cleanup_staged_manifest_files(&object_store, &staged_data_files, &[]) .await; diff --git a/rust/lance-table/src/transaction.rs b/rust/lance-table/src/transaction.rs index ad97cc205f5..d1c768804d1 100644 --- a/rust/lance-table/src/transaction.rs +++ b/rust/lance-table/src/transaction.rs @@ -48,7 +48,7 @@ pub use update_map::{ UpdateMap, UpdateMapEntry, translate_config_updates, translate_schema_metadata_updates, }; pub use validate::{ - SchemaInputKind, canonicalize_stable_field_ids, validate_detached_stable_field_ids, + canonicalize_stable_field_ids, resolve_arrow_field_ids, validate_detached_stable_field_ids, validate_operation, validate_stable_field_id_manifest, validate_stable_field_id_transition, }; diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 064236c670f..25c8e0671c9 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -24,51 +24,35 @@ struct FieldIdRemap { raw_source_ids: HashMap, } -/// How to interpret field IDs in the schema supplied to a commit. +/// Resolve an Arrow-derived operation against the dataset it was read from. /// -/// This is input context, not part of the persisted schema or transaction. -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub enum SchemaInputKind { - /// IDs describe Lance fields, including fields preserved across renames. - #[default] - Lance, - /// IDs come from Arrow conversion or caller metadata. Overwrite and Merge - /// match existing fields by name and type, not by these input IDs. Project - /// may use explicit IDs to rename existing fields; missing IDs must remain - /// unassigned rather than receiving positional IDs during conversion. - Arrow, -} - -/// Assign field IDs for an operation before validation. +/// Overwrite and Merge match fields by name and type, not positional Arrow IDs. +/// Project may use explicit IDs for renames; its conversion must leave missing +/// IDs unassigned. New file mappings follow the resolved schema; retained files +/// are unchanged. Call this at the conversion boundary, before committing the +/// resulting Lance operation. Commit-time allocation still handles retries. /// -/// - New datasets receive field IDs starting at zero. -/// - Stable datasets preserve IDs for compatible existing fields and assign new -/// fields IDs above the persisted high-water mark. -/// - Arrow field-ID metadata cannot choose the ID of a new field. -/// - Raw Arrow Project operations must match existing fields: they write no data -/// and cannot add fields. -/// - New files' field mappings follow the assigned IDs; files retained by a merge -/// are never remapped. -pub fn canonicalize_stable_field_ids( +/// ```no_run +/// # use lance_table::{format::Manifest, transaction::{Operation, resolve_arrow_field_ids}}; +/// # fn convert(manifest: &Manifest, operation: &mut Operation) -> lance_core::Result<()> { +/// resolve_arrow_field_ids(Some(manifest), operation)?; +/// # Ok(()) +/// # } +/// ``` +pub fn resolve_arrow_field_ids( manifest: Option<&Manifest>, operation: &mut Operation, - input_kind: SchemaInputKind, ) -> Result<()> { - let raw_arrow_schema = input_kind == SchemaInputKind::Arrow; if manifest.is_some_and(|manifest| !manifest.uses_stable_field_ids()) { - if raw_arrow_schema { - match operation { - Operation::Overwrite { schema, .. } - | Operation::Project { schema, .. } - | Operation::Merge { schema, .. } => { - // Legacy datasets retain the standalone Arrow conversion - // contract. Missing IDs still need to be assigned. - schema.try_set_field_id(None)?; - schema.validate()?; - schema.verify_primary_key()?; - } - _ => {} + match operation { + Operation::Overwrite { schema, .. } + | Operation::Project { schema, .. } + | Operation::Merge { schema, .. } => { + schema.try_set_field_id(None)?; + schema.validate()?; + schema.verify_primary_key()?; } + _ => {} } return Ok(()); } @@ -77,11 +61,10 @@ pub fn canonicalize_stable_field_ids( Operation::Overwrite { schema, fragments, .. } => { - let field_id_remap = - canonicalize_schema(manifest, schema, !raw_arrow_schema, raw_arrow_schema)?; + let field_id_remap = canonicalize_schema(manifest, schema, None, true)?; remap_fragment_field_ids(fragments, &field_id_remap, &HashSet::new())?; } - Operation::Project { schema, .. } if raw_arrow_schema => { + Operation::Project { schema, .. } => { let Some(manifest) = manifest else { return Ok(()); }; @@ -99,11 +82,63 @@ pub fn canonicalize_stable_field_ids( .flat_map(|fragment| fragment.referenced_lance_files()) .map(|file| (file.base_id, file.path.clone())) .collect(); - if raw_arrow_schema { - let field_id_remap = canonicalize_schema(Some(manifest), schema, false, true)?; - remap_fragment_field_ids(fragments, &field_id_remap, &retained_files)?; + let field_id_remap = canonicalize_schema(Some(manifest), schema, None, true)?; + remap_fragment_field_ids(fragments, &field_id_remap, &retained_files)?; + } + _ => {} + } + Ok(()) +} + +/// Assign IDs and update new file mappings before committing a Lance operation. +/// +/// `manifest` is the latest version. On a retry, `read_schema` identifies which +/// input IDs referred to existing fields when the transaction was prepared; +/// `None` uses the manifest's schema. New fields are allocated above the latest +/// high-water mark, without binding their provisional IDs to fields introduced +/// by a concurrent commit. Arrow inputs must first use [`resolve_arrow_field_ids`]. +/// +/// ```no_run +/// # use lance_table::{format::Manifest, transaction::{Operation, canonicalize_stable_field_ids}}; +/// # fn commit(read: &Manifest, latest: &Manifest, operation: &mut Operation) -> lance_core::Result<()> { +/// canonicalize_stable_field_ids(Some(latest), operation, Some(&read.schema))?; +/// # Ok(()) +/// # } +/// ``` +pub fn canonicalize_stable_field_ids( + manifest: Option<&Manifest>, + operation: &mut Operation, + read_schema: Option<&Schema>, +) -> Result<()> { + if manifest.is_some_and(|manifest| !manifest.uses_stable_field_ids()) { + return Ok(()); + } + match operation { + Operation::Overwrite { + schema, fragments, .. + } => { + let identity_schema = read_schema.or_else(|| manifest.map(|manifest| &manifest.schema)); + let remap = canonicalize_schema(manifest, schema, identity_schema, false)?; + remap_fragment_field_ids(fragments, &remap, &HashSet::new())?; + } + Operation::Merge { + schema, fragments, .. + } => { + if let Some(manifest) = manifest { + let retained_files = manifest + .fragments + .iter() + .flat_map(|fragment| fragment.referenced_lance_files()) + .map(|file| (file.base_id, file.path.clone())) + .collect(); + canonicalize_merge_replacements( + manifest, + schema, + fragments, + &retained_files, + read_schema.unwrap_or(&manifest.schema), + )?; } - canonicalize_merge_replacements(manifest, schema, fragments, &retained_files)?; } _ => {} } @@ -115,6 +150,7 @@ fn canonicalize_merge_replacements( schema: &mut Schema, fragments: &mut [Fragment], retained_files: &HashSet, + read_schema: &Schema, ) -> Result<()> { let mut replaced_field_ids = HashSet::new(); for fragment in fragments.iter() { @@ -131,16 +167,11 @@ fn canonicalize_merge_replacements( .filter(|field_id| retained_field_ids.contains(field_id)), ); } - if replaced_field_ids.is_empty() { - return Ok(()); - } - let original = schema.clone(); - let max_field_id = manifest.max_field_id(); for field in &mut schema.fields { - clear_replaced_and_new_field_ids(field, max_field_id, &replaced_field_ids); + clear_replaced_and_new_field_ids(field, read_schema, &replaced_field_ids); } - schema.try_set_field_id(Some(max_field_id))?; + schema.try_set_field_id(Some(manifest.max_field_id()))?; schema.validate()?; schema.verify_primary_key()?; @@ -159,22 +190,22 @@ fn canonicalize_merge_replacements( fn clear_replaced_and_new_field_ids( field: &mut Field, - max_field_id: i32, + base_schema: &Schema, replaced_field_ids: &HashSet, ) { - if field.id > max_field_id || replaced_field_ids.contains(&field.id) { + if base_schema.field_by_id(field.id).is_none() || replaced_field_ids.contains(&field.id) { clear_field_ids(field); return; } for child in &mut field.children { - clear_replaced_and_new_field_ids(child, max_field_id, replaced_field_ids); + clear_replaced_and_new_field_ids(child, base_schema, replaced_field_ids); } } fn canonicalize_raw_project_schema(manifest: &Manifest, schema: &mut Schema) -> Result<()> { let mut unmatched_fields = Vec::new(); for field in &mut schema.fields { - if !canonicalize_field(field, -1, &manifest.schema, None, true) { + if !canonicalize_field(field, -1, &manifest.schema, None, Some(&manifest.schema)) { unmatched_fields.push(field.name.clone()); } } @@ -192,7 +223,7 @@ fn canonicalize_raw_project_schema(manifest: &Manifest, schema: &mut Schema) -> fn canonicalize_schema( manifest: Option<&Manifest>, schema: &mut Schema, - allow_id_binding: bool, + identity_schema: Option<&Schema>, remap_raw_source_ids: bool, ) -> Result { let original = schema.clone(); @@ -209,7 +240,7 @@ fn canonicalize_schema( schema.try_reassign_field_ids(max_existing_id)?; } else if let Some(manifest) = manifest { for field in &mut schema.fields { - canonicalize_field(field, -1, &manifest.schema, None, allow_id_binding); + canonicalize_field(field, -1, &manifest.schema, None, identity_schema); } schema.try_set_field_id(max_existing_id)?; } @@ -241,7 +272,7 @@ fn canonicalize_field( parent_id: i32, base_schema: &Schema, base_parent: Option<&Field>, - allow_id_binding: bool, + identity_schema: Option<&Schema>, ) -> bool { let same_name = match base_parent { Some(parent) => parent.children.iter().find(|base| base.name == field.name), @@ -250,6 +281,8 @@ fn canonicalize_field( .iter() .find(|base| base.name == field.name), }; + let allow_id_binding = + identity_schema.is_some_and(|schema| schema.field_by_id(field.id).is_some()); let by_id = (allow_id_binding && field.id >= 0) .then(|| base_schema.field_by_id(field.id)) .flatten() @@ -276,7 +309,7 @@ fn canonicalize_field( field.id, base_schema, Some(base_field), - allow_id_binding, + identity_schema, ) { all_children_match = false; } @@ -1119,8 +1152,7 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Lance) - .unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, None).unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1144,8 +1176,7 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap(); + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1157,6 +1188,55 @@ mod tests { assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]); } + #[rstest::rstest] + #[case::concurrent_add(false, vec![0, 2])] + #[case::concurrent_replace(true, vec![2, 3])] + fn canonicalize_overwrite_retry_preserves_only_live_read_identities( + #[case] replace: bool, + #[case] expected_ids: Vec, + ) { + let read = activated_manifest(); + let mut staged = read.schema.clone(); + let mut new_field = Field::new_arrow("new_column", DataType::Int32, true).unwrap(); + new_field.id = 1; + staged.fields.push(new_field.clone()); + let mut latest_schema = read.schema.clone(); + if replace { + latest_schema.fields.clear(); + } + new_field.name = "concurrent_column".to_string(); + latest_schema.fields.push(new_field); + let latest_ids = latest_schema.field_ids().into_iter().collect(); + let mut latest = manifest_with_file_fields(latest_schema, latest_ids); + latest.activate_stable_field_ids(); + let mut operation = Operation::Overwrite { + schema: staged, + fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0, 1])], + config_upsert_values: None, + initial_bases: None, + }; + + canonicalize_stable_field_ids(Some(&latest), &mut operation, Some(&read.schema)).unwrap(); + validate_operation(Some(&latest), &operation).unwrap(); + + let Operation::Overwrite { + schema, fragments, .. + } = operation + else { + unreachable!(); + }; + assert_eq!( + schema + .fields + .iter() + .map(|field| field.id) + .collect::>(), + expected_ids + ); + assert_eq!(fragments[0].files[0].fields.as_ref(), expected_ids); + assert_eq!(schema.fields[1].name, "new_column"); + } + #[test] fn canonicalize_raw_arrow_overwrite_matches_reordered_fields_by_name() { let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ @@ -1181,8 +1261,7 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap(); + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); let Operation::Overwrite { schema, .. } = operation else { unreachable!(); @@ -1209,8 +1288,7 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap(); + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1250,8 +1328,7 @@ mod tests { initial_bases: None, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap(); + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); let Operation::Overwrite { schema, fragments, .. @@ -1287,9 +1364,7 @@ mod tests { initial_bases: None, }; - let err = - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap_err(); + let err = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err(); assert!(err.to_string().contains("ambiguous raw Arrow field IDs")); } @@ -1305,9 +1380,7 @@ mod tests { preserves_nullability: true, }; - let err = - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap_err(); + let err = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err(); assert!(err.to_string().contains("writes no data"), "{err}"); } @@ -1322,8 +1395,7 @@ mod tests { preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap(); + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); let Operation::Project { schema, .. } = operation else { unreachable!(); @@ -1342,8 +1414,7 @@ mod tests { preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap(); + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); let Operation::Project { schema, .. } = operation else { unreachable!(); @@ -1367,8 +1438,7 @@ mod tests { preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Lance) - .unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, None).unwrap(); let Operation::Merge { schema, fragments, .. @@ -1405,8 +1475,7 @@ mod tests { preserves_nullability: true, }; - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Lance) - .unwrap(); + canonicalize_stable_field_ids(Some(&manifest), &mut operation, None).unwrap(); let Operation::Merge { schema, fragments, .. @@ -1443,9 +1512,7 @@ mod tests { preserves_nullability: true, }; - let err = - canonicalize_stable_field_ids(Some(&manifest), &mut operation, SchemaInputKind::Arrow) - .unwrap_err(); + let err = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err(); assert!(err.to_string().contains("ambiguous raw Arrow field IDs")); } diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 26fc2970cbc..a081c1dbdc7 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -103,9 +103,9 @@ mod take; pub mod transaction { pub use lance_table::transaction::{ DataOverlayGroup, DataReplacementGroup, Operation, ReadVersionState, RewriteGroup, - RewrittenIndex, SchemaInputKind, Transaction, TransactionBuilder, UpdateMap, - UpdateMapEntry, UpdateMode, UpdatedFragmentOffsets, translate_config_updates, - translate_schema_metadata_updates, validate_operation, + RewrittenIndex, Transaction, TransactionBuilder, UpdateMap, UpdateMapEntry, UpdateMode, + UpdatedFragmentOffsets, translate_config_updates, translate_schema_metadata_updates, + validate_operation, }; } pub mod udtf; @@ -4106,8 +4106,6 @@ pub(crate) struct ManifestWriteConfig { migration_next_row_id: Option, // default None /// Whether this commit activates stable field IDs. activate_stable_field_ids: bool, - /// How to interpret the schema supplied to this commit; never persisted. - pub(crate) schema_input_kind: lance_table::transaction::SchemaInputKind, } impl Default for ManifestWriteConfig { @@ -4121,7 +4119,6 @@ impl Default for ManifestWriteConfig { storage_format: None, migration_next_row_id: None, activate_stable_field_ids: false, - schema_input_kind: Default::default(), } } } diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index db06f3f76e7..ff32459025a 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1371,7 +1371,6 @@ async fn test_write_manifest( disable_transaction_file: false, migration_next_row_id: None, activate_stable_field_ids: false, - schema_input_kind: Default::default(), }, dataset.manifest_location.naming_scheme, None, @@ -3537,7 +3536,6 @@ async fn write_manifest_file_rejects_a_nullable_primary_key() { disable_transaction_file: false, migration_next_row_id: None, activate_stable_field_ids: false, - schema_input_kind: Default::default(), }, dataset.manifest_location.naming_scheme, None, diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index 97085e54ada..f4fac512823 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -29,7 +29,7 @@ use lance_io::traits::Writer; use lance_table::format::{BasePath, DataFile, Fragment, IndexMetadata}; use lance_table::io::commit::{CommitHandler, commit_handler_from_url}; use lance_table::io::manifest::ManifestDescribing; -use lance_table::transaction::{Operation, SchemaInputKind, canonicalize_stable_field_ids}; +use lance_table::transaction::{Operation, resolve_arrow_field_ids}; use object_store::path::Path; use std::borrow::Cow; use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; @@ -1530,19 +1530,14 @@ pub(super) fn prepare_write_schema( { // Uncommitted fragment APIs return files without the schema used to // write them, so their mappings must already use commit-time IDs. - // The converted Arrow schema carries positional IDs, not trusted Lance - // field IDs, so explicitly select Arrow input semantics. + // Resolve positional Arrow IDs before the schema enters the writer. let mut operation = Operation::Overwrite { fragments: Vec::new(), schema: normalized_converted_schema, config_upsert_values: None, initial_bases: None, }; - canonicalize_stable_field_ids( - Some(&dataset.manifest), - &mut operation, - SchemaInputKind::Arrow, - )?; + resolve_arrow_field_ids(Some(&dataset.manifest), &mut operation)?; match operation { Operation::Overwrite { schema, .. } => schema, _ => { diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index cf0d6399aad..95a36a9d228 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -11,7 +11,6 @@ use lance_select::RowAddrTreeMap; use lance_table::{ format::{DataStorageFormat, is_detached_version}, io::commit::{CommitConfig, CommitHandler, ManifestNamingScheme}, - transaction::SchemaInputKind, }; use crate::io::commit::DEFAULT_COMMIT_RETRY_TIMEOUT; @@ -56,7 +55,6 @@ pub struct CommitBuilder<'a> { migration_next_row_id: Option, /// Whether this commit atomically activates stable field IDs. activate_stable_field_ids: bool, - schema_input_kind: SchemaInputKind, } /// Default timeout applied to [`CommitBuilder::execute`] when none is set. @@ -82,27 +80,9 @@ impl<'a> CommitBuilder<'a> { timeout: Some(DEFAULT_COMMIT_TIMEOUT), migration_next_row_id: None, activate_stable_field_ids: false, - schema_input_kind: SchemaInputKind::Lance, } } - /// Interpret transaction schema IDs according to their input source. - /// - /// Defaults to [`SchemaInputKind::Lance`]. Arrow Project inputs must retain - /// explicit IDs and leave missing IDs unassigned. This setting applies to - /// every commit attempt and is not persisted in the transaction. - /// - /// ```no_run - /// # use lance::dataset::CommitBuilder; - /// use lance::dataset::transaction::SchemaInputKind; - /// let builder = CommitBuilder::new("memory://") - /// .with_schema_input_kind(SchemaInputKind::Arrow); - /// ``` - pub fn with_schema_input_kind(mut self, input_kind: SchemaInputKind) -> Self { - self.schema_input_kind = input_kind; - self - } - /// Whether to use stable row ids. This makes the `_rowid` column stable /// after compaction, but not updates. /// @@ -449,7 +429,6 @@ impl<'a> CommitBuilder<'a> { storage_format: self.storage_format.map(DataStorageFormat::new), migration_next_row_id: self.migration_next_row_id, activate_stable_field_ids: self.activate_stable_field_ids, - schema_input_kind: self.schema_input_kind, ..Default::default() }; @@ -628,6 +607,7 @@ mod tests { DataFile, Fragment, IndexMetadata, Manifest, Transaction as TableTransaction, }; use lance_table::io::commit::{CommitError, ManifestLocation, ManifestWriter}; + use lance_table::transaction::resolve_arrow_field_ids; use std::time::Duration; use object_store::throttle::ThrottleConfig; @@ -648,7 +628,7 @@ mod tests { fields: vec![field], metadata: metadata.clone(), }; - let transaction = Transaction::new( + let mut transaction = Transaction::new( 0, Operation::Overwrite { schema, @@ -658,8 +638,8 @@ mod tests { }, None, ); + resolve_arrow_field_ids(None, &mut transaction.operation).unwrap(); let dataset = CommitBuilder::new("memory://") - .with_schema_input_kind(SchemaInputKind::Arrow) .execute(transaction) .await .unwrap(); @@ -675,7 +655,7 @@ mod tests { #[rstest::rstest] #[tokio::test] - async fn raw_arrow_merge_uses_commit_input_kind(#[values(false, true)] detached: bool) { + async fn raw_arrow_merge_resolved_before_commit(#[values(false, true)] detached: bool) { let batch = record_batch!(("a", Int32, [1, 2]), ("b", Int32, [3, 4])).unwrap(); let mut dataset = InsertBuilder::new("memory://") .with_params(&WriteParams { @@ -690,7 +670,7 @@ mod tests { for field in &mut raw_schema.fields { field.id += 10; } - let transaction = Transaction::new( + let mut transaction = Transaction::new( dataset.version().version, Operation::Merge { schema: raw_schema, @@ -699,8 +679,8 @@ mod tests { }, None, ); + resolve_arrow_field_ids(Some(&dataset.manifest), &mut transaction.operation).unwrap(); let committed = CommitBuilder::new(Arc::new(dataset.clone())) - .with_schema_input_kind(SchemaInputKind::Arrow) .with_detached(detached) .execute(transaction) .await diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 81a049e2acd..129d84b3eed 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -361,11 +361,7 @@ async fn do_commit_new_dataset( store_registry: Arc, ) -> Result<(Manifest, ManifestLocation)> { let mut transaction = transaction.clone(); - canonicalize_stable_field_ids( - None, - &mut transaction.operation, - write_config.schema_input_kind, - )?; + canonicalize_stable_field_ids(None, &mut transaction.operation, None)?; let transaction = &transaction; validate_operation(None, &transaction.operation)?; let pb_transaction = pb::Transaction::from(transaction); @@ -1097,11 +1093,7 @@ pub(crate) async fn do_commit_detached_transaction( ) -> Result<(Manifest, ManifestLocation)> { ensure_can_write_manifest(&dataset.manifest)?; let mut transaction = transaction.clone(); - canonicalize_stable_field_ids( - Some(&dataset.manifest), - &mut transaction.operation, - write_config.schema_input_kind, - )?; + canonicalize_stable_field_ids(Some(&dataset.manifest), &mut transaction.operation, None)?; let transaction = &transaction; validate_detached_stable_field_ids(&dataset.manifest, &transaction.operation)?; validate_operation(Some(&dataset.manifest), &transaction.operation)?; @@ -1463,13 +1455,15 @@ pub(crate) async fn commit_transaction( ensure_can_write_manifest(&dataset.manifest)?; } - // Assign IDs on an attempt-local copy: a retry must bind the original - // input IDs and file mappings against the new manifest's high-water mark. + // Preserve identities from the read version, but allocate new IDs above + // the latest high-water mark. Each attempt remaps its own copy of the + // staged files, so a retry never mistakes a provisional ID for a field + // introduced by a concurrent commit. let mut attempt_transaction = transaction.clone(); canonicalize_stable_field_ids( Some(&dataset.manifest), &mut attempt_transaction.operation, - write_config.schema_input_kind, + Some(read_version_dataset.schema()), )?; validate_operation(Some(&dataset.manifest), &attempt_transaction.operation)?; @@ -1739,7 +1733,7 @@ mod tests { CommitLease, CommitLock, ManifestWriter, RenameCommitHandler, UnsafeCommitHandler, commit_handler_from_url, }; - use lance_table::transaction::SchemaInputKind; + use lance_table::transaction::resolve_arrow_field_ids; use lance_testing::datagen::generate_random_array; use super::*; @@ -1966,12 +1960,8 @@ mod tests { schema: raw_schema.clone(), preserves_nullability: true, }; - canonicalize_stable_field_ids( - Some(&foreign_manifest), - &mut expected, - SchemaInputKind::Arrow, - ) - .unwrap(); + resolve_arrow_field_ids(Some(&dataset.manifest), &mut expected).unwrap(); + let operation = expected.clone(); let Operation::Project { schema: expected_schema, .. @@ -1982,15 +1972,7 @@ mod tests { let committed = CommitBuilder::new(Arc::new(dataset.clone())) .with_commit_handler(handler) - .with_schema_input_kind(SchemaInputKind::Arrow) - .execute(Transaction::new( - dataset.version().version, - Operation::Project { - schema: raw_schema, - preserves_nullability: true, - }, - None, - )) + .execute(Transaction::new(dataset.version().version, operation, None)) .await .unwrap(); @@ -2001,8 +1983,11 @@ mod tests { ); } + #[rstest::rstest] #[tokio::test] - async fn raw_arrow_merge_retry_rebinds_after_allocator_advance() { + async fn raw_arrow_retry_rebinds_after_allocator_advance( + #[values(false, true)] overwrite: bool, + ) { let tmp = TempStrDir::default(); let uri = tmp.as_str(); let mut dataset = Dataset::write( @@ -2039,31 +2024,45 @@ mod tests { ], metadata: HashMap::new(), }; - let operation = Operation::Merge { - fragments: vec![merged_fragment], - schema: raw_schema, - preserves_nullability: true, + let mut operation = if overwrite { + Operation::Overwrite { + fragments: vec![merged_fragment], + schema: raw_schema, + config_upsert_values: None, + initial_bases: None, + } + } else { + Operation::Merge { + fragments: vec![merged_fragment], + schema: raw_schema, + preserves_nullability: true, + } }; + resolve_arrow_field_ids(Some(&dataset.manifest), &mut operation).unwrap(); let mut expected = operation.clone(); canonicalize_stable_field_ids( Some(&foreign_manifest), &mut expected, - SchemaInputKind::Arrow, + Some(dataset.schema()), ) .unwrap(); let committed = CommitBuilder::new(Arc::new(dataset.clone())) .with_commit_handler(handler) - .with_schema_input_kind(SchemaInputKind::Arrow) .execute(Transaction::new(dataset.version().version, operation, None)) .await .unwrap(); - let Operation::Merge { + let (Operation::Merge { schema: expected_schema, fragments: expected_fragments, .. - } = expected + } + | Operation::Overwrite { + schema: expected_schema, + fragments: expected_fragments, + .. + }) = expected else { unreachable!(); }; From f73688a093011793632a834facb8c5716c7a3580 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 21 Sep 2026 22:03:06 +0800 Subject: [PATCH 20/24] fix: reconcile Arrow schema conversion after merging main --- python/python/tests/test_dataset.py | 30 ++++++---- python/src/dataset.rs | 83 +++++++++++++++++++++++---- python/src/transaction.rs | 11 ++-- rust/lance-table/src/feature_flags.rs | 4 +- 4 files changed, 99 insertions(+), 29 deletions(-) diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index c6272ce8abc..ff2714e1905 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -1984,15 +1984,25 @@ def test_cleanup_with_rate_limit(tmp_path): assert (finished - start) >= 2_000_000_000 # 2s -def test_create_from_commit(tmp_path: Path): - table = pa.Table.from_pydict({"a": range(100), "b": range(100)}) +@pytest.mark.parametrize("as_transaction", [False, True]) +def test_create_from_commit(tmp_path: Path, as_transaction: bool): + metadata = {b"source": b"user metadata"} + table = pa.Table.from_pydict( + {"a": range(100), "b": range(100)} + ).replace_schema_metadata(metadata) base_dir = tmp_path / "test" - fragment = lance.fragment.LanceFragment.create(base_dir, table) + fragments = [ + lance.fragment.LanceFragment.create(base_dir, table.slice(offset, 50)) + for offset in (0, 50) + ] - operation = lance.LanceOperation.Overwrite(table.schema, [fragment]) - dataset = lance.LanceDataset.commit(base_dir, operation) + operation = lance.LanceOperation.Overwrite(table.schema, fragments) + transaction = lance.Transaction(0, operation) if as_transaction else operation + dataset = lance.LanceDataset.commit(base_dir, transaction) tbl = dataset.to_table() assert tbl == table + assert len(dataset.get_fragments()) == 2 + assert dataset.schema.metadata == metadata def test_strict_overwrite(tmp_path: Path): @@ -5790,8 +5800,10 @@ def test_detached_commits(tmp_path: Path): assert detached2.to_table() == pa.table({"x": [0, 1, 3]}) -def test_detached_raw_arrow_merge_strips_transaction_metadata(tmp_path: Path): - dataset = lance.write_dataset(pa.table({"x": [0, 1]}), tmp_path) +def test_detached_raw_arrow_merge_preserves_schema_metadata(tmp_path: Path): + metadata = {b"source": b"user metadata"} + table = pa.table({"x": [0, 1]}).replace_schema_metadata(metadata) + dataset = lance.write_dataset(table, tmp_path) fragment = dataset.get_fragments()[0].metadata with pytest.deprecated_call(): operation = lance.LanceOperation.Merge([fragment], dataset.schema, True) @@ -5804,9 +5816,7 @@ def test_detached_raw_arrow_merge_strips_transaction_metadata(tmp_path: Path): ) assert detached.to_table() == dataset.to_table() - assert b"lance:transaction_schema_source_raw_arrow" not in ( - detached.schema.metadata or {} - ) + assert detached.schema.metadata == metadata def test_dataset_drop(tmp_path: Path): diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 6b442fb3027..80b7727430d 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -95,6 +95,7 @@ use lance_linalg::distance::MetricType; use lance_table::format::{BasePath, Fragment, IndexMetadata}; use lance_table::io::commit::CommitHandler; use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler; +use lance_table::transaction::resolve_arrow_field_ids; use crate::error::PythonErrorExt; use crate::file::object_store_from_uri_or_path; @@ -3048,7 +3049,7 @@ impl Dataset { #[pyo3(signature = (dest, operation, read_version = None, commit_lock = None, storage_options = None, enable_v2_manifest_paths = None, detached = None, max_retries = None, commit_message = None, enable_stable_row_ids = None, namespace_client = None, table_id = None, namespace_client_managed_versioning = false, commit_timeout = None))] fn commit( dest: PyWriteDest, - operation: PyLance, + operation: &Bound<'_, PyAny>, read_version: Option, commit_lock: Option<&Bound<'_, PyAny>>, storage_options: Option>, @@ -3062,18 +3063,22 @@ impl Dataset { namespace_client_managed_versioning: bool, commit_timeout: Option, ) -> PyResult { - let mut transaction = Transaction::new(read_version.unwrap_or_default(), operation.0, None); + let transaction = operation + .py() + .import(intern!(operation.py(), "lance"))? + .getattr("Transaction")? + .call1((read_version.unwrap_or_default(), operation))?; if let Some(commit_message) = commit_message { - transaction.transaction_properties = Some(Arc::new(HashMap::from([( - LANCE_COMMIT_MESSAGE_KEY.to_string(), - commit_message, - )]))); + transaction.setattr( + "transaction_properties", + HashMap::from([(LANCE_COMMIT_MESSAGE_KEY.to_string(), commit_message)]), + )?; } Self::commit_transaction( dest, - PyLance(transaction), + &transaction, commit_lock, storage_options, enable_v2_manifest_paths, @@ -3093,7 +3098,7 @@ impl Dataset { #[pyo3(signature = (dest, transaction, commit_lock = None, storage_options = None, enable_v2_manifest_paths = None, detached = None, max_retries = None, enable_stable_row_ids = None, namespace_client = None, table_id = None, namespace_client_managed_versioning = false, commit_timeout = None))] fn commit_transaction( dest: PyWriteDest, - transaction: PyLance, + transaction: &Bound<'_, PyAny>, commit_lock: Option<&Bound<'_, PyAny>>, storage_options: Option>, enable_v2_manifest_paths: Option, @@ -3105,6 +3110,16 @@ impl Dataset { namespace_client_managed_versioning: bool, commit_timeout: Option, ) -> PyResult { + let mut rust_transaction = transaction.extract::>()?.0; + let operation = transaction.getattr("operation")?; + let input_schema = match &rust_transaction.operation { + Operation::Overwrite { .. } => Some(operation.getattr("new_schema")?), + Operation::Merge { .. } | Operation::Project { .. } => { + Some(operation.getattr("schema")?) + } + _ => None, + }; + let is_arrow = input_schema.is_some_and(|schema| !schema.is_instance_of::()); let accessor = crate::storage_options::create_accessor_from_storage_options(storage_options.clone())?; @@ -3148,7 +3163,55 @@ impl Dataset { None }; - let mut builder = CommitBuilder::new(dest.as_dest()) + // Resolve Arrow IDs while the Python input type is still available. + // The core commit path receives only ordinary Lance operations. + let read_dataset = if is_arrow { + rt().block_on(Some(transaction.py()), async { + let dataset = match &dest { + PyWriteDest::Dataset(dataset) => Some(dataset.ds.clone()), + PyWriteDest::Uri(uri) => { + match DatasetBuilder::from_uri(&**uri) + .with_read_params(ReadParams { + store_options: object_store_params.clone(), + commit_handler: commit_handler.clone(), + ..Default::default() + }) + .load() + .await + { + Ok(dataset) => Some(Arc::new(dataset)), + Err(Error::DatasetNotFound { .. } | Error::NotFound { .. }) => None, + Err(error) => return Err(error), + } + } + }; + let dataset = match dataset { + Some(dataset) + if rust_transaction.read_version != 0 + && dataset.version().version != rust_transaction.read_version => + { + Some(Arc::new( + dataset + .checkout_version(rust_transaction.read_version) + .await?, + )) + } + dataset => dataset, + }; + resolve_arrow_field_ids( + dataset.as_ref().map(|dataset| dataset.manifest()), + &mut rust_transaction.operation, + )?; + Ok::<_, Error>(dataset) + })? + .infer_error()? + } else { + None + }; + let destination = read_dataset + .map(WriteDestination::Dataset) + .unwrap_or_else(|| dest.as_dest()); + let mut builder = CommitBuilder::new(destination) .enable_v2_manifest_paths(enable_v2_manifest_paths.unwrap_or(true)) .with_detached(detached.unwrap_or(false)) .with_max_retries(max_retries.unwrap_or(20)) @@ -3169,7 +3232,7 @@ impl Dataset { let ds = rt() .block_on( commit_lock.map(|cl| cl.py()), - builder.execute(transaction.0), + builder.execute(rust_transaction), )? .io_or_timeout_error()?; diff --git a/python/src/transaction.rs b/python/src/transaction.rs index 82a4ed07de0..6e63d2b8d86 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -14,7 +14,6 @@ use lance::dataset::transaction::{ use lance::datatypes::{Field, Schema}; use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; use lance_table::format::{BasePath, DataFile, Fragment, IndexFile, IndexMetadata}; -use lance_table::transaction::TRANSACTION_SCHEMA_SOURCE_RAW_ARROW; use pyo3::exceptions::PyValueError; use pyo3::{Bound, FromPyObject, PyAny, PyResult, Python}; use pyo3::{intern, prelude::*}; @@ -1101,10 +1100,8 @@ fn convert_schema(arrow_schema: &ArrowSchema) -> PyResult { "Failed to convert Arrow schema to Lance schema: {e}" )) })?; - let mut metadata = arrow_schema.metadata.clone(); - metadata.insert( - TRANSACTION_SCHEMA_SOURCE_RAW_ARROW.to_string(), - String::new(), - ); - Ok(Schema { fields, metadata }) + Ok(Schema { + fields, + metadata: arrow_schema.metadata.clone(), + }) } diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index ac1257b016e..20d71f54142 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -402,7 +402,7 @@ mod tests { assert!(can_read_dataset(super::FLAG_BASE_PATHS)); assert!(can_read_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); assert!(can_read_dataset(super::FLAG_STABLE_FIELD_IDS)); - assert!(!can_read_dataset(super::FLAG_MIXED_DATA_FILE_VERSIONS)); + assert!(can_read_dataset(super::FLAG_MIXED_DATA_FILE_VERSIONS)); // Overlay support is gated on the build profile / env opt-in, so the // flag is readable exactly when overlays are enabled (see // test_data_overlay_flag_release_gating for the full policy). @@ -480,7 +480,7 @@ mod tests { assert!(can_write_dataset(super::FLAG_BASE_PATHS)); assert!(can_write_dataset(super::FLAG_DISABLE_TRANSACTION_FILE)); assert!(can_write_dataset(super::FLAG_STABLE_FIELD_IDS)); - assert!(!can_write_dataset(super::FLAG_MIXED_DATA_FILE_VERSIONS)); + assert!(can_write_dataset(super::FLAG_MIXED_DATA_FILE_VERSIONS)); // Overlay support is gated on the build profile / env opt-in, so the // flag is writable exactly when overlays are enabled (see // test_data_overlay_flag_release_gating for the full policy). From f4afe817fa414a040fdcf588133d380c2e4a7a1d Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 21 Sep 2026 22:56:02 +0800 Subject: [PATCH 21/24] refactor: simplify stable field ID allocation and tests --- java/lance-jni/src/transaction.rs | 140 +++++-------- .../java/org/lance/operation/MergeTest.java | 17 -- .../lance/operation/OperationTestBase.java | 18 ++ .../java/org/lance/operation/ProjectTest.java | 17 -- .../python/tests/compat/compat_decorator.py | 6 +- .../python/tests/compat/test_file_formats.py | 3 +- .../tests/compat/test_scalar_indices.py | 21 +- .../tests/compat/test_vector_indices.py | 13 +- .../lance-namespace-impls/src/dir/manifest.rs | 31 +-- rust/lance-table/src/transaction.rs | 2 +- .../src/transaction/manifest_build.rs | 30 +-- rust/lance-table/src/transaction/validate.rs | 184 +++++++++--------- rust/lance/src/dataset/schema_evolution.rs | 40 +--- .../src/dataset/tests/dataset_migrations.rs | 48 ++--- rust/lance/src/io/commit.rs | 65 ++----- 15 files changed, 219 insertions(+), 416 deletions(-) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index da92c02b41d..44a27ea08d0 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -1513,26 +1513,6 @@ fn convert_to_rust_transaction( .build()) } -#[derive(Clone, Copy)] -enum RawArrowFieldIdMode { - AssignMissing, - // Preserve the distinction between explicit IDs and fields without identity metadata. - ExplicitOnly, -} - -#[derive(Clone, Copy)] -enum LegacyFieldIdMode { - Inherit, - Standalone, -} - -struct SchemaConversionOptions { - raw_field_id_mode: RawArrowFieldIdMode, - legacy_field_id_mode: LegacyFieldIdMode, -} - -type SchemaReadContext = (LanceSchema, i32, bool); - struct ConvertedSchema { schema: LanceSchema, field_id_remap: HashMap, @@ -1540,11 +1520,11 @@ struct ConvertedSchema { fn convert_arrow_schema( arrow_schema: &Schema, - read_context: Option, - options: SchemaConversionOptions, + manifest: Option<&Manifest>, + operation_name: &str, ) -> Result { - let original_schema = if matches!(options.raw_field_id_mode, RawArrowFieldIdMode::ExplicitOnly) - { + // Project can rename by explicit ID but must not treat positional IDs as identity. + let original_schema = if operation_name == "Project" { LanceSchema { fields: arrow_schema .fields @@ -1562,26 +1542,19 @@ fn convert_arrow_schema( })? }; - if read_context - .as_ref() - .is_none_or(|(_, _, stable_field_ids)| *stable_field_ids) - { - return Ok(ConvertedSchema { - schema: original_schema, - field_id_remap: HashMap::new(), - }); - } - - if matches!(options.legacy_field_id_mode, LegacyFieldIdMode::Standalone) { + let Some(manifest) = manifest + .filter(|manifest| !manifest.uses_stable_field_ids() && operation_name != "Overwrite") + else { return Ok(ConvertedSchema { schema: original_schema, field_id_remap: HashMap::new(), }); - } - - let (read_schema, max_field_id, _) = read_context.expect("legacy dataset context"); - let schema = - LanceSchema::from_arrow_schema(arrow_schema, Some(read_schema), Some(max_field_id))?; + }; + let schema = LanceSchema::from_arrow_schema( + arrow_schema, + Some(manifest.schema.clone()), + Some(manifest.max_field_id()), + )?; let field_id_remap = original_schema .fields_pre_order() @@ -1601,7 +1574,7 @@ fn convert_schema_from_operation( java_operation: &JObject, java_allocator: &JObject, manifest: Option<&Manifest>, - options: SchemaConversionOptions, + operation_name: &str, ) -> Result { let schema_ptr = env .call_method( @@ -1615,15 +1588,7 @@ fn convert_schema_from_operation( let c_schema = unsafe { FFI_ArrowSchema::from_raw(c_schema_ptr) }; let arrow_schema = Schema::try_from(&c_schema)?; - let read_context = manifest.map(|manifest| { - ( - manifest.schema.clone(), - manifest.max_field_id(), - manifest.uses_stable_field_ids(), - ) - }); - - convert_arrow_schema(&arrow_schema, read_context, options) + convert_arrow_schema(&arrow_schema, manifest, operation_name) } type DataFileIdentity = (Option, String); @@ -1801,10 +1766,7 @@ fn convert_to_rust_operation( ) })?, manifest, - SchemaConversionOptions { - raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, - legacy_field_id_mode: LegacyFieldIdMode::Inherit, - }, + &op_name, )?; Operation::Project { preserves_nullability: env @@ -1916,7 +1878,7 @@ fn convert_to_rust_operation( } } "Overwrite" => { - let mut fragments: Vec = + let fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; @@ -1935,10 +1897,7 @@ fn convert_to_rust_operation( base.extract_object(env) }) })?; - let ConvertedSchema { - schema, - field_id_remap, - } = convert_schema_from_operation( + let ConvertedSchema { schema, .. } = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1947,12 +1906,8 @@ fn convert_to_rust_operation( ) })?, manifest, - SchemaConversionOptions { - raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, - legacy_field_id_mode: LegacyFieldIdMode::Standalone, - }, + &op_name, )?; - remap_fragment_field_ids(&mut fragments, &field_id_remap, &HashSet::new()); Operation::Overwrite { fragments, schema, @@ -2136,10 +2091,7 @@ fn convert_to_rust_operation( ) })?, manifest, - SchemaConversionOptions { - raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, - legacy_field_id_mode: LegacyFieldIdMode::Inherit, - }, + &op_name, )?; let retained_files = if field_id_remap.is_empty() { HashSet::new() @@ -2549,18 +2501,18 @@ mod tests { let arrow_schema = ArrowSchema::new(vec![ArrowField::new("a", ArrowDataType::Int32, false)]); + let mut manifest = Manifest::new( + base_schema, + Arc::default(), + Default::default(), + HashMap::new(), + ); + manifest.activate_stable_field_ids(); + let ConvertedSchema { schema, field_id_remap, - } = convert_arrow_schema( - &arrow_schema, - Some((base_schema, 0, true)), - SchemaConversionOptions { - raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, - legacy_field_id_mode: LegacyFieldIdMode::Inherit, - }, - ) - .unwrap(); + } = convert_arrow_schema(&arrow_schema, Some(&manifest), "Project").unwrap(); assert!(schema.metadata.is_empty()); assert_eq!(schema.field("a").unwrap().id, -1); @@ -2582,15 +2534,14 @@ mod tests { ])), ]) .with_metadata(metadata.clone()); - let converted = convert_arrow_schema( - &arrow_schema, - Some((base_schema, 0, true)), - SchemaConversionOptions { - raw_field_id_mode: RawArrowFieldIdMode::ExplicitOnly, - legacy_field_id_mode: LegacyFieldIdMode::Inherit, - }, - ) - .unwrap(); + let mut manifest = Manifest::new( + base_schema, + Arc::default(), + Default::default(), + HashMap::new(), + ); + manifest.activate_stable_field_ids(); + let converted = convert_arrow_schema(&arrow_schema, Some(&manifest), "Project").unwrap(); assert_eq!(converted.schema.field("renamed").unwrap().id, 0); assert_eq!(converted.schema.metadata, metadata); assert!(converted.field_id_remap.is_empty()); @@ -2606,18 +2557,17 @@ mod tests { }; let arrow_schema = ArrowSchema::new(vec![ArrowField::new("a", ArrowDataType::Utf8, false)]); + let manifest = Manifest::new( + base_schema, + Arc::default(), + Default::default(), + HashMap::new(), + ); + let ConvertedSchema { schema, field_id_remap, - } = convert_arrow_schema( - &arrow_schema, - Some((base_schema, 0, false)), - SchemaConversionOptions { - raw_field_id_mode: RawArrowFieldIdMode::AssignMissing, - legacy_field_id_mode: LegacyFieldIdMode::Standalone, - }, - ) - .unwrap(); + } = convert_arrow_schema(&arrow_schema, Some(&manifest), "Overwrite").unwrap(); assert_eq!(schema.field("a").unwrap().data_type(), ArrowDataType::Utf8); assert!(schema.metadata.is_empty()); diff --git a/java/src/test/java/org/lance/operation/MergeTest.java b/java/src/test/java/org/lance/operation/MergeTest.java index 1fb210b987e..d57a08ca495 100644 --- a/java/src/test/java/org/lance/operation/MergeTest.java +++ b/java/src/test/java/org/lance/operation/MergeTest.java @@ -38,18 +38,14 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; -import java.io.IOException; import java.nio.charset.StandardCharsets; -import java.nio.file.Files; import java.nio.file.Path; -import java.nio.file.StandardCopyOption; import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; -import java.util.stream.Stream; public class MergeTest extends OperationTestBase { @@ -405,19 +401,6 @@ void testLegacyMergeInheritsNonContiguousFieldIds(@TempDir Path tempDir) throws } } - private void copyDirectory(Path source, Path target) throws IOException { - try (Stream paths = Files.walk(source)) { - for (Path path : (Iterable) paths::iterator) { - Path destination = target.resolve(source.relativize(path)); - if (Files.isDirectory(path)) { - Files.createDirectories(destination); - } else { - Files.copy(path, destination, StandardCopyOption.REPLACE_EXISTING); - } - } - } - } - private Map fieldMeta(int fieldId) { Map idMeta = new HashMap<>(); idMeta.put("lance:field_id", String.valueOf(fieldId)); diff --git a/java/src/test/java/org/lance/operation/OperationTestBase.java b/java/src/test/java/org/lance/operation/OperationTestBase.java index df6ddb2830d..1d14a3962e9 100644 --- a/java/src/test/java/org/lance/operation/OperationTestBase.java +++ b/java/src/test/java/org/lance/operation/OperationTestBase.java @@ -28,8 +28,13 @@ import org.junit.jupiter.api.TestInstance; import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; import java.util.Collections; import java.util.UUID; +import java.util.stream.Stream; @TestInstance(TestInstance.Lifecycle.PER_CLASS) public class OperationTestBase { @@ -49,6 +54,19 @@ void tearDown() { } } + protected void copyDirectory(Path source, Path target) throws IOException { + try (Stream paths = Files.walk(source)) { + for (Path path : (Iterable) paths::iterator) { + Path destination = target.resolve(source.relativize(path)); + if (Files.isDirectory(path)) { + Files.createDirectories(destination); + } else { + Files.copy(path, destination, StandardCopyOption.REPLACE_EXISTING); + } + } + } + } + /** Helper method to append simple data to a dataset. */ protected Dataset createAndAppendRows(TestUtils.SimpleTestDataset suite, int rowCount) { dataset = suite.createEmptyDataset(); diff --git a/java/src/test/java/org/lance/operation/ProjectTest.java b/java/src/test/java/org/lance/operation/ProjectTest.java index aaf59fe88a7..c28503f4503 100644 --- a/java/src/test/java/org/lance/operation/ProjectTest.java +++ b/java/src/test/java/org/lance/operation/ProjectTest.java @@ -29,16 +29,12 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; -import java.io.IOException; -import java.nio.file.Files; import java.nio.file.Path; -import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; @@ -197,19 +193,6 @@ void testLegacyProjectPreservesNonContiguousFieldIds(@TempDir Path tempDir) thro } } - private void copyDirectory(Path source, Path target) throws IOException { - try (Stream paths = Files.walk(source)) { - for (Path path : (Iterable) paths::iterator) { - Path destination = target.resolve(source.relativize(path)); - if (Files.isDirectory(path)) { - Files.createDirectories(destination); - } else { - Files.copy(path, destination, StandardCopyOption.REPLACE_EXISTING); - } - } - } - } - private LanceField findField(List fields, String fieldName) { return fields.stream() .filter(field -> field.getName().equals(fieldName)) diff --git a/python/python/tests/compat/compat_decorator.py b/python/python/tests/compat/compat_decorator.py index 1ebdba59a6a..3b6c039cea3 100644 --- a/python/python/tests/compat/compat_decorator.py +++ b/python/python/tests/compat/compat_decorator.py @@ -201,10 +201,6 @@ def compat_env(self, version: str, method_name: str) -> Dict[str, str]: return {} -class DatasetUpgradeDowngradeTest(UpgradeDowngradeTest): - """Compatibility contract for datasets.""" - - @contextmanager def _temporary_env(overrides: Optional[Dict[str, str]]): if not overrides: @@ -369,7 +365,7 @@ def _make_test_function(cls, param_names, test_type): if test_type == "downgrade": func_body = f''' def test_func({sig_params}): - """Test the old-version contract for data written by the current version.""" + """Test that old Lance version can read data written by current version.""" from pathlib import Path obj = cls(tmp_path / "data.lance", {init_params}) obj.compat_version = version diff --git a/python/python/tests/compat/test_file_formats.py b/python/python/tests/compat/test_file_formats.py index 679968019d3..98967ccc561 100644 --- a/python/python/tests/compat/test_file_formats.py +++ b/python/python/tests/compat/test_file_formats.py @@ -15,7 +15,6 @@ from lance.file import LanceFileReader, LanceFileWriter from .compat_decorator import ( - DatasetUpgradeDowngradeTest, UpgradeDowngradeTest, compat_test, ) @@ -131,7 +130,7 @@ def check_write(self): @compat_test(min_version="0.16.0") -class BasicTypesLegacy(DatasetUpgradeDowngradeTest): +class BasicTypesLegacy(UpgradeDowngradeTest): """Test legacy data storage version 0.1 compatibility.""" def __init__(self, path: Path): diff --git a/python/python/tests/compat/test_scalar_indices.py b/python/python/tests/compat/test_scalar_indices.py index b9fcb9def27..850c1212532 100644 --- a/python/python/tests/compat/test_scalar_indices.py +++ b/python/python/tests/compat/test_scalar_indices.py @@ -16,12 +16,15 @@ import lance import pyarrow as pa -from .compat_decorator import DatasetUpgradeDowngradeTest, compat_test +from .compat_decorator import ( + UpgradeDowngradeTest, + compat_test, +) from .util import safe_data_storage_version @compat_test(min_version="0.30.0") -class BTreeIndex(DatasetUpgradeDowngradeTest): +class BTreeIndex(UpgradeDowngradeTest): """Test BTREE scalar index compatibility (introduced in 0.20.0). Started fully working in 0.30.0 with various fixes. @@ -77,7 +80,7 @@ def check_write(self): @compat_test(min_version="0.22.0") -class BitmapLabelListIndex(DatasetUpgradeDowngradeTest): +class BitmapLabelListIndex(UpgradeDowngradeTest): """Test BITMAP and LABEL_LIST scalar index compatibility (introduced in 0.20.0). Started fully working in 0.22.0 with fixes to LABEL_LIST index. @@ -135,7 +138,7 @@ def check_write(self): @compat_test(min_version="0.36.0") -class NgramIndex(DatasetUpgradeDowngradeTest): +class NgramIndex(UpgradeDowngradeTest): """Test NGRAM index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -184,7 +187,7 @@ def check_write(self): @compat_test(min_version="0.36.0") -class ZonemapBloomfilterIndex(DatasetUpgradeDowngradeTest): +class ZonemapBloomfilterIndex(UpgradeDowngradeTest): """Test ZONEMAP and BLOOMFILTER index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -261,11 +264,11 @@ def check_write(self): def skip_downgrade(self, version: str) -> bool: # In 0.X the zonemap index did not properly handle NULL in filters - return super().skip_downgrade(version) or version.startswith("0.") + return version.startswith("0.") @compat_test(min_version="0.36.0") -class JsonIndex(DatasetUpgradeDowngradeTest): +class JsonIndex(UpgradeDowngradeTest): """Test JSON index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -321,7 +324,7 @@ def check_write(self): @compat_test(min_version="0.36.0") -class FtsIndex(DatasetUpgradeDowngradeTest): +class FtsIndex(UpgradeDowngradeTest): """Test FTS (full-text search) index compatibility (introduced in 0.36.0).""" def __init__(self, path: Path): @@ -376,7 +379,7 @@ def check_write(self): ds.optimize.compact_files() def skip_downgrade(self, version: str) -> bool: - return super().skip_downgrade(version) or version.startswith("0.") + return version.startswith("0.") def current_env(self, method_name: str) -> dict[str, str]: if method_name == "create": diff --git a/python/python/tests/compat/test_vector_indices.py b/python/python/tests/compat/test_vector_indices.py index 0b7239ec5fb..e97d6be8bf6 100644 --- a/python/python/tests/compat/test_vector_indices.py +++ b/python/python/tests/compat/test_vector_indices.py @@ -17,12 +17,15 @@ import pyarrow as pa import pyarrow.compute as pc -from .compat_decorator import DatasetUpgradeDowngradeTest, compat_test +from .compat_decorator import ( + UpgradeDowngradeTest, + compat_test, +) from .util import safe_data_storage_version @compat_test(min_version="0.29.1.beta2") -class PqVectorIndex(DatasetUpgradeDowngradeTest): +class PqVectorIndex(UpgradeDowngradeTest): """Test PQ (Product Quantization) vector index compatibility.""" def __init__(self, path: Path): @@ -98,7 +101,7 @@ def check_write(self): @compat_test(min_version="0.39.0") -class HnswPqVectorIndex(DatasetUpgradeDowngradeTest): +class HnswPqVectorIndex(UpgradeDowngradeTest): """Test IVF_HNSW_PQ vector index compatibility. Note: Only tests versions >= 0.39.0 because earlier versions don't support @@ -179,7 +182,7 @@ def check_write(self): @compat_test(min_version="0.39.0") -class HnswSqVectorIndex(DatasetUpgradeDowngradeTest): +class HnswSqVectorIndex(UpgradeDowngradeTest): """Test IVF_HNSW_SQ vector index compatibility. Note: Only tests versions >= 0.39.0 because earlier versions don't support @@ -260,7 +263,7 @@ def check_write(self): @compat_test(min_version="4.0.0-beta.8") -class IvfRqVectorIndex(DatasetUpgradeDowngradeTest): +class IvfRqVectorIndex(UpgradeDowngradeTest): """Test IVF_RQ vector index compatibility. V2 was introduced in v4.0.0-beta.8""" def __init__(self, path: Path): diff --git a/rust/lance-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index ad03219ab5b..39325ba67de 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -60,9 +60,7 @@ use lance_table::format::{Fragment, IndexMetadata, Manifest}; use lance_table::io::commit::{ CommitError, CommitHandler, commit_handler_from_url, write_manifest_file_to_path, }; -use lance_table::transaction::{ - canonicalize_stable_field_ids, validate_stable_field_id_transition, -}; +use lance_table::transaction::validate_stable_field_id_transition; use object_store::{Error as ObjectStoreError, path::Path}; use roaring::RoaringBitmap; use std::io::Cursor; @@ -1984,7 +1982,7 @@ impl ManifestNamespace { ..WriteParams::default() }; - let mut transaction = match InsertBuilder::new(dataset.clone()) + let transaction = match InsertBuilder::new(dataset.clone()) .with_params(&write_params) .execute_uncommitted_stream(output_stream) .await @@ -2000,7 +1998,10 @@ impl ManifestNamespace { let (mutation, index_data) = Self::take_manifest_rewrite_result(&shared)?; - let Operation::Overwrite { fragments, .. } = &transaction.operation else { + let Operation::Overwrite { + fragments, schema, .. + } = &transaction.operation + else { return Err(NamespaceError::Internal { message: "Manifest rewrite transaction is not an overwrite".to_string(), } @@ -2021,26 +2022,6 @@ impl ManifestNamespace { return Ok(mutation.result); } - if let Err(err) = canonicalize_stable_field_ids( - Some(dataset.manifest()), - &mut transaction.operation, - None, - ) { - self.cleanup_staged_manifest_files(&object_store, &staged_data_files, &[]) - .await; - return Err(err); - } - - let Operation::Overwrite { - fragments, schema, .. - } = &transaction.operation - else { - return Err(NamespaceError::Internal { - message: "Manifest rewrite transaction is not an overwrite".to_string(), - } - .into()); - }; - let mut manifest = Self::manifest_from_overwrite_transaction( dataset.manifest(), schema.clone(), diff --git a/rust/lance-table/src/transaction.rs b/rust/lance-table/src/transaction.rs index ffda5e93649..1fc095cfed9 100644 --- a/rust/lance-table/src/transaction.rs +++ b/rust/lance-table/src/transaction.rs @@ -50,7 +50,7 @@ pub use update_map::{ }; pub use validate::{ canonicalize_stable_field_ids, resolve_arrow_field_ids, validate_detached_stable_field_ids, - validate_operation, validate_stable_field_id_manifest, validate_stable_field_id_transition, + validate_operation, validate_stable_field_id_transition, }; use crate::format::{IndexMetadata, Manifest}; diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index 249cf15d393..02b2a4a63ba 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -35,9 +35,7 @@ use crate::system_index::mem_wal::{ use crate::transaction::UpdateMode::{RewriteColumns, RewriteRows}; use crate::transaction::row_version::resolve_update_version_metadata; use crate::transaction::update_map::apply_update_map; -use crate::transaction::validate::{ - merge_fragment_physically_rewritten, validate_stable_field_id_manifest, -}; +use crate::transaction::validate::merge_fragment_physically_rewritten; use crate::transaction::{ CoverageIdentity, DataReplacementGroup, LogicalIndexSegments, Operation, ReadVersionState, RewriteGroup, Transaction, UpdatedFragmentOffsets, @@ -1677,7 +1675,6 @@ impl Transaction { } validate_stable_field_id_flags(&manifest)?; - validate_stable_field_id_manifest(&manifest)?; Ok((manifest, final_indices)) } @@ -1736,31 +1733,6 @@ mod tests { ) } - #[test] - fn new_dataset_preserves_legacy_field_id_allocation_by_default() { - let arrow_schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); - let mut schema = LanceSchema::try_from(&arrow_schema).unwrap(); - schema.try_set_field_id(None).unwrap(); - let transaction = Transaction::new( - 0, - Operation::Overwrite { - fragments: vec![], - schema, - config_upsert_values: None, - initial_bases: None, - }, - None, - ); - - let (manifest, _) = transaction - .build_manifest(None, vec![], "txn", &default_build_config()) - .unwrap(); - - assert_eq!(manifest.max_allocated_field_id, None); - assert_eq!(manifest.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); - assert_eq!(manifest.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); - } - #[test] fn activation_sets_writer_gate_when_auto_flags_are_disabled() { let manifest = sample_manifest(); diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 25c8e0671c9..15587a4af33 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -17,13 +17,6 @@ use std::collections::{HashMap, HashSet}; type DataFileIdentity = (Option, String); -#[derive(Default)] -struct FieldIdRemap { - canonical_ids: HashSet, - explicit_ids: HashMap, - raw_source_ids: HashMap, -} - /// Resolve an Arrow-derived operation against the dataset it was read from. /// /// Overwrite and Merge match fields by name and type, not positional Arrow IDs. @@ -62,7 +55,7 @@ pub fn resolve_arrow_field_ids( schema, fragments, .. } => { let field_id_remap = canonicalize_schema(manifest, schema, None, true)?; - remap_fragment_field_ids(fragments, &field_id_remap, &HashSet::new())?; + resolve_fragment_field_ids(fragments, &field_id_remap, &HashSet::new(), schema)?; } Operation::Project { schema, .. } => { let Some(manifest) = manifest else { @@ -83,7 +76,7 @@ pub fn resolve_arrow_field_ids( .map(|file| (file.base_id, file.path.clone())) .collect(); let field_id_remap = canonicalize_schema(Some(manifest), schema, None, true)?; - remap_fragment_field_ids(fragments, &field_id_remap, &retained_files)?; + resolve_fragment_field_ids(fragments, &field_id_remap, &retained_files, schema)?; } _ => {} } @@ -119,7 +112,7 @@ pub fn canonicalize_stable_field_ids( } => { let identity_schema = read_schema.or_else(|| manifest.map(|manifest| &manifest.schema)); let remap = canonicalize_schema(manifest, schema, identity_schema, false)?; - remap_fragment_field_ids(fragments, &remap, &HashSet::new())?; + remap_fragment_field_ids(fragments, &remap, &HashSet::new()); } Operation::Merge { schema, fragments, .. @@ -175,16 +168,13 @@ fn canonicalize_merge_replacements( schema.validate()?; schema.verify_primary_key()?; - let mut field_id_remap = FieldIdRemap::default(); - for (original, canonical) in original.fields_pre_order().zip(schema.fields_pre_order()) { - field_id_remap.canonical_ids.insert(canonical.id); - if original.id >= 0 { - field_id_remap - .explicit_ids - .insert(original.id, canonical.id); - } - } - remap_fragment_field_ids(fragments, &field_id_remap, retained_files)?; + let field_id_remap = original + .fields_pre_order() + .zip(schema.fields_pre_order()) + .filter(|(original, _)| original.id >= 0) + .map(|(original, canonical)| (original.id, canonical.id)) + .collect(); + remap_fragment_field_ids(fragments, &field_id_remap, retained_files); Ok(()) } @@ -225,15 +215,11 @@ fn canonicalize_schema( schema: &mut Schema, identity_schema: Option<&Schema>, remap_raw_source_ids: bool, -) -> Result { - let original = schema.clone(); - let raw_source = if remap_raw_source_ids { - let mut raw_source = original.clone(); - raw_source.try_set_field_id(None)?; - Some(raw_source) - } else { - None - }; +) -> Result> { + let mut original = schema.clone(); + if remap_raw_source_ids { + original.try_set_field_id(None)?; + } let max_existing_id = manifest.map(Manifest::max_field_id); if manifest.is_none() { @@ -247,24 +233,12 @@ fn canonicalize_schema( schema.validate()?; schema.verify_primary_key()?; - let mut field_id_remap = FieldIdRemap::default(); - for (original, canonical) in original.fields_pre_order().zip(schema.fields_pre_order()) { - field_id_remap.canonical_ids.insert(canonical.id); - if original.id >= 0 { - field_id_remap - .explicit_ids - .insert(original.id, canonical.id); - } - } - if let Some(raw_source) = raw_source { - field_id_remap.raw_source_ids.extend( - raw_source - .fields_pre_order() - .zip(schema.fields_pre_order()) - .map(|(source, canonical)| (source.id, canonical.id)), - ); - } - Ok(field_id_remap) + Ok(original + .fields_pre_order() + .zip(schema.fields_pre_order()) + .filter(|(original, _)| original.id >= 0) + .map(|(original, canonical)| (original.id, canonical.id)) + .collect()) } fn canonicalize_field( @@ -325,66 +299,78 @@ fn clear_field_ids(field: &mut Field) { } } -fn remap_fragment_field_ids( +fn resolve_fragment_field_ids( fragments: &mut [Fragment], - field_id_remap: &FieldIdRemap, + field_id_remap: &HashMap, retained_files: &HashSet, + schema: &Schema, ) -> Result<()> { + let canonical_ids = schema.field_ids(); for fragment in fragments { - let source_ids = if field_id_remap.raw_source_ids.is_empty() { - &field_id_remap.explicit_ids - } else { - // Raw Arrow overwrite fragments may have been written either by a - // standalone writer using the source schema IDs or by a - // dataset-aware writer using canonical IDs. Resolve that namespace - // once for the whole fragment so split files cannot disagree. If - // both interpretations are possible and produce different - // identities then there is no safe mapping without provenance. - let fragment_field_ids = fragment - .referenced_lance_files() - .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone()))) - .flat_map(|file| file.fields.iter().copied()) - .filter(|field_id| *field_id >= 0) - .collect::>(); - let canonical_source = fragment_field_ids - .iter() - .all(|field_id| field_id_remap.canonical_ids.contains(field_id)); - let raw_source = fragment_field_ids - .iter() - .all(|field_id| field_id_remap.raw_source_ids.contains_key(field_id)); - let raw_changes_identity = fragment_field_ids - .iter() - .any(|field_id| field_id_remap.raw_source_ids.get(field_id) != Some(field_id)); - - match (canonical_source, raw_source, raw_changes_identity) { - (true, true, true) => { - return Err(Error::invalid_input(format!( - "Fragment {} has ambiguous raw Arrow field IDs; its file mappings can be interpreted as either source or canonical identities", - fragment.id - ))); - } - (true, _, _) => continue, - (false, true, _) => &field_id_remap.raw_source_ids, - (false, false, _) => { - return Err(Error::invalid_input(format!( - "Fragment {} field IDs do not match either the raw Arrow source schema or the canonical replacement schema", - fragment.id - ))); - } + // Raw Arrow overwrite fragments may have been written either by a + // standalone writer using the source schema IDs or by a + // dataset-aware writer using canonical IDs. Resolve that namespace + // once for the whole fragment so split files cannot disagree. If + // both interpretations are possible and produce different + // identities then there is no safe mapping without provenance. + let fragment_field_ids = fragment + .referenced_lance_files() + .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone()))) + .flat_map(|file| file.fields.iter().copied()) + .filter(|field_id| *field_id >= 0) + .collect::>(); + let canonical_source = fragment_field_ids + .iter() + .all(|field_id| canonical_ids.contains(field_id)); + let raw_source = fragment_field_ids + .iter() + .all(|field_id| field_id_remap.contains_key(field_id)); + let raw_changes_identity = fragment_field_ids + .iter() + .any(|field_id| field_id_remap.get(field_id) != Some(field_id)); + + match (canonical_source, raw_source, raw_changes_identity) { + (true, true, true) => { + return Err(Error::invalid_input(format!( + "Fragment {} has ambiguous raw Arrow field IDs; its file mappings can be interpreted as either source or canonical identities", + fragment.id + ))); } - }; + (true, _, _) => continue, + (false, true, _) => {} + (false, false, _) => { + return Err(Error::invalid_input(format!( + "Fragment {} field IDs do not match either the raw Arrow source schema or the canonical replacement schema", + fragment.id + ))); + } + } + remap_fragment_field_ids( + std::slice::from_mut(fragment), + field_id_remap, + retained_files, + ); + } + Ok(()) +} + +fn remap_fragment_field_ids( + fragments: &mut [Fragment], + field_id_remap: &HashMap, + retained_files: &HashSet, +) { + for fragment in fragments { for file in fragment.referenced_lance_files_mut() { if retained_files.contains(&(file.base_id, file.path.clone())) { continue; } for field_id in std::sync::Arc::make_mut(&mut file.fields) { - if let Some(canonical_id) = source_ids.get(field_id) { + if let Some(canonical_id) = field_id_remap.get(field_id) { *field_id = *canonical_id; } } } } - Ok(()) } /// Validate the operation is valid for the given manifest. @@ -473,7 +459,7 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> } /// Validate stable-field-ID invariants that are independent of one operation. -pub fn validate_stable_field_id_manifest(manifest: &Manifest) -> Result<()> { +fn validate_stable_field_id_manifest(manifest: &Manifest) -> Result<()> { let Some(max_allocated_field_id) = manifest.max_allocated_field_id else { return Ok(()); }; @@ -1271,16 +1257,20 @@ mod tests { assert_eq!(schema.metadata.get("source").unwrap(), "user metadata"); } - #[test] - fn canonicalize_overwrite_preserves_already_canonical_fragment_field_ids() { + #[rstest::rstest] + #[case::canonical_ids([1, 2])] + #[case::duplicate_arrow_metadata([42, 42])] + fn canonicalize_overwrite_preserves_already_canonical_fragment_field_ids( + #[case] input_ids: [i32; 2], + ) { let manifest = activated_manifest(); let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![ ArrowField::new("b", DataType::Int32, true), ArrowField::new("c", DataType::Int32, true), ])) .unwrap(); - schema.fields[0].id = 1; - schema.fields[1].id = 2; + schema.fields[0].id = input_ids[0]; + schema.fields[1].id = input_ids[1]; let mut operation = Operation::Overwrite { fragments: vec![fragment_with_file_fields(0, "new.lance", vec![1, 2])], schema, diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 4bd22c79407..8f1dd75cb02 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -742,9 +742,10 @@ pub(super) async fn alter_columns( let mut new_schema = dataset.schema().clone(); // Mapping of old to new fields that need to be casted. - let mut cast_sources: Vec = Vec::new(); + let mut cast_fields: Vec<(Field, Field)> = Vec::new(); let mut tightens_nullability = false; + let mut next_field_id = i64::from(dataset.manifest.max_field_id()) + 1; let fallback_version = dataset.manifest.data_storage_format.lance_file_format(); for alteration in alterations { @@ -796,45 +797,12 @@ pub(super) async fn alter_columns( field_dest.nullable, ); *field_dest = Field::try_from(&arrow_field)?; - // Keep the old id temporarily so the replacement can be located - // after every alteration has been applied. Fresh ids are assigned - // below in canonical schema order, independent of request order. - field_dest.id = field_src.id; - field_dest.parent_id = field_src.parent_id; + field_dest.try_set_id(field_src.parent_id, &mut next_field_id)?; - cast_sources.push(field_src.clone()); + cast_fields.push((field_src.clone(), field_dest.clone())); } } - let mut cast_fields = Vec::with_capacity(cast_sources.len()); - if !cast_sources.is_empty() { - let destination_paths = cast_sources - .iter() - .map(|source| new_schema.field_path(source.id)) - .collect::>>()?; - - for source in &cast_sources { - new_schema - .mut_field_by_id(source.id) - .expect("cast source must still identify its replacement") - .id = -1; - } - new_schema.try_set_field_id(Some(dataset.manifest.max_field_id()))?; - - cast_fields = cast_sources - .into_iter() - .zip(destination_paths) - .map(|(source, path)| { - let destination = new_schema.field(&path).ok_or_else(|| { - Error::internal(format!( - "cast replacement field '{path}' disappeared while assigning field ids" - )) - })?; - Ok((source, destination.clone())) - }) - .collect::>>()?; - } - new_schema.validate()?; new_schema.verify_primary_key()?; diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index 6a7068b8215..2de0d53dd10 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -927,18 +927,8 @@ async fn test_overwrite_preserves_compatible_stable_field_identities() { let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; dataset.migrate_to_stable_field_ids().await.unwrap(); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("id", DataType::Int64, false), - ArrowField::new("replacement", DataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int64Array::from_iter_values(0..10)), - Arc::new(Int64Array::from_iter_values(10..20)), - ], - ) - .unwrap(); + let batch = record_batch!(("id", Int64, [0, 1]), ("replacement", Int64, [10, 11])).unwrap(); + let schema = batch.schema(); let overwritten = Dataset::write( RecordBatchIterator::new(vec![Ok(batch)], schema), source_uri.as_str(), @@ -1019,26 +1009,15 @@ async fn test_stable_field_id_rename_and_nullability_preserve_identity() { #[tokio::test] async fn test_stable_field_id_multi_cast_uses_schema_order() { - let source_uri = TempStrDir::default(); - let schema = Arc::new(ArrowSchema::new(vec![ - ArrowField::new("a", DataType::Int32, false), - ArrowField::new("b", DataType::Int32, false), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int32Array::from(vec![1, 2])), - Arc::new(Int32Array::from(vec![3, 4])), - ], - ) - .unwrap(); - let mut dataset = Dataset::write( - RecordBatchIterator::new(vec![Ok(batch)], schema), - source_uri.as_str(), - None, - ) - .await - .unwrap(); + let batch = record_batch!(("a", Int32, [1, 2]), ("b", Int32, [3, 4])).unwrap(); + let mut dataset = InsertBuilder::new("memory://") + .with_params(&WriteParams { + max_rows_per_file: 1, + ..Default::default() + }) + .execute(vec![batch]) + .await + .unwrap(); dataset.migrate_to_stable_field_ids().await.unwrap(); dataset @@ -1053,6 +1032,11 @@ async fn test_stable_field_id_multi_cast_uses_schema_order() { assert_eq!(dataset.schema().field("b").unwrap().id, 3); assert_eq!(dataset.manifest.max_allocated_field_id, Some(3)); dataset.validate().await.unwrap(); + assert_eq!(dataset.manifest.fragments.len(), 2); + assert_eq!( + dataset.scan().try_into_batch().await.unwrap(), + record_batch!(("a", Int64, [1, 2]), ("b", Int64, [3, 4])).unwrap() + ); } #[tokio::test] diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 4bb5386dd16..95f73a403bb 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -45,7 +45,7 @@ use lance_table::io::commit::{ use lance_table::io::manifest::read_manifest; use lance_table::transaction::{ canonicalize_stable_field_ids, validate_detached_stable_field_ids, validate_operation, - validate_stable_field_id_manifest, validate_stable_field_id_transition, + validate_stable_field_id_transition, }; use rand::{Rng, rng}; use roaring::RoaringBitmap; @@ -1252,10 +1252,8 @@ pub(crate) async fn do_commit_detached_transaction( // diagnostics. Finalization repeats this at the manifest write boundary. fix_schema(&mut manifest)?; validate_stable_field_id_transition(&dataset.manifest, &manifest, &transaction.operation)?; - manifest.update_max_field_id(); crate::dataset::versions::check_manifest_storage_version_for_commit(&mut manifest)?; check_fragment_ids(&manifest)?; - validate_stable_field_id_manifest(&manifest)?; // Runs after the coverage derivation and can replace a fragment bitmap // while keeping its UUID, so anything it narrowed loses its position. let recovered_coverage = migrate_indices(dataset, &mut indices).await?; @@ -1628,11 +1626,9 @@ pub(crate) async fn commit_transaction( &manifest, &attempt_transaction.operation, )?; - manifest.update_max_field_id(); crate::dataset::versions::check_manifest_storage_version_for_commit(&mut manifest)?; check_fragment_ids(&manifest)?; - validate_stable_field_id_manifest(&manifest)?; // Runs after the coverage derivation and can replace a fragment bitmap // while keeping its UUID, so anything it narrowed loses its position. @@ -2043,7 +2039,7 @@ mod tests { } #[tokio::test] - async fn raw_arrow_project_retry_matches_single_attempt_ids() { + async fn raw_arrow_project_retry_preserves_identity_and_metadata() { let tmp = TempStrDir::default(); let uri = tmp.as_str(); let mut dataset = Dataset::write( @@ -2065,25 +2061,19 @@ mod tests { Operation::ReserveFragments { num_fragments: 1 }, None, ); - let handler = inject_foreign_commit_handler(foreign_manifest.clone(), &foreign_transaction); + let handler = inject_foreign_commit_handler(foreign_manifest, &foreign_transaction); let raw_schema = Schema { fields: vec![Field::new_arrow("x", DataType::Int32, false).unwrap()], metadata: HashMap::from([("input-note".to_string(), "project".to_string())]), }; - let mut expected = Operation::Project { - schema: raw_schema.clone(), + let mut expected_schema = raw_schema.clone(); + expected_schema.fields[0].id = 0; + let mut operation = Operation::Project { + schema: raw_schema, preserves_nullability: true, }; - resolve_arrow_field_ids(Some(&dataset.manifest), &mut expected).unwrap(); - let operation = expected.clone(); - let Operation::Project { - schema: expected_schema, - .. - } = expected - else { - unreachable!(); - }; + resolve_arrow_field_ids(Some(&dataset.manifest), &mut operation).unwrap(); let committed = CommitBuilder::new(Arc::new(dataset.clone())) .with_commit_handler(handler) @@ -2092,10 +2082,6 @@ mod tests { .unwrap(); assert_eq!(committed.schema(), &expected_schema); - assert_eq!( - committed.schema().metadata.get("input-note").unwrap(), - "project" - ); } #[rstest::rstest] @@ -2125,7 +2111,7 @@ mod tests { Operation::ReserveFragments { num_fragments: 1 }, None, ); - let handler = inject_foreign_commit_handler(foreign_manifest.clone(), &foreign_transaction); + let handler = inject_foreign_commit_handler(foreign_manifest, &foreign_transaction); let mut merged_fragment = dataset.manifest.fragments[0].clone(); let mut new_file = merged_fragment.files[0].clone(); @@ -2139,6 +2125,9 @@ mod tests { ], metadata: HashMap::new(), }; + let mut expected_schema = raw_schema.clone(); + expected_schema.fields[0].id = 0; + expected_schema.fields[1].id = 2; let mut operation = if overwrite { Operation::Overwrite { fragments: vec![merged_fragment], @@ -2154,13 +2143,6 @@ mod tests { } }; resolve_arrow_field_ids(Some(&dataset.manifest), &mut operation).unwrap(); - let mut expected = operation.clone(); - canonicalize_stable_field_ids( - Some(&foreign_manifest), - &mut expected, - Some(dataset.schema()), - ) - .unwrap(); let committed = CommitBuilder::new(Arc::new(dataset.clone())) .with_commit_handler(handler) @@ -2168,24 +2150,15 @@ mod tests { .await .unwrap(); - let (Operation::Merge { - schema: expected_schema, - fragments: expected_fragments, - .. - } - | Operation::Overwrite { - schema: expected_schema, - fragments: expected_fragments, - .. - }) = expected - else { - unreachable!(); - }; assert_eq!(committed.schema(), &expected_schema); - assert_eq!(committed.schema().field("new_column").unwrap().id, 2); + assert_eq!(committed.manifest.max_allocated_field_id, Some(2)); + assert_eq!( + committed.manifest.fragments[0].files[0].fields.as_ref(), + &[0] + ); assert_eq!( - committed.manifest.fragments[0].files[1].fields, - expected_fragments[0].files[1].fields + committed.manifest.fragments[0].files[1].fields.as_ref(), + &[2] ); } From 477cea23dca7ab16481f20315ed30b8808839333 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 2 Oct 2026 17:50:54 +0800 Subject: [PATCH 22/24] docs: clarify stable field ID migration requirements --- docs/src/format/table/schema.md | 5 ++- docs/src/format/table/versioning.md | 2 +- rust/lance-table/src/format/manifest.rs | 59 ++++++++++++++++++++----- rust/lance/src/dataset.rs | 6 +++ 4 files changed, 58 insertions(+), 14 deletions(-) diff --git a/docs/src/format/table/schema.md b/docs/src/format/table/schema.md index c0e0ee5d2eb..f7916cc1112 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -244,7 +244,10 @@ must honor that state. A legacy manifest sets neither value. A stable manifest s that sets only one is invalid. The reader flag for stable field IDs must remain unset because the feature does not change read behavior. -A dataset changes to stable field IDs only through an explicit migration commit. +A dataset changes to stable field IDs only through an explicit migration commit. Before activation, +operators must ensure that all clients that can write to the dataset enforce writer feature flags, +rejecting writes when they do not support a required flag. Clients that ignore these flags must no +longer write to the dataset: they may discard the high-water mark and allow field IDs to be reused. A dataset cannot return to the legacy behavior. After activation, a restore must fail if it targets a version that does not set `max_allocated_field_id`. Before activation, different fields may have diff --git a/docs/src/format/table/versioning.md b/docs/src/format/table/versioning.md index a95eeab3777..6a605a5f2ec 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -34,7 +34,7 @@ they should return an "unsupported" error on any read or write operation. | 256 | `FLAG_MIXED_DATA_FILE_VERSIONS` | Yes | Yes | The snapshot may reference recognized V2 data files with different exact versions. Both bits must be set and remain set on later versions. | | 512 | `FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS` | Yes | Yes | The table uses stable row IDs and carries a [Fragment Reuse Index](../index/system/frag_reuse.md). | | 1024 | `FLAG_FRAGMENT_REUSE_INDEX` | Yes | Yes | The fragment reuse index records tagged transitions (`IndexMetadata.index_version >= 1`). Readers must translate row addresses through them; writers must preserve them. An implementation without this flag would decode the details as the legacy format and silently drop the transitions when it next rewrites the fragment reuse index. See [FRI index versions](../index/system/frag_reuse.md#fri-index-versions). | -| 2048 | `FLAG_STABLE_FIELD_IDS` | No | Yes | The manifest sets `max_allocated_field_id`, and a writer must assign new field IDs above it. See [Field IDs](schema.md#field-ids). | +| 2048 | `FLAG_STABLE_FIELD_IDS` | No | Yes | The manifest sets `max_allocated_field_id`, and a writer must assign new field IDs above it. Before activation, all clients that can write to the dataset must enforce writer feature flags. See [Field IDs](schema.md#field-ids). | diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index ac971176955..9a2f25c0c01 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -1212,7 +1212,7 @@ impl SelfDescribingFileReader for V1FileReader { #[cfg(test)] mod tests { use crate::feature_flags::{FLAG_STABLE_FIELD_IDS, FLAG_USE_V2_FORMAT_DEPRECATED}; - use crate::format::overlay::{DataOverlayFile, OverlayCoverage}; + use crate::format::overlay::{DataOverlayFile, OverlayCoverage, TOMBSTONE_FIELD_ID}; use crate::format::{DataFile, DeletionFile, DeletionFileType}; use std::num::NonZero; @@ -1712,30 +1712,65 @@ mod tests { #[test] fn stable_field_id_high_water_mark_survives_dropped_references_and_round_trip() { - let arrow_schema = ArrowSchema::new(vec![ArrowField::new( - "a", - arrow_schema::DataType::Int64, - false, - )]); + let arrow_schema = ArrowSchema::new(vec![ + ArrowField::new("a", arrow_schema::DataType::Int64, false), + ArrowField::new("b", arrow_schema::DataType::Int64, false), + ]); let schema = Schema::try_from(&arrow_schema).unwrap(); + let mut fragment = Fragment::new(0); + fragment.files.push(DataFile::new( + "ab.lance", + vec![0, 1], + vec![0, 1], + ConcreteFileVersion::V2_0, + None, + None, + )); let mut manifest = Manifest::new( schema, - Arc::new(vec![]), + Arc::new(vec![fragment]), DataStorageFormat::default(), HashMap::new(), ); manifest.activate_stable_field_ids(); - manifest.max_allocated_field_id = Some(43); manifest.writer_feature_flags |= FLAG_STABLE_FIELD_IDS; + assert_eq!(manifest.max_allocated_field_id, Some(1)); + + manifest.schema.fields.pop(); + let file = &mut Arc::make_mut(&mut manifest.fragments)[0].files[0]; + Arc::make_mut(&mut file.fields)[1] = TOMBSTONE_FIELD_ID; + manifest.update_max_field_id(); assert_eq!(manifest.max_referenced_field_id(), 0); - assert_eq!(manifest.max_field_id(), 43); + assert_eq!(manifest.max_field_id(), 1); - let recovered = Manifest::try_from(pb::Manifest::from(&manifest)).unwrap(); - assert_eq!(recovered.max_allocated_field_id, Some(43)); - assert_eq!(recovered.max_field_id(), 43); + let encoded = manifest.serialized(); + let mut recovered = + Manifest::try_from(pb::Manifest::decode(encoded.as_slice()).unwrap()).unwrap(); + assert_eq!(recovered.max_allocated_field_id, Some(1)); + assert_eq!(recovered.max_field_id(), 1); + assert_eq!( + recovered.fragments[0].files[0].fields.as_ref(), + &[0, TOMBSTONE_FIELD_ID] + ); + assert_eq!( + recovered.fragments[0].files[0].column_indices.as_ref(), + &[0, 1] + ); assert_eq!(recovered.reader_feature_flags & FLAG_STABLE_FIELD_IDS, 0); assert_ne!(recovered.writer_feature_flags & FLAG_STABLE_FIELD_IDS, 0); + + recovered.schema.fields.push( + Field::try_from(ArrowField::new("c", arrow_schema::DataType::Int64, false)).unwrap(), + ); + let max_field_id = recovered.max_field_id(); + recovered + .schema + .try_set_field_id(Some(max_field_id)) + .unwrap(); + assert_eq!(recovered.schema.field("c").unwrap().id, 2); + recovered.update_max_field_id(); + assert_eq!(recovered.max_allocated_field_id, Some(2)); } #[test] diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index f5be0401ea5..7e4a953b078 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -3267,6 +3267,12 @@ impl Dataset { /// after fields and their files are dropped. Activation is one-way and /// idempotent. /// + /// Before calling this method, ensure that all clients that can write to the + /// dataset enforce writer feature flags, rejecting writes when they do not + /// support a required flag. Clients that ignore these flags must no longer + /// write to the dataset: they may discard the high-water mark and allow + /// field IDs to be reused. This method cannot enforce that client policy. + /// /// ``` /// # use lance::{Dataset, Result}; /// # async fn activate(dataset: &mut Dataset) -> Result<()> { From e35633e661525e802ee78dc141044efb7b6aab1a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 2 Oct 2026 18:19:17 +0800 Subject: [PATCH 23/24] fix: normalize new dataset IDs before blob promotion --- rust/lance/src/dataset/versions/mod.rs | 6 +++++- rust/lance/src/dataset/write.rs | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/rust/lance/src/dataset/versions/mod.rs b/rust/lance/src/dataset/versions/mod.rs index ce57af0e13e..d831432abc1 100644 --- a/rust/lance/src/dataset/versions/mod.rs +++ b/rust/lance/src/dataset/versions/mod.rs @@ -151,7 +151,11 @@ pub async fn write_fragments( // set them aside before the schema is checked against the dataset's and // put them back on the schema that is written. This has to come before // the blob promotion, which gives every negative field id a new one. - let (normalized_schema, lineage_fields) = split_row_lineage_fields(normalized_schema)?; + let (mut normalized_schema, lineage_fields) = split_row_lineage_fields(normalized_schema)?; + if dataset.is_none() { + // Input Arrow IDs must not seed allocation for a new dataset's blob children. + normalized_schema.try_reassign_field_ids(None)?; + } let normalized_schema = match version { ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 => { write::promote_legacy_blob_schema(&normalized_schema)? diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index b11dbc907d8..3db012aa6ef 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -1840,7 +1840,7 @@ pub(super) fn promote_legacy_blob_schema(schema: &Schema) -> Result { for field in &mut schema.fields { field.promote_blob_v2()?; } - schema.set_field_id(schema.max_field_id()); + schema.try_set_field_id(schema.max_field_id())?; Ok(schema) } From 2620abec34ea99a8c7d5f6079f12aa1956b5b9ac Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 7 Oct 2026 12:59:15 +0800 Subject: [PATCH 24/24] test: stabilize multivector recall and bound compaction futures --- .../src/dataset/optimize/tests/binary_copy.rs | 2 +- rust/lance/src/index/vector/ivf/v2.rs | 23 +++++++++++++++---- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/rust/lance/src/dataset/optimize/tests/binary_copy.rs b/rust/lance/src/dataset/optimize/tests/binary_copy.rs index 9dba0de02f0..e54198d81d3 100644 --- a/rust/lance/src/dataset/optimize/tests/binary_copy.rs +++ b/rust/lance/src/dataset/optimize/tests/binary_copy.rs @@ -14,7 +14,7 @@ const NON_LEGACY_VERSIONS: [LanceFileVersion; 4] = [ #[tokio::test] async fn test_binary_copy_merge_small_files() { for version in NON_LEGACY_VERSIONS { - do_test_binary_copy_merge_small_files(version).await; + Box::pin(do_test_binary_copy_merge_small_files(version)).await; } } diff --git a/rust/lance/src/index/vector/ivf/v2.rs b/rust/lance/src/index/vector/ivf/v2.rs index a791540f81c..91d5e7beacf 100644 --- a/rust/lance/src/index/vector/ivf/v2.rs +++ b/rust/lance/src/index/vector/ivf/v2.rs @@ -3144,6 +3144,7 @@ mod tests { use uuid::Uuid; const NUM_ROWS: usize = 512; + const MULTIVEC_VECTORS_PER_ROW: usize = 3; const DIM: usize = 32; // 8-bit PQ needs at least 256 training vectors; 320 leaves a stable margin // while 20 neighbors provide a useful recall oracle. @@ -3617,13 +3618,12 @@ mod tests { where T::Native: SampleUniform, { - const VECTOR_NUM_PER_ROW: usize = 3; let start_id = start_id.unwrap_or(0); let ids = Arc::new(UInt64Array::from_iter_values( start_id..start_id + num_rows as u64, )); let total_floats = match is_multivector { - true => num_rows * VECTOR_NUM_PER_ROW * DIM, + true => num_rows * MULTIVEC_VECTORS_PER_ROW * DIM, false => num_rows * DIM, }; let vectors = generate_random_array_with_range::(total_floats, range); @@ -3647,7 +3647,7 @@ mod tests { )); let array = Arc::new(ListArray::new( vector_field, - OffsetBuffer::from_lengths(std::iter::repeat_n(VECTOR_NUM_PER_ROW, num_rows)), + OffsetBuffer::from_lengths(std::iter::repeat_n(MULTIVEC_VECTORS_PER_ROW, num_rows)), Arc::new(fsl), None, )); @@ -6497,7 +6497,22 @@ mod tests { #[tokio::test] async fn test_legacy_ivf_pq_cosine_multivec_smoke() { - let params = pq_matrix_params(1, DistanceType::Cosine, IndexFileVersion::Legacy); + // Train on all vectors: the smaller single-vector matrix's sample budget + // makes this fixture's recall depend on which rows training samples. + let num_vectors = NUM_ROWS * MULTIVEC_VECTORS_PER_ROW; + let mut ivf_params = IvfBuildParams::new(1); + ivf_params.max_iters = 2; + ivf_params.sample_rate = num_vectors; + let pq_params = PQBuildParams { + num_sub_vectors: 4, + num_bits: 8, + max_iters: 2, + sample_rate: num_vectors.div_ceil(1 << 8), + ..Default::default() + }; + let mut params = + VectorIndexParams::with_ivf_pq_params(DistanceType::Cosine, ivf_params, pq_params); + params.version(IndexFileVersion::Legacy); test_index_multivec_impl::(params, 1, 0.5, 0.0..1.0).await; }