1use crate::format::{Fragment, Manifest};
11use crate::io::deletion::relative_deletion_file_path;
12use crate::transaction::{Operation, UpdateMode, UpdatedFragmentOffsets};
13use lance_core::datatypes::{Field, Schema};
14use lance_core::{Error, Result};
15use lance_file::version::ConcreteFileVersion;
16use std::collections::{HashMap, HashSet};
17
18type DataFileIdentity = (Option<u32>, String);
19
20pub fn resolve_arrow_field_ids(
37 manifest: Option<&Manifest>,
38 operation: &mut Operation,
39) -> Result<()> {
40 if manifest.is_some_and(|manifest| !manifest.uses_non_reusable_field_ids()) {
41 match operation {
42 Operation::Overwrite { schema, .. }
43 | Operation::Project { schema, .. }
44 | Operation::Merge { schema, .. } => {
45 schema.try_set_field_id(None)?;
46 schema.validate()?;
47 schema.verify_primary_key()?;
48 }
49 _ => {}
50 }
51 return Ok(());
52 }
53
54 match operation {
55 Operation::Overwrite {
56 schema, fragments, ..
57 } => {
58 let field_id_remap = canonicalize_schema(manifest, schema, None, true)?;
59 resolve_fragment_field_ids(fragments, &field_id_remap, &HashSet::new(), schema)?;
60 }
61 Operation::Project { schema, .. } => {
62 let Some(manifest) = manifest else {
63 return Ok(());
64 };
65 canonicalize_raw_project_schema(manifest, schema)?;
66 }
67 Operation::Merge {
68 schema, fragments, ..
69 } => {
70 let Some(manifest) = manifest else {
71 return Ok(());
72 };
73 let retained_files = manifest
74 .fragments
75 .iter()
76 .flat_map(|fragment| fragment.referenced_lance_files())
77 .map(|file| (file.base_id, file.path.clone()))
78 .collect();
79 let field_id_remap =
80 canonicalize_schema(Some(manifest), schema, Some(&manifest.schema), true)?;
81 resolve_fragment_field_ids(fragments, &field_id_remap, &retained_files, schema)?;
82 }
83 _ => {}
84 }
85 Ok(())
86}
87
88pub fn canonicalize_non_reusable_field_ids(
104 manifest: Option<&Manifest>,
105 operation: &mut Operation,
106 read_schema: Option<&Schema>,
107) -> Result<()> {
108 if manifest.is_some_and(|manifest| !manifest.uses_non_reusable_field_ids()) {
109 return Ok(());
110 }
111 match operation {
112 Operation::Overwrite {
113 schema, fragments, ..
114 } => {
115 let remap = canonicalize_schema(manifest, schema, None, false)?;
116 remap_fragment_field_ids(fragments, &remap, &HashSet::new());
117 }
118 Operation::Merge {
119 schema, fragments, ..
120 } => {
121 if let Some(manifest) = manifest {
122 let retained_files = manifest
123 .fragments
124 .iter()
125 .flat_map(|fragment| fragment.referenced_lance_files())
126 .map(|file| (file.base_id, file.path.clone()))
127 .collect();
128 canonicalize_merge_replacements(
129 manifest,
130 schema,
131 fragments,
132 &retained_files,
133 read_schema.unwrap_or(&manifest.schema),
134 )?;
135 }
136 }
137 _ => {}
138 }
139 Ok(())
140}
141
142fn canonicalize_merge_replacements(
143 manifest: &Manifest,
144 schema: &mut Schema,
145 fragments: &mut [Fragment],
146 retained_files: &HashSet<DataFileIdentity>,
147 read_schema: &Schema,
148) -> Result<()> {
149 let mut replaced_field_ids = HashSet::new();
150 for fragment in fragments.iter() {
151 let retained_field_ids = fragment
152 .referenced_lance_files()
153 .filter(|file| retained_files.contains(&(file.base_id, file.path.clone())))
154 .flat_map(|file| file.fields.iter().copied())
155 .collect::<HashSet<_>>();
156 replaced_field_ids.extend(
157 fragment
158 .referenced_lance_files()
159 .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone())))
160 .flat_map(|file| file.fields.iter().copied())
161 .filter(|field_id| retained_field_ids.contains(field_id)),
162 );
163 }
164 let original = schema.clone();
165 for field in &mut schema.fields {
166 clear_replaced_and_new_field_ids(field, read_schema, &replaced_field_ids);
167 }
168 schema.try_set_field_id(Some(manifest.max_field_id()))?;
169 schema.validate()?;
170 schema.verify_primary_key()?;
171
172 let field_id_remap = original
173 .fields_pre_order()
174 .zip(schema.fields_pre_order())
175 .filter(|(original, _)| original.id >= 0)
176 .map(|(original, canonical)| (original.id, canonical.id))
177 .collect();
178 remap_fragment_field_ids(fragments, &field_id_remap, retained_files);
179 Ok(())
180}
181
182fn clear_replaced_and_new_field_ids(
183 field: &mut Field,
184 base_schema: &Schema,
185 replaced_field_ids: &HashSet<i32>,
186) {
187 if base_schema.field_by_id(field.id).is_none() || replaced_field_ids.contains(&field.id) {
188 clear_field_ids(field);
189 return;
190 }
191 for child in &mut field.children {
192 clear_replaced_and_new_field_ids(child, base_schema, replaced_field_ids);
193 }
194}
195
196fn canonicalize_raw_project_schema(manifest: &Manifest, schema: &mut Schema) -> Result<()> {
197 let mut unmatched_fields = Vec::new();
198 for field in &mut schema.fields {
199 if !canonicalize_field(field, -1, &manifest.schema, None, Some(&manifest.schema)) {
200 unmatched_fields.push(field.name.clone());
201 }
202 }
203 if !unmatched_fields.is_empty() {
204 return Err(Error::invalid_input(format!(
205 "Raw Arrow Project fields [{}] do not match existing field identities; Project cannot allocate new identities because it writes no data",
206 unmatched_fields.join(", ")
207 )));
208 }
209 schema.validate()?;
210 schema.verify_primary_key()?;
211 Ok(())
212}
213
214fn canonicalize_schema(
215 manifest: Option<&Manifest>,
216 schema: &mut Schema,
217 matching_schema: Option<&Schema>,
218 remap_raw_source_ids: bool,
219) -> Result<HashMap<i32, i32>> {
220 let mut original = schema.clone();
221 if remap_raw_source_ids {
222 original.try_set_field_id(None)?;
223 }
224
225 let max_existing_id = manifest.map(Manifest::max_field_id);
226 if let Some(matching_schema) = matching_schema {
227 for field in &mut schema.fields {
228 canonicalize_field(field, -1, matching_schema, None, None);
229 }
230 schema.try_set_field_id(max_existing_id)?;
231 } else {
232 schema.try_reassign_field_ids(max_existing_id)?;
233 }
234 schema.validate()?;
235 schema.verify_primary_key()?;
236
237 Ok(original
238 .fields_pre_order()
239 .zip(schema.fields_pre_order())
240 .filter(|(original, _)| original.id >= 0)
241 .map(|(original, canonical)| (original.id, canonical.id))
242 .collect())
243}
244
245fn canonicalize_field(
246 field: &mut Field,
247 parent_id: i32,
248 base_schema: &Schema,
249 base_parent: Option<&Field>,
250 identity_schema: Option<&Schema>,
251) -> bool {
252 let same_name = match base_parent {
253 Some(parent) => parent.children.iter().find(|base| base.name == field.name),
254 None => base_schema
255 .fields
256 .iter()
257 .find(|base| base.name == field.name),
258 };
259 let allow_id_binding =
260 identity_schema.is_some_and(|schema| schema.field_by_id(field.id).is_some());
261 let by_id = (allow_id_binding && field.id >= 0)
262 .then(|| base_schema.field_by_id(field.id))
263 .flatten()
264 .filter(|base| base.parent_id == parent_id);
265 let base_field = if allow_id_binding && field.id >= 0 {
266 by_id
267 .filter(|base| base.logical_type == field.logical_type)
268 .or_else(|| same_name.filter(|base| base.logical_type == field.logical_type))
269 } else {
270 same_name.filter(|base| base.logical_type == field.logical_type)
271 };
272
273 let Some(base_field) = base_field else {
274 clear_field_ids(field);
275 return false;
276 };
277
278 field.id = base_field.id;
279 field.parent_id = parent_id;
280 let mut all_children_match = true;
281 for child in &mut field.children {
282 if !canonicalize_field(
283 child,
284 field.id,
285 base_schema,
286 Some(base_field),
287 identity_schema,
288 ) {
289 all_children_match = false;
290 }
291 }
292 all_children_match
293}
294
295fn clear_field_ids(field: &mut Field) {
296 field.id = -1;
297 field.parent_id = -1;
298 for child in &mut field.children {
299 clear_field_ids(child);
300 }
301}
302
303fn resolve_fragment_field_ids(
304 fragments: &mut [Fragment],
305 field_id_remap: &HashMap<i32, i32>,
306 retained_files: &HashSet<DataFileIdentity>,
307 schema: &Schema,
308) -> Result<()> {
309 let canonical_ids = schema.field_ids();
310 for fragment in fragments {
311 let fragment_field_ids = fragment
318 .referenced_lance_files()
319 .filter(|file| !retained_files.contains(&(file.base_id, file.path.clone())))
320 .flat_map(|file| file.fields.iter().copied())
321 .filter(|field_id| *field_id >= 0)
322 .collect::<HashSet<_>>();
323 let canonical_source = fragment_field_ids
324 .iter()
325 .all(|field_id| canonical_ids.contains(field_id));
326 let raw_source = fragment_field_ids
327 .iter()
328 .all(|field_id| field_id_remap.contains_key(field_id));
329 let raw_changes_identity = fragment_field_ids
330 .iter()
331 .any(|field_id| field_id_remap.get(field_id) != Some(field_id));
332
333 match (canonical_source, raw_source, raw_changes_identity) {
334 (true, true, true) => {
335 return Err(Error::invalid_input(format!(
336 "Fragment {} has ambiguous raw Arrow field IDs; its file mappings can be interpreted as either source or canonical identities",
337 fragment.id
338 )));
339 }
340 (true, _, _) => continue,
341 (false, true, _) => {}
342 (false, false, _) => {
343 return Err(Error::invalid_input(format!(
344 "Fragment {} field IDs do not match either the raw Arrow source schema or the canonical replacement schema",
345 fragment.id
346 )));
347 }
348 }
349 remap_fragment_field_ids(
350 std::slice::from_mut(fragment),
351 field_id_remap,
352 retained_files,
353 );
354 }
355 Ok(())
356}
357
358fn remap_fragment_field_ids(
359 fragments: &mut [Fragment],
360 field_id_remap: &HashMap<i32, i32>,
361 retained_files: &HashSet<DataFileIdentity>,
362) {
363 for fragment in fragments {
364 for file in fragment.referenced_lance_files_mut() {
365 if retained_files.contains(&(file.base_id, file.path.clone())) {
366 continue;
367 }
368 for field_id in std::sync::Arc::make_mut(&mut file.fields) {
369 if let Some(canonical_id) = field_id_remap.get(field_id) {
370 *field_id = *canonical_id;
371 }
372 }
373 }
374 }
375}
376
377pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> Result<()> {
379 let manifest = match (manifest, operation) {
380 (
381 None,
382 Operation::Overwrite {
383 fragments, schema, ..
384 },
385 ) => {
386 overwrite_fragments_valid(fragments)?;
388 schema_fragments_valid(None, schema, fragments)?;
389
390 return Ok(());
391 }
392 (None, Operation::Clone { .. }) => return Ok(()),
393 (Some(manifest), _) => manifest,
394 (None, _) => {
395 return Err(Error::invalid_input(format!(
396 "Cannot apply operation {} to non-existent dataset",
397 operation.name()
398 )));
399 }
400 };
401
402 let result = match operation {
403 Operation::Append { fragments } => {
404 schema_fragments_valid(Some(manifest), &manifest.schema, fragments)
406 }
407 Operation::Project { schema, .. } => {
408 schema_fragments_valid(Some(manifest), schema, manifest.fragments.as_ref())
409 }
410 Operation::Merge {
411 fragments, schema, ..
412 } => {
413 merge_fragments_valid(manifest, fragments)?;
414 merge_schema_valid(manifest, schema, fragments)?;
415 schema_fragments_valid(Some(manifest), schema, fragments)
416 }
417 Operation::Overwrite {
418 fragments, schema, ..
419 } => {
420 overwrite_fragments_valid(fragments)?;
421 schema_fragments_valid(None, schema, fragments)
425 }
426 Operation::Update {
427 updated_fragments,
428 new_fragments,
429 updated_fragment_offsets,
430 update_mode,
431 ..
432 } => {
433 schema_fragments_valid(Some(manifest), &manifest.schema, updated_fragments)?;
434 schema_fragments_valid(Some(manifest), &manifest.schema, new_fragments)?;
435 if matches!(update_mode, Some(UpdateMode::RewriteColumns))
441 && let Some(UpdatedFragmentOffsets(off_map)) = updated_fragment_offsets
442 {
443 let updated_ids: HashSet<u64> = updated_fragments.iter().map(|f| f.id).collect();
444 for &frag_id in off_map.keys() {
445 if !updated_ids.contains(&frag_id) {
446 return Err(Error::invalid_input(format!(
447 "updatedFragmentOffsets key {} is not in updated_fragments; \
448 offsets must reference only fragments being rewritten",
449 frag_id
450 )));
451 }
452 }
453 }
454 Ok(())
455 }
456 _ => Ok(()),
457 };
458 result?;
459 validate_non_reusable_field_id_operation(manifest, operation)
460}
461
462fn validate_non_reusable_field_id_manifest(manifest: &Manifest) -> Result<()> {
464 let Some(max_allocated_field_id) = manifest.max_allocated_field_id else {
465 return Ok(());
466 };
467 let max_referenced_field_id = manifest.max_referenced_field_id();
468 if max_allocated_field_id < max_referenced_field_id {
469 return Err(Error::invalid_input(format!(
470 "Non-reusable field-ID high-water mark {} is below referenced field ID {}",
471 max_allocated_field_id, max_referenced_field_id
472 )));
473 }
474 Ok(())
475}
476
477pub fn validate_non_reusable_field_id_transition(
485 parent: &Manifest,
486 successor: &Manifest,
487 operation: &Operation,
488) -> Result<()> {
489 let Some(parent_max_field_id) = parent.max_allocated_field_id else {
490 return Ok(());
491 };
492 let Some(successor_max_field_id) = successor.max_allocated_field_id else {
493 return Err(Error::invalid_input(
494 "Non-reusable field-ID activation marker is missing from the successor manifest",
495 ));
496 };
497 if successor_max_field_id < parent_max_field_id {
498 return Err(Error::invalid_input(format!(
499 "Non-reusable field-ID high-water mark decreases from {parent_max_field_id} to {successor_max_field_id}"
500 )));
501 }
502 let successor_schema_ids = successor
503 .schema
504 .fields_pre_order()
505 .map(|field| field.id)
506 .collect::<HashSet<_>>();
507
508 if !matches!(operation, Operation::Restore { .. }) {
509 validate_new_field_ids(
510 parent,
511 successor.schema.fields_pre_order().filter(|field| {
512 matches!(operation, Operation::Overwrite { .. })
513 || parent.schema.field_by_id(field.id).is_none()
514 }),
515 )?;
516 }
517
518 for field_id in successor
519 .fragments
520 .iter()
521 .flat_map(|fragment| fragment.referenced_lance_files())
522 .flat_map(|file| file.fields.iter())
523 .copied()
524 .filter(|field_id| *field_id >= 0)
525 {
526 if matches!(operation, Operation::Overwrite { .. })
527 && !successor_schema_ids.contains(&field_id)
528 {
529 return Err(Error::invalid_input(format!(
530 "Overwrite references field ID {field_id}, which is not in its replacement schema"
531 )));
532 }
533 if field_id > parent_max_field_id && !successor_schema_ids.contains(&field_id) {
534 return Err(Error::invalid_input(format!(
535 "Data file or overlay references new field ID {field_id}, but the canonical successor schema does not contain that identity"
536 )));
537 }
538 }
539 Ok(())
540}
541
542fn validate_new_field_ids<'a>(
543 manifest: &Manifest,
544 new_fields: impl Iterator<Item = &'a Field>,
545) -> Result<()> {
546 let max_allocated_field_id = manifest.max_field_id();
547 for field in new_fields {
548 if field.id <= max_allocated_field_id {
549 return Err(Error::invalid_input(format!(
550 "New field '{}' has ID {}, but non-reusable field IDs must be greater than the high-water mark {}",
551 field.name, field.id, max_allocated_field_id
552 )));
553 }
554 }
555 Ok(())
556}
557
558fn validate_non_reusable_field_id_operation(
559 manifest: &Manifest,
560 operation: &Operation,
561) -> Result<()> {
562 if !manifest.uses_non_reusable_field_ids() {
563 return Ok(());
564 }
565 validate_non_reusable_field_id_manifest(manifest)?;
566
567 let (Operation::Overwrite { schema, .. }
568 | Operation::Merge { schema, .. }
569 | Operation::Project { schema, .. }) = operation
570 else {
571 return Ok(());
572 };
573 schema.validate()?;
574
575 if matches!(operation, Operation::Overwrite { .. }) {
576 return validate_new_field_ids(manifest, schema.fields_pre_order());
577 }
578
579 for field in schema.fields_pre_order() {
580 let Some(prior_field) = manifest.schema.field_by_id(field.id) else {
581 continue;
582 };
583 if field.parent_id != prior_field.parent_id {
584 return Err(Error::invalid_input(format!(
585 "Field ID {} moves from parent {} to parent {}; non-reusable field identity cannot move between parents",
586 field.id, prior_field.parent_id, field.parent_id
587 )));
588 }
589 if field.logical_type != prior_field.logical_type {
590 return Err(Error::invalid_input(format!(
591 "Field ID {} changes logical type from {} to {}; type replacement must allocate a new field ID",
592 field.id, prior_field.logical_type, field.logical_type
593 )));
594 }
595 }
596
597 validate_new_field_ids(
598 manifest,
599 schema
600 .fields_pre_order()
601 .filter(|field| manifest.schema.field_by_id(field.id).is_none()),
602 )
603}
604
605pub fn validate_detached_non_reusable_field_ids(
607 manifest: &Manifest,
608 operation: &Operation,
609) -> Result<()> {
610 if !manifest.uses_non_reusable_field_ids() {
611 return Ok(());
612 }
613 match operation {
614 Operation::Merge { schema, .. } if schema == &manifest.schema => Ok(()),
615 Operation::Merge { .. }
616 | Operation::Project { .. }
617 | Operation::Overwrite { .. }
618 | Operation::Restore { .. } => Err(Error::invalid_input(
619 "Detached commits cannot change schema after non-reusable field IDs are activated",
620 )),
621 _ => Ok(()),
622 }
623}
624
625fn overwrite_fragments_valid(fragments: &[Fragment]) -> Result<()> {
630 for fragment in fragments {
631 if let Some(deletion_file) = &fragment.deletion_file {
632 return Err(Error::invalid_input(format!(
633 "Overwrite fragments must be newly written, but fragment {} carries \
634 deletion file {}. Use Delete to commit deletions against existing \
635 fragments, or Merge to change their schema.",
636 fragment.id,
637 relative_deletion_file_path(fragment.id, deletion_file)
638 )));
639 }
640 }
641 Ok(())
642}
643
644fn schema_fragments_valid(
645 manifest: Option<&Manifest>,
646 schema: &Schema,
647 fragments: &[Fragment],
648) -> Result<()> {
649 if let Some(manifest) = manifest {
650 return match manifest.data_storage_format.lance_file_format() {
651 ConcreteFileVersion::V1 => schema_fragments_legacy_valid(schema, fragments),
652 ConcreteFileVersion::V2_0
653 | ConcreteFileVersion::V2_1
654 | ConcreteFileVersion::V2_2
655 | ConcreteFileVersion::V2_3 => schema_fragments_modern_valid(schema, fragments),
656 };
657 }
658 schema_fragments_modern_valid(schema, fragments)
659}
660
661pub fn schema_fragments_modern_valid(_schema: &Schema, fragments: &[Fragment]) -> Result<()> {
662 for fragment in fragments {
664 for data_file in &fragment.files {
665 if data_file.fields.iter().len() == 0 {
666 return Err(Error::invalid_input(format!(
667 "Datafile {} does not contain any fields",
668 data_file.path
669 )));
670 }
671 }
672 }
673 Ok(())
674}
675
676pub fn schema_fragments_legacy_valid(schema: &Schema, fragments: &[Fragment]) -> Result<()> {
680 for fragment in fragments {
683 for field in schema.fields_pre_order() {
684 if !fragment
685 .files
686 .iter()
687 .flat_map(|f| f.fields.iter())
688 .any(|f_id| f_id == &field.id)
689 {
690 return Err(Error::invalid_input(format!(
691 "Fragment {} does not contain field {:?}",
692 fragment.id, field
693 )));
694 }
695 }
696 }
697 Ok(())
698}
699
700#[inline]
708pub(super) fn merge_fragment_physically_rewritten(prev: &Fragment, merged: &Fragment) -> bool {
709 debug_assert_eq!(prev.id, merged.id);
710 if prev.files.len() != merged.files.len() {
711 return true;
712 }
713 prev.files.iter().zip(merged.files.iter()).any(|(p, m)| {
717 p.path != m.path
718 || p.fields != m.fields
719 || p.column_indices != m.column_indices
720 || p.file_major_version != m.file_major_version
721 || p.file_minor_version != m.file_minor_version
722 || p.base_id != m.base_id
723 })
724}
725
726fn merge_fragments_valid(manifest: &Manifest, new_fragments: &[Fragment]) -> Result<()> {
730 let original_fragments = manifest.fragments.as_ref();
731
732 if new_fragments.len() < original_fragments.len() {
734 return Err(Error::invalid_input(format!(
735 "Merge operation reduced fragment count from {} to {}. \
736 Merge operations should only add columns, not reduce fragments.",
737 original_fragments.len(),
738 new_fragments.len()
739 )));
740 }
741
742 let mut new_fragment_map: HashMap<u64, &Fragment> = HashMap::new();
747 for fragment in new_fragments {
748 if fragment.id != 0 && new_fragment_map.contains_key(&fragment.id) {
749 return Err(Error::invalid_input(format!(
750 "Merge operation contains duplicate fragment id {}. \
751 New fragments must use id 0 (assigned at commit time) or a unique reserved id.",
752 fragment.id
753 )));
754 }
755 new_fragment_map.entry(fragment.id).or_insert(fragment);
756 }
757
758 let previous_ids: HashSet<u64> = original_fragments.iter().map(|f| f.id).collect();
765 let mut seen_previous: HashSet<u64> = HashSet::new();
766 let mut first_new_id: Option<u64> = None;
767 for fragment in new_fragments {
768 let is_existing = previous_ids.contains(&fragment.id) && seen_previous.insert(fragment.id);
769 if !is_existing {
770 first_new_id.get_or_insert(fragment.id);
771 } else if let Some(new_id) = first_new_id {
772 return Err(Error::invalid_input(format!(
773 "Merge operation lists existing fragment {} after a new fragment \
774 (id {}). New fragments must be listed after every existing fragment.",
775 fragment.id, new_id
776 )));
777 }
778 }
779
780 let mut missing_fragments: Vec<u64> = Vec::new();
783 for original_fragment in original_fragments {
784 if let Some(new_fragment) = new_fragment_map.get(&original_fragment.id) {
785 if original_fragment.physical_rows != new_fragment.physical_rows {
787 return Err(Error::invalid_input(format!(
788 "Merge operation changed row count for fragment {}. \
789 Original: {:?}, New: {:?}. \
790 Merge operations should preserve fragment row counts and only add new columns.",
791 original_fragment.id,
792 original_fragment.physical_rows,
793 new_fragment.physical_rows
794 )));
795 }
796 if manifest.uses_stable_row_ids() && new_fragment.row_id_meta.is_none() {
799 return Err(Error::invalid_input(format!(
800 "Merge operation dropped row id metadata for existing fragment {}. \
801 New fragments (id 0) must be listed after the existing fragments; \
802 they are assigned row ids at commit time.",
803 original_fragment.id
804 )));
805 }
806 } else {
807 missing_fragments.push(original_fragment.id);
808 }
809 }
810
811 if !missing_fragments.is_empty() {
812 return Err(Error::invalid_input(format!(
813 "Merge operation is missing original fragments: {:?}. \
814 Merge operations should preserve all original fragments and only add new columns. \
815 Expected fragments: {:?}, but got: {:?}",
816 missing_fragments,
817 original_fragments.iter().map(|f| f.id).collect::<Vec<_>>(),
818 new_fragment_map.keys().copied().collect::<Vec<_>>()
819 )));
820 }
821
822 Ok(())
823}
824
825fn merge_schema_valid(
840 manifest: &Manifest,
841 new_schema: &Schema,
842 fragments: &[Fragment],
843) -> Result<()> {
844 let prior_schema = &manifest.schema;
845 let mut new_fragment_map: HashMap<u64, &Fragment> = HashMap::new();
848 for fragment in fragments {
849 new_fragment_map.entry(fragment.id).or_insert(fragment);
850 }
851
852 for field in new_schema.fields_pre_order() {
855 let Some(prior_field) = prior_schema.field_by_id(field.id) else {
856 continue;
857 };
858 let prior_parent = prior_schema
860 .field_ancestry_by_id(field.id)
861 .ok_or_else(|| {
862 Error::internal(format!(
863 "Missing prior schema ancestry for field id {}",
864 field.id
865 ))
866 })?
867 .into_iter()
868 .rev()
869 .nth(1);
870 let new_parent = new_schema
871 .field_ancestry_by_id(field.id)
872 .ok_or_else(|| {
873 Error::internal(format!(
874 "Missing proposed schema ancestry for field id {}",
875 field.id
876 ))
877 })?
878 .into_iter()
879 .rev()
880 .nth(1);
881 let prior_parent_id = prior_parent.map_or(-1, |parent| parent.id);
882 let new_parent_id = new_parent.map_or(-1, |parent| parent.id);
883 if prior_field.parent_id != prior_parent_id || field.parent_id != new_parent_id {
884 return Err(Error::invalid_input(format!(
885 "Merge operation has inconsistent parent ids for field id {} (\"{}\"): \
886 prior field declares {} but its schema ancestry has {}; proposed field \
887 declares {} but its schema ancestry has {}.",
888 field.id,
889 new_schema.field_path(field.id)?,
890 prior_field.parent_id,
891 prior_parent_id,
892 field.parent_id,
893 new_parent_id
894 )));
895 }
896 let prior_siblings = prior_parent.map_or(prior_schema.fields.as_slice(), |parent| {
898 parent.children.as_slice()
899 });
900 if new_parent_id != prior_parent_id
901 || prior_siblings
902 .iter()
903 .any(|sibling| sibling.name == field.name && sibling.id != field.id)
904 {
905 return Err(Error::invalid_input(format!(
906 "Merge operation remaps field id {} from \"{}\" to \"{}\". \
907 Merge must preserve the dataset's field ids: derive the new schema \
908 from the dataset's current schema instead of renumbering fields.",
909 field.id,
910 prior_schema.field_path(field.id)?,
911 new_schema.field_path(field.id)?
912 )));
913 }
914 if let Some(changes) = shared_field_binding_changes(prior_field, field)
915 && !is_field_binding_fully_rewritten(manifest, &new_fragment_map, field.id)
916 {
917 return Err(Error::invalid_input(format!(
918 "Merge operation changes field id {} (\"{}\") without rewriting it in \
919 every existing fragment: {}. Merge must preserve each existing field's \
920 logical type, nullability, storage encoding, and dictionary unless all \
921 existing base and overlay files carrying that field are replaced.",
922 field.id,
923 new_schema.field_path(field.id)?,
924 changes
925 )));
926 }
927 }
928
929 let max_field_id = manifest.max_field_id();
930 for field in new_schema.fields_pre_order() {
931 if prior_schema.field_by_id(field.id).is_none() && field.id <= max_field_id {
932 let next_id_msg = match max_field_id.checked_add(1) {
933 Some(next_id) => format!("New fields must use ids of at least {}.", next_id),
934 None => {
935 "No further field id can be allocated because ids are exhausted.".to_string()
936 }
937 };
938 return Err(Error::invalid_input(format!(
939 "Merge operation assigns id {} to new field \"{}\", but ids up to {} are \
940 already used by current or dropped fields. {}",
941 field.id,
942 new_schema.field_path(field.id)?,
943 max_field_id,
944 next_id_msg
945 )));
946 }
947 }
948
949 let mut prior_paths = HashMap::with_capacity(prior_schema.fields_pre_order().count());
950 for field in prior_schema.fields_pre_order() {
951 prior_paths.insert(prior_schema.field_path(field.id)?, field);
952 }
953 for field in new_schema.fields_pre_order() {
954 if prior_schema.field_by_id(field.id).is_some() {
955 continue;
956 }
957 let new_path = new_schema.field_path(field.id)?;
958 let Some(prior_field) = prior_paths.get(&new_path) else {
959 continue;
960 };
961 let materialized = fragments.iter().all(|fragment| {
962 fragment
963 .files
964 .iter()
965 .any(|file| file_materializes_field(file.fields.as_ref(), field))
966 });
967 if !materialized {
968 return Err(Error::invalid_input(format!(
969 "Merge operation remaps existing field \"{}\" from id {} to id {} without \
970 rewriting its data. Every proposed fragment must materialize the new field \
971 id in a base data file.",
972 new_path, prior_field.id, field.id
973 )));
974 }
975 }
976
977 Ok(())
978}
979
980fn file_materializes_field(file_field_ids: &[i32], field: &Field) -> bool {
981 file_field_ids.contains(&field.id)
982 || field
983 .children
984 .iter()
985 .any(|child| file_materializes_field(file_field_ids, child))
986}
987
988fn is_field_binding_fully_rewritten(
989 manifest: &Manifest,
990 new_fragment_map: &HashMap<u64, &Fragment>,
991 field_id: i32,
992) -> bool {
993 manifest.fragments.iter().all(|prior_fragment| {
994 let Some(new_fragment) = new_fragment_map.get(&prior_fragment.id) else {
995 return false;
996 };
997
998 let is_materialized = new_fragment
999 .files
1000 .iter()
1001 .any(|file| file.fields.contains(&field_id));
1002 if !is_materialized {
1003 return false;
1004 }
1005
1006 prior_fragment
1007 .referenced_lance_files()
1008 .filter(|file| file.fields.contains(&field_id))
1009 .all(|prior_file| {
1010 !new_fragment.referenced_lance_files().any(|new_file| {
1011 new_file.fields.contains(&field_id)
1012 && new_file.base_id == prior_file.base_id
1013 && new_file.path == prior_file.path
1014 })
1015 })
1016 })
1017}
1018
1019fn shared_field_binding_changes(prior: &Field, new: &Field) -> Option<String> {
1020 let mut changes = Vec::with_capacity(4);
1021 if prior.logical_type != new.logical_type {
1022 changes.push(format!(
1023 "logical type {} -> {}",
1024 prior.logical_type, new.logical_type
1025 ));
1026 }
1027 if prior.nullable != new.nullable {
1028 changes.push(format!("nullable {} -> {}", prior.nullable, new.nullable));
1029 }
1030 if prior.encoding != new.encoding {
1031 changes.push(format!(
1032 "storage encoding {:?} -> {:?}",
1033 prior.encoding, new.encoding
1034 ));
1035 }
1036 if prior.dictionary != new.dictionary {
1037 changes.push("dictionary".to_string());
1038 }
1039 if changes.is_empty() {
1040 None
1041 } else {
1042 Some(changes.join(", "))
1043 }
1044}
1045
1046#[cfg(test)]
1047mod tests {
1048 use super::*;
1049 use crate::format::overlay::{DataOverlayFile, OverlayCoverage};
1050 use crate::format::{DataFile, DataStorageFormat};
1051 use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema};
1052 use lance_core::datatypes::{Field as LanceCoreField, LogicalType, Schema as LanceSchema};
1053 use roaring::RoaringBitmap;
1054 use std::collections::HashMap;
1055 use std::sync::Arc;
1056
1057 #[test]
1058 fn test_merge_fragments_valid() {
1059 let schema = ArrowSchema::new(vec![
1061 ArrowField::new("id", DataType::Int32, false),
1062 ArrowField::new("name", DataType::Utf8, false),
1063 ]);
1064
1065 let original_fragments = vec![Fragment::new(1), Fragment::new(2), Fragment::new(3)];
1067
1068 let manifest = Manifest::new(
1070 LanceSchema::try_from(&schema).unwrap(),
1071 Arc::new(original_fragments),
1072 DataStorageFormat::new(ConcreteFileVersion::V2_0),
1073 HashMap::new(),
1074 );
1075
1076 let empty_fragments = vec![];
1078 let result = merge_fragments_valid(&manifest, &empty_fragments);
1079 assert!(result.is_err());
1080 assert!(
1081 result
1082 .unwrap_err()
1083 .to_string()
1084 .contains("reduced fragment count")
1085 );
1086
1087 let missing_fragments = vec![
1089 Fragment::new(1),
1090 Fragment::new(2),
1091 Fragment::new(4), ];
1094 let result = merge_fragments_valid(&manifest, &missing_fragments);
1095 assert!(result.is_err());
1096 assert!(
1097 result
1098 .unwrap_err()
1099 .to_string()
1100 .contains("missing original fragments")
1101 );
1102
1103 let reduced_fragments = vec![
1105 Fragment::new(1),
1106 Fragment::new(2),
1107 ];
1109 let result = merge_fragments_valid(&manifest, &reduced_fragments);
1110 assert!(result.is_err());
1111 assert!(
1112 result
1113 .unwrap_err()
1114 .to_string()
1115 .contains("reduced fragment count")
1116 );
1117
1118 let valid_fragments = vec![
1120 Fragment::new(1),
1121 Fragment::new(2),
1122 Fragment::new(3),
1123 Fragment::new(4), Fragment::new(5), ];
1126 let result = merge_fragments_valid(&manifest, &valid_fragments);
1127 assert!(result.is_ok());
1128
1129 let same_fragments = vec![Fragment::new(1), Fragment::new(2), Fragment::new(3)];
1131 let result = merge_fragments_valid(&manifest, &same_fragments);
1132 assert!(result.is_ok());
1133 }
1134
1135 fn one_field_schema() -> LanceSchema {
1136 LanceSchema::try_from(&ArrowSchema::new(vec![ArrowField::new(
1137 "a",
1138 DataType::Int32,
1139 true,
1140 )]))
1141 .unwrap()
1142 }
1143
1144 fn fragment_with_file_fields(id: u64, path: &str, fields: Vec<i32>) -> Fragment {
1145 let mut fragment = Fragment::new(id);
1146 fragment
1147 .files
1148 .push(DataFile::new_legacy_from_fields(path, fields, None));
1149 fragment
1150 }
1151
1152 fn manifest_with_file_fields(schema: LanceSchema, fields: Vec<i32>) -> Manifest {
1153 Manifest::new(
1154 schema,
1155 Arc::new(vec![fragment_with_file_fields(0, "f.lance", fields)]),
1156 DataStorageFormat::new(ConcreteFileVersion::V2_0),
1157 HashMap::new(),
1158 )
1159 }
1160
1161 fn activated_manifest() -> Manifest {
1162 let schema = one_field_schema();
1163 let mut manifest = manifest_with_file_fields(schema, vec![0]);
1164 manifest.activate_non_reusable_field_ids();
1165 manifest
1166 }
1167
1168 #[test]
1169 fn non_reusable_field_ids_allow_reserved_ids_above_high_water_mark() {
1170 let mut manifest = activated_manifest();
1171 manifest.max_allocated_field_id = Some(5);
1172 let mut schema = manifest.schema.clone();
1173 let mut new_field =
1174 LanceCoreField::try_from(&ArrowField::new("b", DataType::Int32, true)).unwrap();
1175 new_field.id = 6;
1176 schema.fields.push(new_field);
1177 let valid = Operation::Project {
1178 schema: schema.clone(),
1179 preserves_nullability: true,
1180 };
1181 validate_operation(Some(&manifest), &valid).unwrap();
1182
1183 schema.fields.last_mut().unwrap().id = 7;
1184 let reserved = Operation::Project {
1185 schema: schema.clone(),
1186 preserves_nullability: true,
1187 };
1188 validate_operation(Some(&manifest), &reserved).unwrap();
1189
1190 schema.fields.last_mut().unwrap().id = 5;
1191 let reused = Operation::Project {
1192 schema,
1193 preserves_nullability: true,
1194 };
1195 let err = validate_operation(Some(&manifest), &reused).unwrap_err();
1196 assert!(matches!(err, Error::InvalidInput { .. }), "{err}");
1197 assert!(
1198 err.to_string()
1199 .contains("greater than the high-water mark 5"),
1200 "{err}"
1201 );
1202 }
1203
1204 #[test]
1205 fn non_reusable_field_ids_require_fresh_identity_for_type_replacement() {
1206 let manifest = activated_manifest();
1207 let mut schema = manifest.schema.clone();
1208 schema.fields[0].logical_type = LogicalType::try_from(&DataType::Float32).unwrap();
1209 let operation = Operation::Project {
1210 schema,
1211 preserves_nullability: true,
1212 };
1213
1214 let err = validate_operation(Some(&manifest), &operation).unwrap_err();
1215
1216 assert!(err.to_string().contains("type replacement"), "{err}");
1217 }
1218
1219 #[test]
1220 fn non_reusable_field_ids_reject_overwrite_with_existing_ids() {
1221 let manifest = activated_manifest();
1222 let schema = manifest.schema.clone();
1223 let operation = Operation::Overwrite {
1224 fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0])],
1225 schema,
1226 config_upsert_values: None,
1227 initial_bases: None,
1228 };
1229
1230 for err in [
1231 validate_operation(Some(&manifest), &operation).unwrap_err(),
1232 validate_non_reusable_field_id_transition(&manifest, &manifest, &operation)
1233 .unwrap_err(),
1234 ] {
1235 assert!(matches!(err, Error::InvalidInput { .. }));
1236 assert!(
1237 err.to_string()
1238 .contains("greater than the high-water mark 0")
1239 );
1240 }
1241 }
1242
1243 #[test]
1244 fn canonicalize_overwrite_remaps_hostile_arrow_field_ids() {
1245 let manifest = activated_manifest();
1246 let mut schema = one_field_schema();
1247 schema.fields[0].id = i32::MAX;
1248 let mut operation = Operation::Overwrite {
1249 fragments: vec![fragment_with_file_fields(0, "new.lance", vec![i32::MAX])],
1250 schema,
1251 config_upsert_values: None,
1252 initial_bases: None,
1253 };
1254
1255 canonicalize_non_reusable_field_ids(Some(&manifest), &mut operation, None).unwrap();
1256
1257 let Operation::Overwrite {
1258 schema, fragments, ..
1259 } = operation
1260 else {
1261 unreachable!();
1262 };
1263 assert_eq!(schema.fields[0].id, 1);
1264 assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]);
1265 }
1266
1267 #[rstest::rstest]
1268 #[case::last_id(i32::MAX - 1)]
1269 #[case::exhausted(i32::MAX)]
1270 fn overwrite_cannot_keep_existing_ids_to_avoid_exhaustion(#[case] max_id: i32) {
1271 let mut manifest = activated_manifest();
1272 manifest.max_allocated_field_id = Some(max_id);
1273 let mut operation = Operation::Overwrite {
1274 schema: manifest.schema.clone(),
1275 fragments: vec![],
1276 config_upsert_values: None,
1277 initial_bases: None,
1278 };
1279 let result = canonicalize_non_reusable_field_ids(Some(&manifest), &mut operation, None);
1280 if max_id == i32::MAX {
1281 let err = result.unwrap_err();
1282 assert!(matches!(err, Error::InvalidInput { .. }));
1283 assert!(err.to_string().contains("IDs are exhausted"));
1284 } else {
1285 result.unwrap();
1286 let Operation::Overwrite { schema, .. } = operation else {
1287 unreachable!();
1288 };
1289 assert_eq!(schema.fields[0].id, i32::MAX);
1290 }
1291 }
1292
1293 #[test]
1294 fn canonicalize_overwrite_remaps_standalone_fragment_field_ids() {
1295 let manifest = activated_manifest();
1296 let mut schema = one_field_schema();
1297 schema.fields[0].id = -1;
1298 let mut operation = Operation::Overwrite {
1299 fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0])],
1300 schema,
1301 config_upsert_values: None,
1302 initial_bases: None,
1303 };
1304
1305 resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap();
1306
1307 let Operation::Overwrite {
1308 schema, fragments, ..
1309 } = operation
1310 else {
1311 unreachable!();
1312 };
1313 assert_eq!(schema.fields[0].id, 1);
1314 assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]);
1315 }
1316
1317 #[rstest::rstest]
1318 fn canonicalize_overwrite_retry_assigns_all_ids_above_latest_mark(
1319 #[values(false, true)] replace: bool,
1320 ) {
1321 let read = activated_manifest();
1322 let mut staged = read.schema.clone();
1323 let mut new_field = Field::new_arrow("new_column", DataType::Int32, true).unwrap();
1324 new_field.id = 1;
1325 staged.fields.push(new_field.clone());
1326 let mut latest_schema = read.schema.clone();
1327 if replace {
1328 latest_schema.fields.clear();
1329 }
1330 new_field.name = "concurrent_column".to_string();
1331 latest_schema.fields.push(new_field);
1332 let latest_ids = latest_schema.field_ids().into_iter().collect();
1333 let mut latest = manifest_with_file_fields(latest_schema, latest_ids);
1334 latest.activate_non_reusable_field_ids();
1335 let mut operation = Operation::Overwrite {
1336 schema: staged,
1337 fragments: vec![fragment_with_file_fields(0, "new.lance", vec![0, 1])],
1338 config_upsert_values: None,
1339 initial_bases: None,
1340 };
1341
1342 canonicalize_non_reusable_field_ids(Some(&latest), &mut operation, Some(&read.schema))
1343 .unwrap();
1344 validate_operation(Some(&latest), &operation).unwrap();
1345
1346 let Operation::Overwrite {
1347 schema, fragments, ..
1348 } = operation
1349 else {
1350 unreachable!();
1351 };
1352 assert_eq!(
1353 schema
1354 .fields
1355 .iter()
1356 .map(|field| field.id)
1357 .collect::<Vec<_>>(),
1358 vec![2, 3]
1359 );
1360 assert_eq!(fragments[0].files[0].fields.as_ref(), &[2, 3]);
1361 assert_eq!(schema.fields[1].name, "new_column");
1362 }
1363
1364 #[test]
1365 fn canonicalize_raw_arrow_overwrite_assigns_new_ids_in_schema_order() {
1366 let schema = LanceSchema::try_from(&ArrowSchema::new(vec![
1367 ArrowField::new("a", DataType::Int32, true),
1368 ArrowField::new("b", DataType::Int32, true),
1369 ]))
1370 .unwrap();
1371 let mut manifest = manifest_with_file_fields(schema, vec![0, 1]);
1372 manifest.activate_non_reusable_field_ids();
1373 let mut raw_schema = LanceSchema::try_from(&ArrowSchema::new(vec![
1374 ArrowField::new("b", DataType::Int32, true),
1375 ArrowField::new("a", DataType::Int32, true),
1376 ]))
1377 .unwrap();
1378 raw_schema
1379 .metadata
1380 .insert("source".to_string(), "user metadata".to_string());
1381 let mut operation = Operation::Overwrite {
1382 fragments: vec![],
1383 schema: raw_schema,
1384 config_upsert_values: None,
1385 initial_bases: None,
1386 };
1387
1388 resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap();
1389
1390 let Operation::Overwrite { schema, .. } = operation else {
1391 unreachable!();
1392 };
1393 assert_eq!(schema.field("b").unwrap().id, 2);
1394 assert_eq!(schema.field("a").unwrap().id, 3);
1395 assert_eq!(schema.metadata.get("source").unwrap(), "user metadata");
1396 }
1397
1398 #[rstest::rstest]
1399 #[case::canonical_ids([1, 2])]
1400 #[case::duplicate_arrow_metadata([42, 42])]
1401 fn canonicalize_overwrite_preserves_already_canonical_fragment_field_ids(
1402 #[case] input_ids: [i32; 2],
1403 ) {
1404 let manifest = activated_manifest();
1405 let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![
1406 ArrowField::new("b", DataType::Int32, true),
1407 ArrowField::new("c", DataType::Int32, true),
1408 ]))
1409 .unwrap();
1410 schema.fields[0].id = input_ids[0];
1411 schema.fields[1].id = input_ids[1];
1412 let mut operation = Operation::Overwrite {
1413 fragments: vec![fragment_with_file_fields(0, "new.lance", vec![1, 2])],
1414 schema,
1415 config_upsert_values: None,
1416 initial_bases: None,
1417 };
1418
1419 resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap();
1420
1421 let Operation::Overwrite {
1422 schema, fragments, ..
1423 } = operation
1424 else {
1425 unreachable!();
1426 };
1427 assert_eq!(
1428 schema
1429 .fields_pre_order()
1430 .map(|field| field.id)
1431 .collect::<Vec<_>>(),
1432 vec![1, 2]
1433 );
1434 assert_eq!(fragments[0].files[0].fields.as_ref(), &[1, 2]);
1435 }
1436
1437 #[test]
1438 fn canonicalize_overwrite_uses_one_raw_namespace_for_split_files() {
1439 let manifest = activated_manifest();
1440 let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![
1441 ArrowField::new("b", DataType::Int32, true),
1442 ArrowField::new("c", DataType::Int32, true),
1443 ]))
1444 .unwrap();
1445 for field in &mut schema.fields {
1446 field.id = -1;
1447 }
1448 let mut fragment = fragment_with_file_fields(0, "b.lance", vec![0]);
1449 fragment
1450 .files
1451 .push(DataFile::new_legacy_from_fields("c.lance", vec![1], None));
1452 let mut operation = Operation::Overwrite {
1453 fragments: vec![fragment],
1454 schema,
1455 config_upsert_values: None,
1456 initial_bases: None,
1457 };
1458
1459 resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap();
1460
1461 let Operation::Overwrite {
1462 schema, fragments, ..
1463 } = operation
1464 else {
1465 unreachable!();
1466 };
1467 assert_eq!(
1468 schema
1469 .fields_pre_order()
1470 .map(|field| field.id)
1471 .collect::<Vec<_>>(),
1472 vec![1, 2]
1473 );
1474 assert_eq!(fragments[0].files[0].fields.as_ref(), &[1]);
1475 assert_eq!(fragments[0].files[1].fields.as_ref(), &[2]);
1476 }
1477
1478 #[test]
1479 fn canonicalize_overwrite_rejects_ambiguous_raw_field_ids() {
1480 let manifest = activated_manifest();
1481 let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![
1482 ArrowField::new("b", DataType::Int32, true),
1483 ArrowField::new("c", DataType::Int32, true),
1484 ]))
1485 .unwrap();
1486 schema.fields[0].id = 2;
1487 schema.fields[1].id = 1;
1488 let mut operation = Operation::Overwrite {
1489 fragments: vec![fragment_with_file_fields(0, "new.lance", vec![2, 1])],
1490 schema,
1491 config_upsert_values: None,
1492 initial_bases: None,
1493 };
1494
1495 let err = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err();
1496
1497 assert!(err.to_string().contains("ambiguous raw Arrow field IDs"));
1498 }
1499
1500 #[test]
1501 fn canonicalize_raw_arrow_project_rejects_unmatched_field() {
1502 let manifest = activated_manifest();
1503 let mut schema = one_field_schema();
1504 schema.fields[0].name = "renamed".to_string();
1505 schema.fields[0].id = -1;
1506 let mut operation = Operation::Project {
1507 schema,
1508 preserves_nullability: true,
1509 };
1510
1511 let err = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err();
1512
1513 assert!(err.to_string().contains("writes no data"), "{err}");
1514 }
1515
1516 #[test]
1517 fn canonicalize_raw_arrow_project_preserves_explicit_existing_identity() {
1518 let manifest = activated_manifest();
1519 let mut schema = one_field_schema();
1520 schema.fields[0].name = "renamed".to_string();
1521 let mut operation = Operation::Project {
1522 schema,
1523 preserves_nullability: true,
1524 };
1525
1526 resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap();
1527
1528 let Operation::Project { schema, .. } = operation else {
1529 unreachable!();
1530 };
1531 assert_eq!(schema.fields[0].name, "renamed");
1532 assert_eq!(schema.fields[0].id, 0);
1533 }
1534
1535 #[test]
1536 fn canonicalize_raw_arrow_schema_for_legacy_dataset() {
1537 let manifest = manifest_with_file_fields(one_field_schema(), vec![0]);
1538 let mut schema = one_field_schema();
1539 schema.fields[0].id = -1;
1540 let mut operation = Operation::Project {
1541 schema,
1542 preserves_nullability: true,
1543 };
1544
1545 resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap();
1546
1547 let Operation::Project { schema, .. } = operation else {
1548 unreachable!();
1549 };
1550 assert_eq!(schema.fields[0].id, 0);
1551 assert!(schema.metadata.is_empty());
1552 }
1553
1554 #[test]
1555 fn canonicalize_merge_preserves_identity_for_physical_rewrite() {
1556 let mut manifest = activated_manifest();
1557 Arc::make_mut(&mut manifest.fragments).push(fragment_with_file_fields(
1558 1,
1559 "retained.lance",
1560 vec![0],
1561 ));
1562 let rewritten_fragment = fragment_with_file_fields(0, "replacement.lance", vec![0]);
1563 let mut operation = Operation::Merge {
1564 fragments: vec![rewritten_fragment, manifest.fragments[1].clone()],
1565 schema: manifest.schema.clone(),
1566 preserves_nullability: true,
1567 };
1568
1569 canonicalize_non_reusable_field_ids(Some(&manifest), &mut operation, None).unwrap();
1570
1571 let Operation::Merge {
1572 schema, fragments, ..
1573 } = operation
1574 else {
1575 unreachable!();
1576 };
1577 assert_eq!(schema.fields[0].id, 0);
1578 assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]);
1579 assert_eq!(fragments[1].files[0].fields.as_ref(), &[0]);
1580 validate_operation(
1581 Some(&manifest),
1582 &Operation::Merge {
1583 fragments,
1584 schema,
1585 preserves_nullability: true,
1586 },
1587 )
1588 .unwrap();
1589 }
1590
1591 #[test]
1592 fn canonicalize_merge_allocates_fresh_identity_for_overlaid_column() {
1593 let manifest = activated_manifest();
1594 let mut merged_fragment = manifest.fragments[0].clone();
1595 merged_fragment.files.push(DataFile::new_legacy_from_fields(
1596 "replacement.lance",
1597 vec![0],
1598 None,
1599 ));
1600 let mut operation = Operation::Merge {
1601 fragments: vec![merged_fragment],
1602 schema: manifest.schema.clone(),
1603 preserves_nullability: true,
1604 };
1605
1606 canonicalize_non_reusable_field_ids(Some(&manifest), &mut operation, None).unwrap();
1607
1608 let Operation::Merge {
1609 schema, fragments, ..
1610 } = operation
1611 else {
1612 unreachable!();
1613 };
1614 assert_eq!(schema.fields[0].id, 1);
1615 assert_eq!(fragments[0].files[0].fields.as_ref(), &[0]);
1616 assert_eq!(fragments[0].files[1].fields.as_ref(), &[1]);
1617 }
1618
1619 #[test]
1620 fn canonicalize_merge_rejects_ambiguous_raw_field_ids() {
1621 let manifest = activated_manifest();
1622 let mut schema = LanceSchema::try_from(&ArrowSchema::new(vec![
1623 ArrowField::new("a", DataType::Int32, true),
1624 ArrowField::new("b", DataType::Int32, true),
1625 ArrowField::new("c", DataType::Int32, true),
1626 ]))
1627 .unwrap();
1628 schema.fields[0].id = 0;
1629 schema.fields[1].id = 2;
1630 schema.fields[2].id = 1;
1631 let mut merged_fragment = manifest.fragments[0].clone();
1632 merged_fragment.files.push(DataFile::new_legacy_from_fields(
1633 "new.lance",
1634 vec![1, 2],
1635 None,
1636 ));
1637 let mut operation = Operation::Merge {
1638 fragments: vec![merged_fragment],
1639 schema,
1640 preserves_nullability: true,
1641 };
1642
1643 let err = resolve_arrow_field_ids(Some(&manifest), &mut operation).unwrap_err();
1644
1645 assert!(err.to_string().contains("ambiguous raw Arrow field IDs"));
1646 }
1647
1648 #[test]
1649 fn non_reusable_field_id_manifest_rejects_high_water_mark_below_overlay_reference() {
1650 let mut manifest = activated_manifest();
1651 Arc::make_mut(&mut manifest.fragments)[0]
1652 .overlays
1653 .push(DataOverlayFile {
1654 data_file: DataFile::new_legacy_from_fields("overlay.lance", vec![7], None),
1655 coverage: OverlayCoverage::Shared(Arc::new(RoaringBitmap::from_iter([0_u32]))),
1656 committed_version: 1,
1657 });
1658
1659 let err = validate_non_reusable_field_id_manifest(&manifest).unwrap_err();
1660
1661 assert!(
1662 err.to_string().contains("below referenced field ID 7"),
1663 "{err}"
1664 );
1665 }
1666
1667 #[test]
1668 fn non_reusable_field_id_transition_rejects_file_only_allocator_advance() {
1669 let manifest = activated_manifest();
1670 let mut successor = Manifest::new_from_previous(
1671 &manifest,
1672 manifest.schema.clone(),
1673 Arc::new(vec![fragment_with_file_fields(0, "new.lance", vec![0, 1])]),
1674 );
1675 successor.max_allocated_field_id = manifest.max_allocated_field_id;
1676 let operation = Operation::UpdateConfig {
1677 config_updates: None,
1678 table_metadata_updates: None,
1679 schema_metadata_updates: None,
1680 field_metadata_updates: HashMap::new(),
1681 };
1682
1683 let err = validate_non_reusable_field_id_transition(&manifest, &successor, &operation)
1684 .unwrap_err();
1685
1686 assert!(
1687 err.to_string().contains("canonical successor schema"),
1688 "{err}"
1689 );
1690 }
1691
1692 #[test]
1693 fn non_reusable_field_id_transition_rejects_decreasing_high_water_mark() {
1694 let manifest = activated_manifest();
1695 let mut successor = Manifest::new_from_previous(
1696 &manifest,
1697 manifest.schema.clone(),
1698 manifest.fragments.clone(),
1699 );
1700 successor.max_allocated_field_id = Some(manifest.max_field_id() - 1);
1701 let operation = Operation::UpdateConfig {
1702 config_updates: None,
1703 table_metadata_updates: None,
1704 schema_metadata_updates: None,
1705 field_metadata_updates: HashMap::new(),
1706 };
1707
1708 let err = validate_non_reusable_field_id_transition(&manifest, &successor, &operation)
1709 .unwrap_err();
1710
1711 assert!(
1712 err.to_string().contains("high-water mark decreases"),
1713 "{err}"
1714 );
1715 }
1716
1717 #[test]
1718 fn detached_non_reusable_field_ids_allow_data_only_merge_and_reject_schema_change() {
1719 let manifest = activated_manifest();
1720 let data_only = Operation::Merge {
1721 fragments: manifest.fragments.as_ref().clone(),
1722 schema: manifest.schema.clone(),
1723 preserves_nullability: true,
1724 };
1725 validate_detached_non_reusable_field_ids(&manifest, &data_only).unwrap();
1726
1727 let mut changed_schema = manifest.schema.clone();
1728 changed_schema.fields[0].name = "renamed".to_string();
1729 let schema_change = Operation::Project {
1730 schema: changed_schema,
1731 preserves_nullability: true,
1732 };
1733 let err = validate_detached_non_reusable_field_ids(&manifest, &schema_change).unwrap_err();
1734 assert!(
1735 err.to_string()
1736 .contains("Detached commits cannot change schema"),
1737 "{err}"
1738 );
1739 }
1740
1741 #[rstest::rstest]
1742 #[case::logical_type(DataType::Float32, true)]
1743 #[case::nullability(DataType::Int32, false)]
1744 #[test]
1745 fn test_merge_shared_id_change_requires_full_rewrite(
1746 #[case] data_type: DataType,
1747 #[case] nullable: bool,
1748 ) {
1749 let schema = one_field_schema();
1750 let prior_fragments = vec![
1751 fragment_with_file_fields(0, "old-0.lance", vec![0]),
1752 fragment_with_file_fields(1, "old-1.lance", vec![0]),
1753 ];
1754 let manifest = Manifest::new(
1755 schema.clone(),
1756 Arc::new(prior_fragments.clone()),
1757 DataStorageFormat::new(ConcreteFileVersion::V2_0),
1758 HashMap::new(),
1759 );
1760 let mut new_schema = schema;
1761 new_schema.fields[0].logical_type = LogicalType::try_from(&data_type).unwrap();
1762 new_schema.fields[0].nullable = nullable;
1763
1764 let rewritten_fragments = vec![
1765 fragment_with_file_fields(0, "new-0.lance", vec![0]),
1766 fragment_with_file_fields(1, "new-1.lance", vec![0]),
1767 ];
1768 merge_schema_valid(&manifest, &new_schema, &rewritten_fragments).unwrap();
1769
1770 let partially_rewritten = vec![rewritten_fragments[0].clone(), prior_fragments[1].clone()];
1771 let err = merge_schema_valid(&manifest, &new_schema, &partially_rewritten).unwrap_err();
1772 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
1773 assert!(
1774 err.to_string()
1775 .contains("without rewriting it in every existing fragment"),
1776 "unexpected error: {}",
1777 err
1778 );
1779 }
1780
1781 #[test]
1782 fn test_merge_preserved_dictionary_not_flagged() {
1783 let dictionary = lance_core::datatypes::Dictionary {
1789 offset: 100,
1790 length: 5,
1791 values: None,
1792 };
1793 let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ArrowField::new(
1794 "d",
1795 DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)),
1796 true,
1797 )]))
1798 .unwrap();
1799 let mut prior_schema = schema.clone();
1800 prior_schema.fields[0].dictionary = Some(dictionary.clone());
1801 let manifest = Manifest::new(
1802 prior_schema,
1803 Arc::new(vec![fragment_with_file_fields(0, "old.lance", vec![0])]),
1804 DataStorageFormat::new(ConcreteFileVersion::V2_0),
1805 HashMap::new(),
1806 );
1807
1808 let mut unchanged = schema.clone();
1809 unchanged.fields[0].dictionary = Some(dictionary);
1810 let retained = vec![fragment_with_file_fields(0, "old.lance", vec![0])];
1811 assert_eq!(
1812 shared_field_binding_changes(&manifest.schema.fields[0], &unchanged.fields[0]),
1813 None
1814 );
1815 merge_schema_valid(&manifest, &unchanged, &retained).unwrap();
1816
1817 let mut moved = schema;
1818 moved.fields[0].dictionary = Some(lance_core::datatypes::Dictionary {
1819 offset: 200,
1820 length: 5,
1821 values: None,
1822 });
1823 assert_eq!(
1824 shared_field_binding_changes(&manifest.schema.fields[0], &moved.fields[0]),
1825 Some("dictionary".to_string())
1826 );
1827 let err = merge_schema_valid(&manifest, &moved, &retained).unwrap_err();
1828 assert!(matches!(err, Error::InvalidInput { .. }), "got {err:?}");
1829 assert!(
1830 err.to_string()
1831 .contains(r#"changes field id 0 ("d") without rewriting it in every existing fragment: dictionary."#),
1832 "unexpected error: {err}"
1833 );
1834 }
1835
1836 #[test]
1837 fn test_merge_shared_id_change_rejects_retained_overlay() {
1838 let schema = one_field_schema();
1839 let mut prior_fragment = fragment_with_file_fields(0, "old.lance", vec![0]);
1840 prior_fragment.overlays.push(DataOverlayFile {
1841 data_file: DataFile::new_legacy_from_fields("old-overlay.lance", vec![0], None),
1842 coverage: OverlayCoverage::Shared(Arc::new(RoaringBitmap::from_iter([0_u32]))),
1843 committed_version: 1,
1844 });
1845 let manifest = Manifest::new(
1846 schema.clone(),
1847 Arc::new(vec![prior_fragment.clone()]),
1848 DataStorageFormat::new(ConcreteFileVersion::V2_0),
1849 HashMap::new(),
1850 );
1851 let mut new_schema = schema;
1852 new_schema.fields[0].nullable = false;
1853
1854 let mut rewritten = fragment_with_file_fields(0, "new.lance", vec![0]);
1855 rewritten.overlays = prior_fragment.overlays.clone();
1856 let err = merge_schema_valid(&manifest, &new_schema, &[rewritten]).unwrap_err();
1857 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
1858 assert!(
1859 err.to_string()
1860 .contains("without rewriting it in every existing fragment"),
1861 "unexpected error: {}",
1862 err
1863 );
1864 }
1865
1866 #[test]
1867 fn test_merge_allows_rewritten_fresh_field_id() {
1868 let schema = one_field_schema();
1869 let manifest = manifest_with_file_fields(schema.clone(), vec![0]);
1870 let mut rewritten_schema = schema;
1871 rewritten_schema.fields[0].id = 1;
1872 let mut rewritten = manifest.fragments[0].clone();
1873 rewritten.files[0] = DataFile::new_legacy_from_fields("rewritten.lance", vec![1], None);
1874 merge_schema_valid(&manifest, &rewritten_schema, &[rewritten]).unwrap();
1875 }
1876
1877 #[rstest::rstest]
1878 #[case::parent("s")]
1879 #[case::child("s.x")]
1880 #[test]
1881 fn test_merge_allows_id_preserving_rename(#[case] path: &str) {
1882 let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ArrowField::new(
1883 "s",
1884 DataType::Struct(
1885 vec![
1886 ArrowField::new("x", DataType::Int32, true),
1887 ArrowField::new("y", DataType::Int32, true),
1888 ]
1889 .into(),
1890 ),
1891 true,
1892 )]))
1893 .unwrap();
1894 let manifest = manifest_with_file_fields(schema.clone(), vec![1, 2]);
1895 let mut renamed = schema;
1896 let field_id = renamed.field(path).unwrap().id;
1897 renamed.mut_field_by_id(field_id).unwrap().name = "renamed".into();
1898 merge_schema_valid(&manifest, &renamed, &manifest.fragments).unwrap();
1899
1900 renamed.fields[0].children.swap(0, 1);
1902 renamed.fields[0].children[0].id = 1;
1903 renamed.fields[0].children[1].id = 2;
1904 let err = merge_schema_valid(&manifest, &renamed, &manifest.fragments).unwrap_err();
1905 assert!(matches!(err, Error::InvalidInput { .. }), "{err}");
1906 assert!(err.to_string().contains("remaps field id 1"), "{err}");
1907 }
1908
1909 #[rstest::rstest]
1910 #[case::dropped_parent_stale(false, false)]
1911 #[case::retained_parent_stale(true, false)]
1912 #[case::dropped_parent_updated(false, true)]
1913 #[case::retained_parent_updated(true, true)]
1914 #[test]
1915 fn test_merge_rejects_moving_field_to_unused_path(
1916 #[case] retain_parent: bool,
1917 #[case] update_parent_id: bool,
1918 ) {
1919 let schema = LanceSchema::try_from(&ArrowSchema::new(vec![ArrowField::new(
1920 "s",
1921 DataType::Struct(
1922 vec![
1923 ArrowField::new("x", DataType::Int32, true),
1924 ArrowField::new("y", DataType::Int32, true),
1925 ]
1926 .into(),
1927 ),
1928 true,
1929 )]))
1930 .unwrap();
1931 let manifest = manifest_with_file_fields(schema.clone(), vec![1, 2]);
1932 let mut moved = schema;
1933 let mut child = moved.fields[0].children.remove(0);
1934 if !retain_parent {
1935 moved.fields.clear();
1936 }
1937 child.name = "renamed".into();
1938 if update_parent_id {
1939 child.parent_id = -1;
1940 }
1941 moved.fields.push(child);
1942 moved.validate().unwrap();
1943
1944 let operation = Operation::Merge {
1945 schema: moved,
1946 fragments: manifest.fragments.as_ref().clone(),
1947 preserves_nullability: true,
1948 };
1949 let err = validate_operation(Some(&manifest), &operation).unwrap_err();
1950 assert!(matches!(err, Error::InvalidInput { .. }), "{err}");
1951 let expected = if update_parent_id {
1952 "remaps field id 1 from \"s.x\" to \"renamed\""
1953 } else {
1954 "inconsistent parent ids for field id 1"
1955 };
1956 assert!(err.to_string().contains(expected), "{err}");
1957 }
1958
1959 #[test]
1960 fn test_merge_rejects_max_field_id_overflow() {
1961 let schema = one_field_schema();
1962 let manifest = manifest_with_file_fields(schema.clone(), vec![0, i32::MAX]);
1963 assert_eq!(manifest.max_field_id(), i32::MAX);
1964
1965 let mut new_schema = schema;
1966 let mut extra =
1967 LanceCoreField::try_from(&ArrowField::new("b", DataType::Int32, true)).unwrap();
1968 extra.id = 1;
1969 new_schema.fields.push(extra);
1970
1971 let err = merge_schema_valid(&manifest, &new_schema, &manifest.fragments).unwrap_err();
1972 assert!(matches!(err, Error::InvalidInput { .. }), "got {:?}", err);
1973 let message = err.to_string();
1974 assert!(
1975 message.contains("assigns id 1 to new field \"b\"") && message.contains("exhausted"),
1976 "unexpected error: {}",
1977 message
1978 );
1979 }
1980
1981 #[test]
1987 fn test_overwrite_legacy_to_stable_with_struct_fields() {
1988 use arrow_schema::Fields;
1989
1990 let arrow_schema = ArrowSchema::new(vec![
1993 ArrowField::new("id", DataType::Int32, false),
1994 ArrowField::new("name", DataType::Utf8, false),
1995 ArrowField::new(
1996 "address",
1997 DataType::Struct(Fields::from(vec![
1998 ArrowField::new("city", DataType::Utf8, false),
1999 ArrowField::new("country", DataType::Utf8, false),
2000 ])),
2001 false,
2002 ),
2003 ]);
2004 let schema = LanceSchema::try_from(&arrow_schema).unwrap();
2005
2006 let legacy_manifest = Manifest::new(
2008 schema.clone(),
2009 Arc::new(vec![Fragment::new(0)]),
2010 DataStorageFormat::new(ConcreteFileVersion::V1),
2011 HashMap::new(),
2012 );
2013
2014 let stable_fragment = Fragment {
2017 id: 0,
2018 files: vec![DataFile::new(
2019 "data.lance",
2020 vec![0, 1, 3, 4], vec![0, 1, 2, 3],
2022 ConcreteFileVersion::V1,
2023 None,
2024 None,
2025 )],
2026 physical_rows: Some(10),
2027 overlays: vec![],
2028 deletion_file: None,
2029 row_id_meta: None,
2030 last_updated_at_version_meta: None,
2031 created_at_version_meta: None,
2032 };
2033
2034 let operation = Operation::Overwrite {
2035 fragments: vec![stable_fragment],
2036 schema,
2037 config_upsert_values: None,
2038 initial_bases: None,
2039 };
2040
2041 validate_operation(Some(&legacy_manifest), &operation).unwrap();
2044 }
2045}