Replace Parquet strategy pattern with unified parallel approach

Author: Artem

migrator/src/main/scala/com/scylladb/migrator/config/Savepoints.scala

@@ -3,36 +3,7 @@ package com.scylladb.migrator.config
 import io.circe.{ Decoder, Encoder }
 import io.circe.generic.semiauto.{ deriveDecoder, deriveEncoder }
 
-sealed trait ParquetProcessingMode
-object ParquetProcessingMode {
-  case object Parallel extends ParquetProcessingMode
-  case object Sequential extends ParquetProcessingMode
-
-  implicit val encoder: Encoder[ParquetProcessingMode] = Encoder.encodeString.contramap {
-    case Parallel   => "parallel"
-    case Sequential => "sequential"
-  }
-
-  implicit val decoder: Decoder[ParquetProcessingMode] = Decoder.decodeString.emap {
-    case "parallel"   => Right(Parallel)
-    case "sequential" => Right(Sequential)
-    case other =>
-      Left(s"Unknown parquet processing mode: $other. Valid values: parallel, sequential")
-  }
-}
-
-case class Savepoints(intervalSeconds: Int,
-                      path: String,
-                      parquetProcessingMode: Option[ParquetProcessingMode]) {
-
-  /**
-    * Returns the configured Parquet processing mode.
-    * Defaults to [[ParquetProcessingMode.Parallel]] if not specified.
-    * This default affects migration semantics.
-    */
-  def getParquetProcessingMode: ParquetProcessingMode =
-    parquetProcessingMode.getOrElse(ParquetProcessingMode.Parallel)
-}
+case class Savepoints(intervalSeconds: Int, path: String)
 
 object Savepoints {
   implicit val encoder: Encoder[Savepoints] = deriveEncoder[Savepoints]

migrator/src/main/scala/com/scylladb/migrator/readers/FileCompletionListener.scala

@@ -0,0 +1,93 @@
+package com.scylladb.migrator.readers
+
+import org.apache.log4j.LogManager
+import org.apache.spark.scheduler.{ SparkListener, SparkListenerTaskEnd }
+import org.apache.spark.Success
+
+import scala.collection.concurrent.TrieMap
+
+/**
+  * SparkListener that tracks partition completion and aggregates it to file-level completion.
+  *
+  * This listener monitors Spark task completion events and maintains mappings between
+  * partitions and files. When all partitions belonging to a file have been successfully
+  * completed, it marks the file as processed via the ParquetSavepointsManager.
+  *
+  * @param partitionToFile Mapping from Spark partition ID to source file paths
+  * @param fileToPartitions Mapping from file path to the set of partition IDs reading from it
+  * @param savepointsManager Manager to notify when files are completed
+  */
+class FileCompletionListener(
+  partitionToFiles: Map[Int, Set[String]],
+  fileToPartitions: Map[String, Set[Int]],
+  savepointsManager: ParquetSavepointsManager
+) extends SparkListener {
+
+  private val log = LogManager.getLogger("com.scylladb.migrator.readers.FileCompletionListener")
+
+  private val completedPartitions = TrieMap.empty[Int, Boolean]
+
+  private val completedFiles = TrieMap.empty[String, Boolean]
+
+  log.info(
+    s"FileCompletionListener initialized: tracking ${fileToPartitions.size} files " +
+      s"across ${partitionToFiles.size} partitions")
+
+  override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit =
+    if (taskEnd.reason == Success) {
+      val partitionId = taskEnd.taskInfo.partitionId
+
+      partitionToFiles.get(partitionId) match {
+        case Some(filenames) =>
+          if (completedPartitions.putIfAbsent(partitionId, true).isEmpty) {
+            filenames.foreach { filename =>
+              log.debug(s"Partition $partitionId completed (file: $filename)")
+              checkFileCompletion(filename)
+            }
+          }
+
+        case None =>
+          log.trace(s"Task completed for untracked partition $partitionId")
+      }
+    } else {
+      log.debug(
+        s"Task for partition ${taskEnd.taskInfo.partitionId} did not complete successfully: ${taskEnd.reason}")
+    }
+
+  private def checkFileCompletion(filename: String): Unit = {
+    if (completedFiles.contains(filename)) {
+      return
+    }
+
+    fileToPartitions.get(filename) match {
+      case Some(allPartitions) =>
+        val allComplete = allPartitions.forall(completedPartitions.contains)
+
+        if (allComplete) {
+          if (completedFiles.putIfAbsent(filename, true).isEmpty) {
+            savepointsManager.markFileAsProcessed(filename)
+
+            val progress = s"${completedFiles.size}/${fileToPartitions.size}"
+            log.info(s"File completed: $filename (progress: $progress)")
+          }
+        } else {
+          val completedCount = allPartitions.count(completedPartitions.contains)
+          log.trace(s"File $filename: $completedCount/${allPartitions.size} partitions complete")
+        }
+
+      case None =>
+        log.warn(s"File $filename not found in fileToPartitions map (this shouldn't happen)")
+    }
+  }
+
+  def getCompletedFilesCount: Int = completedFiles.size
+
+  def getTotalFilesCount: Int = fileToPartitions.size
+
+  def getProgressReport: String = {
+    val filesCompleted = getCompletedFilesCount
+    val totalFiles = getTotalFilesCount
+
+    s"Progress: $filesCompleted/$totalFiles files"
+  }
+}

migrator/src/main/scala/com/scylladb/migrator/readers/ParallelParquetStrategy.scala

@@ -1,61 +1,73 @@
 package com.scylladb.migrator.readers
 
-import com.scylladb.migrator.config.{
-  MigratorConfig,
-  ParquetProcessingMode,
-  SourceSettings,
-  TargetSettings
-}
-import com.scylladb.migrator.scylla.SourceDataFrame
-import com.scylladb.migrator.scylla
+import com.scylladb.migrator.config.{ MigratorConfig, SourceSettings, TargetSettings }
+import com.scylladb.migrator.scylla.{ ScyllaParquetMigrator, SourceDataFrame }
 import org.apache.log4j.LogManager
 import org.apache.spark.sql.{ AnalysisException, SparkSession }
 import scala.util.Using
 
-case class ParquetReaderWithSavepoints(source: SourceSettings.Parquet,
-                                       allFiles: Seq[String],
-                                       skipFiles: Set[String]) {
-
-  val filesToProcess: Seq[String] = allFiles.filterNot(skipFiles.contains)
-
-  def configureHadoop(spark: SparkSession): Unit =
-    Parquet.configureHadoopCredentials(spark, source)
-}
-
 object Parquet {
   val log = LogManager.getLogger("com.scylladb.migrator.readers.Parquet")
 
   def migrateToScylla(config: MigratorConfig,
                       source: SourceSettings.Parquet,
                       target: TargetSettings.Scylla)(implicit spark: SparkSession): Unit = {
-    val processingMode = config.savepoints.getParquetProcessingMode
-
-    val strategy: ParquetProcessingStrategy = processingMode match {
-      case ParquetProcessingMode.Parallel =>
-        log.info("Selected PARALLEL processing mode (default)")
-        new ParallelParquetStrategy()
-      case ParquetProcessingMode.Sequential =>
-        log.info("Selected SEQUENTIAL processing mode (with savepoints)")
-        new SequentialParquetStrategy()
-    }
-
-    strategy.migrate(config, source, target)
-  }
-
-  def prepareParquetReader(spark: SparkSession,
-                           source: SourceSettings.Parquet,
-                           skipFiles: Set[String] = Set.empty): ParquetReaderWithSavepoints = {
+    log.info("Starting Parquet migration with parallel processing and file-level savepoints")
 
     configureHadoopCredentials(spark, source)
 
     val allFiles = listParquetFiles(spark, source.path)
-    log.info(s"Found ${allFiles.size} Parquet files in ${source.path}")
+    val skipFiles = config.getSkipParquetFilesOrEmptySet
+    val filesToProcess = allFiles.filterNot(skipFiles.contains)
+
+    if (filesToProcess.isEmpty) {
+      log.info("No Parquet files to process. Migration is complete.")
+      return
+    }
+
+    log.info(s"Processing ${filesToProcess.size} Parquet files")
 
-    if (skipFiles.nonEmpty) {
-      log.info(s"Skipping ${skipFiles.size} already processed files")
+    val df = if (skipFiles.isEmpty) {
+      spark.read.parquet(source.path)
+    } else {
+      spark.read.parquet(filesToProcess: _*)
     }
 
-    ParquetReaderWithSavepoints(source, allFiles, skipFiles)
+    log.info("Reading partition metadata for file tracking...")
+    val metadata = PartitionMetadataReader.readMetadataFromDataFrame(df)
+
+    val partitionToFiles = PartitionMetadataReader.buildPartitionToFileMap(metadata)
+    val fileToPartitions = PartitionMetadataReader.buildFileToPartitionsMap(metadata)
+
+    log.info(
+      s"Discovered ${fileToPartitions.size} files with ${metadata.size} total partitions to process")
+
+    Using.resource(ParquetSavepointsManager(config, spark.sparkContext)) { savepointsManager =>
+      val listener = new FileCompletionListener(
+        partitionToFiles,
+        fileToPartitions,
+        savepointsManager
+      )
+      spark.sparkContext.addSparkListener(listener)
+
+      try {
+        val sourceDF = SourceDataFrame(df, None, savepointsSupported = false)
+
+        log.info("Created DataFrame from Parquet source")
+
+        ScyllaParquetMigrator.migrate(config, target, sourceDF, savepointsManager)
+
+        savepointsManager.dumpMigrationState("completed")
+
+        log.info(
+          s"Parquet migration completed successfully: " +
+            s"${listener.getCompletedFilesCount}/${listener.getTotalFilesCount} files processed")
+
+      } finally {
+        spark.sparkContext.removeSparkListener(listener)
+        log.info(s"Final progress: ${listener.getProgressReport}")
+      }
+    }
   }
 
   def listParquetFiles(spark: SparkSession, path: String): Seq[String] = {
@@ -88,7 +100,7 @@ object Parquet {
     * This method sets the necessary Hadoop configuration properties for AWS access key, secret key,
     * and optionally a session token. When a session token is present, it sets the credentials provider
     * to TemporaryAWSCredentialsProvider as required by Hadoop.
-    * 
+    *
     * If a region is specified in the source configuration, this method also sets the S3A endpoint region
     * via the `fs.s3a.endpoint.region` property.
     *

migrator/src/main/scala/com/scylladb/migrator/readers/Parquet.scala

@@ -0,0 +1,85 @@
+package com.scylladb.migrator.readers
+
+import org.apache.log4j.LogManager
+import org.apache.spark.sql.{ DataFrame, SparkSession }
+import org.apache.spark.sql.functions._
+
+case class PartitionMetadata(
+  partitionId: Int,
+  filename: String
+)
+
+/**
+  * This reader uses Spark's internal partition information to build mappings
+  * between partition IDs and file paths. This allows us to track when all
+  * partitions of a file have been processed, enabling file-level savepointse.
+  */
+object PartitionMetadataReader {
+  private val logger = LogManager.getLogger("com.scylladb.migrator.readers.PartitionMetadataReader")
+
+  def readMetadata(spark: SparkSession, filePaths: Seq[String]): Seq[PartitionMetadata] = {
+    logger.info(s"Reading partition metadata from ${filePaths.size} file(s)")
+    val df = spark.read.parquet(filePaths: _*)
+    readMetadataFromDataFrame(df)
+  }
+
+  def readMetadataFromDataFrame(df: DataFrame): Seq[PartitionMetadata] =
+    try {
+
+      val partitionInfo = df
+        .select(input_file_name().as("filename"))
+        .rdd
+        .mapPartitionsWithIndex { (partitionId, iter) =>
+          val files = iter.map(row => row.getString(0)).toSet
+          files.map(filename => (partitionId, filename)).iterator
+        }
+        .collect()
+        .toSeq
+
+      val metadata = partitionInfo.zipWithIndex.map {
+        case ((partitionId, filename), idx) =>
+          PartitionMetadata(
+            partitionId = partitionId,
+            filename    = filename
+          )
+      }
+
+      logger.info(s"Discovered ${metadata.size} partition-to-file mappings")
+
+      val fileStats = metadata.groupBy(_.filename).view.mapValues(_.size)
+      logger.info(s"Files distribution: ${fileStats.size} unique files")
+      fileStats.foreach {
+        case (file, partCount) =>
+          logger.debug(s"  File: $file -> $partCount partition(s)")
+      }
+
+      metadata
+
+    } catch {
+      case e: Exception =>
+        logger.error(s"Failed to read partition metadata", e)
+        throw new RuntimeException(s"Could not read partition metadata: ${e.getMessage}", e)
+    }
+
+  def buildFileToPartitionsMap(metadata: Seq[PartitionMetadata]): Map[String, Set[Int]] = {
+    val result = metadata
+      .groupBy(_.filename)
+      .view
+      .mapValues(_.map(_.partitionId).toSet)
+      .toMap
+
+    logger.debug(s"Built file-to-partitions map with ${result.size} files")
+    result
+  }
+
+  def buildPartitionToFileMap(metadata: Seq[PartitionMetadata]): Map[Int, Set[String]] = {
+    val result = metadata
+      .groupBy(_.partitionId)
+      .view
+      .mapValues(_.map(_.filename).toSet)
+      .toMap
+
+    logger.debug(s"Built partition-to-file map with ${result.size} partitions")
+    result
+  }
+}