Construct iceberg data files during commit step (#4140)

Author: thisisArjit

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/IcebergOverwritePartitionsStep.java

@@ -19,7 +19,6 @@
 
 import java.io.IOException;
 import java.time.Duration;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import java.util.Properties;
@@ -28,7 +27,6 @@
 
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.catalog.TableIdentifier;
-import org.apache.iceberg.util.SerializationUtil;
 
 import com.github.rholder.retry.Attempt;
 import com.github.rholder.retry.RetryException;
@@ -38,6 +36,7 @@
 import com.typesafe.config.Config;
 import com.typesafe.config.ConfigFactory;
 
+import lombok.Setter;
 import lombok.extern.slf4j.Slf4j;
 
 import org.apache.gobblin.commit.CommitStep;
@@ -59,8 +58,8 @@
 public class IcebergOverwritePartitionsStep implements CommitStep {
   private final String destTableIdStr;
   private final Properties properties;
-  // Data files are kept as a list of base64 encoded strings for optimised de-serialization.
-  private final List<String> base64EncodedDataFiles;
+  // data files are populated once all the copy tasks are done. Each IcebergPartitionCopyableFile has a serialized data file
+  @Setter private List<DataFile> dataFiles;
   private final String partitionColName;
   private final String partitionValue;
   public static final String OVERWRITE_PARTITIONS_RETRYER_CONFIG_PREFIX = IcebergDatasetFinder.ICEBERG_DATASET_PREFIX +
@@ -74,14 +73,12 @@ public class IcebergOverwritePartitionsStep implements CommitStep {
    * Constructs an {@code IcebergReplacePartitionsStep} with the specified parameters.
    *
    * @param destTableIdStr the identifier of the destination table as a string
-   * @param base64EncodedDataFiles [from List<DataFiles>] the serialized data files to be used for replacing partitions
    * @param properties the properties containing configuration
    */
-  public IcebergOverwritePartitionsStep(String destTableIdStr, String partitionColName, String partitionValue, List<String> base64EncodedDataFiles, Properties properties) {
+  public IcebergOverwritePartitionsStep(String destTableIdStr, String partitionColName, String partitionValue, Properties properties) {
     this.destTableIdStr = destTableIdStr;
     this.partitionColName = partitionColName;
     this.partitionValue = partitionValue;
-    this.base64EncodedDataFiles = base64EncodedDataFiles;
     this.properties = properties;
   }
 
@@ -103,7 +100,6 @@ public void execute() throws IOException {
     // our copying. any new data written in the meanwhile to THE SAME partitions we are about to overwrite will be
     // clobbered and replaced by the copy entities from our execution.
     IcebergTable destTable = createDestinationCatalog().openTable(TableIdentifier.parse(this.destTableIdStr));
-    List<DataFile> dataFiles = getDataFiles();
     try {
       log.info("~{}~ Starting partition overwrite - partition: {}; value: {}; numDataFiles: {}; path[0]: {}",
           this.destTableIdStr,
@@ -140,14 +136,6 @@ public void execute() throws IOException {
     }
   }
 
-  private List<DataFile> getDataFiles() {
-    List<DataFile> dataFiles = new ArrayList<>(base64EncodedDataFiles.size());
-    for (String base64EncodedDataFile : base64EncodedDataFiles) {
-      dataFiles.add(SerializationUtil.deserializeFromBase64(base64EncodedDataFile));
-    }
-    return dataFiles;
-  }
-
   protected IcebergCatalog createDestinationCatalog() throws IOException {
     return IcebergDatasetFinder.createIcebergCatalog(this.properties, IcebergDatasetFinder.CatalogLocation.DESTINATION);
   }

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/IcebergPartitionCopyableFile.java

@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.gobblin.data.management.copy.iceberg;
+
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.util.SerializationUtil;
+
+import lombok.AccessLevel;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+import lombok.Setter;
+import lombok.extern.slf4j.Slf4j;
+
+import org.apache.gobblin.data.management.copy.CopyableFile;
+
+
+/**
+ * An extension of {@link CopyableFile} that includes a base64-encoded Iceberg {@link DataFile}.
+ */
+@Getter
+@Setter
+@NoArgsConstructor(access = AccessLevel.PROTECTED)
+@EqualsAndHashCode(callSuper = true)
+@Slf4j
+public class IcebergPartitionCopyableFile extends CopyableFile {
+
+  /**
+   * Base64-encoded Iceberg {@link DataFile} associated with this copyable file.
+   */
+  private String base64EncodedDataFile;
+
+  public IcebergPartitionCopyableFile(CopyableFile copyableFile, DataFile dataFile) {
+    super(copyableFile.getOrigin(), copyableFile.getDestination(), copyableFile.getDestinationOwnerAndPermission(),
+        copyableFile.getAncestorsOwnerAndPermission(), copyableFile.getChecksum(), copyableFile.getPreserve(),
+        copyableFile.getFileSet(), copyableFile.getOriginTimestamp(), copyableFile.getUpstreamTimestamp(),
+        copyableFile.getAdditionalMetadata(), copyableFile.datasetOutputPath, copyableFile.getDataFileVersionStrategy());
+    this.base64EncodedDataFile = SerializationUtil.serializeToBase64(dataFile);
+  }
+
+  public DataFile getDataFile() {
+    return SerializationUtil.deserializeFromBase64(base64EncodedDataFile);
+  }
+}

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/IcebergPartitionDataset.java

@@ -20,15 +20,13 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.UUID;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
 import java.util.function.Predicate;
-import java.util.stream.Collectors;
 
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang3.StringUtils;
@@ -42,9 +40,7 @@
 import org.apache.iceberg.StructLike;
 import org.apache.iceberg.TableMetadata;
 import org.apache.iceberg.TableProperties;
-import org.apache.iceberg.util.SerializationUtil;
 
-import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.ImmutableList;
 import com.google.common.base.Preconditions;
@@ -101,48 +97,16 @@ Collection<CopyEntity> generateCopyEntities(FileSystem targetFs, CopyConfigurati
     // TODO: Refactor the IcebergDataset::generateCopyEntities to avoid code duplication
     //  Differences are getting data files, copying ancestor permission and adding post publish steps
     String fileSet = this.getFileSetId();
-    List<CopyEntity> copyEntities = Lists.newArrayList();
     IcebergTable srcIcebergTable = getSrcIcebergTable();
     List<DataFile> srcDataFiles = srcIcebergTable.getPartitionSpecificDataFiles(this.partitionFilterPredicate);
-    Map<Path, DataFile> destDataFileBySrcPath = calcDestDataFileBySrcPath(srcDataFiles);
-    Configuration defaultHadoopConfiguration = new Configuration();
-
-    for (Map.Entry<Path, FileStatus> entry : calcSrcFileStatusByDestFilePath(destDataFileBySrcPath).entrySet()) {
-      Path destPath = entry.getKey();
-      FileStatus srcFileStatus = entry.getValue();
-      // TODO: should be the same FS each time; try creating once, reusing thereafter, to not recreate wastefully
-      FileSystem actualSourceFs = getSourceFileSystemFromFileStatus(srcFileStatus, defaultHadoopConfiguration);
-
-      CopyableFile fileEntity = CopyableFile.fromOriginAndDestination(
-              actualSourceFs, srcFileStatus, targetFs.makeQualified(destPath), copyConfig)
-          .fileSet(fileSet)
-          .datasetOutputPath(targetFs.getUri().getPath())
-          .build();
-
-      fileEntity.setSourceData(getSourceDataset(this.sourceFs));
-      fileEntity.setDestinationData(getDestinationDataset(targetFs));
-      copyEntities.add(fileEntity);
-    }
-
-    // Adding this check to avoid adding post publish step when there are no files to copy.
-    List<DataFile> destDataFiles = new ArrayList<>(destDataFileBySrcPath.values());
-    if (CollectionUtils.isNotEmpty(destDataFiles)) {
-      copyEntities.add(createOverwritePostPublishStep(destDataFiles));
-    }
-
-    log.info("~{}~ generated {} copy entities", fileSet, copyEntities.size());
-    return copyEntities;
-  }
 
-  private Map<Path, DataFile> calcDestDataFileBySrcPath(List<DataFile> srcDataFiles)
-      throws IcebergTable.TableNotFoundException {
-    String fileSet = this.getFileSetId();
-    Map<Path, DataFile> destDataFileBySrcPath = new ConcurrentHashMap<>(srcDataFiles.size());
     if (srcDataFiles.isEmpty()) {
       log.warn("~{}~ found no data files for partition col : {} with partition value : {} to copy", fileSet,
           this.partitionColumnName, this.partitionColValue);
-      return destDataFileBySrcPath;
+      return new ArrayList<>(0);
     }
+
+    // get source & destination write data locations to update data file paths
     TableMetadata srcTableMetadata = getSrcIcebergTable().accessTableMetadata();
     TableMetadata destTableMetadata = getDestIcebergTable().accessTableMetadata();
     PartitionSpec partitionSpec = destTableMetadata.spec();
@@ -160,17 +124,58 @@ private Map<Path, DataFile> calcDestDataFileBySrcPath(List<DataFile> srcDataFile
           destWriteDataLocation
       );
     }
-    srcDataFiles.forEach(dataFile -> {
+
+    List<CopyEntity> copyEntities = getIcebergParitionCopyEntities(targetFs, srcDataFiles, srcWriteDataLocation, destWriteDataLocation, partitionSpec, copyConfig);
+    // Adding this check to avoid adding post publish step when there are no files to copy.
+    if (CollectionUtils.isNotEmpty(copyEntities)) {
+      copyEntities.add(createOverwritePostPublishStep());
+    }
+
+    log.info("~{}~ generated {} copy entities", fileSet, copyEntities.size());
+    return copyEntities;
+  }
+
+  private List<CopyEntity> getIcebergParitionCopyEntities(
+      FileSystem targetFs,
+      List<DataFile> srcDataFiles,
+      String srcWriteDataLocation,
+      String destWriteDataLocation,
+      PartitionSpec partitionSpec,
+      CopyConfiguration copyConfig) {
+    String fileSet = this.getFileSetId();
+    Configuration defaultHadoopConfiguration = new Configuration();
+    List<CopyEntity> copyEntities = Collections.synchronizedList(new ArrayList<>(srcDataFiles.size()));
+    Function<Path, FileStatus> getFileStatus = CheckedExceptionFunction.wrapToTunneled(this.sourceFs::getFileStatus);
+
+    srcDataFiles.parallelStream().forEach(dataFile -> {
+      // create destination data file from source data file by replacing the source path with destination path
       String srcFilePath = dataFile.path().toString();
       Path updatedDestFilePath = relocateDestPath(srcFilePath, srcWriteDataLocation, destWriteDataLocation);
       log.debug("~{}~ Path changed from Src : {} to Dest : {}", fileSet, srcFilePath, updatedDestFilePath);
-      destDataFileBySrcPath.put(new Path(srcFilePath), DataFiles.builder(partitionSpec)
+      DataFile destDataFile = DataFiles.builder(partitionSpec)
           .copy(dataFile)
           .withPath(updatedDestFilePath.toString())
-          .build());
+          .build();
+
+      // get file status of source file
+      FileStatus srcFileStatus = getFileStatus.apply(new Path(srcFilePath));
+      try {
+        // TODO: should be the same FS each time; try creating once, reusing thereafter, to not recreate wastefully
+        FileSystem actualSourceFs = getSourceFileSystemFromFileStatus(srcFileStatus, defaultHadoopConfiguration);
+        // create copyable file entity
+        CopyableFile fileEntity = CopyableFile.fromOriginAndDestination(actualSourceFs, srcFileStatus,
+                targetFs.makeQualified(updatedDestFilePath), copyConfig).fileSet(fileSet)
+            .datasetOutputPath(targetFs.getUri().getPath()).build();
+        fileEntity.setSourceData(getSourceDataset(this.sourceFs));
+        fileEntity.setDestinationData(getDestinationDataset(targetFs));
+        // add corresponding data file to each copyable iceberg partition file
+        IcebergPartitionCopyableFile icebergPartitionCopyableFile = new IcebergPartitionCopyableFile(fileEntity, destDataFile);
+        copyEntities.add(icebergPartitionCopyableFile);
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
     });
-    log.info("~{}~ created {} destination data files", fileSet, destDataFileBySrcPath.size());
-    return destDataFileBySrcPath;
+    return copyEntities;
   }
 
   private Path relocateDestPath(String curPathStr, String prefixToBeReplaced, String prefixToReplaceWith) {
@@ -186,43 +191,17 @@ private Path addUUIDToPath(String filePathStr) {
     return new Path(fileDir, newFileName);
   }
 
-  private Map<Path, FileStatus> calcSrcFileStatusByDestFilePath(Map<Path, DataFile> destDataFileBySrcPath)
-      throws IOException {
-    Function<Path, FileStatus> getFileStatus = CheckedExceptionFunction.wrapToTunneled(this.sourceFs::getFileStatus);
-    Map<Path, FileStatus> srcFileStatusByDestFilePath = new ConcurrentHashMap<>();
-    try {
-      srcFileStatusByDestFilePath = destDataFileBySrcPath.entrySet()
-          .parallelStream()
-          .collect(Collectors.toConcurrentMap(entry -> new Path(entry.getValue().path().toString()),
-              entry -> getFileStatus.apply(entry.getKey())));
-    } catch (CheckedExceptionFunction.WrappedIOException wrapper) {
-      wrapper.rethrowWrapped();
-    }
-    return srcFileStatusByDestFilePath;
-  }
-
-  private PostPublishStep createOverwritePostPublishStep(List<DataFile> destDataFiles) {
-    List<String> serializedDataFiles = getBase64EncodedDataFiles(destDataFiles);
-
+  private PostPublishStep createOverwritePostPublishStep() {
     IcebergOverwritePartitionsStep icebergOverwritePartitionStep = new IcebergOverwritePartitionsStep(
         this.getDestIcebergTable().getTableId().toString(),
         this.partitionColumnName,
         this.partitionColValue,
-        serializedDataFiles,
         this.properties
     );
 
     return new PostPublishStep(this.getFileSetId(), Maps.newHashMap(), icebergOverwritePartitionStep, 0);
   }
 
-  private List<String> getBase64EncodedDataFiles(List<DataFile> destDataFiles) {
-    List<String> base64EncodedDataFiles = new ArrayList<>(destDataFiles.size());
-    for (DataFile dataFile : destDataFiles) {
-      base64EncodedDataFiles.add(SerializationUtil.serializeToBase64(dataFile));
-    }
-    return base64EncodedDataFiles;
-  }
-
   private Predicate<StructLike> createPartitionFilterPredicate() throws IOException {
     //TODO: Refactor it later using factory or other way to support different types of filter predicate
     // Also take into consideration creation of Expression Filter to be used in overwrite api