diff --git a/docs/src/format/table/index.md b/docs/src/format/table/index.md index a20d4c564c2..548449f6061 100644 --- a/docs/src/format/table/index.md +++ b/docs/src/format/table/index.md @@ -28,6 +28,20 @@ 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. If a manifest sets it, the dataset uses non-reusable 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 must not assign a previously allocated field ID to another field. +- 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 15e8946c708..96e482f72e1 100644 --- a/docs/src/format/table/schema.md +++ b/docs/src/format/table/schema.md @@ -221,15 +221,61 @@ 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. + +`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 non-reusable 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 non-reusable 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_NON_REUSABLE_FIELD_IDS` tells writers that they +must honor that state. A legacy manifest sets neither value. An activated manifest sets both. A manifest +that sets only one is invalid. The reader flag for non-reusable field IDs must remain unset because the +feature does not change read behavior. + +A dataset changes to non-reusable 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 +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 -- **Immutable**: Once assigned, a field's ID never changes -- **Unique**: Each field within a table has a unique ID -- **Stable**: IDs are preserved across schema evolution operations -- **Sparse**: Field IDs may not form a contiguous sequence after schema evolution +- **Preserved**: 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: + +- **Not reassigned**: After activation, a writer must not assign a field's ID to another field, + even after the original field is dropped or replaced. +- **Increasing**: Every new ID is greater than the activation high-water mark and every ID assigned + after activation. + +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 @@ -296,13 +342,71 @@ 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 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 -- **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**: 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**: When `max_allocated_field_id` is set, replace all fields and assign every field, + including nested fields, a new ID above the previous high-water mark. This applies even when + names and types are unchanged. References to old field IDs do not identify the replacement fields The use of field IDs ensures that data files can be correctly interpreted even as the schema changes over time. +### Blob Identity Namespace + +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 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: + +```python +pa.schema([pa.field("image", pa.large_binary())]) +``` + +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 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 f35afe88b5b..540a0ac9d1d 100644 --- a/docs/src/format/table/versioning.md +++ b/docs/src/format/table/versioning.md @@ -37,7 +37,9 @@ they should return an "unsupported" error on any read or write operation. | 2048 | `FLAG_UNSTABLE_SPILLED_ROW_LINEAGE` | Yes | Yes | Some fragment stores its row ids or row version sequences as hidden columns of a data file rather than inline. A reader without this flag would see the fragment as having no row ids. Unstable: release builds reject it unless explicitly opted in. | | 4096 | `FLAG_FRAGMENT_TREE` | Yes | Yes | Fragment records live in a [fragment tree](fragment_metadata.md). `Manifest.fragments` is empty. | | 8192 | `FLAG_INDEPENDENT_COVERING_FIELDS` | Yes | Yes | Requires `FLAG_COVERED_INDEX_METADATA` and is retained together with it. `IndexMetadata.fields` contains only key fields, while `covering_fields` independently declares carried fields and may overlap `fields`. Implementations that only support the legacy suffix contract must reject the dataset. | +| 16384 | `FLAG_MANAGED_BLOBS` | Yes | Yes | Blob v2 descriptors may independently address Lance-owned objects. Both flags remain set across subsequent commits and restore. | +| 32768 | `FLAG_NON_REUSABLE_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). | -Flags with bit values 16384 and above are unknown; unknown flags cause implementations to reject the dataset with an "unsupported" error. The paired mixed-version reader and writer bits must either both be set or both be clear; a half-set manifest is invalid. +Flags with bit values 65536 and above are unknown; unknown flags cause implementations to reject the dataset with an "unsupported" error. The paired mixed-version reader and writer bits must either both be set or both be clear; a half-set manifest is invalid. diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index e1ca961af72..c87efb44d56 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -30,13 +30,15 @@ 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::resolve_arrow_field_ids; use prost::Message; use prost_types::Any; use roaring::RoaringBitmap; use std::cell::Cell; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use uuid::Uuid; @@ -1514,13 +1516,69 @@ fn convert_to_rust_transaction( .build()) } +struct ConvertedSchema { + schema: LanceSchema, + field_id_remap: HashMap, +} + +fn convert_arrow_schema( + arrow_schema: &Schema, + manifest: Option<&Manifest>, + operation_name: &str, +) -> Result { + // 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 + .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 Some(manifest) = manifest.filter(|manifest| { + !manifest.uses_non_reusable_field_ids() && operation_name != "Overwrite" + }) else { + return Ok(ConvertedSchema { + schema: original_schema, + field_id_remap: HashMap::new(), + }); + }; + 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() + .zip(schema.fields_pre_order()) + .filter_map(|(original, canonical)| { + (original.id >= 0 && original.id != canonical.id).then_some((original.id, canonical.id)) + }) + .collect(); + Ok(ConvertedSchema { + 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, -) -> Result { + manifest: Option<&Manifest>, + operation_name: &str, +) -> Result { let schema_ptr = env .call_method( java_operation, @@ -1531,32 +1589,32 @@ fn convert_schema_from_operation( .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)?; - if let Some(dataset) = dataset { - let arrow_schema = Schema::try_from(&c_schema)?; + convert_arrow_schema(&arrow_schema, manifest, operation_name) +} - // 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() - } - }; +type DataFileIdentity = (Option, String); - 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) - } 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 - )) - }) +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; + } + } + } } } @@ -1581,10 +1639,7 @@ 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` using the legacy Java conversion rules. fn from_arrow_schema( arrow_schema: &Schema, base_schema: Option, @@ -1598,7 +1653,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) { @@ -1612,7 +1666,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(()) } @@ -1690,11 +1744,23 @@ fn convert_to_rust_operation( read_version: u64, ) -> 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( + 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, .. } = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1702,10 +1768,15 @@ fn convert_to_rust_operation( "BufferAllocator is required for Project operations".to_string(), ) })?, - dataset, - read_version, - )?, - }, + manifest, + &op_name, + )?; + Operation::Project { + preserves_nullability: env + .get_boolean_from_method(java_operation, "preservesNullability")?, + schema, + } + } "UpdateConfig" => { let config_updates_obj = env .call_method( @@ -1829,10 +1900,7 @@ fn convert_to_rust_operation( base.extract_object(env) }) })?; - // 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 ConvertedSchema { schema, .. } = convert_schema_from_operation( env, java_operation, allocator.ok_or_else(|| { @@ -1840,8 +1908,8 @@ fn convert_to_rust_operation( "BufferAllocator is required for Overwrite operations".to_string(), ) })?, - None, - read_version, + manifest, + &op_name, )?; Operation::Overwrite { fragments, @@ -2010,25 +2078,40 @@ fn convert_to_rust_operation( Operation::DataOverlay { groups } } "Merge" => { - let fragments: Vec = + let mut fragments: Vec = import_vec_from_method(env, java_operation, "fragments", |env, fragment| { fragment.extract_object(env) })?; + let ConvertedSchema { + 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(), + ) + })?, + manifest, + &op_name, + )?; + let retained_files = if field_id_remap.is_empty() { + HashSet::new() + } else { + 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 { 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" => { @@ -2037,7 +2120,7 @@ fn convert_to_rust_operation( env.call_method(java_operation, "version", "()J", &[])? .j()?, )?; - return Ok(Operation::Restore { version }); + Operation::Restore { version } } "ReserveFragments" => { let java_num_fragments = env @@ -2048,7 +2131,7 @@ fn convert_to_rust_operation( "reserveFragments.numFragments must be non-negative, got {java_num_fragments}" )) })?; - return Ok(Operation::ReserveFragments { num_fragments }); + Operation::ReserveFragments { num_fragments } } "CreateIndex" => { let new_indices = @@ -2059,10 +2142,10 @@ 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, - }); + } } "UpdateMemWalState" => { let compacted_sstables = import_vec_from_method( @@ -2097,6 +2180,7 @@ fn convert_to_rust_operation( ))); } }; + resolve_arrow_field_ids(manifest, &mut op)?; Ok(op) } @@ -2383,6 +2467,116 @@ mod tests { pub const LANCE_FIELD_ID_KEY: &str = "lance:field_id"; + #[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)).unwrap(); + + assert_eq!(schema.field("a").unwrap().id, 5); + assert_eq!(schema.field("b").unwrap().id, 9); + } + + #[test] + fn non_reusable_java_schema_conversion_preserves_unassigned_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)]); + + let mut manifest = Manifest::new( + base_schema, + Arc::default(), + Default::default(), + HashMap::new(), + ); + manifest.activate_non_reusable_field_ids(); + + let ConvertedSchema { + schema, + field_id_remap, + } = convert_arrow_schema(&arrow_schema, Some(&manifest), "Project").unwrap(); + + assert!(schema.metadata.is_empty()); + assert_eq!(schema.field("a").unwrap().id, -1); + assert!(field_id_remap.is_empty()); + } + + #[test] + fn non_reusable_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 mut manifest = Manifest::new( + base_schema, + Arc::default(), + Default::default(), + HashMap::new(), + ); + manifest.activate_non_reusable_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()); + } + + #[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 manifest = Manifest::new( + base_schema, + Arc::default(), + Default::default(), + HashMap::new(), + ); + + let ConvertedSchema { + schema, + field_id_remap, + } = convert_arrow_schema(&arrow_schema, Some(&manifest), "Overwrite").unwrap(); + + assert_eq!(schema.field("a").unwrap().data_type(), ArrowDataType::Utf8); + assert!(schema.metadata.is_empty()); + assert!(field_id_remap.is_empty()); + } + #[test] fn test_create_schema_from_arrow() { // base_schema has an existing field id @@ -2542,7 +2736,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: parent manual, entries/key/value max_field_id let map_entries = ArrowField::new( "entries", ArrowDataType::Struct(ArrowFields::from(vec![ diff --git a/java/src/main/java/org/lance/operation/SchemaOperation.java b/java/src/main/java/org/lance/operation/SchemaOperation.java index 509492c852d..ca0e470db95 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. *
+ * + *

Datasets use non-reusable field IDs only after explicit migration. On those datasets, + * Overwrite assigns new IDs to every field, even if names and types are unchanged. Other operations + * preserve existing field IDs unless they replace the field. New IDs are assigned by the dataset + * allocator, and field mappings in newly written fragments are updated to match. */ 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..d57a08ca495 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; @@ -44,6 +45,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Optional; public class MergeTest extends OperationTestBase { @@ -359,6 +361,46 @@ void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Ex } } + @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 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/OverwriteTest.java b/java/src/test/java/org/lance/operation/OverwriteTest.java index c1def711edc..24317271a93 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[] {0, 1}, + committedOverwrite.fragments().get(0).getFiles().get(0).getFields()); + } } } } diff --git a/java/src/test/java/org/lance/operation/ProjectTest.java b/java/src/test/java/org/lance/operation/ProjectTest.java index bd3dd7d2960..c28503f4503 100644 --- a/java/src/test/java/org/lance/operation/ProjectTest.java +++ b/java/src/test/java/org/lance/operation/ProjectTest.java @@ -17,9 +17,14 @@ 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; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; @@ -27,7 +32,9 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; @@ -103,4 +110,94 @@ void testPreservesNullabilityEqualityAndRoundTrip(@TempDir Path tempDir) { } } } + + @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()); + + 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)); + } + } + } + } + + 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/protos/table.proto b/protos/table.proto index 1d0febbf215..2370324a346 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -167,6 +167,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 << 15: 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; @@ -274,6 +277,19 @@ message Manifest { // The branch of the dataset. None means main branch. optional string branch = 20; + + /* The non-reusable field-ID allocator high-water mark. + * + * 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. + * An overwrite assigns every field a new ID above the previous high-water + * mark, even when field names and types are unchanged. + * + * This field and FLAG_NON_REUSABLE_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 = 24; } // Manifest // external dataset base path diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 2e5dd63b658..a6416ed2dce 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -6261,8 +6261,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 @@ -6504,7 +6502,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/test_dataset.py b/python/python/tests/test_dataset.py index 585fb9aedc5..c23bc323e1a 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): @@ -5840,6 +5850,25 @@ def test_detached_commits(tmp_path: Path): assert detached2.to_table() == pa.table({"x": [0, 1, 3]}) +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) + + detached = lance.LanceDataset.commit( + dataset, + operation, + read_version=dataset.version, + detached=True, + ) + + assert detached.to_table() == dataset.to_table() + assert detached.schema.metadata == metadata + + def test_dataset_drop(tmp_path: Path): table = pa.table({"x": [0]}) lance.write_dataset(table, tmp_path) diff --git a/python/python/tests/test_multi_base.py b/python/python/tests/test_multi_base.py index 7e045e9b9dc..fc283d63a79 100644 --- a/python/python/tests/test_multi_base.py +++ b/python/python/tests/test_multi_base.py @@ -1290,6 +1290,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/python/python/tests/test_schema.py b/python/python/tests/test_schema.py index 67ba402a116..c885183419e 100644 --- a/python/python/tests/test_schema.py +++ b/python/python/tests/test_schema.py @@ -97,6 +97,16 @@ def test_lance_schema_from_protos_rejects_missing_parent(): 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 + + def test_lance_schema_field_lookup(tmp_path: Path): dataset = lance.write_dataset( pa.table({"x": range(2), "s": [{"a": 1}, {"a": 2}]}), tmp_path diff --git a/python/src/dataset.rs b/python/src/dataset.rs index eca9c2e756d..abbac0d5a2c 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; @@ -3090,7 +3091,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>, @@ -3104,18 +3105,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, @@ -3135,7 +3140,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, @@ -3147,6 +3152,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())?; @@ -3190,7 +3205,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)) @@ -3211,7 +3274,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/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 b3f2311b3f4..6e63d2b8d86 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -11,7 +11,7 @@ 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 pyo3::exceptions::PyValueError; @@ -1090,11 +1090,18 @@ 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| { - PyValueError::new_err(format!( - "Failed to convert Arrow schema to Lance schema: {}", - e - )) + 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}" + )) + })?; + Ok(Schema { + fields, + metadata: arrow_schema.metadata.clone(), }) } diff --git a/rust/lance-core/src/datatypes/field.rs b/rust/lance-core/src/datatypes/field.rs index 17545a760c7..27d357256c6 100644 --- a/rust/lance-core/src/datatypes/field.rs +++ b/rust/lance-core/src/datatypes/field.rs @@ -1104,6 +1104,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 0a132875ff5..d4e3f4fe8db 100644 --- a/rust/lance-core/src/datatypes/schema.rs +++ b/rust/lance-core/src/datatypes/schema.rs @@ -704,6 +704,31 @@ 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`. + /// 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; + 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`. 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<()> { + 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()); } @@ -899,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()?; @@ -1759,11 +1784,53 @@ 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::*; + #[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 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] fn test_resolve_with_quoted_fields() { // Create a schema with fields containing dots @@ -2461,8 +2528,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-namespace-impls/src/dir/manifest.rs b/rust/lance-namespace-impls/src/dir/manifest.rs index 5ca1ab036e6..9f7f410a4bd 100644 --- a/rust/lance-namespace-impls/src/dir/manifest.rs +++ b/rust/lance-namespace-impls/src/dir/manifest.rs @@ -60,6 +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::validate_non_reusable_field_id_transition; use object_store::{Error as ObjectStoreError, path::Path}; use roaring::RoaringBitmap; use std::io::Cursor; @@ -2065,6 +2066,16 @@ impl ManifestNamespace { schema.clone(), fragments, ); + if let Err(err) = validate_non_reusable_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()]; diff --git a/rust/lance-table/src/feature_flags.rs b/rust/lance-table/src/feature_flags.rs index c1b4fa10110..360acee0fd9 100644 --- a/rust/lance-table/src/feature_flags.rs +++ b/rust/lance-table/src/feature_flags.rs @@ -99,15 +99,19 @@ pub const FLAG_INDEPENDENT_COVERING_FIELDS: u64 = 1 << 13; /// must resolve their explicit bases and writers/GC must preserve those references. /// This capability is sticky, including across restore, and requires both words. pub const FLAG_MANAGED_BLOBS: u64 = 1 << 14; - +/// Field IDs are allocated from a persistent high-water mark and are never reused. +/// Writers must understand this allocation contract. It does not change how +/// readers interpret the schema or data files. +pub const FLAG_NON_REUSABLE_FIELD_IDS: u64 = 1 << 15; /// The first bit that is unknown as a feature flag -pub const FLAG_UNKNOWN: u64 = 1 << 15; +pub const FLAG_UNKNOWN: u64 = 1 << 16; const _: () = assert!(FLAG_COVERED_INDEX_METADATA < FLAG_UNKNOWN); // The fence needs a bit the current released build already refuses, which means // at or above the boundary that build shipped with (bit 7). const _: () = assert!(FLAG_COVERED_INDEX_METADATA >= 1 << 7); const _: () = assert!(FLAG_MIXED_DATA_FILE_VERSIONS < FLAG_UNKNOWN); +const _: () = assert!(FLAG_NON_REUSABLE_FIELD_IDS < FLAG_UNKNOWN); // Same fence for this bit: v12.0.0 refuses bit 9 and up. const _: () = assert!(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS >= 1 << 9); const _: () = assert!(FLAG_FRAG_REUSE_WITH_STABLE_ROW_IDS < FLAG_UNKNOWN); @@ -120,6 +124,8 @@ const _: () = assert!(FLAG_MANAGED_BLOBS < FLAG_UNKNOWN); pub(crate) const STICKY_PAIRED_FLAGS: u64 = FLAG_MIXED_DATA_FILE_VERSIONS | FLAG_FRAGMENT_REUSE_INDEX | FLAG_MANAGED_BLOBS; +pub(crate) const STICKY_READER_FLAGS: u64 = STICKY_PAIRED_FLAGS; +pub(crate) const STICKY_WRITER_FLAGS: u64 = STICKY_PAIRED_FLAGS | FLAG_NON_REUSABLE_FIELD_IDS; /// Environment variable that opts a release build into reading and writing data /// overlay files before the feature is generally released. @@ -138,14 +144,19 @@ pub fn apply_feature_flags( disable_transaction_file: bool, ) -> Result<()> { // Carried across the reset: a `Manifest` only points at its index section, - // so whether any index declares covering columns is not visible here. `build_manifest` decides it from the index list it is - // committing and sets the bit after calling this; without the carry the - // second call, from `write_manifest_file`, would clear that decision - // immediately before the write. + // so whether any index declares covering columns is not visible here. + // `build_manifest` decides it from the index list it is committing and sets + // the bit after calling this; without the carry the second call, from + // `write_manifest_file`, would clear that decision immediately before the + // write. 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 non_reusable_field_ids = manifest.max_allocated_field_id.is_some(); + if non_reusable_field_ids { + manifest.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + } + validate_non_reusable_field_id_flags(manifest)?; // Reset flags manifest.reader_feature_flags = 0; manifest.writer_feature_flags = 0; @@ -213,6 +224,10 @@ pub fn apply_feature_flags( manifest.writer_feature_flags |= FLAG_DISABLE_TRANSACTION_FILE; } + if non_reusable_field_ids { + manifest.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + } + manifest.reader_feature_flags |= covered_index_metadata; manifest.writer_feature_flags |= covered_index_metadata; manifest.reader_feature_flags |= sticky_paired_flags; @@ -221,20 +236,20 @@ 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. +/// Non-reusable 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_non_reusable_field_id_flags(source)?; destination.reader_feature_flags |= sticky_flags; - destination.writer_feature_flags |= sticky_flags; + destination.writer_feature_flags |= + sticky_flags | (source.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS); Ok(()) } @@ -317,6 +332,7 @@ pub fn can_write_dataset(writer_flags: u64) -> bool { /// not support or whose paired capabilities are inconsistent. pub fn ensure_can_read_manifest(manifest: &Manifest) -> Result<()> { validate_paired_feature_flags(manifest)?; + validate_non_reusable_field_id_flags(manifest)?; if !can_read_dataset(manifest.reader_feature_flags) { return Err(Error::not_supported_source( format!( @@ -334,6 +350,7 @@ pub fn ensure_can_read_manifest(manifest: &Manifest) -> Result<()> { /// not support or whose paired capabilities are inconsistent. pub fn ensure_can_write_manifest(manifest: &Manifest) -> Result<()> { validate_paired_feature_flags(manifest)?; + validate_non_reusable_field_id_flags(manifest)?; if !can_write_dataset(manifest.writer_feature_flags) { return Err(Error::not_supported_source( format!( @@ -385,6 +402,30 @@ pub fn validate_paired_feature_flags(manifest: &Manifest) -> Result<()> { Ok(()) } +/// Refuse a manifest whose non-reusable-field-ID marker and required flags disagree. +/// +/// The high-water mark is the activation marker and always requires the writer +/// bit. Non-reusable field IDs do not change read semantics, so the reader bit is not +/// a valid activation mode. +pub fn validate_non_reusable_field_id_flags(manifest: &Manifest) -> Result<()> { + let activated = manifest.max_allocated_field_id.is_some(); + let reader = manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS != 0; + let writer = manifest.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS != 0; + if activated != writer { + return Err(Error::corrupt_file_named( + "manifest", + "Manifest non-reusable-field-ID high-water mark and writer feature flag disagree", + )); + } + if reader { + return Err(Error::corrupt_file_named( + "manifest", + "Manifest has a non-reusable-field-ID reader feature flag, but non-reusable field IDs only require writer support", + )); + } + Ok(()) +} + fn validated_sticky_paired_flags(manifest: &Manifest) -> Result { validate_paired_feature_flags(manifest)?; Ok(manifest.reader_feature_flags & STICKY_PAIRED_FLAGS) @@ -484,6 +525,8 @@ 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_NON_REUSABLE_FIELD_IDS)); + 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). @@ -643,6 +686,8 @@ 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_NON_REUSABLE_FIELD_IDS)); + 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). @@ -731,6 +776,25 @@ mod tests { ); } + #[test] + fn inheriting_preserves_non_reusable_field_id_writer_gate() { + let mut source = empty_manifest(); + source.activate_non_reusable_field_ids(); + source.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + let mut destination = empty_manifest(); + + inherit_sticky_feature_flags(&mut destination, &source).unwrap(); + + assert_eq!( + destination.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_ne!( + destination.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + } + #[test] fn inheriting_refuses_a_half_set_source() { for (reader, writer) in [ @@ -806,6 +870,75 @@ mod tests { assert!(err.to_string().contains("cannot be written"), "{err}"); } + #[rstest::rstest] + fn apply_feature_flags_sets_writer_gate_for_explicit_non_reusable_field_id_activation( + #[values(false, true)] managed_blobs: bool, + ) { + let mut manifest = empty_manifest(); + let managed_blob_flags = if managed_blobs { FLAG_MANAGED_BLOBS } else { 0 }; + manifest.reader_feature_flags = managed_blob_flags; + manifest.writer_feature_flags = managed_blob_flags; + manifest.activate_non_reusable_field_ids(); + + apply_feature_flags(&mut manifest, false, false).unwrap(); + + assert_eq!( + manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_ne!( + manifest.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_eq!( + manifest.reader_feature_flags & FLAG_MANAGED_BLOBS, + managed_blob_flags + ); + assert_eq!( + manifest.writer_feature_flags & FLAG_MANAGED_BLOBS, + managed_blob_flags + ); + } + + #[test] + fn apply_feature_flags_rejects_non_reusable_field_id_reader_flag() { + let mut manifest = empty_manifest(); + manifest.activate_non_reusable_field_ids(); + manifest.reader_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + manifest.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + + let err = apply_feature_flags(&mut manifest, false, false).unwrap_err(); + + assert!( + err.to_string().contains("only require writer support"), + "{err}" + ); + } + + #[test] + fn non_reusable_field_id_marker_and_writer_gate_must_agree() { + let mut activated_without_gate = empty_manifest(); + activated_without_gate.activate_non_reusable_field_ids(); + assert!(validate_non_reusable_field_id_flags(&activated_without_gate).is_err()); + + let mut gate_without_marker = empty_manifest(); + gate_without_marker.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + assert!(validate_non_reusable_field_id_flags(&gate_without_marker).is_err()); + + let mut writer_only = empty_manifest(); + writer_only.activate_non_reusable_field_ids(); + writer_only.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + validate_non_reusable_field_id_flags(&writer_only).unwrap(); + + let mut paired = writer_only.clone(); + paired.reader_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + assert!(validate_non_reusable_field_id_flags(&paired).is_err()); + + let mut reader_without_activation = empty_manifest(); + reader_without_activation.reader_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + assert!(validate_non_reusable_field_id_flags(&reader_without_activation).is_err()); + } + #[rstest::rstest] #[case::reader_only(true, false)] #[case::writer_only(false, true)] diff --git a/rust/lance-table/src/format/manifest.rs b/rust/lance-table/src/format/manifest.rs index 9862d242f57..97162ee2e11 100644 --- a/rust/lance-table/src/format/manifest.rs +++ b/rust/lance-table/src/format/manifest.rs @@ -19,7 +19,7 @@ use std::ops::Range; use std::sync::Arc; use super::{Fragment, InlineRowIds, RowIdMeta}; -use crate::feature_flags::{FLAG_COVERED_INDEX_METADATA, 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; @@ -77,6 +77,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 non-reusable 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, @@ -197,6 +201,7 @@ impl Manifest { reader_feature_flags: 0, // These will be set on commit writer_feature_flags: 0, // These will be set on commit max_fragment_id: None, + max_allocated_field_id: None, transaction_file: None, transaction_section: None, fragment_offsets, @@ -225,9 +230,10 @@ 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, + 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, transaction_section: None, fragment_offsets, @@ -292,10 +298,11 @@ 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_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), transaction_section: None, fragment_offsets: self.fragment_offsets.clone(), @@ -468,12 +475,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 @@ -489,6 +501,26 @@ impl Manifest { schema_max_id.max(fragment_max_id) } + /// Whether the non-reusable-field-ID allocation contract is active. + pub fn uses_non_reusable_field_ids(&self) -> bool { + self.max_allocated_field_id.is_some() + } + + /// Activate non-reusable field IDs at the maximum ID visible in this snapshot. + pub fn activate_non_reusable_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> { @@ -774,6 +806,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, + /// Whether this commit atomically activates non-reusable field IDs. + pub activate_non_reusable_field_ids: bool, /// Row lineage sequences of the current manifest's fragments that live /// outside the manifest, read ahead of the build. An update that rewrites /// rows needs the existing row ids and created-at versions to carry each @@ -1058,6 +1092,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 @@ -1124,6 +1159,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 { @@ -1211,8 +1247,8 @@ impl SelfDescribingFileReader for V1FileReader { #[cfg(test)] mod tests { - use crate::feature_flags::FLAG_USE_V2_FORMAT_DEPRECATED; - use crate::format::overlay::{DataOverlayFile, OverlayCoverage}; + use crate::feature_flags::{FLAG_NON_REUSABLE_FIELD_IDS, FLAG_USE_V2_FORMAT_DEPRECATED}; + use crate::format::overlay::{DataOverlayFile, OverlayCoverage, TOMBSTONE_FIELD_ID}; use crate::format::{DataFile, DeletionFile, DeletionFileType}; use std::num::NonZero; @@ -1710,6 +1746,105 @@ mod tests { assert_eq!(manifest.max_field_id(), 43); } + #[test] + fn non_reusable_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), + 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![fragment]), + DataStorageFormat::default(), + HashMap::new(), + ); + manifest.activate_non_reusable_field_ids(); + manifest.writer_feature_flags |= FLAG_NON_REUSABLE_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(), 1); + + 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_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_ne!( + recovered.writer_feature_flags & FLAG_NON_REUSABLE_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] + fn shallow_clone_preserves_non_reusable_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.writer_feature_flags |= FLAG_NON_REUSABLE_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_NON_REUSABLE_FIELD_IDS, 0); + assert_ne!(cloned.writer_feature_flags & FLAG_NON_REUSABLE_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 9a000afb8d1..17b5a187f98 100644 --- a/rust/lance-table/src/transaction.rs +++ b/rust/lance-table/src/transaction.rs @@ -52,7 +52,11 @@ pub use row_version::has_writer_placed_lineage; pub use update_map::{ UpdateMap, UpdateMapEntry, translate_config_updates, translate_schema_metadata_updates, }; -pub use validate::validate_operation; +pub use validate::{ + canonicalize_non_reusable_field_ids, resolve_arrow_field_ids, + validate_detached_non_reusable_field_ids, validate_non_reusable_field_id_transition, + validate_operation, +}; 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 5c262965304..09bec42f285 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -12,8 +12,9 @@ use crate::feature_flags::{ FLAG_COVERED_INDEX_METADATA, FLAG_FRAGMENT_REUSE_INDEX, FLAG_MANAGED_BLOBS, - FLAG_STABLE_ROW_IDS, apply_feature_flags, ensure_can_read_manifest, ensure_can_write_manifest, - inherit_sticky_feature_flags, + FLAG_NON_REUSABLE_FIELD_IDS, FLAG_STABLE_ROW_IDS, apply_feature_flags, + ensure_can_read_manifest, ensure_can_write_manifest, inherit_sticky_feature_flags, + validate_non_reusable_field_id_flags, }; use crate::format::overlay::{OverlayCoverage, TOMBSTONE_FIELD_ID}; use crate::format::{ @@ -189,6 +190,18 @@ impl Transaction { manifest.max_fragment_id = manifest .max_fragment_id .max(current_manifest.max_fragment_id); + if current_manifest.uses_non_reusable_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}: non-reusable field IDs were activated after that version" + ))); + }; + manifest.max_allocated_field_id = + Some(restored_max_field_id.max(current_manifest.max_field_id())); + manifest.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + } // Row ids are a high-water mark like fragment ids: rewinding hands old ids to new rows. manifest.next_row_id = manifest.next_row_id.max(current_manifest.next_row_id); // Turning stable row ids off would revert `_rowid` to row addresses, whose @@ -1724,6 +1737,16 @@ impl Transaction { ) }; + if config.activate_non_reusable_field_ids { + let already_active = current_manifest + .map(|manifest| manifest.uses_non_reusable_field_ids()) + .unwrap_or(false); + if !already_active { + manifest.activate_non_reusable_field_ids(); + manifest.writer_feature_flags |= FLAG_NON_REUSABLE_FIELD_IDS; + } + } + // Only newly published Blob data files activate the capability. Comparing // physical files also covers column rewrites and overlays while leaving // metadata-only changes and deletion vectors on old tables alone. @@ -1799,6 +1822,7 @@ impl Transaction { manifest.set_timestamp(config.timestamp_nanos); manifest.update_max_fragment_id(); + manifest.update_max_field_id(); match &self.operation { Operation::Overwrite { @@ -1996,6 +2020,7 @@ impl Transaction { manifest.writer_feature_flags |= FLAG_FRAGMENT_REUSE_INDEX; } + validate_non_reusable_field_id_flags(&manifest)?; Ok((manifest, final_indices)) } @@ -3408,6 +3433,38 @@ mod tests { assert_eq!(overlay_paths, ["kept-mixed.lance", "kept-live.lance"]); } + #[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.activate_non_reusable_field_ids = true; + + let (activated, _) = transaction + .build_manifest(Some(&manifest), vec![], "txn", &config) + .unwrap(); + + assert!(activated.uses_non_reusable_field_ids()); + assert_eq!( + activated.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_ne!( + activated.writer_feature_flags & FLAG_NON_REUSABLE_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 e14260066fb..164feb2cbf2 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, + activate_non_reusable_field_ids: false, spilled_row_lineage: Default::default(), } } diff --git a/rust/lance-table/src/transaction/validate.rs b/rust/lance-table/src/transaction/validate.rs index 0e223095732..d42189f5b7d 100644 --- a/rust/lance-table/src/transaction/validate.rs +++ b/rust/lance-table/src/transaction/validate.rs @@ -15,6 +15,365 @@ use lance_core::{Error, Result}; use lance_file::version::ConcreteFileVersion; use std::collections::{HashMap, HashSet}; +type DataFileIdentity = (Option, String); + +/// Resolve an Arrow-derived operation against the dataset it was read from. +/// +/// Overwrite assigns new IDs to every field. Merge matches 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. +/// +/// ```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, +) -> Result<()> { + if manifest.is_some_and(|manifest| !manifest.uses_non_reusable_field_ids()) { + 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(()); + } + + match operation { + Operation::Overwrite { + schema, fragments, .. + } => { + let field_id_remap = canonicalize_schema(manifest, schema, None, true)?; + resolve_fragment_field_ids(fragments, &field_id_remap, &HashSet::new(), schema)?; + } + Operation::Project { schema, .. } => { + let Some(manifest) = manifest else { + return Ok(()); + }; + canonicalize_raw_project_schema(manifest, schema)?; + } + 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(); + let field_id_remap = + canonicalize_schema(Some(manifest), schema, Some(&manifest.schema), true)?; + resolve_fragment_field_ids(fragments, &field_id_remap, &retained_files, schema)?; + } + _ => {} + } + Ok(()) +} + +/// Assign IDs and update new file mappings before committing a Lance operation. +/// +/// `manifest` is the latest version. Overwrite assigns new IDs to every field. +/// For Merge, `read_schema` identifies 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_non_reusable_field_ids}}; +/// # fn commit(read: &Manifest, latest: &Manifest, operation: &mut Operation) -> lance_core::Result<()> { +/// canonicalize_non_reusable_field_ids(Some(latest), operation, Some(&read.schema))?; +/// # Ok(()) +/// # } +/// ``` +pub fn canonicalize_non_reusable_field_ids( + manifest: Option<&Manifest>, + operation: &mut Operation, + read_schema: Option<&Schema>, +) -> Result<()> { + if manifest.is_some_and(|manifest| !manifest.uses_non_reusable_field_ids()) { + return Ok(()); + } + match operation { + Operation::Overwrite { + schema, fragments, .. + } => { + let remap = canonicalize_schema(manifest, schema, None, 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), + )?; + } + } + _ => {} + } + Ok(()) +} + +fn canonicalize_merge_replacements( + manifest: &Manifest, + 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() { + 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)), + ); + } + let original = schema.clone(); + for field in &mut schema.fields { + clear_replaced_and_new_field_ids(field, read_schema, &replaced_field_ids); + } + schema.try_set_field_id(Some(manifest.max_field_id()))?; + schema.validate()?; + schema.verify_primary_key()?; + + 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(()) +} + +fn clear_replaced_and_new_field_ids( + field: &mut Field, + base_schema: &Schema, + replaced_field_ids: &HashSet, +) { + 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, 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, Some(&manifest.schema)) { + 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, + matching_schema: Option<&Schema>, + remap_raw_source_ids: bool, +) -> 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 let Some(matching_schema) = matching_schema { + for field in &mut schema.fields { + canonicalize_field(field, -1, matching_schema, None, None); + } + schema.try_set_field_id(max_existing_id)?; + } else { + schema.try_reassign_field_ids(max_existing_id)?; + } + schema.validate()?; + schema.verify_primary_key()?; + + 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( + field: &mut Field, + parent_id: i32, + base_schema: &Schema, + base_parent: Option<&Field>, + identity_schema: Option<&Schema>, +) -> 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 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() + .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) + .or_else(|| same_name.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 false; + }; + + field.id = base_field.id; + field.parent_id = parent_id; + let mut all_children_match = true; + for child in &mut field.children { + if !canonicalize_field( + child, + field.id, + base_schema, + Some(base_field), + identity_schema, + ) { + all_children_match = false; + } + } + all_children_match +} + +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 resolve_fragment_field_ids( + fragments: &mut [Fragment], + field_id_remap: &HashMap, + retained_files: &HashSet, + schema: &Schema, +) -> Result<()> { + let canonical_ids = schema.field_ids(); + for fragment in fragments { + // 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) = 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,7 +399,7 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &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) @@ -95,6 +454,171 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> Ok(()) } _ => Ok(()), + }; + result?; + validate_non_reusable_field_id_operation(manifest, operation) +} + +/// Validate non-reusable-field-ID invariants that are independent of one operation. +fn validate_non_reusable_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!( + "Non-reusable 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_non_reusable_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( + "Non-reusable 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!( + "Non-reusable 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_new_field_ids( + parent, + successor.schema.fields_pre_order().filter(|field| { + matches!(operation, Operation::Overwrite { .. }) + || 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_new_field_ids<'a>( + manifest: &Manifest, + new_fields: impl Iterator, +) -> Result<()> { + 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 non-reusable field IDs must be greater than the high-water mark {}", + field.name, field.id, max_allocated_field_id + ))); + } + } + Ok(()) +} + +fn validate_non_reusable_field_id_operation( + manifest: &Manifest, + operation: &Operation, +) -> Result<()> { + if !manifest.uses_non_reusable_field_ids() { + return Ok(()); + } + validate_non_reusable_field_id_manifest(manifest)?; + + let (Operation::Overwrite { schema, .. } + | Operation::Merge { schema, .. } + | Operation::Project { schema, .. }) = operation + else { + return Ok(()); + }; + schema.validate()?; + + if matches!(operation, Operation::Overwrite { .. }) { + return validate_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 {}; non-reusable 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_new_field_ids( + manifest, + schema + .fields_pre_order() + .filter(|field| manifest.schema.field_by_id(field.id).is_none()), + ) +} + +/// Reject detached schema changes once non-reusable field identity is active. +pub fn validate_detached_non_reusable_field_ids( + manifest: &Manifest, + operation: &Operation, +) -> Result<()> { + if !manifest.uses_non_reusable_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 non-reusable field IDs are activated", + )), + _ => Ok(()), } } @@ -388,7 +912,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!( @@ -403,6 +927,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, @@ -576,6 +1108,586 @@ mod tests { ) } + fn activated_manifest() -> Manifest { + let schema = one_field_schema(); + let mut manifest = manifest_with_file_fields(schema, vec![0]); + manifest.activate_non_reusable_field_ids(); + manifest + } + + #[test] + fn non_reusable_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(); + 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 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), &reused).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!( + err.to_string() + .contains("greater than the high-water mark 5"), + "{err}" + ); + } + + #[test] + fn non_reusable_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 non_reusable_field_ids_reject_overwrite_with_existing_ids() { + 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, + }; + + for err in [ + validate_operation(Some(&manifest), &operation).unwrap_err(), + validate_non_reusable_field_id_transition(&manifest, &manifest, &operation) + .unwrap_err(), + ] { + assert!(matches!(err, Error::InvalidInput { .. })); + assert!( + err.to_string() + .contains("greater than the high-water mark 0") + ); + } + } + + #[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_non_reusable_field_ids(Some(&manifest), &mut operation, None).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]); + } + + #[rstest::rstest] + #[case::last_id(i32::MAX - 1)] + #[case::exhausted(i32::MAX)] + fn overwrite_cannot_keep_existing_ids_to_avoid_exhaustion(#[case] max_id: i32) { + let mut manifest = activated_manifest(); + manifest.max_allocated_field_id = Some(max_id); + let mut operation = Operation::Overwrite { + schema: manifest.schema.clone(), + fragments: vec![], + config_upsert_values: None, + initial_bases: None, + }; + let result = canonicalize_non_reusable_field_ids(Some(&manifest), &mut operation, None); + if max_id == i32::MAX { + let err = result.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + assert!(err.to_string().contains("IDs are exhausted")); + } else { + result.unwrap(); + let Operation::Overwrite { schema, .. } = operation else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, i32::MAX); + } + } + + #[test] + fn canonicalize_overwrite_remaps_standalone_fragment_field_ids() { + let manifest = activated_manifest(); + let mut schema = one_field_schema(); + schema.fields[0].id = -1; + let mut operation = Operation::Overwrite { + fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0])], + schema, + config_upsert_values: None, + initial_bases: None, + }; + + resolve_arrow_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]); + } + + #[rstest::rstest] + fn canonicalize_overwrite_retry_assigns_all_ids_above_latest_mark( + #[values(false, true)] replace: bool, + ) { + 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_non_reusable_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_non_reusable_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::>(), + vec![2, 3] + ); + assert_eq!(fragments[0].files[0].fields.as_ref(), &[2, 3]); + assert_eq!(schema.fields[1].name, "new_column"); + } + + #[test] + fn canonicalize_raw_arrow_overwrite_assigns_new_ids_in_schema_order() { + 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_non_reusable_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("source".to_string(), "user metadata".to_string()); + let mut operation = Operation::Overwrite { + fragments: vec![], + schema: raw_schema, + config_upsert_values: None, + initial_bases: None, + }; + + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Overwrite { schema, .. } = operation else { + unreachable!(); + }; + assert_eq!(schema.field("b").unwrap().id, 2); + assert_eq!(schema.field("a").unwrap().id, 3); + assert_eq!(schema.metadata.get("source").unwrap(), "user metadata"); + } + + #[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 = 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, + config_upsert_values: None, + initial_bases: None, + }; + + resolve_arrow_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, 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; + } + 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, + }; + + resolve_arrow_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; + 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 = resolve_arrow_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(); + let mut schema = one_field_schema(); + schema.fields[0].name = "renamed".to_string(); + schema.fields[0].id = -1; + let mut operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + let err = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err(); + + 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(); + let mut operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + resolve_arrow_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]); + let mut schema = one_field_schema(); + schema.fields[0].id = -1; + let mut operation = Operation::Project { + schema, + preserves_nullability: true, + }; + + resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap(); + + let Operation::Project { schema, .. } = operation else { + unreachable!(); + }; + assert_eq!(schema.fields[0].id, 0); + assert!(schema.metadata.is_empty()); + } + + #[test] + 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_non_reusable_field_ids(Some(&manifest), &mut operation, None).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( + "replacement.lance", + vec![0], + None, + )); + let mut operation = Operation::Merge { + fragments: vec![merged_fragment], + schema: manifest.schema.clone(), + preserves_nullability: true, + }; + + canonicalize_non_reusable_field_ids(Some(&manifest), &mut operation, None).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 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; + 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 = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err(); + + assert!(err.to_string().contains("ambiguous raw Arrow field IDs")); + } + + #[test] + fn non_reusable_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_non_reusable_field_id_manifest(&manifest).unwrap_err(); + + assert!( + err.to_string().contains("below referenced field ID 7"), + "{err}" + ); + } + + #[test] + fn non_reusable_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_non_reusable_field_id_transition(&manifest, &successor, &operation) + .unwrap_err(); + + assert!( + err.to_string().contains("canonical successor schema"), + "{err}" + ); + } + + #[test] + fn non_reusable_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_non_reusable_field_id_transition(&manifest, &successor, &operation) + .unwrap_err(); + + assert!( + err.to_string().contains("high-water mark decreases"), + "{err}" + ); + } + + #[test] + fn detached_non_reusable_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_non_reusable_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_non_reusable_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 c5d8fd45b5e..4fc3f6301c8 100644 --- a/rust/lance/src/blob.rs +++ b/rust/lance/src/blob.rs @@ -1734,6 +1734,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, 2); + assert_eq!(logical.fields[0].children[1].id, 3); + assert_eq!(logical.max_field_id(), Some(3)); + + // 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)); + } + #[test] fn test_prepared_blob_schema_normalizes_by_semantic_child_name() { let mut metadata = HashMap::new(); diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index 27c48b45b32..a0b14253776 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -136,7 +136,7 @@ use crate::dataset::sql::SqlQueryBuilder; use crate::datatypes::Schema; use crate::io::commit::{ commit_detached_transaction, commit_new_dataset, commit_transaction, - default_commit_retry_timeout, detect_overlapping_fragments, + default_commit_retry_timeout, detect_overlapping_fragments, fix_schema, }; use crate::session::Session; use crate::utils::temporal::{SystemTime, timestamp_to_nanos, utc_now}; @@ -152,7 +152,7 @@ use lance_index::scalar::lance_format::LanceIndexStore; use lance_namespace::models::{DeclareTableRequest, DescribeTableRequest}; use lance_table::feature_flags::{ apply_feature_flags, ensure_can_read_manifest, ensure_can_write_manifest, - validate_paired_feature_flags, + validate_non_reusable_field_id_flags, validate_paired_feature_flags, }; use lance_table::io::deletion::{DELETIONS_DIR, relative_deletion_file_path}; use lance_table::rowids::{RowIdSequence, write_row_ids}; @@ -3296,6 +3296,51 @@ 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. 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<()> { + /// dataset.migrate_to_non_reusable_field_ids().await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn migrate_to_non_reusable_field_ids(&mut self) -> Result<()> { + if self.manifest.uses_non_reusable_field_ids() { + 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: repaired_manifest.fragments.as_ref().clone(), + schema: repaired_manifest.schema, + preserves_nullability: true, + }, + None, + ); + let new_ds = CommitBuilder::new(Arc::new(self.clone())) + .with_max_retries(0) + .with_non_reusable_field_id_migration_activation() + .execute(transaction) + .await?; + *self = new_ds; + Ok(()) + } + /// Shared clone-target preflight for `shallow_clone` and `deep_clone`: /// permit the clone only when the target definitively holds no dataset. /// Only the codebase-wide "dataset absent" pair passes: the built-in @@ -3865,7 +3910,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. @@ -4201,6 +4246,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 + /// Whether this commit activates non-reusable field IDs. + activate_non_reusable_field_ids: bool, /// This commit is a tagged fragment-reuse-index trim derived by /// `cleanup_frag_reuse_index` against the current manifest entry; see /// `ManifestBuildConfig::tagged_frag_reuse_trim`. @@ -4217,6 +4264,7 @@ impl Default for ManifestWriteConfig { use_legacy_format: None, storage_format: None, migration_next_row_id: None, + activate_non_reusable_field_ids: false, tagged_frag_reuse_trim: false, } } @@ -4259,6 +4307,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, + activate_non_reusable_field_ids: self.activate_non_reusable_field_ids, spilled_row_lineage: Default::default(), } } @@ -4302,6 +4351,7 @@ pub(crate) async fn write_manifest_file( transaction: Option, may_change_schema: bool, ) -> std::result::Result { + manifest.update_max_field_id(); validate_paired_feature_flags(manifest)?; if let Some(indices) = &indices { lance_table::system_index::frag_reuse::metadata::validate_flags(manifest, indices)?; @@ -4332,7 +4382,6 @@ pub(crate) async fn write_manifest_file( blob::validate_blob_threshold_metadata(&manifest.schema) .map_err(CommitError::OtherError)?; } - 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 @@ -4355,6 +4404,7 @@ pub(crate) async fn write_manifest_file( indices.as_deref().unwrap_or_default(), )?; + validate_non_reusable_field_id_flags(manifest).map_err(CommitError::OtherError)?; versions::finalize_manifest_storage_version(manifest)?; manifest.set_timestamp(timestamp_to_nanos(config.timestamp)); diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 4231d4f7928..5c941adc911 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -2350,7 +2350,7 @@ impl FileFragment { // the right_on key. let mut new_schema: Schema = self.schema().merge(joiner.out_schema().as_ref())?; // Use the same starting id as the updater so schema and data file ids match. - new_schema.set_field_id(Some(self.dataset.manifest.max_field_id())); + new_schema.try_set_field_id(Some(self.dataset.manifest.max_field_id()))?; let new_fragment = self .clone() 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/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 1301ff274c9..18c178498f5 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -390,7 +390,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); @@ -864,7 +864,7 @@ pub(super) async fn alter_columns( let mut cast_fields: Vec<(Field, Field)> = Vec::new(); let mut tightens_nullability = false; - let mut next_field_id = dataset.manifest.max_field_id() + 1; + 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 { @@ -916,7 +916,7 @@ 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); + field_dest.try_set_id(field_src.parent_id, &mut next_field_id)?; cast_fields.push((field_src.clone(), field_dest.clone())); } @@ -2791,13 +2791,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 { @@ -2826,16 +2854,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) } @@ -2850,7 +2869,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", @@ -4748,6 +4773,8 @@ mod test { }), ) .await?; + dataset.migrate_to_non_reusable_field_ids().await?; + assert!(dataset.manifest.uses_non_reusable_field_ids()); assert_eq!(dataset.manifest.max_field_id(), 0); // Test we can add 1 column, drop it, then add another column. Validate @@ -4762,7 +4789,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( @@ -4771,7 +4798,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( @@ -4783,7 +4810,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. @@ -4797,12 +4824,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( @@ -4811,7 +4838,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 9791358e2a1..afb2e50c2d7 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -1356,7 +1356,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 retain legacy field-ID allocation until explicitly migrated. let manifest = read_manifest( dataset.object_store.as_ref(), &dataset @@ -1379,6 +1379,7 @@ async fn test_write_manifest( "stable" | "next" )); assert_eq!(manifest.reader_feature_flags, 0); + assert_eq!(manifest.writer_feature_flags, 0); // Create one with deletions dataset.delete("i < 10").await.unwrap(); @@ -1424,6 +1425,7 @@ async fn test_write_manifest( storage_format: None, disable_transaction_file: false, migration_next_row_id: None, + activate_non_reusable_field_ids: false, tagged_frag_reuse_trim: false, }, dataset.manifest_location.naming_scheme, @@ -1461,6 +1463,59 @@ async fn test_write_manifest( assert!(matches!(write_result, Err(Error::NotSupported { .. }))); } +#[tokio::test] +async fn test_clone_rejects_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 mut unknown_writer = dataset.manifest.as_ref().clone(); + unknown_writer.version += 1; + unknown_writer.writer_feature_flags |= feature_flags::FLAG_UNKNOWN; + write_manifest_file( + dataset.object_store.as_ref(), + dataset.commit_handler.as_ref(), + &dataset.base, + &mut unknown_writer, + None, + &ManifestWriteConfig { + auto_set_feature_flags: false, + ..Default::default() + }, + dataset.manifest_location.naming_scheme, + None, + true, + ) + .await + .unwrap(); + + let mut source = Dataset::open(&source_uri).await.unwrap(); + let error = source + .shallow_clone(shallow_clone_uri.as_str(), source.version().version, None) + .await + .unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. }), "{error}"); + let error = source + .deep_clone(deep_clone_uri.as_str(), source.version().version, None) + .await + .unwrap_err(); + assert!(matches!(error, Error::NotSupported { .. }), "{error}"); +} + #[tokio::test] async fn open_rejects_mixed_file_versions_without_capability() { let uri = TempStdDir::default(); @@ -3987,6 +4042,7 @@ async fn write_manifest_file_rejects_a_nullable_primary_key() { storage_format: None, disable_transaction_file: false, migration_next_row_id: None, + activate_non_reusable_field_ids: false, tagged_frag_reuse_trim: false, }, dataset.manifest_location.naming_scheme, diff --git a/rust/lance/src/dataset/tests/dataset_migrations.rs b/rust/lance/src/dataset/tests/dataset_migrations.rs index 4615d1559d7..4b2e346bf44 100644 --- a/rust/lance/src/dataset/tests/dataset_migrations.rs +++ b/rust/lance/src/dataset/tests/dataset_migrations.rs @@ -1,24 +1,29 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use std::collections::HashMap; use std::sync::Arc; use std::vec; -use crate::dataset::InsertBuilder; use crate::dataset::optimize::{CompactionOptions, compact_files}; +use crate::dataset::{ColumnAlteration, InsertBuilder, NewColumnTransform}; 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::frag_reuse::FRAG_REUSE_INDEX_NAME; use lance_index::{IndexCriteria, IndexType, scalar::ScalarIndexParams}; -use lance_table::feature_flags::FLAG_STABLE_ROW_IDS; +use lance_table::feature_flags::{FLAG_NON_REUSABLE_FIELD_IDS, FLAG_STABLE_ROW_IDS}; use lance_table::format::{Fragment, IndexMetadata, RowIdMeta}; 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, Int64Array, ListArray, RecordBatchIterator, UInt32Array}; +use arrow_array::{ + Array, ArrayRef, Float32Array, Int32Array, Int64Array, ListArray, RecordBatchIterator, + StructArray, UInt32Array, +}; +use arrow_array::{RecordBatch, record_batch}; use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use lance_file::version::LanceFileVersion; @@ -377,6 +382,53 @@ async fn test_fix_v0_10_5_corrupt_schema() { ); } +#[tokio::test] +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(); + + 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_non_reusable_field_ids()); + assert_eq!( + cloned.manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_eq!( + cloned.manifest.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); +} + +#[tokio::test] +async fn test_non_reusable_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_non_reusable_field_ids().await.unwrap(); + + dataset.validate().await.unwrap(); + assert!(dataset.manifest.uses_non_reusable_field_ids()); + assert_eq!( + dataset.manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + + let activation_version = dataset.version().version; + dataset.migrate_to_non_reusable_field_ids().await.unwrap(); + assert_eq!(dataset.version().version, activation_version); +} + #[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 @@ -735,6 +787,346 @@ async fn make_simple_dataset(uri: &str, n: i64) -> Dataset { .unwrap() } +#[tokio::test] +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_non_reusable_field_ids()); + assert_eq!(dataset.manifest.max_allocated_field_id, None); + assert_eq!( + dataset.manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_eq!( + dataset.manifest.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + let created_version = dataset.version().version; + + dataset.migrate_to_non_reusable_field_ids().await.unwrap(); + assert_eq!(dataset.version().version, created_version + 1); + assert!(dataset.manifest.uses_non_reusable_field_ids()); + assert_eq!(dataset.manifest.max_allocated_field_id, Some(0)); + assert_eq!( + dataset.manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_ne!( + dataset.manifest.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + let activation_version = dataset.version().version; + + dataset.migrate_to_non_reusable_field_ids().await.unwrap(); + assert_eq!(dataset.version().version, activation_version); +} + +#[tokio::test] +async fn test_non_reusable_field_id_restore_boundary_and_high_water_mark() { + 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_non_reusable_field_ids().await.unwrap(); + 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_eq!( + activation_snapshot.manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 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("non-reusable 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] +async fn test_shallow_clone_preserves_non_reusable_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_non_reusable_field_ids().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_eq!( + cloned.manifest.reader_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); + assert_ne!( + cloned.manifest.writer_feature_flags & FLAG_NON_REUSABLE_FIELD_IDS, + 0 + ); +} + +#[tokio::test] +async fn test_overwrite_assigns_all_new_field_ids() { + let source_uri = TempStrDir::default(); + let mut dataset = make_simple_dataset(source_uri.as_str(), 10).await; + dataset.migrate_to_non_reusable_field_ids().await.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(), + 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)); +} + +#[tokio::test] +async fn test_raw_arrow_overwrite_assigns_new_ids_to_reordered_fields() { + 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_non_reusable_field_ids().await.unwrap(); + + 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() + }) + .execute(vec![reordered_batch]) + .await + .unwrap(); + + assert_eq!(overwritten.schema().field("b").unwrap().id, 2); + assert_eq!(overwritten.schema().field("a").unwrap().id, 3); + assert_eq!(overwritten.manifest.fragments.len(), 2); + assert_eq!(overwritten.scan().try_into_batch().await.unwrap(), expected); + assert!( + overwritten + .manifest + .fragments + .iter() + .flat_map(|fragment| &fragment.files) + .all(|file| file.fields.as_ref() == [2, 3]) + ); +} + +#[rstest] +#[tokio::test] +async fn test_repeated_overwrite_assigns_new_nested_field_ids( + #[values(false, true)] activate: bool, +) { + let nested = StructArray::from(record_batch!(("value", Int32, [1, 2])).unwrap()); + let batch = RecordBatch::try_from_iter([("s", Arc::new(nested) as ArrayRef)]).unwrap(); + let dir = TempStrDir::default(); + let mut dataset = InsertBuilder::new(&dir) + .execute(vec![batch.clone()]) + .await + .unwrap(); + if activate { + dataset.migrate_to_non_reusable_field_ids().await.unwrap(); + } + let original = dataset.clone(); + for first_id in [2, 4] { + InsertBuilder::new(Arc::new(dataset)) + .with_params(&WriteParams { + mode: WriteMode::Overwrite, + max_rows_per_file: 1, + ..Default::default() + }) + .execute(vec![batch.clone()]) + .await + .unwrap(); + dataset = Dataset::open(dir.as_str()).await.unwrap(); + let expected_ids = if activate { + vec![first_id, first_id + 1] + } else { + vec![0, 1] + }; + assert_eq!( + dataset + .schema() + .fields_pre_order() + .map(|field| field.id) + .collect::>(), + expected_ids + ); + assert_eq!( + dataset.manifest.max_allocated_field_id, + activate.then_some(first_id + 1) + ); + assert_eq!(dataset.manifest.fragments.len(), 2); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), batch); + } + assert_eq!(original.scan().try_into_batch().await.unwrap(), batch); +} + +#[tokio::test] +async fn test_non_reusable_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_non_reusable_field_ids().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_non_reusable_field_id_multi_cast_uses_schema_order() { + 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_non_reusable_field_ids().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(); + 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] +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, None); + 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/updater.rs b/rust/lance/src/dataset/updater.rs index 5f2edb96477..db070aa8fbd 100644 --- a/rust/lance/src/dataset/updater.rs +++ b/rust/lance/src/dataset/updater.rs @@ -224,7 +224,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/versions/mod.rs b/rust/lance/src/dataset/versions/mod.rs index 75c84bfe569..daf780a27a1 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 20047b11a61..6376277f2c4 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -34,6 +34,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, resolve_arrow_field_ids}; use object_store::{ObjectStoreExt, path::Path}; use std::borrow::Cow; use std::collections::{BTreeSet, HashMap, HashSet, VecDeque}; @@ -1848,7 +1849,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) } @@ -1858,7 +1859,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; @@ -1917,6 +1922,28 @@ pub(super) fn prepare_write_schema( } } projected + } else if let Some(dataset) = dataset + && matches!(params.mode, WriteMode::Overwrite) + && dataset.manifest.uses_non_reusable_field_ids() + { + // Uncommitted fragment APIs return files without the schema used to + // write them, so their mappings must already use commit-time IDs. + // 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, + }; + resolve_arrow_field_ids(Some(&dataset.manifest), &mut operation)?; + match operation { + Operation::Overwrite { schema, .. } => schema, + _ => { + return Err(Error::internal( + "Non-reusable field-ID canonicalization changed an Overwrite operation", + )); + } + } } else { normalized_converted_schema }; @@ -3352,7 +3379,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(), @@ -3368,7 +3395,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() @@ -3379,16 +3415,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 4af9a12fe93..ff13b114d6b 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -28,7 +28,6 @@ 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 tracing::info; @@ -54,6 +53,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, + /// Whether this commit atomically activates non-reusable field IDs. + activate_non_reusable_field_ids: bool, /// Set only by `Dataset::deep_clone`, after it has copied the source files. deep_clone_files_copied: bool, } @@ -80,6 +81,7 @@ impl<'a> CommitBuilder<'a> { transaction_properties: None, timeout: Some(DEFAULT_COMMIT_TIMEOUT), migration_next_row_id: None, + activate_non_reusable_field_ids: false, deep_clone_files_copied: false, } } @@ -280,6 +282,11 @@ impl<'a> CommitBuilder<'a> { self } + pub(crate) fn with_non_reusable_field_id_migration_activation(mut self) -> Self { + self.activate_non_reusable_field_ids = true; + self + } + /// Mark this commit as the last step of [`Dataset::deep_clone`], which has /// already copied the source's data, deletion and index files to the /// destination. @@ -416,14 +423,6 @@ impl<'a> CommitBuilder<'a> { )); } - // 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) => ( @@ -454,6 +453,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, + activate_non_reusable_field_ids: self.activate_non_reusable_field_ids, ..Default::default() }; @@ -624,6 +624,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_core::utils::tempfile::TempStrDir; @@ -633,6 +634,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; @@ -643,6 +645,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 mut transaction = Transaction::new( + 0, + Operation::Overwrite { + schema, + fragments: vec![], + config_upsert_values: None, + initial_bases: None, + }, + None, + ); + resolve_arrow_field_ids(None, &mut transaction.operation).unwrap(); + let dataset = CommitBuilder::new("memory://") + .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_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 { + max_rows_per_file: 1, + ..Default::default() + }) + .execute(vec![batch.clone()]) + .await + .unwrap(); + dataset.migrate_to_non_reusable_field_ids().await.unwrap(); + let mut raw_schema = dataset.schema().clone(); + for field in &mut raw_schema.fields { + field.id += 10; + } + let mut transaction = Transaction::new( + dataset.version().version, + Operation::Merge { + schema: raw_schema, + fragments: dataset.manifest.fragments.as_ref().clone(), + preserves_nullability: true, + }, + None, + ); + resolve_arrow_field_ids(Some(&dataset.manifest), &mut transaction.operation).unwrap(); + let committed = CommitBuilder::new(Arc::new(dataset.clone())) + .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/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; } diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index 2a45637995e..675ebb4a07c 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -44,7 +44,11 @@ use lance_table::io::commit::{ CommitConfig, CommitError, CommitHandler, ManifestLocation, ManifestNamingScheme, }; use lance_table::io::manifest::read_manifest; -use lance_table::transaction::{FragReuseUpdate, PreparedIndices, has_writer_placed_lineage}; +use lance_table::transaction::{ + FragReuseUpdate, PreparedIndices, canonicalize_non_reusable_field_ids, + has_writer_placed_lineage, validate_detached_non_reusable_field_ids, + validate_non_reusable_field_id_transition, validate_operation, +}; use rand::{Rng, rng}; use roaring::RoaringBitmap; @@ -432,6 +436,10 @@ async fn do_commit_new_dataset( metadata_cache: &DSMetadataCache, session: Arc, ) -> Result<(Manifest, ManifestLocation)> { + let mut transaction = transaction.clone(); + canonicalize_non_reusable_field_ids(None, &mut transaction.operation, None)?; + let transaction = &transaction; + validate_operation(None, &transaction.operation)?; let pb_transaction = pb::Transaction::try_from(transaction)?; let inline_transaction = pb_transaction.encoded_len() <= MAX_INLINE_TRANSACTION_BYTES; // Classified from the operation itself. Reading it back off the inline @@ -623,6 +631,10 @@ async fn do_commit_new_dataset( (manifest, indices) }; + if !manifest.uses_non_reusable_field_ids() { + fix_schema(&mut manifest)?; + } + let result = write_manifest_file( object_store, commit_handler.as_ref(), @@ -864,7 +876,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(()); @@ -887,13 +899,24 @@ fn fix_schema(manifest: &mut Manifest) -> Result<()> { return Ok(()); } + if manifest.uses_non_reusable_field_ids() { + return Err(Error::invalid_input( + "Cannot repair duplicate field IDs after non-reusable 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(); @@ -1266,6 +1289,11 @@ pub(crate) async fn do_commit_detached_transaction( retry_timeout: Duration, ) -> Result<(Manifest, ManifestLocation)> { ensure_can_write_manifest(&dataset.manifest)?; + let mut transaction = transaction.clone(); + canonicalize_non_reusable_field_ids(Some(&dataset.manifest), &mut transaction.operation, None)?; + let transaction = &transaction; + validate_detached_non_reusable_field_ids(&dataset.manifest, &transaction.operation)?; + validate_operation(Some(&dataset.manifest), &transaction.operation)?; // Detached commits skip the rebase pipeline, so a rewrite's transition // intent would never be assembled or validated (a dummy intent plus a // hand-built entry would satisfy the manifest chokepoint unvalidated). @@ -1348,6 +1376,11 @@ pub(crate) async fn do_commit_detached_transaction( // Validate before the fragment-id check to preserve legacy migration // diagnostics. Finalization repeats this at the manifest write boundary. fix_schema(&mut manifest)?; + validate_non_reusable_field_id_transition( + &dataset.manifest, + &manifest, + &transaction.operation, + )?; crate::dataset::versions::check_manifest_storage_version_for_commit(&mut manifest)?; check_fragment_ids(&manifest)?; // Runs after the coverage derivation and can replace a fragment bitmap @@ -1749,9 +1782,22 @@ pub(crate) async fn commit_transaction( let attempt_start = Instant::now(); let retry_start = *retry_start.get_or_insert(attempt_start); + // Merge preserves fields from the read version; Overwrite replaces all + // fields. 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_non_reusable_field_ids( + Some(&dataset.manifest), + &mut attempt_transaction.operation, + Some(read_version_dataset.schema()), + )?; + validate_operation(Some(&dataset.manifest), &attempt_transaction.operation)?; + // Recomputed every attempt: the rebase above may have rewritten the // transaction. - let pb_transaction = pb::Transaction::try_from(&transaction)?; + let pb_transaction = pb::Transaction::try_from(&attempt_transaction)?; let inline_transaction = pb_transaction.encoded_len() <= MAX_INLINE_TRANSACTION_BYTES; // Classified from the operation itself. Reading it back off the inline // copy would tie the verdict to the payload size instead. @@ -1773,8 +1819,9 @@ pub(crate) async fn commit_transaction( // Build an up-to-date manifest from the transaction and current // manifest: prepare the index list against this attempt's manifest // first, then build from the prepared result. - let build_config = build_config_for_attempt(&dataset, &transaction, write_config).await?; - let (mut manifest, mut indices) = match transaction.operation { + let build_config = + build_config_for_attempt(&dataset, &attempt_transaction, write_config).await?; + let (mut manifest, mut indices) = match attempt_transaction.operation { Operation::Restore { version } => { // A restore reinstates its snapshot's index list and entry // whole; nothing is prepared or withdrawn a second time. @@ -1796,8 +1843,9 @@ pub(crate) async fn commit_transaction( None => FragReuseUpdate::None, }; let prepared = - prepare_attempt(&dataset, &transaction, &build_config, frag_reuse).await?; - transaction.build_manifest_prepared( + prepare_attempt(&dataset, &attempt_transaction, &build_config, frag_reuse) + .await?; + attempt_transaction.build_manifest_prepared( Some(dataset.manifest.as_ref()), prepared, transaction_file, @@ -1819,6 +1867,11 @@ pub(crate) async fn commit_transaction( migrate_manifest(&dataset, &mut manifest, recompute_stats).await?; fix_schema(&mut manifest)?; + validate_non_reusable_field_id_transition( + &dataset.manifest, + &manifest, + &attempt_transaction.operation, + )?; crate::dataset::versions::check_manifest_storage_version_for_commit(&mut manifest)?; check_fragment_ids(&manifest)?; @@ -1854,7 +1907,7 @@ pub(crate) async fn commit_transaction( Ok(manifest_location) => { record_successful_commit( &dataset, - &transaction, + &attempt_transaction, &manifest, &manifest_location, indices, @@ -1875,7 +1928,7 @@ pub(crate) async fn commit_transaction( commit_handler, &dataset.base, target_version, - &transaction, + &attempt_transaction, ) .await { @@ -1886,7 +1939,7 @@ pub(crate) async fn commit_transaction( let committed_manifest = *committed_manifest; record_successful_commit( &dataset, - &transaction, + &attempt_transaction, &committed_manifest, &location, indices, @@ -1948,7 +2001,7 @@ pub(crate) async fn commit_transaction( commit_handler, &dataset.base, target_version, - &transaction, + &attempt_transaction, ) .await { @@ -1959,7 +2012,7 @@ pub(crate) async fn commit_transaction( let committed_manifest = *committed_manifest; record_successful_commit( &dataset, - &transaction, + &attempt_transaction, &committed_manifest, &location, indices, @@ -2026,12 +2079,13 @@ mod tests { CommitLease, CommitLock, ManifestWriter, RenameCommitHandler, UnsafeCommitHandler, commit_handler_from_url, }; + use lance_table::transaction::resolve_arrow_field_ids; 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}; @@ -2197,6 +2251,195 @@ 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::try_from(transaction).unwrap().into(); + Arc::new(InjectForeignCommitHandler { + foreign: Mutex::new(Some((manifest, transaction))), + }) + } + + #[tokio::test] + async fn raw_arrow_project_retry_preserves_identity_and_metadata() { + let tmp = TempStrDir::default(); + let uri = tmp.as_str(); + let mut dataset = Dataset::write( + RecordBatchIterator::new( + vec![Ok(simple_batch(&simple_schema(), vec![1, 2, 3]))], + simple_schema(), + ), + uri, + None, + ) + .await + .unwrap(); + dataset.migrate_to_non_reusable_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); + let foreign_transaction = Transaction::new( + dataset.version().version, + Operation::ReserveFragments { num_fragments: 1 }, + None, + ); + 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_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 operation).unwrap(); + + let committed = CommitBuilder::new(Arc::new(dataset.clone())) + .with_commit_handler(handler) + .execute(Transaction::new(dataset.version().version, operation, None)) + .await + .unwrap(); + + assert_eq!(committed.schema(), &expected_schema); + } + + #[rstest::rstest] + #[tokio::test] + 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( + RecordBatchIterator::new( + vec![Ok(simple_batch(&simple_schema(), vec![1, 2, 3]))], + simple_schema(), + ), + uri, + None, + ) + .await + .unwrap(); + dataset.migrate_to_non_reusable_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); + 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 }, + None, + ); + 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(); + 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::new(), + }; + let mut expected_schema = raw_schema.clone(); + expected_schema.fields[0].id = if overwrite { 2 } else { 0 }; + expected_schema.fields[1].id = if overwrite { 3 } else { 2 }; + 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 committed = CommitBuilder::new(Arc::new(dataset.clone())) + .with_commit_handler(handler) + .execute(Transaction::new(dataset.version().version, operation, None)) + .await + .unwrap(); + + assert_eq!(committed.schema(), &expected_schema); + assert_eq!( + committed.manifest.max_allocated_field_id, + Some(if overwrite { 3 } else { 2 }) + ); + assert_eq!( + committed.manifest.fragments[0].files[0].fields.as_ref(), + &[expected_schema.fields[0].id] + ); + assert_eq!( + committed.manifest.fragments[0].files[1].fields.as_ref(), + &[expected_schema.fields[1].id] + ); + } + #[tokio::test] async fn test_unsafe_commit_handler() { let handler = Arc::new(UnsafeCommitHandler); @@ -2582,6 +2825,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; @@ -2881,6 +3141,63 @@ mod tests { } } + #[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_non_reusable_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_non_reusable_field_ids(); + + let err = fix_schema(&mut manifest).unwrap_err(); + + assert!( + err.to_string().contains("non-reusable 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.