From 8542e306dff7e9fb6588fa16bbe5b7c36cc7f916 Mon Sep 17 00:00:00 2001 From: Prabhu Shankar Date: Fri, 6 Mar 2026 03:40:54 -0500 Subject: [PATCH 1/4] migratebucketprocedurefile --- .../MigrateTableBucketProcedure.java | 218 ++++++++++++++++++ 1 file changed, 218 insertions(+) create mode 100644 presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java new file mode 100644 index 0000000000000..a7ae2e1813e4d --- /dev/null +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java @@ -0,0 +1,218 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.iceberg.procedure; + +import com.facebook.presto.common.type.TypeManager; +import com.facebook.presto.iceberg.IcebergDistributedProcedureHandle; +import com.facebook.presto.iceberg.IcebergProcedureContext; +import com.facebook.presto.iceberg.IcebergTableHandle; +import com.facebook.presto.iceberg.IcebergTableLayoutHandle; +import com.facebook.presto.spi.ConnectorDistributedProcedureHandle; +import com.facebook.presto.spi.ConnectorSession; +import com.facebook.presto.spi.ConnectorTableLayoutHandle; +import com.facebook.presto.spi.connector.ConnectorProcedureContext; +import com.facebook.presto.spi.procedure.DistributedProcedure; +import com.facebook.presto.spi.procedure.DistributedProcedure.Argument; +import com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import io.airlift.slice.Slice; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.RewriteFiles; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableScan; +import org.apache.iceberg.Transaction; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.io.PositionOutputStream; +import org.apache.iceberg.io.SeekableInputStream; + +import javax.inject.Inject; +import javax.inject.Provider; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Collection; +import java.util.HashSet; +import java.util.Map; +import java.util.OptionalInt; +import java.util.Set; + +import static com.facebook.presto.common.Utils.checkArgument; +import static com.facebook.presto.common.type.StandardTypes.VARCHAR; +import static com.facebook.presto.iceberg.IcebergSessionProperties.getCompressionCodec; +import static com.facebook.presto.iceberg.IcebergUtil.getColumns; +import static com.facebook.presto.iceberg.IcebergUtil.getFileFormat; +import static com.facebook.presto.iceberg.PartitionSpecConverter.toPrestoPartitionSpec; +import static com.facebook.presto.iceberg.SchemaConverter.toPrestoSchema; +import static com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure.SCHEMA; +import static com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure.TABLE_NAME; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class MigrateTableBucketProcedure + implements Provider +{ + private final TypeManager typeManager; + + @Inject + public MigrateTableBucketProcedure(TypeManager typeManager) + { + this.typeManager = requireNonNull(typeManager, "typeManager is null"); + } + + @Override + public DistributedProcedure get() + { + return new TableDataRewriteDistributedProcedure( + "system", + "migrate_table_bucket", + ImmutableList.of( + new Argument(SCHEMA, VARCHAR), + new Argument(TABLE_NAME, VARCHAR), + new Argument("new_base_path", VARCHAR)), + this::beginCallDistributedProcedure, + this::finishCallDistributedProcedure, + arguments -> { + checkArgument(arguments.length == 2, + format( + "invalid number of arguments: %s (should have %s)", + arguments.length, + 2)); + checkArgument( + arguments[0] instanceof Table && arguments[1] instanceof Transaction, + "Invalid arguments, required: [Table, Transaction]"); + + Table table = (Table) arguments[0]; + Transaction transaction = (Transaction) arguments[1]; + + return new IcebergProcedureContext(table, transaction); + }); + } + + private ConnectorDistributedProcedureHandle beginCallDistributedProcedure( + ConnectorSession session, + ConnectorProcedureContext procedureContext, + ConnectorTableLayoutHandle tableLayoutHandle, + Object[] arguments, + OptionalInt sortOrderIndex) + { + IcebergProcedureContext icebergContext = (IcebergProcedureContext) procedureContext; + IcebergTableLayoutHandle layoutHandle = (IcebergTableLayoutHandle) tableLayoutHandle; + String newBasePath = (String) arguments[2]; + Map relevantData = ImmutableMap.of("new_base_path", newBasePath); + IcebergTableHandle tableHandle = layoutHandle.getTable(); + Table icebergTable = icebergContext.getTable(); + + return new IcebergDistributedProcedureHandle( + tableHandle.getSchemaName(), + tableHandle.getIcebergTableName(), + toPrestoSchema(icebergTable.schema(), typeManager), + toPrestoPartitionSpec(icebergTable.spec(), typeManager), + getColumns(icebergTable.schema(), icebergTable.spec(), typeManager), + icebergTable.location(), + getFileFormat(icebergTable), + getCompressionCodec(session), + icebergTable.properties(), + layoutHandle, + ImmutableList.of(), + relevantData); } + + private void finishCallDistributedProcedure( + ConnectorSession session, + ConnectorProcedureContext procedureContext, + ConnectorDistributedProcedureHandle handle, + Collection fragments) + { + IcebergProcedureContext icebergContext = (IcebergProcedureContext) procedureContext; + IcebergDistributedProcedureHandle icebergHandle = (IcebergDistributedProcedureHandle) handle; + Table icebergTable = icebergContext.getTransaction().table(); + String newBasePath = icebergHandle.getRelevantData().get("new_base_path"); + Set existingDataFiles = new HashSet<>(); + TableScan tableScan = icebergTable.newScan().useSnapshot(icebergTable.currentSnapshot().snapshotId()); + + try (CloseableIterable tasks = tableScan.planFiles()) { + for (FileScanTask task : tasks) { + existingDataFiles.add(task.file()); + } + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + + Set deleteFiles = new HashSet<>(); + TableScan deleteScan = icebergTable.newScan().useSnapshot(icebergTable.currentSnapshot().snapshotId()); + + try (CloseableIterable tasks = deleteScan.planFiles()) { + for (FileScanTask task : tasks) { + if (!task.deletes().isEmpty()) { + deleteFiles.addAll(task.deletes()); + } + } + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + + Set newDataFiles = new HashSet<>(); + + 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); + OutputFile outputFile = fileIO.newOutputFile(newPath); + + 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); + } + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + + DataFiles.Builder builder = DataFiles.builder(icebergTable.spec()).copy(oldFile).withPath(newPath); + + if (oldFile.partition() != null) { + builder.withPartition(oldFile.partition()); + } + + DataFile newFile = builder.build(); + + 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()); + } + + rewrite.commit(); + } +} From 78621d3bc290cbdfd91aab8bab8983b51e8037fb Mon Sep 17 00:00:00 2001 From: Prabhu Shankar Date: Mon, 9 Mar 2026 11:57:51 -0400 Subject: [PATCH 2/4] migratebucketprocedurefile updates --- .../MigrateTableBucketProcedure.java | 175 ++++++++++++------ 1 file changed, 119 insertions(+), 56 deletions(-) diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java index a7ae2e1813e4d..551be1107955c 100644 --- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java @@ -21,6 +21,7 @@ import com.facebook.presto.spi.ConnectorDistributedProcedureHandle; import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.ConnectorTableLayoutHandle; +import com.facebook.presto.spi.PrestoException; import com.facebook.presto.spi.connector.ConnectorProcedureContext; import com.facebook.presto.spi.procedure.DistributedProcedure; import com.facebook.presto.spi.procedure.DistributedProcedure.Argument; @@ -29,9 +30,9 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import io.airlift.slice.Slice; +import io.airlift.slice.Slices; import org.apache.iceberg.DataFile; import org.apache.iceberg.DataFiles; -import org.apache.iceberg.DeleteFile; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.RewriteFiles; import org.apache.iceberg.Snapshot; @@ -50,6 +51,7 @@ import java.io.IOException; import java.io.UncheckedIOException; +import java.nio.charset.StandardCharsets; import java.util.Collection; import java.util.HashSet; import java.util.Map; @@ -63,6 +65,7 @@ import static com.facebook.presto.iceberg.IcebergUtil.getFileFormat; import static com.facebook.presto.iceberg.PartitionSpecConverter.toPrestoPartitionSpec; import static com.facebook.presto.iceberg.SchemaConverter.toPrestoSchema; +import static com.facebook.presto.spi.StandardErrorCode.NOT_SUPPORTED; import static com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure.SCHEMA; import static com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure.TABLE_NAME; import static java.lang.String.format; @@ -71,6 +74,8 @@ public class MigrateTableBucketProcedure implements Provider { + private static final String FRAGMENT_DELIMITER = "\u0000"; + private final TypeManager typeManager; @Inject @@ -118,9 +123,31 @@ private ConnectorDistributedProcedureHandle beginCallDistributedProcedure( IcebergProcedureContext icebergContext = (IcebergProcedureContext) procedureContext; IcebergTableLayoutHandle layoutHandle = (IcebergTableLayoutHandle) tableLayoutHandle; String newBasePath = (String) arguments[2]; - Map relevantData = ImmutableMap.of("new_base_path", newBasePath); IcebergTableHandle tableHandle = layoutHandle.getTable(); Table icebergTable = icebergContext.getTable(); + Snapshot currentSnapshot = icebergTable.currentSnapshot(); + if (currentSnapshot != null) { + TableScan deleteScan = icebergTable.newScan().useSnapshot(currentSnapshot.snapshotId()); + try (CloseableIterable tasks = deleteScan.planFiles()) { + for (FileScanTask task : tasks) { + if (!task.deletes().isEmpty()) { + throw new PrestoException( + NOT_SUPPORTED, + "migrate_table_bucket does not support tables with delete files. " + + "Please compact or merge delete files before migrating."); + } + } + } + catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + String normalizedNewBasePath = stripTrailingSlash(newBasePath); + + Map relevantData = ImmutableMap.of( + "new_base_path", normalizedNewBasePath, + "table_location", stripTrailingSlash(icebergTable.location())); return new IcebergDistributedProcedureHandle( tableHandle.getSchemaName(), @@ -134,7 +161,46 @@ private ConnectorDistributedProcedureHandle beginCallDistributedProcedure( icebergTable.properties(), layoutHandle, ImmutableList.of(), - relevantData); } + relevantData); + } + + public static Slice copyFileAndBuildFragment( + FileIO fileIO, + String tableLocation, + String normalizedNewBasePath, + DataFile dataFile) + { + String oldPath = dataFile.path().toString(); + String normalizedTableLocation = stripTrailingSlash(tableLocation); + String relativePath = oldPath.startsWith(normalizedTableLocation) + ? oldPath.substring(normalizedTableLocation.length() + 1) + : oldPath.substring(oldPath.lastIndexOf('/') + 1); + String newPath = normalizedNewBasePath + "/" + relativePath; + + if (newPath.equals(oldPath)) { + String payload = oldPath + FRAGMENT_DELIMITER + oldPath; + return Slices.wrappedBuffer(payload.getBytes(StandardCharsets.UTF_8)); + } + + InputFile inputFile = fileIO.newInputFile(oldPath); + OutputFile outputFile = fileIO.newOutputFile(newPath); + + try (SeekableInputStream in = inputFile.newStream(); + PositionOutputStream out = outputFile.create()) { + byte[] buffer = new byte[64 * 1024]; + int bytesRead; + while ((bytesRead = in.read(buffer)) > 0) { + out.write(buffer, 0, bytesRead); + } + } + catch (IOException e) { + throw new UncheckedIOException( + new IOException(format("Failed to copy file from %s to %s", oldPath, newPath), e)); + } + + String payload = oldPath + FRAGMENT_DELIMITER + newPath; + return Slices.wrappedBuffer(payload.getBytes(StandardCharsets.UTF_8)); + } private void finishCallDistributedProcedure( ConnectorSession session, @@ -145,74 +211,71 @@ private void finishCallDistributedProcedure( IcebergProcedureContext icebergContext = (IcebergProcedureContext) procedureContext; IcebergDistributedProcedureHandle icebergHandle = (IcebergDistributedProcedureHandle) handle; Table icebergTable = icebergContext.getTransaction().table(); - String newBasePath = icebergHandle.getRelevantData().get("new_base_path"); + Map pathRemapping = decodeFragments(fragments); Set existingDataFiles = new HashSet<>(); - TableScan tableScan = icebergTable.newScan().useSnapshot(icebergTable.currentSnapshot().snapshotId()); - - try (CloseableIterable tasks = tableScan.planFiles()) { - for (FileScanTask task : tasks) { - existingDataFiles.add(task.file()); - } - } - catch (IOException e) { - throw new UncheckedIOException(e); + Set newDataFiles = new HashSet<>(); + Snapshot currentSnapshot = icebergTable.currentSnapshot(); + if (currentSnapshot == null) { + return; } + long validationSnapshotId = currentSnapshot.snapshotId(); - Set deleteFiles = new HashSet<>(); - TableScan deleteScan = icebergTable.newScan().useSnapshot(icebergTable.currentSnapshot().snapshotId()); + TableScan tableScan = icebergTable.newScan().useSnapshot(currentSnapshot.snapshotId()); - try (CloseableIterable tasks = deleteScan.planFiles()) { + try (CloseableIterable tasks = tableScan.planFiles()) { for (FileScanTask task : tasks) { - if (!task.deletes().isEmpty()) { - deleteFiles.addAll(task.deletes()); - } - } - } - catch (IOException e) { - throw new UncheckedIOException(e); - } + DataFile oldFile = task.file(); + String oldPath = oldFile.path().toString(); + String newPath = pathRemapping.get(oldPath); - Set newDataFiles = new HashSet<>(); - - FileIO fileIO = icebergTable.io(); + if (newPath == null) { + throw new IllegalStateException( + format("No fragment received for data file: %s. " + + "The distributed copy may be incomplete.", oldPath)); + } - for (DataFile oldFile : existingDataFiles) { - String oldPath = oldFile.path().toString(); - String fileName = oldPath.substring(oldPath.lastIndexOf('/') + 1); - String newPath = newBasePath + "/" + fileName; + existingDataFiles.add(oldFile); - InputFile inputFile = fileIO.newInputFile(oldPath); - OutputFile outputFile = fileIO.newOutputFile(newPath); + DataFiles.Builder builder = DataFiles.builder(icebergTable.spec()) + .copy(oldFile) + .withPath(newPath); - 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); + if (oldFile.partition() != null) { + builder.withPartition(oldFile.partition()); } - } - catch (IOException e) { - throw new UncheckedIOException(e); - } - - DataFiles.Builder builder = DataFiles.builder(icebergTable.spec()).copy(oldFile).withPath(newPath); - if (oldFile.partition() != null) { - builder.withPartition(oldFile.partition()); + newDataFiles.add(builder.build()); } - - DataFile newFile = builder.build(); - - newDataFiles.add(newFile); } - RewriteFiles rewrite = icebergContext.getTransaction().newRewrite().rewriteFiles(existingDataFiles, deleteFiles, newDataFiles, ImmutableSet.of()); - + catch (IOException e) { + throw new UncheckedIOException(e); + } + RewriteFiles rewrite = icebergContext.getTransaction().newRewrite().rewriteFiles(existingDataFiles, ImmutableSet.of(), newDataFiles, ImmutableSet.of()); + rewrite.validateFromSnapshot(currentSnapshot.snapshotId()); Snapshot snapshot = icebergTable.currentSnapshot(); - if (snapshot != null) { - rewrite.validateFromSnapshot(snapshot.snapshotId()); + rewrite.commit(); + } + + private static Map decodeFragments(Collection fragments) + { + ImmutableMap.Builder map = ImmutableMap.builder(); + for (Slice fragment : fragments) { + String payload = new String(fragment.getBytes(), StandardCharsets.UTF_8); + int delimIndex = payload.indexOf(FRAGMENT_DELIMITER); + if (delimIndex < 0) { + throw new IllegalArgumentException( + format("Malformed fragment (missing delimiter): %s", payload)); + } + String oldPath = payload.substring(0, delimIndex); + String newPath = payload.substring(delimIndex + 1); + map.put(oldPath, newPath); } + return map.build(); + } - rewrite.commit(); + private static String stripTrailingSlash(String path) + { + requireNonNull(path, "path is null"); + return path.endsWith("/") ? path.substring(0, path.length() - 1) : path; } } From d627bd8fdf506e06b81d91bcccc37fa0f3736933 Mon Sep 17 00:00:00 2001 From: Prabhu Shankar <157943490+prabhushankar-ps@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:19:16 -0400 Subject: [PATCH 3/4] Add IcebergProcedureContext for procedure lifecycle --- .../procedure/IcebergProcedureContext.java | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/IcebergProcedureContext.java diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/IcebergProcedureContext.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/IcebergProcedureContext.java new file mode 100644 index 0000000000000..6b26ad50492b3 --- /dev/null +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/IcebergProcedureContext.java @@ -0,0 +1,36 @@ +package com.facebook.presto.iceberg.procedure; + +import com.facebook.presto.spi.connector.ConnectorProcedureContext; +import org.apache.iceberg.Table; +import org.apache.iceberg.Transaction; + +import static java.util.Objects.requireNonNull; + +/** + * Simple context wrapper for Iceberg distributed procedures. + * + * Carries Iceberg Table + Transaction through the procedure lifecycle: + * begin → worker → finish + */ +public class IcebergProcedureContext + implements ConnectorProcedureContext +{ + private final Table table; + private final Transaction transaction; + + public IcebergProcedureContext(Table table, Transaction transaction) + { + this.table = requireNonNull(table, "table is null"); + this.transaction = requireNonNull(transaction, "transaction is null"); + } + + public Table getTable() + { + return table; + } + + public Transaction getTransaction() + { + return transaction; + } +} From 3736520dc188423dc9f416703caf4be1d598207d Mon Sep 17 00:00:00 2001 From: Prabhu Shankar <157943490+prabhushankar-ps@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:22:49 -0400 Subject: [PATCH 4/4] Refactor import and snapshot validation logic --- .../procedure/MigrateTableBucketProcedure.java | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java index 551be1107955c..4373bc73b5850 100644 --- a/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java +++ b/presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/MigrateTableBucketProcedure.java @@ -15,7 +15,7 @@ import com.facebook.presto.common.type.TypeManager; import com.facebook.presto.iceberg.IcebergDistributedProcedureHandle; -import com.facebook.presto.iceberg.IcebergProcedureContext; +import com.facebook.presto.iceberg.procedure.IcebergProcedureContext; import com.facebook.presto.iceberg.IcebergTableHandle; import com.facebook.presto.iceberg.IcebergTableLayoutHandle; import com.facebook.presto.spi.ConnectorDistributedProcedureHandle; @@ -218,7 +218,6 @@ private void finishCallDistributedProcedure( if (currentSnapshot == null) { return; } - long validationSnapshotId = currentSnapshot.snapshotId(); TableScan tableScan = icebergTable.newScan().useSnapshot(currentSnapshot.snapshotId()); @@ -250,9 +249,11 @@ private void finishCallDistributedProcedure( catch (IOException e) { throw new UncheckedIOException(e); } - RewriteFiles rewrite = icebergContext.getTransaction().newRewrite().rewriteFiles(existingDataFiles, ImmutableSet.of(), newDataFiles, ImmutableSet.of()); - rewrite.validateFromSnapshot(currentSnapshot.snapshotId()); - Snapshot snapshot = icebergTable.currentSnapshot(); + long validationSnapshotId = currentSnapshot.snapshotId(); + RewriteFiles rewrite = icebergContext.getTransaction() + .newRewrite() + .rewriteFiles(existingDataFiles, ImmutableSet.of(), newDataFiles, ImmutableSet.of()); + rewrite.validateFromSnapshot(validationSnapshotId); rewrite.commit(); }