diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index de14803cafa..606e3d8a94f 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -624,6 +624,10 @@ impl Transaction { for fragment in updated_fragments { updated_by_id.entry(fragment.id).or_insert(fragment); } + let live_field_ids = schema + .fields_pre_order() + .map(|field| field.id) + .collect::>(); let updated_frags: Vec = existing_fragments .iter() .filter_map(|f| { @@ -638,6 +642,15 @@ impl Transaction { // fields it rewrote, since the fresh base values // supersede them. updated.overlays = f.overlays.clone(); + // A concurrent Project may have pruned files from the + // current fragment after this post-image was staged. + // Match Project's rule: retain a file if any field in + // it remains live, including mixed live/dropped files. + updated.files.retain(|file| { + file.fields + .iter() + .any(|field_id| live_field_ids.contains(field_id)) + }); if matches!(update_mode, Some(RewriteColumns)) { crate::format::overlay::tombstone_overlay_fields( &mut updated.overlays, @@ -728,6 +741,14 @@ impl Transaction { let mut new_fragments = Self::fragments_with_ids(new_fragments.clone(), &mut fragment_id) .collect::>(); + // New fragments were staged against the same pre-Project schema. + for fragment in &mut new_fragments { + fragment.files.retain(|file| { + file.fields + .iter() + .any(|field_id| live_field_ids.contains(field_id)) + }); + } // Assign row IDs to any fragments that don't have them yet // (e.g., inserted rows from merge_insert operations) @@ -1755,11 +1776,70 @@ mod tests { assert_eq!(rows, vec![None, Some(42), None, Some(43)]); } + #[rstest::rstest] + #[case::unspecified(None)] + #[case::rewrite_rows(Some(UpdateMode::RewriteRows))] + #[case::rewrite_columns(Some(UpdateMode::RewriteColumns))] + fn test_update_build_manifest_does_not_restore_projected_files( + #[case] update_mode: Option, + ) { + let mut manifest = sample_manifest_with_fragments(0..3); + let projected_file = DataFile::new_legacy_from_fields("projected.lance", vec![0], None); + Arc::make_mut(&mut manifest.fragments)[1].files = vec![projected_file.clone()]; + + // Model an update staged before a projection removed field 1's file. + let mut updated = manifest.fragments[1].clone(); + updated.files.push(DataFile::new_legacy_from_fields( + "dropped.lance", + vec![1], + None, + )); + updated.physical_rows = Some(42); + let inserted_projected_file = + DataFile::new_legacy_from_fields("inserted-projected.lance", vec![0], None); + let mut inserted = Fragment::new(0); + inserted.files = vec![ + inserted_projected_file.clone(), + DataFile::new_legacy_from_fields("inserted-dropped.lance", vec![1], None), + ]; + + let transaction = Transaction::new( + manifest.version, + Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: vec![updated], + new_fragments: vec![inserted], + fields_modified: vec![1], + compacted_sstables: vec![], + fields_for_preserving_frag_bitmap: vec![], + update_mode, + inserted_rows_filter: None, + updated_fragment_offsets: None, + }, + None, + ); + + let (new_manifest, _) = transaction + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap(); + + let fragment = &new_manifest.fragments[1]; + assert_eq!(fragment.files, vec![projected_file]); + assert_eq!(fragment.physical_rows, Some(42)); + assert_eq!( + new_manifest.fragments[3].files, + vec![inserted_projected_file] + ); + assert_eq!(new_manifest.max_field_id(), 0); + } + #[test] fn test_delete_build_manifest_applies_deletion_to_current_fragment() { let mut manifest = sample_manifest_with_fragments(0..5); manifest.version = 2; + let projected_file = DataFile::new_legacy_from_fields("projected.lance", vec![0], None); let current_fragment = &mut Arc::make_mut(&mut manifest.fragments)[2]; + current_fragment.files = vec![projected_file.clone()]; current_fragment.physical_rows = Some(42); current_fragment.overlays = vec![overlay_with_field(0, 2)]; current_fragment.last_updated_at_version_meta = Some( @@ -1769,8 +1849,13 @@ mod tests { .unwrap(), ); + // Model a delete staged before a projection removed field 1's file. let mut updated2 = Fragment::new(2); updated2.physical_rows = Some(42); + updated2.files = vec![ + projected_file.clone(), + DataFile::new_legacy_from_fields("dropped.lance", vec![1], None), + ]; let deletion_file = DeletionFile { read_version: 1, id: 10, @@ -1796,16 +1881,17 @@ mod tests { let ids: Vec = new_manifest.fragments.iter().map(|f| f.id).collect(); assert_eq!(ids, vec![0, 2, 4]); - let rows: Vec> = new_manifest + let fragment2 = new_manifest .fragments .iter() - .map(|f| f.physical_rows) - .collect(); - assert_eq!(rows, vec![None, Some(42), None]); - let overlays = &new_manifest.fragments[1].overlays; - assert_eq!(overlays.len(), 1); - assert_eq!(overlays[0].committed_version, 2); - assert_eq!(new_manifest.fragments[1].deletion_file, Some(deletion_file)); + .find(|fragment| fragment.id == 2) + .unwrap(); + assert_eq!(fragment2.files, vec![projected_file]); + assert_eq!(fragment2.deletion_file, Some(deletion_file)); + assert_eq!(fragment2.overlays.len(), 1); + assert_eq!(fragment2.overlays[0].committed_version, 2); + assert_eq!(fragment2.physical_rows, Some(42)); + assert_eq!(new_manifest.max_field_id(), 0); assert_eq!(last_updated_at_versions(&new_manifest, 2), vec![2; 42]); }