feat(plugin-iceberg): Add MigrateTableBucket distributed procedure - #27278
Open
prabhushankar-ps wants to merge 8 commits into
Open
feat(plugin-iceberg): Add MigrateTableBucket distributed procedure#27278prabhushankar-ps wants to merge 8 commits into
prabhushankar-ps wants to merge 8 commits into
Conversation
prabhushankar-ps
requested review from
a team,
ZacBlanco and
hantangwangd
as code owners
March 6, 2026 09:46
Contributor
Reviewer's GuideIntroduces a new Iceberg distributed system procedure Sequence diagram for system.migrate_table_bucket distributed rewritesequenceDiagram
actor User
participant PrestoCoordinator
participant IcebergConnector
participant MigrateTableBucketProcedure
participant TableDataRewriteDistributedProcedure
participant IcebergTable
participant IcebergTransaction
participant IcebergFileIO
User->>PrestoCoordinator: CALL iceberg.system.migrate_table_bucket(schema_name, table_name, new_base_path)
PrestoCoordinator->>IcebergConnector: Execute system.migrate_table_bucket
IcebergConnector->>MigrateTableBucketProcedure: Provider.get()
MigrateTableBucketProcedure-->>IcebergConnector: DistributedProcedure(TableDataRewriteDistributedProcedure)
PrestoCoordinator->>TableDataRewriteDistributedProcedure: beginCallDistributedProcedure
TableDataRewriteDistributedProcedure->>MigrateTableBucketProcedure: beginCallDistributedProcedure(session, context, layout, args, sortOrderIndex)
MigrateTableBucketProcedure->>IcebergTable: load table
MigrateTableBucketProcedure->>IcebergDistributedProcedureHandle: create(handle with relevantData new_base_path)
IcebergDistributedProcedureHandle-->>TableDataRewriteDistributedProcedure: handle
loop Distributed scan and rewrite planning
TableDataRewriteDistributedProcedure->>IcebergTable: newScan().useSnapshot(currentSnapshot)
IcebergTable-->>TableDataRewriteDistributedProcedure: FileScanTasks
end
PrestoCoordinator->>TableDataRewriteDistributedProcedure: finishCallDistributedProcedure
TableDataRewriteDistributedProcedure->>MigrateTableBucketProcedure: finishCallDistributedProcedure(session, context, handle, fragments)
MigrateTableBucketProcedure->>IcebergTable: newScan().useSnapshot(currentSnapshot)
IcebergTable-->>MigrateTableBucketProcedure: existingDataFiles
MigrateTableBucketProcedure->>IcebergTable: newScan().useSnapshot(currentSnapshot)
IcebergTable-->>MigrateTableBucketProcedure: deleteFiles
loop For each existing DataFile
MigrateTableBucketProcedure->>IcebergFileIO: newInputFile(oldPath)
MigrateTableBucketProcedure->>IcebergFileIO: newOutputFile(newBasePath/fileName)
IcebergFileIO-->>MigrateTableBucketProcedure: InputFile, OutputFile
MigrateTableBucketProcedure->>IcebergFileIO: copy bytes from InputFile to OutputFile
MigrateTableBucketProcedure->>MigrateTableBucketProcedure: build new DataFile with newPath and copied metadata
end
MigrateTableBucketProcedure->>IcebergTransaction: newRewrite()
IcebergTransaction-->>MigrateTableBucketProcedure: RewriteFiles
MigrateTableBucketProcedure->>IcebergTransaction: rewriteFiles(existingDataFiles, deleteFiles, newDataFiles, emptySet)
MigrateTableBucketProcedure->>IcebergTable: currentSnapshot
IcebergTable-->>MigrateTableBucketProcedure: Snapshot
MigrateTableBucketProcedure->>IcebergTransaction: validateFromSnapshot(snapshotId)
MigrateTableBucketProcedure->>IcebergTransaction: commit()
IcebergTransaction-->>PrestoCoordinator: Rewrite committed
PrestoCoordinator-->>User: CALL completed successfully
Class diagram for MigrateTableBucketProcedure and related typesclassDiagram
class Provider {
<<interface>>
+get() DistributedProcedure
}
class DistributedProcedure {
<<interface>>
}
class TableDataRewriteDistributedProcedure {
+TableDataRewriteDistributedProcedure(schemaName, procedureName, arguments, beginCallback, finishCallback, contextFactory)
}
class MigrateTableBucketProcedure {
-TypeManager typeManager
+MigrateTableBucketProcedure(TypeManager typeManager)
+get() DistributedProcedure
-beginCallDistributedProcedure(ConnectorSession session, ConnectorProcedureContext procedureContext, ConnectorTableLayoutHandle tableLayoutHandle, Object[] arguments, OptionalInt sortOrderIndex) ConnectorDistributedProcedureHandle
-finishCallDistributedProcedure(ConnectorSession session, ConnectorProcedureContext procedureContext, ConnectorDistributedProcedureHandle handle, Collection fragments) void
}
class TypeManager
class ConnectorSession
class ConnectorProcedureContext
class ConnectorTableLayoutHandle
class ConnectorDistributedProcedureHandle
class IcebergProcedureContext {
+getTable() Table
+getTransaction() Transaction
}
class IcebergTableLayoutHandle {
+getTable() IcebergTableHandle
}
class IcebergTableHandle {
+getSchemaName() String
+getIcebergTableName() String
}
class IcebergDistributedProcedureHandle {
+IcebergDistributedProcedureHandle(String schemaName, String tableName, Object prestoSchema, Object prestoPartitionSpec, Object columns, String location, Object fileFormat, Object compressionCodec, Map properties, IcebergTableLayoutHandle layoutHandle, List tableColumns, Map relevantData)
+getRelevantData() Map
}
class Table {
+schema() Object
+spec() Object
+location() String
+properties() Map
+io() FileIO
+newScan() TableScan
+currentSnapshot() Snapshot
}
class Transaction {
+table() Table
+newRewrite() RewriteFiles
}
class TableScan {
+useSnapshot(long snapshotId) TableScan
+planFiles() CloseableIterable
}
class FileScanTask {
+file() DataFile
+deletes() List
}
class FileIO {
+newInputFile(String path) InputFile
+newOutputFile(String path) OutputFile
}
class DataFile {
+path() Object
+partition() Object
}
class DeleteFile
class Snapshot {
+snapshotId() long
}
class RewriteFiles {
+rewriteFiles(Set existingDataFiles, Set deleteFiles, Set newDataFiles, Set additionalFiles) RewriteFiles
+validateFromSnapshot(long snapshotId) RewriteFiles
+commit() void
}
class InputFile {
+newStream() SeekableInputStream
}
class OutputFile {
+create() PositionOutputStream
}
class SeekableInputStream {
+read(byte[] buffer) int
+close() void
}
class PositionOutputStream {
+write(byte[] buffer, int offset, int length) void
+close() void
}
Provider <|.. MigrateTableBucketProcedure
DistributedProcedure <|.. TableDataRewriteDistributedProcedure
MigrateTableBucketProcedure --> TypeManager
MigrateTableBucketProcedure ..> TableDataRewriteDistributedProcedure
MigrateTableBucketProcedure ..> IcebergProcedureContext
MigrateTableBucketProcedure ..> IcebergTableLayoutHandle
MigrateTableBucketProcedure ..> IcebergTableHandle
MigrateTableBucketProcedure ..> IcebergDistributedProcedureHandle
MigrateTableBucketProcedure ..> Table
MigrateTableBucketProcedure ..> Transaction
MigrateTableBucketProcedure ..> TableScan
MigrateTableBucketProcedure ..> FileScanTask
MigrateTableBucketProcedure ..> DataFile
MigrateTableBucketProcedure ..> DeleteFile
MigrateTableBucketProcedure ..> FileIO
MigrateTableBucketProcedure ..> RewriteFiles
MigrateTableBucketProcedure ..> Snapshot
MigrateTableBucketProcedure ..> InputFile
MigrateTableBucketProcedure ..> OutputFile
MigrateTableBucketProcedure ..> SeekableInputStream
MigrateTableBucketProcedure ..> PositionOutputStream
IcebergProcedureContext --> Table
IcebergProcedureContext --> Transaction
IcebergTableLayoutHandle --> IcebergTableHandle
IcebergDistributedProcedureHandle --> IcebergTableLayoutHandle
Table --> FileIO
Table --> TableScan
Table --> Snapshot
Transaction --> Table
Transaction --> RewriteFiles
TableScan --> FileScanTask
FileScanTask --> DataFile
FileScanTask --> DeleteFile
FileIO --> InputFile
FileIO --> OutputFile
InputFile --> SeekableInputStream
OutputFile --> PositionOutputStream
File-Level Changes
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey - I've found 4 issues, and left some high level feedback:
- The actual data copy happens entirely in finishCallDistributedProcedure on the coordinator and the fragments argument is ignored, which defeats the purpose of using TableDataRewriteDistributedProcedure for distributed rewrites; consider moving the file copy into distributed tasks and using fragments to report new/old files.
- The handling of delete files is inconsistent with the documented limitation that delete files are not supported: you collect deleteFiles and pass them to rewriteFiles without providing replacement delete files, which will effectively drop them and can corrupt tables with deletes; either hard-fail when delete files are present or implement a delete-file-preserving strategy.
- Building new paths as new_base_path + "/" + fileName loses the original directory/partition structure and can also introduce double slashes if new_base_path already ends with "/"; consider preserving the relative path under the table root or normalizing paths to avoid layout and partitioning issues.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The actual data copy happens entirely in finishCallDistributedProcedure on the coordinator and the fragments argument is ignored, which defeats the purpose of using TableDataRewriteDistributedProcedure for distributed rewrites; consider moving the file copy into distributed tasks and using fragments to report new/old files.
- The handling of delete files is inconsistent with the documented limitation that delete files are not supported: you collect deleteFiles and pass them to rewriteFiles without providing replacement delete files, which will effectively drop them and can corrupt tables with deletes; either hard-fail when delete files are present or implement a delete-file-preserving strategy.
- Building new paths as new_base_path + "/" + fileName loses the original directory/partition structure and can also introduce double slashes if new_base_path already ends with "/"; consider preserving the relative path under the table root or normalizing paths to avoid layout and partitioning issues.
## Individual Comments
### Comment 1
<location path="presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java" line_range="150-159" />
<code_context>
+ Table icebergTable = icebergContext.getTransaction().table();
+ String newBasePath = icebergHandle.getRelevantData().get("new_base_path");
+ Set<DataFile> existingDataFiles = new HashSet<>();
+ TableScan tableScan = icebergTable.newScan().useSnapshot(icebergTable.currentSnapshot().snapshotId());
+
+ try (CloseableIterable<FileScanTask> tasks = tableScan.planFiles()) {
+ for (FileScanTask task : tasks) {
+ existingDataFiles.add(task.file());
+ }
+ }
+ catch (IOException e) {
+ throw new UncheckedIOException(e);
+ }
+
+ Set<DeleteFile> deleteFiles = new HashSet<>();
+ TableScan deleteScan = icebergTable.newScan().useSnapshot(icebergTable.currentSnapshot().snapshotId());
+
+ try (CloseableIterable<FileScanTask> tasks = deleteScan.planFiles()) {
</code_context>
<issue_to_address>
**issue (bug_risk):** Handle empty tables where currentSnapshot() may be null before calling useSnapshot
`tableScan` and `deleteScan` both call `useSnapshot(icebergTable.currentSnapshot().snapshotId())` assuming a non-null snapshot. For empty tables or when the current snapshot was expired, `currentSnapshot()` may be null and cause an NPE. Please either no-op when there is no snapshot, or validate upfront and fail with a clear error instead of relying on an NPE.
</issue_to_address>
### Comment 2
<location path="presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java" line_range="180-182" />
<code_context>
+ FileIO fileIO = icebergTable.io();
+
+ for (DataFile oldFile : existingDataFiles) {
+ String oldPath = oldFile.path().toString();
+ String fileName = oldPath.substring(oldPath.lastIndexOf('/') + 1);
+ String newPath = newBasePath + "/" + fileName;
+
+ InputFile inputFile = fileIO.newInputFile(oldPath);
</code_context>
<issue_to_address>
**issue:** Guard against copying a file onto itself when new_base_path matches the existing directory
If `newBasePath` resolves to the same directory as `oldPath`'s parent, `newPath` may equal `oldPath`. That would mean reading and writing the same physical file, risking corruption. Please either fail fast when `newPath.equals(oldPath)` or detect and skip the copy, reusing the existing `DataFile` in that case.
</issue_to_address>
### Comment 3
<location path="presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java" line_range="187-193" />
<code_context>
+
+ try (SeekableInputStream in = inputFile.newStream();
+ PositionOutputStream out = outputFile.create()) {
+ byte[] buffer = new byte[8192];
+ int bytesRead;
+ while ((bytesRead = in.read(buffer)) > 0) {
+ out.write(buffer, 0, bytesRead);
+ }
</code_context>
<issue_to_address>
**suggestion (performance):** Consider using a larger buffer or a streaming utility for potentially large data files
The current 8KB buffer is correct but likely suboptimal for typical large Iceberg files. Using a larger buffer (e.g., 64KB+), `ByteStreams.copy`, or an Iceberg `FileIO` copy/streaming helper would reduce syscall overhead and improve migration throughput on large tables.
```suggestion
try (SeekableInputStream in = inputFile.newStream();
PositionOutputStream out = outputFile.create()) {
// Use a larger buffer to reduce syscall overhead when copying large Iceberg data files
byte[] buffer = new byte[64 * 1024];
int bytesRead;
while ((bytesRead = in.read(buffer)) > 0) {
out.write(buffer, 0, bytesRead);
}
```
</issue_to_address>
### Comment 4
<location path="presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java" line_range="209-211" />
<code_context>
+
+ newDataFiles.add(newFile);
+ }
+ RewriteFiles rewrite = icebergContext.getTransaction().newRewrite().rewriteFiles(existingDataFiles, deleteFiles, newDataFiles, ImmutableSet.of());
+
+ Snapshot snapshot = icebergTable.currentSnapshot();
+ if (snapshot != null) {
+ rewrite.validateFromSnapshot(snapshot.snapshotId());
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Clarify snapshot consistency between scans and validation to avoid subtle race conditions
`existingDataFiles`/`deleteFiles` are scanned using whatever `currentSnapshot()` was when the scans were created, but `validateFromSnapshot` uses `icebergTable.currentSnapshot()` again at validation time. With concurrent commits, these can diverge and cause validation to fail or misbehave. Capture the snapshot ID once at the start of `finishCallDistributedProcedure` and reuse it for both the scans and `validateFromSnapshot` to ensure consistency.
Suggested implementation:
```java
RewriteFiles rewrite = icebergContext.getTransaction()
.newRewrite()
.rewriteFiles(existingDataFiles, deleteFiles, newDataFiles, ImmutableSet.of());
if (validationSnapshotId != null) {
rewrite.validateFromSnapshot(validationSnapshotId);
}
rewrite.commit();
```
To fully implement the snapshot consistency guarantee described in your comment, the following additional changes are required elsewhere in `MigrateTableBucketProcedure.java`:
1. At the very beginning of `finishCallDistributedProcedure` (or the method where `existingDataFiles`/`deleteFiles` are prepared), capture the snapshot once:
- `Snapshot snapshot = icebergTable.currentSnapshot();`
- `Long validationSnapshotId = (snapshot != null) ? snapshot.snapshotId() : null;`
- Use a `Long` (or boxed type) so `null` can represent "no snapshot".
2. Ensure that all scans that produce `existingDataFiles` and `deleteFiles` are created against this captured snapshot:
- If you are currently relying on implicit `currentSnapshot()` in the scan builders, switch to explicitly using the captured ID, e.g.:
- `tableScan.useSnapshot(validationSnapshotId)` (or the appropriate Iceberg API in this codebase).
- This guarantees `existingDataFiles` / `deleteFiles` reflect exactly the same snapshot as `validationSnapshotId`.
3. Make `validationSnapshotId` visible where used in the replacement above:
- Either as a local variable in `finishCallDistributedProcedure` that is in scope for the commit section.
- Or as a field on the surrounding class set at the start of the procedure and read here (local variable is preferable if possible).
Once these additional changes are made, `validateFromSnapshot` will be using the same snapshot ID that was used for the scans, avoiding the race condition with concurrent commits.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Contributor
|
Please include documentation for the new property, as described in Designing Your Code in CONTRIBUTING.md: "All new language features, new functions, session and config properties, and major features have documentation added" |
Contributor
Author
|
Hi all, please view the RFC attached with this PR. |
Contributor
|
prabhushankar-ps
changed the title
Contributor
Author
|
Hi all, just checking in to see if there have been any reviews on the propsed RFC. |
hantangwangd
left a comment
hantangwangd
left a comment
Member
There was a problem hiding this comment.
Hi @prabhushankar-ps, thanks for your work. I have some high-level feedback:
Firstly, as I understand, table migration means the entire table root directory – including both data files and metadata files – is migrated. In the current implementation, I only see logic for data files. We may need to rethink the functional scope here and clarify what exactly we aim to achieve.
Secondly, the current code doesn't work properly in the distributed procedure framework. For a distributed procedure to run on Presto, you either need to use an existing subtype that can support this type of task, or extend a new one – which means adding the analysis and planning logic for that subtype. For further details, please refer to this.
Let me know if you see it differently – happy to discuss.
Comment on lines
+205
to
+209
| private void finishCallDistributedProcedure( | ||
| ConnectorSession session, | ||
| ConnectorProcedureContext procedureContext, | ||
| ConnectorDistributedProcedureHandle handle, | ||
| Collection<Slice> fragments) |
Member
There was a problem hiding this comment.
I'm afraid I'm a bit confused here: could you please elaborate on where this method is invoked?
Contributor
Author
There was a problem hiding this comment.
It is passed as a method reference (this::finishCallDistributedProcedure) to TableDataRewriteDistributedProcedure in the get() method. Internally, it's stored as a FinishCallDistributedProcedure functional interface, and invoked when Presto calls DistributedProcedure() after all distributed workers have completed and fragments have been collected.
Contributor
Author
|
Thank you for picking this up @hantangwangd
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR introduces a new distributed Iceberg system procedure:
system.migrate_table_bucket
The procedure enables migration of Iceberg table data files to a new base path (for example, a different S3 bucket or prefix) while preserving table metadata and snapshot consistency.
The implementation leverages TableDataRewriteDistributedProcedure to perform a distributed rewrite, copies underlying data files to the target location, and commits a RewriteFiles transaction to update snapshot metadata atomically.
Motivation and Context
Currently, Presto does not provide a built-in Iceberg procedure to migrate table data files across storage locations while preserving snapshot integrity.
This change introduces a controlled and metadata-consistent way to:
Move Iceberg table data to a new bucket or prefix
Preserve file-level metrics and partition metadata
Maintain snapshot correctness
Use Iceberg’s native RewriteFiles API for atomic commits
This is useful for operational scenarios such as:
Bucket migrations
Storage reorganization
Cross-environment data movement
The implementation follows the distributed procedure framework defined by DistributedProcedure and aligns with the architectural pattern used in RewriteDataFilesProcedure.
#26779
Impact
User-facing change
Adds a new Iceberg system procedure:
CALL iceberg.system.migrate_table_bucket(
schema_name,
table_name,
new_base_path
);
Example:
CALL iceberg.system.migrate_table_bucket(
'default',
'test_table',
's3://new-bucket/path/'
);
Behavior
Current limitations
No public API changes outside Iceberg system procedures.
Test Plan
Due to resource constraints, full end-to-end testing with a custom-built Presto runtime has not yet been completed.
The procedure compiles successfully and integrates with:
IcebergDistributedProcedureHandle
TableDataRewriteDistributedProcedure
Iceberg RewriteFiles commit flow
Next steps for validation include:
I would appreciate guidance from maintainers on a resource-efficient approach for full integration testing.
Contributor checklist
Submission follows code style guidelines. - done
PR description addresses the change accurately and concisely. - done
New SQL procedure is documented above. - yet to
Adequate tests to be added. - yet to
CI expected to validate compilation. - yet to
No new dependencies introduced. - done
Summary by Sourcery
Add a distributed Iceberg procedure to migrate table data files between storage locations with an atomic metadata rewrite.
New Features:
system.migrate_table_bucketdistributed procedure for relocating table data files to a new storage base path while preserving table metadata and snapshot consistency.Enhancements: