Skip to content

Commit 414eb87

Browse files
krisnarukrisnaru
committed
Spark 4.1 only changes
1 parent b88addf commit 414eb87

2 files changed

Lines changed: 198 additions & 134 deletions

File tree

‎core/src/main/java/org/apache/iceberg/RewriteTablePathUtil.java‎

Lines changed: 117 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -100,20 +100,26 @@ public Set<Pair<String, String>> copyPlan() {
100100
* Create a new table metadata object, replacing path references
101101
*
102102
* @param metadata source table metadata
103-
* @param sourcePrefix source prefix that will be replaced
104-
* @param targetPrefix target prefix that will replace it
103+
* @param prefixMappings source prefix amd target prefix mappings
105104
* @return copy of table metadata with paths replaced
106105
*/
107106
public static TableMetadata replacePaths(
108-
TableMetadata metadata, String sourcePrefix, String targetPrefix) {
109-
String newLocation = metadata.location().replaceFirst(sourcePrefix, targetPrefix);
110-
List<Snapshot> newSnapshots = updatePathInSnapshots(metadata, sourcePrefix, targetPrefix);
107+
TableMetadata metadata, Map<String, String> prefixMappings) {
108+
Preconditions.checkArgument(
109+
metadata.partitionStatisticsFiles().isEmpty(),
110+
"Partition statistics files are not supported yet");
111+
Map.Entry<String, String> prefixMap = lookupPrefixMappings(metadata.location(), prefixMappings);
112+
if (prefixMap == null) {
113+
throw new IllegalArgumentException(
114+
"unable to find prefix mapping for path: " + metadata.location());
115+
}
116+
String newLocation = newPath(metadata.location(), prefixMappings);
117+
List<Snapshot> newSnapshots = updatePathInSnapshots(metadata, prefixMappings);
111118
List<TableMetadata.MetadataLogEntry> metadataLogEntries =
112-
updatePathInMetadataLogs(metadata, sourcePrefix, targetPrefix);
119+
updatePathInMetadataLogs(metadata, prefixMappings);
113120
long snapshotId =
114121
metadata.currentSnapshot() == null ? -1 : metadata.currentSnapshot().snapshotId();
115-
Map<String, String> properties =
116-
updateProperties(metadata.properties(), sourcePrefix, targetPrefix);
122+
Map<String, String> properties = updateProperties(metadata.properties(), prefixMappings);
117123

118124
return new TableMetadata(
119125
null,
@@ -137,25 +143,29 @@ public static TableMetadata replacePaths(
137143
metadata.snapshotLog(),
138144
metadataLogEntries,
139145
metadata.refs(),
140-
updatePathInStatisticsFiles(metadata.statisticsFiles(), sourcePrefix, targetPrefix),
141-
updatePathInPartitionStatisticsFiles(
142-
metadata.partitionStatisticsFiles(), sourcePrefix, targetPrefix),
146+
updatePathInStatisticsFiles(metadata.statisticsFiles(), prefixMappings),
147+
updatePathInPartitionStatisticsFiles(metadata.partitionStatisticsFiles(), prefixMappings),
143148
metadata.nextRowId(),
144149
metadata.encryptionKeys(),
145150
metadata.changes());
146151
}
147152

148153
private static Map<String, String> updateProperties(
149-
Map<String, String> tableProperties, String sourcePrefix, String targetPrefix) {
154+
Map<String, String> tableProperties, Map<String, String> prefixMappings) {
150155
Map<String, String> properties = Maps.newHashMap(tableProperties);
151-
updatePathInProperty(properties, sourcePrefix, targetPrefix, TableProperties.OBJECT_STORE_PATH);
152-
updatePathInProperty(
153-
properties, sourcePrefix, targetPrefix, TableProperties.WRITE_FOLDER_STORAGE_LOCATION);
154-
updatePathInProperty(
155-
properties, sourcePrefix, targetPrefix, TableProperties.WRITE_DATA_LOCATION);
156-
updatePathInProperty(
157-
properties, sourcePrefix, targetPrefix, TableProperties.WRITE_METADATA_LOCATION);
158-
156+
for (Map.Entry<String, String> entry : prefixMappings.entrySet()) {
157+
updatePathInProperty(
158+
properties, entry.getKey(), entry.getValue(), TableProperties.OBJECT_STORE_PATH);
159+
updatePathInProperty(
160+
properties,
161+
entry.getKey(),
162+
entry.getValue(),
163+
TableProperties.WRITE_FOLDER_STORAGE_LOCATION);
164+
updatePathInProperty(
165+
properties, entry.getKey(), entry.getValue(), TableProperties.WRITE_DATA_LOCATION);
166+
updatePathInProperty(
167+
properties, entry.getKey(), entry.getValue(), TableProperties.WRITE_METADATA_LOCATION);
168+
}
159169
return properties;
160170
}
161171

@@ -171,13 +181,13 @@ private static void updatePathInProperty(
171181
}
172182

173183
private static List<StatisticsFile> updatePathInStatisticsFiles(
174-
List<StatisticsFile> statisticsFiles, String sourcePrefix, String targetPrefix) {
184+
List<StatisticsFile> statisticsFiles, Map<String, String> prefixMappings) {
175185
return statisticsFiles.stream()
176186
.map(
177187
existing ->
178188
new GenericStatisticsFile(
179189
existing.snapshotId(),
180-
newPath(existing.path(), sourcePrefix, targetPrefix),
190+
newPath(existing.path(), prefixMappings),
181191
existing.fileSizeInBytes(),
182192
existing.fileFooterSizeInBytes(),
183193
existing.blobMetadata()))
@@ -189,46 +199,41 @@ private static List<StatisticsFile> updatePathInStatisticsFiles(
189199
* sourcePrefix in the file paths with the targetPrefix.
190200
*
191201
* @param partitionStatisticsFiles The list of PartitionStatisticsFile to update.
192-
* @param sourcePrefix The prefix to be replaced in the file paths.
193-
* @param targetPrefix The new prefix to replace the sourcePrefix in the file paths.
202+
* @param prefixMappings The mappings between source prefix and destination prefix.
194203
* @return A new list of PartitionStatisticsFile with updated file paths.
195204
*/
196205
private static List<PartitionStatisticsFile> updatePathInPartitionStatisticsFiles(
197-
List<PartitionStatisticsFile> partitionStatisticsFiles,
198-
String sourcePrefix,
199-
String targetPrefix) {
206+
List<PartitionStatisticsFile> partitionStatisticsFiles, Map<String, String> prefixMappings) {
200207

201208
return partitionStatisticsFiles.stream()
202209
.map(
203210
existing ->
204211
ImmutableGenericPartitionStatisticsFile.builder()
205212
.snapshotId(existing.snapshotId())
206-
.path(newPath(existing.path(), sourcePrefix, targetPrefix))
213+
.path(newPath(existing.path(), prefixMappings))
207214
.fileSizeInBytes(existing.fileSizeInBytes())
208215
.build())
209216
.collect(Collectors.toList());
210217
}
211218

212219
private static List<TableMetadata.MetadataLogEntry> updatePathInMetadataLogs(
213-
TableMetadata metadata, String sourcePrefix, String targetPrefix) {
220+
TableMetadata metadata, Map<String, String> prefixMappings) {
214221
List<TableMetadata.MetadataLogEntry> metadataLogEntries =
215222
Lists.newArrayListWithCapacity(metadata.previousFiles().size());
216223
for (TableMetadata.MetadataLogEntry metadataLog : metadata.previousFiles()) {
217224
TableMetadata.MetadataLogEntry newMetadataLog =
218225
new TableMetadata.MetadataLogEntry(
219-
metadataLog.timestampMillis(),
220-
newPath(metadataLog.file(), sourcePrefix, targetPrefix));
226+
metadataLog.timestampMillis(), newPath(metadataLog.file(), prefixMappings));
221227
metadataLogEntries.add(newMetadataLog);
222228
}
223229
return metadataLogEntries;
224230
}
225231

226232
private static List<Snapshot> updatePathInSnapshots(
227-
TableMetadata metadata, String sourcePrefix, String targetPrefix) {
233+
TableMetadata metadata, Map<String, String> prefixMappings) {
228234
List<Snapshot> newSnapshots = Lists.newArrayListWithCapacity(metadata.snapshots().size());
229235
for (Snapshot snapshot : metadata.snapshots()) {
230-
String newManifestListLocation =
231-
newPath(snapshot.manifestListLocation(), sourcePrefix, targetPrefix);
236+
String newManifestListLocation = newPath(snapshot.manifestListLocation(), prefixMappings);
232237
Snapshot newSnapshot =
233238
new BaseSnapshot(
234239
snapshot.sequenceNumber(),
@@ -254,8 +259,7 @@ private static List<Snapshot> updatePathInSnapshots(
254259
* @param io file io
255260
* @param tableMetadata metadata of table
256261
* @param manifestsToRewrite a list of manifest files to filter for rewrite
257-
* @param sourcePrefix source prefix that will be replaced
258-
* @param targetPrefix target prefix that will replace it
262+
* @param prefixMappings map of source prefix to target prefix mappings
259263
* @param stagingDir staging directory
260264
* @param outputPath location to write the manifest list
261265
* @return a copy plan for manifest files whose metadata were contained in the rewritten manifest
@@ -266,21 +270,22 @@ public static RewriteResult<ManifestFile> rewriteManifestList(
266270
FileIO io,
267271
TableMetadata tableMetadata,
268272
Set<String> manifestsToRewrite,
269-
String sourcePrefix,
270-
String targetPrefix,
273+
Map<String, String> prefixMappings,
271274
String stagingDir,
272275
String outputPath) {
273276
RewriteResult<ManifestFile> result = new RewriteResult<>();
274277
OutputFile outputFile = io.newOutputFile(outputPath);
275278

276279
List<ManifestFile> manifestFiles = manifestFilesInSnapshot(io, snapshot);
280+
// Validate that all manifest files can be matched by at least one prefix
277281
manifestFiles.forEach(
278-
mf ->
279-
Preconditions.checkArgument(
280-
mf.path().startsWith(sourcePrefix),
281-
"Encountered manifest file %s not under the source prefix %s",
282-
mf.path(),
283-
sourcePrefix));
282+
mf -> {
283+
Preconditions.checkArgument(
284+
lookupPrefixMappings(mf.path(), prefixMappings) != null,
285+
"Encountered manifest file %s not matching any source prefix in %s",
286+
mf.path(),
287+
prefixMappings.keySet());
288+
});
284289

285290
EncryptionManager encryptionManager =
286291
(io instanceof EncryptingFileIO)
@@ -299,14 +304,16 @@ public static RewriteResult<ManifestFile> rewriteManifestList(
299304

300305
for (ManifestFile file : manifestFiles) {
301306
ManifestFile newFile = file.copy();
302-
((StructLike) newFile).set(0, newPath(newFile.path(), sourcePrefix, targetPrefix));
307+
String newPath = newPath(newFile.path(), prefixMappings);
308+
((StructLike) newFile).set(0, newPath);
303309
writer.add(newFile);
304310

305311
if (manifestsToRewrite.contains(file.path())) {
306312
result.toRewrite().add(file);
313+
String stagingFilePath = stagingPath(file.path(), prefixMappings, stagingDir);
307314
result
308315
.copyPlan()
309-
.add(Pair.of(stagingPath(file.path(), sourcePrefix, stagingDir), newFile.path()));
316+
.add(Pair.of(stagingFilePath != null ? stagingFilePath : file.path(), newPath));
310317
}
311318
}
312319
return result;
@@ -604,28 +611,27 @@ PositionDeleteWriter<Record> writer(
604611
* @param outputFile output file to rewrite delete file to
605612
* @param io file io
606613
* @param spec spec of delete file
607-
* @param sourcePrefix source prefix that will be replaced
608-
* @param targetPrefix target prefix to replace it
614+
* @param prefixMappings source prefix and target prefix mappings
609615
* @param posDeleteReaderWriter class to read and write position delete files
610616
*/
611617
public static void rewritePositionDeleteFile(
612618
DeleteFile deleteFile,
613619
OutputFile outputFile,
614620
FileIO io,
615621
PartitionSpec spec,
616-
String sourcePrefix,
617-
String targetPrefix,
622+
Map<String, String> prefixMappings,
618623
PositionDeleteReaderWriter posDeleteReaderWriter)
619624
throws IOException {
620625
String path = deleteFile.location();
621-
if (!path.startsWith(sourcePrefix)) {
626+
Map.Entry<String, String> matchEntry = lookupPrefixMappings(path, prefixMappings);
627+
if (matchEntry == null) {
622628
throw new UnsupportedOperationException(
623-
String.format("Expected delete file %s to start with prefix: %s", path, sourcePrefix));
629+
String.format("Expected delete file %s to find in prefixes", path));
624630
}
625631

626632
// DV files (Puffin format for v3+) need special handling to rewrite internal blob metadata
627633
if (ContentFileUtil.isDV(deleteFile)) {
628-
rewriteDVFile(deleteFile, outputFile, io, sourcePrefix, targetPrefix);
634+
rewriteDVFile(deleteFile, outputFile, io, matchEntry.getKey(), matchEntry.getValue());
629635
return;
630636
}
631637

@@ -647,12 +653,13 @@ record = recordIt.next();
647653
posDeleteReaderWriter.writer(
648654
outputFile, deleteFile.format(), spec, deleteFile.partition(), rowSchema)) {
649655

650-
writer.write(newPositionDeleteRecord(record, sourcePrefix, targetPrefix));
656+
writer.write(newPositionDeleteRecord(record, matchEntry.getKey(), matchEntry.getValue()));
651657

652658
while (recordIt.hasNext()) {
653659
record = recordIt.next();
654660
if (record != null) {
655-
writer.write(newPositionDeleteRecord(record, sourcePrefix, targetPrefix));
661+
writer.write(
662+
newPositionDeleteRecord(record, matchEntry.getKey(), matchEntry.getValue()));
656663
}
657664
}
658665
}
@@ -728,7 +735,42 @@ private static PositionDelete newPositionDeleteRecord(
728735
}
729736

730737
/**
731-
* Rewrite a path by replacing its source prefix with a target prefix.
738+
* Lookup the longest matching prefix mapping for a given path.
739+
*
740+
* @param path the path to find a prefix mapping for
741+
* @param prefixMappings map of source prefix to target prefix mappings
742+
* @return the Map.Entry with the longest matching source prefix, or null if no match found
743+
*/
744+
public static Map.Entry<String, String> lookupPrefixMappings(
745+
String path, Map<String, String> prefixMappings) {
746+
if (prefixMappings == null || prefixMappings.isEmpty() || path == null) {
747+
return null;
748+
}
749+
750+
String normalizedPath = maybeAppendFileSeparator(path);
751+
return prefixMappings.entrySet().stream()
752+
.filter(entry -> normalizedPath.startsWith(maybeAppendFileSeparator(entry.getKey())))
753+
.max(java.util.Comparator.comparing(entry -> entry.getKey().length()))
754+
.orElse(null);
755+
}
756+
757+
public static String newPath(String path, Map<String, String> prefixMappings) {
758+
if (prefixMappings == null || prefixMappings.isEmpty()) {
759+
return path;
760+
}
761+
762+
Map.Entry<String, String> entry = lookupPrefixMappings(path, prefixMappings);
763+
if (entry != null) {
764+
String sourcePrefix = entry.getKey();
765+
String targetPrefix = entry.getValue();
766+
return combinePaths(targetPrefix, relativize(path, sourcePrefix));
767+
}
768+
769+
return path;
770+
}
771+
772+
/**
773+
* Replace path reference
732774
*
733775
* <p>If the path equals the source prefix (representing a directory location), the result will be
734776
* the target prefix with a trailing separator.
@@ -814,4 +856,23 @@ public static String stagingPath(String originalPath, String sourcePrefix, Strin
814856
String relativePath = relativize(originalPath, sourcePrefix);
815857
return combinePaths(stagingDir, relativePath);
816858
}
859+
860+
/**
861+
* Construct a staging path under a given staging directory, preserving relative directory
862+
* structure to avoid conflicts when multiple files have the same name but different paths.
863+
*
864+
* @param originalPath source path
865+
* @param prefixMappings source prefix and target prefix mappings
866+
* @param stagingDir staging directory
867+
* @return a staging path under the staging directory that preserves the relative path structure
868+
*/
869+
public static String stagingPath(
870+
String originalPath, Map<String, String> prefixMappings, String stagingDir) {
871+
Map.Entry<String, String> entry = lookupPrefixMappings(originalPath, prefixMappings);
872+
if (entry == null) {
873+
throw new IllegalArgumentException("unable to find prefix mapping for path: " + originalPath);
874+
}
875+
String relativePath = relativize(originalPath, entry.getKey());
876+
return combinePaths(stagingDir, relativePath);
877+
}
817878
}

0 commit comments

Comments
 (0)