fix(dataset): resolve Blob v2 external URIs and robustly clean failed writes in add_columns - #7152
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
df5072c to
89bd69a
Compare
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
|
@claude review |
| batch_size: Option<u32>, | ||
| ) -> Result<(Fragment, Schema)> { | ||
| let (fragments, schema) = schema_evolution::add_columns_to_fragments( | ||
| let (fragments, schema, _) = schema_evolution::add_columns_to_fragments( |
There was a problem hiding this comment.
Can we keep fragments_to_cleanup? So that the outer can do some cleanup work.
There was a problem hiding this comment.
Good point. I think this is a broader follow-up than the current blob-v2 add_columns fix.
This PR is scoped to the dataset-level add_columns cleanup path. Keeping
fragments_to_cleanup in Fragment::add_columns would effectively extend the
fragment-level API / ownership boundary so the outer caller can decide how to
clean up on later failures.
I’d prefer to keep this PR focused and handle that as a separate follow-up PR,
unless you feel it is required to land this fix safely.
There was a problem hiding this comment.
Can you file a ticket to track this issue?
There was a problem hiding this comment.
Good point. I filed a follow-up issue to track this: #7231 。
The current PR stays scoped to dataset-level add_columns cleanup, while the follow-up issue tracks the fragment-level cleanup ownership gap for Fragment::add_columns / merge_columns when a later outer merge commit fails.
b74db79 to
d552f82
Compare
yanghua
left a comment
yanghua
left a comment
There was a problem hiding this comment.
Left some new comments.
| struct AddColumnFragments { | ||
| fragments: Vec<Fragment>, | ||
| fragments_to_cleanup: Vec<Fragment>, | ||
| } |
There was a problem hiding this comment.
It would be better to add some doc for this two fields.
There was a problem hiding this comment.
Thanks for the suggestion. I added doc comments for both fields to clarify that fragments contains the fragments returned for the final merge commit, while fragments_to_cleanup tracks uncommitted fragments whose newly written data files should be removed if the operation fails before the merge commit completes.
| } | ||
|
|
||
| /// Clean up any data file and blob sidecars created by the current unfinished writer. | ||
| pub(super) async fn cleanup_unfinished_writer(&mut self) { |
There was a problem hiding this comment.
Shall we clear the self.writer after finishing the cleanup?
There was a problem hiding this comment.
The current callers return immediately after cleanup, so this does not affect the current error path. However, clearing the writer makes the Updater state match the filesystem state after cleanup and avoids leaving a stale writer that points to deleted files. I’ll update cleanup_unfinished_writer to take and drop self.writer after capturing the data file path.
d552f82 to
daa7fb4
Compare
7b2d926 to
f6bc3f9
Compare
|
Since there is no objection, I would merge it at this moment. |
2 participants
Description
This PR addresses two critical functional gaps in the
add_columnspipeline when working with Blob v2 datasets. It ensures that the schema evolution path achieves feature parity with the main write path regarding external URI resolution, and it introduces a robust, leak-free cleanup mechanism for failed mutations.1. Blob v2 External URI Resolution
Context:
In Blob v2, external blobs (where data resides outside the dataset's native storage) rely on an
ExternalBaseResolverto construct the correct absolute URIs from relative paths stored in the dataset manifest.Previously,
add_columnsused a bareopen_writervia theUpdater, which lacked the context to initialize this resolver. As a result, appending or mutating Blob v2 columns withexternalkind would fail to resolve URIs.Fix:
open_update_writerinwrite.rs. When theUpdatercreates a new writer, it now checks ifstorage_version >= V2_2and if the schema containsblob_v2fields.ExternalBaseResolverusing the dataset's registered base paths and passes it to the writer alongside thesource_store_registry.BatchUDForRecordBatchReaderto writeexternalBlob v2 data during schema evolution, achieving strict parity withwrite_fragments_internal.2. Comprehensive Cleanup for Failed Writes
Context:
The
add_columnsoperation is complex and can fail mid-flight due to various reasons: UDF execution panics, stream ingestion errors, checkpoint lookup/insert failures, or schema merge conflicts. Previously, these failures would eagerly propagate via the?operator, leaving behind orphaned.lancedata files and their corresponding Blob v2 sidecar directories.Fix & Architectural Guarantees:
This PR overhauls the error handling and cleanup lifecycle in
add_columns_implandadd_columns_from_stream.Multi-stage Cleanup:
We replaced concurrent
try_collectstreams with sequential processing augmented with scoped error handling. When a failure occurs, the pipeline now executes a two-stage cleanup:updater.cleanup_unfinished_writer(): Cleans up the currently active, unfinished data file that hasn't been finalized into aFragmentyet.cleanup_new_column_data_files(): Cleans up any fully written but uncommitted fragments generated in the current run.Strict Safety Constraint 1: Preservation of External Data
The underlying
cleanup_data_fragmentslogic now strictly checksbase_id.is_none(). This guarantees that cleanup operations never attempt to delete files that belong to an external base, preventing catastrophic deletion of user-managed source data.Strict Safety Constraint 2: Checkpoint Ownership
For long-running UDFs, fragments are incrementally saved to a
UDFCheckpointStore. Once a fragment is successfully inserted into the checkpoint, it is explicitlypop()-ed from the localfragments_to_cleanuplist. This ensures that if a subsequent step fails, we do not physically delete data files that the checkpoint relies on for resumption.Related Issues
Closes #7075
Scope & Follow-up
This PR tightly scopes the fixes to the
add_columnspipeline.Note on
alter_columns: Whilealter_columnsshares the newly improvedadd_columns_implmachinery (and thus benefits from the unfinished writer cleanup), its finalapply_commitstage currently drops thefragments_to_cleanupvector. A commit failure there could still leave orphaned files. To keep this PR focused and reviewable, the commit-failure cleanup foralter_columnswill be addressed in a separate, dedicated follow-up PR.Type of change
Testing Performed
inline,packed,dedicated, andexternalcleanup scenarios. Validated stream errors, UDF panics, and checkpoint lookup/insert anomalies to ensure zero orphan files and zero deleted external files.add_columnswithRecordBatchReaderandBatchUDFacross all Blob v2 kinds. Asserted bit-for-bit data integrity upon reading back the mutated dataset.cargo fmt), linting (cargo clippy), and existing test suites.