diff --git a/pramen/core/src/main/resources/reference.conf b/pramen/core/src/main/resources/reference.conf index 5e1a5330d..ddcc78bd8 100644 --- a/pramen/core/src/main/resources/reference.conf +++ b/pramen/core/src/main/resources/reference.conf @@ -150,8 +150,10 @@ pramen { # When runtime.run.mode = bulk, specifies the size of bulk load #runtime.run.bulk.batch.size = monthly + # You can enable repartitioning for bulk jobs + #runtime.enable.repartitioning = true - #runtime.run.bulk.curreent { + #runtime.run.bulk.current { # date.from = "" # date.to = "" # output.information.date = "" diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/BulkRunConfig.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/BulkRunConfig.scala index daa2d9502..7e4c1a060 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/BulkRunConfig.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/BulkRunConfig.scala @@ -19,7 +19,8 @@ package za.co.absa.pramen.core.app.config import java.time.LocalDate case class BulkRunConfig( - dateFrom: LocalDate, - dateTo: LocalDate, + dataDateFrom: LocalDate, + dataDateTo: LocalDate, + infoDateColumn: Option[String], outputInfoDate: LocalDate ) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/RuntimeConfig.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/RuntimeConfig.scala index 26d6c0ff3..101b08e83 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/RuntimeConfig.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/app/config/RuntimeConfig.scala @@ -49,6 +49,7 @@ case class RuntimeConfig( historicalRunMode: RunMode, bulkBatchSize: BulkBatchSize, bulkLoadCurrent: Option[BulkRunConfig], + enableRepartitioning: Boolean, sparkAppDescriptionTemplate: Option[String], attempt: Int, // Current attempt number for the pipeline run (for auto-retry automation) maxAttempts: Int, // Maximum number of attempts allowed for the pipeline run @@ -65,10 +66,12 @@ object RuntimeConfig { val IS_RERUN = "pramen.runtime.is.rerun" val IS_INVERSE_ORDER = "pramen.runtime.inverse.order" val RUN_MODE = "pramen.runtime.run.mode" - val RUN_BULK_BATCH_SIZE = "pramen.runtime.run.bulk.batch.size" - val BULK_CURRENT_DATE_FROM = "pramen.runtime.run.bulk.curreent.date.from" - val BULK_CURRENT_DATE_TO = "pramen.runtime.run.bulk.curreent.date.to" - val BULK_CURRENT_OUTPUT_INFO_DATE = "pramen.runtime.run.bulk.curreent.output.information.date" + val RUN_BULK_BATCH_SIZE = "pramen.runtime.bulk.batch.size" + val RUN_ENABLE_REPARTITIONING = "pramen.runtime.enable.repartitioning" + val INFO_DATE_COLUMN = "pramen.runtime.info.date.column" + val BULK_CURRENT_DATE_FROM = "pramen.runtime.run.bulk.current.date.from" + val BULK_CURRENT_DATE_TO = "pramen.runtime.run.bulk.current.date.to" + val BULK_CURRENT_OUTPUT_INFO_DATE = "pramen.runtime.run.bulk.current.output.information.date" val RUN_TABLES = "pramen.runtime.run.tables" val UNDERCOVER = "pramen.undercover" val USE_LOCK = "pramen.use.lock" @@ -161,9 +164,11 @@ object RuntimeConfig { val bulkCurrentDateFrom = ConfigUtils.getOptionString(conf, BULK_CURRENT_DATE_FROM).map(getDate) val bulkCurrentDateTo = ConfigUtils.getOptionString(conf, BULK_CURRENT_DATE_TO).map(getDate) val bulkCurrentOutputInfoDate = ConfigUtils.getOptionString(conf, BULK_CURRENT_OUTPUT_INFO_DATE).map(getDate) + val infoDateColumn = ConfigUtils.getOptionString(conf, INFO_DATE_COLUMN) + val enableRepartitioning = ConfigUtils.getOptionBoolean(conf, RUN_ENABLE_REPARTITIONING).getOrElse(false) val bulkLoadCurrent = if (bulkCurrentDateFrom.isDefined && bulkCurrentDateTo.isDefined && bulkCurrentOutputInfoDate.isDefined) { - Some(BulkRunConfig(bulkCurrentDateFrom.get, bulkCurrentDateTo.get, bulkCurrentOutputInfoDate.get)) + Some(BulkRunConfig(bulkCurrentDateFrom.get, bulkCurrentDateTo.get, infoDateColumn, bulkCurrentOutputInfoDate.get)) } else { None } @@ -192,6 +197,7 @@ object RuntimeConfig { attempt = attempt, maxAttempts = maxAttempts, forceReCreateHiveTables = ConfigUtils.getOptionBoolean(conf, FORCE_RECREATE_HIVE_TABLES).getOrElse(false), + enableRepartitioning = enableRepartitioning, executionOptions = executionOptions ) } @@ -217,6 +223,7 @@ object RuntimeConfig { historicalRunMode = RunMode.CheckUpdates, bulkBatchSize = Monthly, bulkLoadCurrent = None, + enableRepartitioning = false, sparkAppDescriptionTemplate = None, attempt = 1, maxAttempts = 1, diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/cmd/CmdLineConfig.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/cmd/CmdLineConfig.scala index e7933c9c0..dc42861a8 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/cmd/CmdLineConfig.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/cmd/CmdLineConfig.scala @@ -44,6 +44,7 @@ case class CmdLineConfig( dateTo: Option[LocalDate] = None, mode: Option[String] = None, bulkSize: Option[String] = None, + infoDateColumn: Option[String] = None, inverseOrder: Option[Boolean] = None, verbose: Option[Boolean] = None, overrideLogLevel: Option[String] = None, @@ -133,6 +134,9 @@ object CmdLineConfig { for (bulkSize <- cmd.bulkSize) accumulatedConfig = accumulatedConfig.withValue(RUN_BULK_BATCH_SIZE, ConfigValueFactory.fromAnyRef(bulkSize)) + for (infoDateColumn <- cmd.infoDateColumn) + accumulatedConfig = accumulatedConfig.withValue(INFO_DATE_COLUMN, ConfigValueFactory.fromAnyRef(infoDateColumn)) + for (logEffectiveConfig <- cmd.logEffectiveConfig) accumulatedConfig = accumulatedConfig.withValue(LOG_EFFECTIVE_CONFIG, ConfigValueFactory.fromAnyRef(logEffectiveConfig)) @@ -213,7 +217,13 @@ object CmdLineConfig { .text("The bulk size for processing date ranges.") .validate(v => if (v == "monthly" || v == "quarterly" || v == "yearly") success - else failure("Invalid bulk size. Must be one of 'monthly', 'quarterly', 'yearly'")) + else failure("Invalid bulk size. Must be one of 'monthly', 'quarterly', 'yearly'")), + opt[String]("info-date-column").optional().action((value, config) => + config.copy(infoDateColumn = Option(value))) + .text("The information date column name to use for repartitioning.") + .validate(v => + if (v.nonEmpty) success + else failure("Invalid information date column name. Must be a non-empty string.")) ) opt[Boolean]("inverse-order").optional().action((value, config) => diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistence.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistence.scala index fb71d9f9c..706944763 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistence.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistence.scala @@ -40,6 +40,12 @@ trait MetastorePersistence { def repairHiveTable(hiveTableName: String, queryExecutor: QueryExecutor, hiveConfig: HiveConfig): Unit + + def isRepartitioningSupported: Boolean + + def repartitionPhase1(infoDateColumn: String, infoDateFrom: LocalDate, infoDateTo: LocalDate, outputInfoDate: LocalDate): Unit = {} + + def repartitionPhase2(infoDateColumn: String, infoDateFrom: LocalDate, infoDateTo: LocalDate, outputInfoDate: LocalDate): Unit = {} } object MetastorePersistence { diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceDelta.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceDelta.scala index 31c9d7b63..f57cc80ff 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceDelta.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceDelta.scala @@ -215,6 +215,8 @@ class MetastorePersistenceDelta(query: Query, throw new UnsupportedOperationException("Delta format does not support Hive tables at the moment.") } + override def isRepartitioningSupported: Boolean = false + def getFilter(infoDateFrom: Option[LocalDate], infoDateTo: Option[LocalDate]): Column = { if (partitionScheme == PartitionScheme.Overwrite) { if (infoDateFrom.isDefined || infoDateTo.isDefined) { diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIceberg.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIceberg.scala index 273ff5f6c..aa27ab150 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIceberg.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIceberg.scala @@ -18,6 +18,7 @@ package za.co.absa.pramen.core.metastore.peristence import org.apache.spark.sql._ import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.DateType import org.slf4j.LoggerFactory import za.co.absa.pramen.api.{CatalogTable, PartitionScheme} import za.co.absa.pramen.core.metastore.MetaTableStats @@ -112,6 +113,29 @@ class MetastorePersistenceIceberg(table: CatalogTable, throw new UnsupportedOperationException("Iceberg only operates on tables in a catalog. Separate Hive options are not supported.") } + override def isRepartitioningSupported: Boolean = true + + override def repartitionPhase1(infoDateDataColumn: String, infoDateFrom: LocalDate, infoDateTo: LocalDate, outputInfoDate: LocalDate): Unit = { + if (infoDateColumn.equalsIgnoreCase(infoDateDataColumn)) + throw new IllegalArgumentException(s"Cannot repartition a table if the metastore info date column is the same as the data info date column ($infoDateDataColumn)") + + val fullTableName = table.getFullTableName + val df = spark.table(fullTableName) + .filter(getFilter(Some(outputInfoDate), Some(outputInfoDate))) + + log.info(s"Running Iceberg repartitioning: UPDATE $fullTableName SET $infoDateColumn = CAST($infoDateDataColumn AS DATE) " + + s"WHERE $infoDateColumn = '$outputInfoDate' AND $infoDateDataColumn >= '$infoDateFrom' AND $infoDateDataColumn <= '$infoDateTo'") + + val dfToWrite = df.withColumn(infoDateColumn, col(infoDateDataColumn).cast(DateType)) + + writeRepartitionedDf(dfToWrite, fullTableName, infoDateColumn, infoDateFrom, infoDateTo, writeOptions) + } + + override def repartitionPhase2(infoDateDataColumn: String, infoDateFrom: LocalDate, infoDateTo: LocalDate, outputInfoDate: LocalDate): Unit = { + if (infoDateColumn.equalsIgnoreCase(infoDateDataColumn)) + throw new IllegalArgumentException(s"Cannot repartition a table if the metastore info date column is the same as the data info date column ($infoDateDataColumn)") + } + def getFilter(infoDateFrom: Option[LocalDate], infoDateTo: Option[LocalDate]): Column = { if (partitionScheme == PartitionScheme.Overwrite) { if (infoDateFrom.isDefined || infoDateTo.isDefined) { diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceNull.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceNull.scala index 996a2189a..c88b76683 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceNull.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceNull.scala @@ -48,4 +48,6 @@ class MetastorePersistenceNull(implicit spark: SparkSession) extends MetastorePe hiveConfig: HiveConfig): Unit = { throw new UnsupportedOperationException("Parquet format does not support Hive tables at the moment.") } + + override def isRepartitioningSupported: Boolean = false } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceParquet.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceParquet.scala index ad77c732d..aa5fab555 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceParquet.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceParquet.scala @@ -145,6 +145,8 @@ class MetastorePersistenceParquet(path: String, throw new UnsupportedOperationException("Parquet format does not support Hive tables at the moment.") } + override def isRepartitioningSupported: Boolean = false + def loadPartitionDirectly(infoDate: LocalDate): DataFrame = { val dateStr = dateFormatter.format(infoDate) val partitionPath = SparkUtils.getPartitionPath(infoDate, infoDateColumn, infoDateFormat, path) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceRaw.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceRaw.scala index 3b7ed0084..91d712bb4 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceRaw.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceRaw.scala @@ -172,6 +172,8 @@ class MetastorePersistenceRaw(path: String, throw new UnsupportedOperationException("Raw format does not support Hive tables.") } + override def isRepartitioningSupported: Boolean = false + private def getListOfFilesRange(infoDateFrom: LocalDate, infoDateTo: LocalDate): Seq[FileStatus] = { if (infoDateFrom.isAfter(infoDateTo)) Seq.empty[FileStatus] diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransient.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransient.scala index 35c1f1adc..dc9528c2a 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransient.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransient.scala @@ -60,4 +60,6 @@ class MetastorePersistenceTransient(tempPath: Option[String], override def repairHiveTable(hiveTableName: String, queryExecutor: QueryExecutor, hiveConfig: HiveConfig): Unit = { throw new UnsupportedOperationException("The 'transient' format does not support Hive tables.") } + + override def isRepartitioningSupported: Boolean = false } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransientEager.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransientEager.scala index 961b1133d..d301cd62d 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransientEager.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceTransientEager.scala @@ -86,4 +86,6 @@ class MetastorePersistenceTransientEager(tempPathOpt: Option[String], hiveConfig: HiveConfig): Unit = { throw new UnsupportedOperationException("Transient format does not support Hive tables.") } + + override def isRepartitioningSupported: Boolean = false } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala index 51b59d285..6b0d283b8 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IncrementalIngestionJob.scala @@ -222,7 +222,7 @@ class IncrementalIngestionJob(operationDef: OperationDef, Reason.Ready case None => log.info(s"Offsets not found for '${outputTable.name}' at '$infoDate'.") - Reason.SkipOnce("No offsets registered") + Reason.SkipOnce(s"Unable to re-run: No offsets registered for $infoDate") } case (true, false) => Reason.Ready diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IngestionJob.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IngestionJob.scala index 7783f8053..3e13c8f84 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IngestionJob.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IngestionJob.scala @@ -87,7 +87,7 @@ class IngestionJob(operationDef: OperationDef, val dataChunk = bookkeeper.getLatestDataChunk(sourceTable.metaTableName, infoDate) val (from, to) = bulkLoadCurrent match { - case Some(bulk) => (bulk.dateFrom, bulk.dateTo) + case Some(bulk) => (bulk.dataDateFrom, bulk.dataDateTo) case None => getInfoDateRange(infoDate, sourceTable.rangeFromExpr, sourceTable.rangeToExpr) } @@ -155,7 +155,7 @@ class IngestionJob(operationDef: OperationDef, val dfTransformed = applyTransformations(df, sourceTable.transformations) val (from, to) = bulkLoadCurrent match { - case Some(bulk) => (bulk.dateFrom, bulk.dateTo) + case Some(bulk) => (bulk.dataDateFrom, bulk.dataDateTo) case None => getInfoDateRange(infoDate, sourceTable.rangeFromExpr, sourceTable.rangeToExpr) } @@ -278,7 +278,7 @@ class IngestionJob(operationDef: OperationDef, private def getSourcingResult(infoDate: LocalDate): SourceResult = { val (from, to) = bulkLoadCurrent match { - case Some(bulk) => (bulk.dateFrom, bulk.dateTo) + case Some(bulk) => (bulk.dataDateFrom, bulk.dataDateTo) case None => getInfoDateRange(infoDate, sourceTable.rangeFromExpr, sourceTable.rangeToExpr) } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/OperationDef.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/OperationDef.scala index 51194e232..83917a58c 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/OperationDef.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/OperationDef.scala @@ -41,6 +41,7 @@ case class OperationDef( ignoreSchemaChange: Boolean, isCritical: Boolean, doNotWriteOutput: Boolean, + enableRepartitioning: Boolean, consumeThreads: Int, dependencies: Seq[MetastoreDependency], outputInfoDateExpression: String, @@ -69,6 +70,7 @@ object OperationDef { val IGNORE_SCHEMA_CHANGE_KEY = "ignore.schema.change" val IS_CRITICAL_KEY = "critical" val DO_NOT_WRITE_OUTPUT_KEY = "do.not.write.output" + val ENABLE_REPARTITIONING_KEY = "enable.repartitioning" val CONSUME_THREADS_KEY = "consume.threads" val DEPENDENCIES_KEY = "dependencies" val STRICT_DEPENDENCY_MANAGEMENT_KEY = "pramen.strict.dependency.management" @@ -109,6 +111,7 @@ object OperationDef { val ignoreSchemaChange = ConfigUtils.getOptionBoolean(conf, IGNORE_SCHEMA_CHANGE_KEY).getOrElse(false) val isCritical = ConfigUtils.getOptionBoolean(conf, IS_CRITICAL_KEY).getOrElse(false) val doNotWriteOutput = ConfigUtils.getOptionBoolean(conf, DO_NOT_WRITE_OUTPUT_KEY).getOrElse(false) + val enableRepartitioning = ConfigUtils.getOptionBoolean(conf, ENABLE_REPARTITIONING_KEY).getOrElse(true) val alwaysAttempt = ConfigUtils.getOptionBoolean(conf, ALWAYS_ATTEMPT_KEY).getOrElse(false) val dependencies = getDependencies(conf, parent, strictDependencyManagement) val outputInfoDateExpressionOpt = ConfigUtils.getOptionString(conf, OUTPUT_INFO_DATE_EXPRESSION_KEY) @@ -158,6 +161,7 @@ object OperationDef { ignoreSchemaChange, isCritical, doNotWriteOutput, + enableRepartitioning, consumeThreads, dependencies, outputInfoDateExpression, diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/AppRunner.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/AppRunner.scala index 6c31c1c84..3aa2487b5 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/AppRunner.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/AppRunner.scala @@ -30,6 +30,7 @@ import za.co.absa.pramen.core.metastore.peristence.{TransientJobManager, Transie import za.co.absa.pramen.core.pipeline._ import za.co.absa.pramen.core.runner.jobrunner.{ConcurrentJobRunner, ConcurrentJobRunnerImpl} import za.co.absa.pramen.core.runner.orchestrator.OrchestratorImpl +import za.co.absa.pramen.core.runner.repartitioner.{JobRepartitioner, JobRepartitionerImpl, JobRepartitionerNull} import za.co.absa.pramen.core.runner.task.{TaskRunner, TaskRunnerMultithreaded} import za.co.absa.pramen.core.state.{PipelineState, PipelineStateImpl, SystemExitCatcherSecurityManager} import za.co.absa.pramen.core.utils.Emoji._ @@ -455,6 +456,11 @@ object AppRunner { taskRunner, spark.sparkContext.applicationId) + implicit val repartitioner: JobRepartitioner = appContext.appConfig.runtimeConfig.bulkLoadCurrent match { + case Some(bulkConfig) => new JobRepartitionerImpl(bulkConfig, appContext.bulkLoadStateManager, appContext.metastore, conf, spark.sparkContext.applicationId, state.getBatchId)(spark) + case None => new JobRepartitionerNull + } + TransientJobManager.setTaskRunner(taskRunner) val orchestrator = new OrchestratorImpl() diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/Orchestrator.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/Orchestrator.scala index 6b3aaecef..83b7fc4e2 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/Orchestrator.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/Orchestrator.scala @@ -21,6 +21,7 @@ import org.apache.spark.sql.SparkSession import za.co.absa.pramen.core.app.AppContext import za.co.absa.pramen.core.pipeline.Job import za.co.absa.pramen.core.runner.jobrunner.ConcurrentJobRunner +import za.co.absa.pramen.core.runner.repartitioner.JobRepartitioner import za.co.absa.pramen.core.state.PipelineState trait Orchestrator { @@ -31,5 +32,6 @@ trait Orchestrator { state: PipelineState, appContext: AppContext, jobRunner: ConcurrentJobRunner, + repartitioner: JobRepartitioner, spark: SparkSession): Unit } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/OrchestratorImpl.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/OrchestratorImpl.scala index 8aad55779..362ed6caf 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/OrchestratorImpl.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/OrchestratorImpl.scala @@ -26,6 +26,7 @@ import za.co.absa.pramen.core.exceptions.{FatalErrorWrapper, ValidationException import za.co.absa.pramen.core.metastore.peristence.TransientJobManager import za.co.absa.pramen.core.pipeline.{Job, JobBase, JobDependency, OperationType} import za.co.absa.pramen.core.runner.jobrunner.ConcurrentJobRunner +import za.co.absa.pramen.core.runner.repartitioner.JobRepartitioner import za.co.absa.pramen.core.runner.splitter.ScheduleStrategyUtils.evaluateRunDate import za.co.absa.pramen.core.state.PipelineState import za.co.absa.pramen.core.utils.Emoji._ @@ -57,6 +58,7 @@ class OrchestratorImpl extends Orchestrator { state: PipelineState, appContext: AppContext, jobRunner: ConcurrentJobRunner, + repartitioner: JobRepartitioner, spark: SparkSession): Unit = { val applicationId = spark.sparkContext.applicationId val allOutputTables = jobs.map(_.outputTable.name) @@ -93,6 +95,7 @@ class OrchestratorImpl extends Orchestrator { jobRunner.startWorkerLoop(runJobChannel) val atLeastOneStarted = sendPendingJobs(runJobChannel, dependencyResolver) + var hasFailures = false var hasFatalErrors = false var hasCriticalJobFailures = false @@ -100,6 +103,7 @@ class OrchestratorImpl extends Orchestrator { completedJobsChannel.foreach { case (finishedJob, taskResults, isSucceeded) => runningJobs.remove(finishedJob) + hasFailures = hasFailures || !isSucceeded hasFatalErrors = hasFatalErrors || taskResults.exists(status => isFatalFailure(status.runStatus)) hasCriticalJobFailures = hasCriticalJobFailures || TransientJobManager.hasCriticalLazyJobFailed || (!isSucceeded && finishedJob.operation.isCritical) @@ -130,6 +134,20 @@ class OrchestratorImpl extends Orchestrator { } } } + + if (!hasFailures && pendingJobs.isEmpty && appContext.appConfig.runtimeConfig.enableRepartitioning) { + log.info("Starting repartitioning of the completed jobs...") + appContext.appConfig.runtimeConfig.bulkLoadCurrent.foreach { bulkConfig => + jobs.filter(job => + !job.taskDef.outputTable.format.isLazy && + !job.taskDef.outputTable.format.isTransient && + job.operation.enableRepartitioning + ).foreach { job => + val taskResults = repartitioner.repartition(job) + state.addTaskCompletion(taskResults) + } + } + } } pendingJobs.foreach(job => { diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitioner.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitioner.scala new file mode 100644 index 000000000..05657c797 --- /dev/null +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitioner.scala @@ -0,0 +1,24 @@ +/* + * Copyright 2022 ABSA Group Limited + * + * 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 za.co.absa.pramen.core.runner.repartitioner + +import za.co.absa.pramen.api.status.TaskResult +import za.co.absa.pramen.core.pipeline.Job + +trait JobRepartitioner { + def repartition(job: Job): Seq[TaskResult] +} diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitionerImpl.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitionerImpl.scala new file mode 100644 index 000000000..1e900f65d --- /dev/null +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitionerImpl.scala @@ -0,0 +1,139 @@ +/* + * Copyright 2022 ABSA Group Limited + * + * 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 za.co.absa.pramen.core.runner.repartitioner + +import com.typesafe.config.Config +import org.apache.spark.sql.SparkSession +import za.co.absa.pramen.api.status.RunStatus.Succeeded +import za.co.absa.pramen.api.status.{RunInfo, RunStatus, TaskDef, TaskResult} +import za.co.absa.pramen.api.status.TaskRunReason.OnRequest +import za.co.absa.pramen.bulkload.BulkLoadStateManager +import za.co.absa.pramen.bulkload.model.BulkLoadPhase +import za.co.absa.pramen.core.app.config.BulkRunConfig +import za.co.absa.pramen.core.metastore.Metastore +import za.co.absa.pramen.core.metastore.peristence.MetastorePersistence +import za.co.absa.pramen.core.pipeline.Job + +import java.time.Instant +import scala.util.control.NonFatal + +class JobRepartitionerImpl(bulkLoadCurrent: BulkRunConfig, + bulkLoadStateManager: BulkLoadStateManager, + metastore: Metastore, + appConfig: Config, + applicationId: String, + batchId: Long)(implicit spark: SparkSession) extends JobRepartitioner { + def repartition(job: Job): Seq[TaskResult] = { + val start = Instant.now() + try { + doRepartition(job) + } catch { + case NonFatal(ex) => + Seq( + TaskResult( + getRepartitionTaskDef(job), + RunStatus.Failed(ex), + Some(RunInfo(bulkLoadCurrent.outputInfoDate, start, Instant.now())), + applicationId, + isTransient = false, + isRawFilesJob = false, + newSchemaRegistered = false, + Seq.empty, + Seq.empty, + Seq.empty, + Map.empty + ) + ) + } + } + + private[core] def doRepartition(job: Job): Seq[TaskResult] = { + val outputTable = job.taskDef.outputTable.name + val bulkLoadStateOpt = bulkLoadStateManager.getState(outputTable, bulkLoadCurrent.outputInfoDate) + + if (bulkLoadStateOpt.isEmpty) return Seq.empty + val bulkLoadState = bulkLoadStateOpt.get + + if (bulkLoadCurrent.infoDateColumn.isEmpty) return Seq.empty + val infoDateColumn = bulkLoadCurrent.infoDateColumn.get + + val metaTable = metastore.getTableDef(outputTable) + val persistence = MetastorePersistence.fromMetaTable(metaTable, appConfig, batchId) + + val start = Instant.now() + val secondStagePhase = if (bulkLoadState.phase == BulkLoadPhase.Processed) { + // Starting repartitioning phase 1 + + if (persistence.isRepartitioningSupported) { + persistence.repartitionPhase1(infoDateColumn, bulkLoadCurrent.outputInfoDate, bulkLoadCurrent.outputInfoDate, bulkLoadCurrent.outputInfoDate) + val updatedState = bulkLoadState.copy(phase = BulkLoadPhase.Repartition1) + bulkLoadStateManager.updatePhase(updatedState) + + BulkLoadPhase.Repartition1 + } else { + BulkLoadPhase.Done + } + } else { + bulkLoadState.phase + } + + val finalPhase = if (secondStagePhase == BulkLoadPhase.Repartition1) { + // Starting repartitioning phase 2 + + if (persistence.isRepartitioningSupported) { + persistence.repartitionPhase2(infoDateColumn, bulkLoadCurrent.outputInfoDate, bulkLoadCurrent.outputInfoDate, bulkLoadCurrent.outputInfoDate) + val updatedState = bulkLoadState.copy(phase = BulkLoadPhase.Done) + bulkLoadStateManager.updatePhase(updatedState) + BulkLoadPhase.Done + } else { + BulkLoadPhase.Done + } + } else { + bulkLoadState.phase + } + + val finish = Instant.now() + + if (finalPhase == BulkLoadPhase.Done) { + val recordCount = metastore.getTable(outputTable, Some(bulkLoadCurrent.dataDateFrom), Some(bulkLoadCurrent.dataDateTo)).count() + Seq( + TaskResult( + getRepartitionTaskDef(job), + Succeeded(None, Some(recordCount), None, None, OnRequest, Seq.empty, Seq.empty, Seq.empty, Seq.empty), + Some(RunInfo(bulkLoadCurrent.outputInfoDate, start, finish)), + applicationId, + isTransient = false, + isRawFilesJob = false, + newSchemaRegistered = false, + Seq.empty, + Seq.empty, + Seq.empty, + Map.empty + ) + ) + + } else { + Seq.empty + } + } + + private[core] def getRepartitionTaskDef(job: Job): TaskDef = { + val originalName = job.taskDef.name + val newName = s"Repartition - $originalName" + job.taskDef.copy(name = newName) + } +} diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitionerNull.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitionerNull.scala new file mode 100644 index 000000000..a68279f1e --- /dev/null +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/repartitioner/JobRepartitionerNull.scala @@ -0,0 +1,24 @@ +/* + * Copyright 2022 ABSA Group Limited + * + * 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 za.co.absa.pramen.core.runner.repartitioner + +import za.co.absa.pramen.api.status.TaskResult +import za.co.absa.pramen.core.pipeline.Job + +class JobRepartitionerNull extends JobRepartitioner { + def repartition(job: Job): Seq[TaskResult] = Seq.empty +} diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleParams.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleParams.scala index f437f3214..590cb5515 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleParams.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleParams.scala @@ -17,7 +17,7 @@ package za.co.absa.pramen.core.runner.splitter import za.co.absa.pramen.api.RunMode -import za.co.absa.pramen.core.app.config.RuntimeConfig +import za.co.absa.pramen.core.app.config.{BulkRunConfig, RuntimeConfig} import java.time.LocalDate @@ -44,19 +44,11 @@ object ScheduleParams { mode: RunMode ) extends ScheduleParams - case class Bulk( - dataDateFrom: LocalDate, - dataDateTo: LocalDate, - outputInfoDate: LocalDate - ) extends ScheduleParams + case class Bulk(bulkRunConfig: BulkRunConfig) extends ScheduleParams def fromRuntimeConfig(conf: RuntimeConfig, backfillDays: Int, trackDays: Int, delayDays: Int): ScheduleParams = { if (conf.bulkLoadCurrent.isDefined) { - ScheduleParams.Bulk( - conf.bulkLoadCurrent.get.dateFrom, - conf.bulkLoadCurrent.get.dateTo, - conf.bulkLoadCurrent.get.outputInfoDate - ) + ScheduleParams.Bulk(conf.bulkLoadCurrent.get) } else if (conf.runDateTo.nonEmpty) { ScheduleParams.Historical( conf.runDate, diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyIncremental.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyIncremental.scala index 6d325e391..2b309f036 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyIncremental.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyIncremental.scala @@ -19,6 +19,7 @@ package za.co.absa.pramen.core.runner.splitter import za.co.absa.pramen.api.jobdef.Schedule import za.co.absa.pramen.api.status.{MetastoreDependency, TaskRunReason} import za.co.absa.pramen.bulkload.BulkLoadStateManager +import za.co.absa.pramen.core.app.config.BulkRunConfig import za.co.absa.pramen.core.bookkeeper.Bookkeeper import za.co.absa.pramen.core.pipeline import za.co.absa.pramen.core.pipeline.TaskPreDef @@ -87,9 +88,9 @@ class ScheduleStrategyIncremental(lastInfoDateProcessedOpt: Option[LocalDate], h case ScheduleParams.Historical(dateFrom, dateTo, inverseDateOrder, mode) => log.info(s"Ranged strategy: from $dateFrom to $dateTo, mode = '${mode.toString}', minimumDate = $minimumDate") getHistorical(outputTable, dateFrom, dateTo, schedule, mode, infoDateExpression, minimumDate, inverseDateOrder, bookkeeper) - case ScheduleParams.Bulk(dataDateFrom, dataDateTo, outputInfoDate) => - log.info(s"Bulk strategy: from $dataDateFrom to $dataDateTo, outputInfoDate = $outputInfoDate") - getBulk(outputTable, dataDateFrom, dataDateTo, outputInfoDate, bulkLoadStateManager) + case ScheduleParams.Bulk(bulkRunConfig: BulkRunConfig) => + log.info(s"Bulk strategy: from $bulkRunConfig.dataDateFrom to $bulkRunConfig.dataDateTo, outputInfoDate = $bulkRunConfig.outputInfoDate") + getBulk(outputTable, bulkRunConfig, bulkLoadStateManager) } filterOutPastMinimumDates(dates, minimumDate) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategySourcing.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategySourcing.scala index 41ae1aa55..5a161165e 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategySourcing.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategySourcing.scala @@ -19,6 +19,7 @@ package za.co.absa.pramen.core.runner.splitter import za.co.absa.pramen.api.jobdef.Schedule import za.co.absa.pramen.api.status.{MetastoreDependency, TaskRunReason} import za.co.absa.pramen.bulkload.BulkLoadStateManager +import za.co.absa.pramen.core.app.config.BulkRunConfig import za.co.absa.pramen.core.bookkeeper.Bookkeeper import za.co.absa.pramen.core.pipeline import za.co.absa.pramen.core.pipeline.TaskPreDef @@ -112,9 +113,9 @@ class ScheduleStrategySourcing(hasInfoDateColumn: Boolean) extends ScheduleStrat case ScheduleParams.Historical(dateFrom, dateTo, inverseDateOrder, mode) => log.info(s"Ranged strategy: from $dateFrom to $dateTo, mode = '${mode.toString}', minimumDate = $minimumDate") getHistorical(outputTable, dateFrom, dateTo, schedule, mode, infoDateExpression, minimumDate, inverseDateOrder, bookkeeper) - case ScheduleParams.Bulk(dataDateFrom, dataDateTo, outputInfoDate) => - log.info(s"Bulk strategy: from $dataDateFrom to $dataDateTo, outputInfoDate = $outputInfoDate") - getBulk(outputTable, dataDateFrom, dataDateTo, outputInfoDate, bulkLoadStateManager) + case ScheduleParams.Bulk(bulkRunConfig: BulkRunConfig) => + log.info(s"Bulk strategy: from ${bulkRunConfig.dataDateFrom} to ${bulkRunConfig.dataDateTo}, outputInfoDate = ${bulkRunConfig.outputInfoDate}") + getBulk(outputTable, bulkRunConfig, bulkLoadStateManager) } filterOutPastMinimumDates(dates, minimumDate) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyUtils.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyUtils.scala index 5faf3d33c..8d5bee431 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyUtils.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/runner/splitter/ScheduleStrategyUtils.scala @@ -22,6 +22,7 @@ import za.co.absa.pramen.api.status.TaskRunReason import za.co.absa.pramen.bulkload.BulkLoadStateManager import za.co.absa.pramen.bulkload.model.BulkLoadPhase.Pending import za.co.absa.pramen.bulkload.model.{BulkLoadPhase, BulkLoadState} +import za.co.absa.pramen.core.app.config.BulkRunConfig import za.co.absa.pramen.core.bookkeeper.Bookkeeper import za.co.absa.pramen.core.expr.DateExprEvaluator import za.co.absa.pramen.core.pipeline @@ -182,27 +183,24 @@ object ScheduleStrategyUtils { } def getBulk(outputTable: String, - dataDateFrom: LocalDate, - dataDateTo: LocalDate, - outputInfoDate: LocalDate, + bulkRunConfig: BulkRunConfig, bulkLoadStateManager: BulkLoadStateManager): List[TaskPreDef] = { - val currentStateOpt = bulkLoadStateManager.getState(outputTable, outputInfoDate) + val currentStateOpt = bulkLoadStateManager.getState(outputTable, bulkRunConfig.outputInfoDate) currentStateOpt match { case Some(state) => - if (!state.dataDateFrom.equals(dataDateFrom) || !state.dataDateTo.equals(dataDateTo)) { - throw new IllegalStateException(s"The job for table '$outputTable' and info date '$outputInfoDate' has different data date range.") + if (!state.dataDateFrom.equals(bulkRunConfig.dataDateFrom) || !state.dataDateTo.equals(bulkRunConfig.dataDateTo)) { + throw new IllegalStateException(s"The job for table '$outputTable' and info date '${bulkRunConfig.outputInfoDate}' has different data date range.") } if (state.phase != Pending) { - return List(TaskPreDef(outputInfoDate, TaskRunReason.Skip("already processed"))) + return List(TaskPreDef(bulkRunConfig.outputInfoDate, TaskRunReason.Skip("already processed"))) } case None => - // ToDo: Propagate info date column for the future repartitioning - val newState = BulkLoadState(outputTable, "", outputInfoDate, dataDateFrom, dataDateTo, Pending) + val newState = BulkLoadState(outputTable, bulkRunConfig.infoDateColumn.getOrElse(""), bulkRunConfig.outputInfoDate, bulkRunConfig.dataDateFrom, bulkRunConfig.dataDateTo, Pending) bulkLoadStateManager.addState(newState) } - List(TaskPreDef(outputInfoDate, TaskRunReason.Rerun)) + List(TaskPreDef(bulkRunConfig.outputInfoDate, TaskRunReason.Rerun)) } def updateBulkLoadCompletion(outputTable: String, diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/state/PipelineStateImpl.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/state/PipelineStateImpl.scala index 1fa9776ff..8067dd69f 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/state/PipelineStateImpl.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/state/PipelineStateImpl.scala @@ -133,7 +133,7 @@ class PipelineStateImpl(implicit conf: Config, notificationBuilder: Notification val (runDateFrom, runDateTo) = runtimeConfig.bulkLoadCurrent match { case Some(bulkLoad) => - (bulkLoad.dateFrom, Option(bulkLoad.dateTo)) + (bulkLoad.dataDateFrom, Option(bulkLoad.dataDateTo)) case None => (runtimeConfig.runDate, runtimeConfig.runDateTo) } diff --git a/pramen/core/src/main/scala_2.11/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala b/pramen/core/src/main/scala_2.11/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala index 6ed6aa398..e81c8a64c 100644 --- a/pramen/core/src/main/scala_2.11/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala +++ b/pramen/core/src/main/scala_2.11/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala @@ -57,4 +57,14 @@ object MetastorePersistenceIcebergOps { writerOptions: Map[String, String]): Unit = { throw new UnsupportedOperationException(s"Iceberg format is not supported in Scala 2.11") } + + def writeRepartitionedDf(df: DataFrame, + table: String, + infoDateColumn: String, + infoDateFrom: LocalDate, + infoDateTo: LocalDate, + writerOptions: Map[String, String]): Unit = { + throw new UnsupportedOperationException(s"Iceberg format is not supported in Scala 2.11") + } + } diff --git a/pramen/core/src/main/scala_2.12/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala b/pramen/core/src/main/scala_2.12/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala index 840245c29..e0e69154b 100644 --- a/pramen/core/src/main/scala_2.12/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala +++ b/pramen/core/src/main/scala_2.12/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala @@ -17,6 +17,7 @@ package za.co.absa.pramen.core.metastore.peristence import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.DateType import org.apache.spark.sql.{DataFrame, SparkSession} import org.slf4j.LoggerFactory import za.co.absa.pramen.api.PartitionScheme @@ -119,4 +120,16 @@ object MetastorePersistenceIcebergOps { .options(writerOptions) .append() } + + def writeRepartitionedDf(df: DataFrame, + table: String, + infoDateColumn: String, + infoDateFrom: LocalDate, + infoDateTo: LocalDate, + writerOptions: Map[String, String]): Unit = { + df.writeTo(table) + .option("check-ordering", "false") + .options(writerOptions) + .overwrite(col(infoDateColumn) >= lit(Date.valueOf(infoDateFrom)) && col(infoDateColumn) <= lit(Date.valueOf(infoDateTo))) + } } diff --git a/pramen/core/src/main/scala_2.13/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala b/pramen/core/src/main/scala_2.13/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala index 840245c29..9ed0c5841 100644 --- a/pramen/core/src/main/scala_2.13/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala +++ b/pramen/core/src/main/scala_2.13/za/co/absa/pramen/core/metastore/peristence/MetastorePersistenceIcebergOps.scala @@ -119,4 +119,16 @@ object MetastorePersistenceIcebergOps { .options(writerOptions) .append() } + + def writeRepartitionedDf(df: DataFrame, + table: String, + infoDateColumn: String, + infoDateFrom: LocalDate, + infoDateTo: LocalDate, + writerOptions: Map[String, String]): Unit = { + df.writeTo(table) + .option("check-ordering", "false") + .options(writerOptions) + .overwrite(col(infoDateColumn) >= lit(Date.valueOf(infoDateFrom)) && col(infoDateColumn) <= lit(Date.valueOf(infoDateTo))) + } } diff --git a/pramen/core/src/test/resources/test/config/integration_bulk_load.conf b/pramen/core/src/test/resources/test/config/integration_bulk_load.conf index ba1918c2f..5fd4a30b2 100644 --- a/pramen/core/src/test/resources/test/config/integration_bulk_load.conf +++ b/pramen/core/src/test/resources/test/config/integration_bulk_load.conf @@ -22,6 +22,8 @@ pramen { temporary.directory = ${base.path}/temp stop.spark.session = false + runtime.enable.repartitioning = true + runtime.info.date.column = "dt" } pramen.metastore { @@ -35,8 +37,8 @@ pramen.metastore { { name = "table2" description = "Table 2" - format = "parquet" - path = ${base.path}/table2 + format = "iceberg" + table = ${iceberg.table.name} } ] } diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/OperationDefFactory.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/OperationDefFactory.scala index b97127822..0666d4708 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/OperationDefFactory.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/OperationDefFactory.scala @@ -33,6 +33,7 @@ object OperationDefFactory { ignoreSchemaChange: Boolean = false, isCritical: Boolean = false, doNotWriteOutput: Boolean = false, + enableRepartitioning: Boolean = true, consumeThreads: Int = 1, dependencies: Seq[MetastoreDependency] = Nil, outputInfoDateExpression: String = "@date", @@ -56,6 +57,7 @@ object OperationDefFactory { ignoreSchemaChange, isCritical, doNotWriteOutput, + enableRepartitioning, consumeThreads, dependencies, outputInfoDateExpression, diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/RuntimeConfigFactory.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/RuntimeConfigFactory.scala index b6f6a738e..47f78bb58 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/RuntimeConfigFactory.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/RuntimeConfigFactory.scala @@ -43,6 +43,7 @@ object RuntimeConfigFactory { historicalRunMode: RunMode = RunMode.CheckUpdates, bulkBatchSize: BulkBatchSize = BulkBatchSize.Monthly, bulkLoadCurrent: Option[BulkRunConfig] = None, + enableRepartitioning: Boolean = false, sparkAppDescriptionTemplate: Option[String] = None, attempt: Int = 1, maxAttempts: Int = 1, @@ -67,6 +68,7 @@ object RuntimeConfigFactory { historicalRunMode, bulkBatchSize, bulkLoadCurrent, + enableRepartitioning, sparkAppDescriptionTemplate, attempt, maxAttempts, diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/cmd/CmdLineConfigSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/cmd/CmdLineConfigSuite.scala index 862874885..702963ece 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/cmd/CmdLineConfigSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/cmd/CmdLineConfigSuite.scala @@ -244,7 +244,7 @@ class CmdLineConfigSuite extends AnyWordSpec { } "return a modified config if bulk mode for date-to override is specified" in { - val cmd = CmdLineConfig.parseCmdLine(Array("--workflow", "dummy.config", "--date-to", "2020-08-15", "--inverse-order", "true", "--run-mode", "bulk", "--bulk-size", "yearly")) + val cmd = CmdLineConfig.parseCmdLine(Array("--workflow", "dummy.config", "--date-to", "2020-08-15", "--inverse-order", "true", "--run-mode", "bulk", "--bulk-size", "yearly", "--info-date-column", "info_date")) val config = CmdLineConfig.applyCmdLineToConfig(emptyConfig, cmd.get) assert(config.hasPath(LOAD_DATE_TO)) @@ -255,6 +255,7 @@ class CmdLineConfigSuite extends AnyWordSpec { assert(config.getBoolean(IS_INVERSE_ORDER)) assert(config.getString(RUN_MODE) == "bulk") assert(config.getString(RUN_BULK_BATCH_SIZE) == "yearly") + assert(config.getString(INFO_DATE_COLUMN) == "info_date") } "return the original config if no cmd line arguments are provided" in { diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/BulkLoadLongSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/BulkLoadLongSuite.scala index a9d3663f8..1516374d2 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/BulkLoadLongSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/integration/BulkLoadLongSuite.scala @@ -18,9 +18,10 @@ package za.co.absa.pramen.core.integration import com.typesafe.config.{Config, ConfigFactory} import org.apache.hadoop.fs.Path -import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll} +import org.apache.spark.sql.functions.col import org.scalatest.wordspec.AnyWordSpec -import za.co.absa.pramen.core.base.SparkTestBase +import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll} +import za.co.absa.pramen.core.base.SparkTestIcebergBase import za.co.absa.pramen.core.fixtures.{RelationalDbFixture, TempDirFixture, TextComparisonFixture} import za.co.absa.pramen.core.rdb.{PramenDb, RdbJdbc} import za.co.absa.pramen.core.reader.model.JdbcConfig @@ -28,10 +29,10 @@ import za.co.absa.pramen.core.runner.AppRunner import za.co.absa.pramen.core.samples.RdbExampleTable import za.co.absa.pramen.core.utils.{FsUtils, ResourceUtils, UsingUtils} -import java.time.LocalDate +import scala.util.Random class BulkLoadLongSuite extends AnyWordSpec - with SparkTestBase + with SparkTestIcebergBase with TempDirFixture with RelationalDbFixture with BeforeAndAfter @@ -41,8 +42,6 @@ class BulkLoadLongSuite extends AnyWordSpec val jdbcConfig: JdbcConfig = JdbcConfig(driver, Some(url), Nil, None, Some(user), Some(password)) var pramenDb: PramenDb = _ - private val infoDate = LocalDate.of(2021, 2, 18) - before { if (pramenDb != null) pramenDb.close() UsingUtils.using(RdbJdbc(jdbcConfig)) { rdb => @@ -126,47 +125,58 @@ class BulkLoadLongSuite extends AnyWordSpec |""".stripMargin "be able to access inner source configuration for strict months" in { + assume(spark.version.split('.').head.toInt >= 3, s"Ignored for too old Delta Lake for Spark ${spark.version}") + withTempDirectory("integration_inner_source") { tempDir => + val tableName = "mt_iceberg_bulktable1" + Math.abs(Random.nextInt()).toString + val fsUtils = new FsUtils(spark.sparkContext.hadoopConfiguration, tempDir) val landingPath = new Path(tempDir, "landing") fsUtils.writeFile(new Path(landingPath, "landing_file1.csv"), csvData) - val conf = getConfig(tempDir) + val conf = getConfig(tempDir, tableName) val exitCode = AppRunner.runBulkPipelines(conf) assert(exitCode == 0) - val table2P = new Path(tempDir, "table2") + val table1P = new Path(tempDir, "table1") - val df = spark.read.parquet(table2P.toString) + val df = spark.read.parquet(table1P.toString) //df.show(1000, truncate = false) - val table2Path1 = new Path(new Path(tempDir, "table2"), s"pramen_info_date=2021-01-01") - val table2Path2 = new Path(new Path(tempDir, "table2"), s"pramen_info_date=2021-02-01") + val table1Path1 = new Path(new Path(tempDir, "table1"), s"pramen_info_date=2021-01-01") + val table1Path2 = new Path(new Path(tempDir, "table1"), s"pramen_info_date=2021-02-01") - assert(fsUtils.exists(table2Path1)) - assert(fsUtils.exists(table2Path2)) + assert(fsUtils.exists(table1Path1)) + assert(fsUtils.exists(table1Path2)) assert(df.count() == 59) + // For now... + assert(!df.filter(col("dt") =!= col("pramen_info_date")).isEmpty) + // Running the job for the second time shouyld not change the output val exitCode2 = AppRunner.runBulkPipelines(conf) assert(exitCode2 == 0) - val df2 = spark.read.parquet(table2P.toString) + val df2 = spark.table(tableName) assert(df2.count() == 59) + assert(df2.filter(col("dt") =!= col("pramen_info_date")).isEmpty) + + spark.sql(s"DELETE FROM $tableName").count() } } } - def getConfig(basePath: String): Config = { + def getConfig(basePath: String, icebergTableName: String): Config = { val configContents = ResourceUtils.getResourceString("/test/config/integration_bulk_load.conf") val basePathEscaped = basePath.replace("\\", "\\\\") val conf = ConfigFactory.parseString( s"""base.path = "$basePathEscaped" + |iceberg.table.name = "$icebergTableName" |pramen { | load.date.from = "2021-01-01" | load.date.to = "2021-02-28" diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/runner/orchestrator/OrchestratorSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/runner/orchestrator/OrchestratorSuite.scala index dc21ddfe4..9143d43c6 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/runner/orchestrator/OrchestratorSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/runner/orchestrator/OrchestratorSuite.scala @@ -28,6 +28,7 @@ import za.co.absa.pramen.core.mocks.runner.ConcurrentJobRunnerSpy import za.co.absa.pramen.core.mocks.state.PipelineStateSpy import za.co.absa.pramen.core.pipeline.OperationDef import za.co.absa.pramen.core.runner.orchestrator.OrchestratorImpl +import za.co.absa.pramen.core.runner.repartitioner.JobRepartitionerNull class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFixture { "runJobs" should { @@ -46,8 +47,9 @@ class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFi val orchestrator = new OrchestratorImpl() val runner = new ConcurrentJobRunnerSpy() + val repartitioner = new JobRepartitionerNull - orchestrator.runJobs(Seq.empty)(conf, null, appContext, runner, spark) + orchestrator.runJobs(Seq.empty)(conf, null, appContext, runner, repartitioner, spark) } } @@ -57,9 +59,10 @@ class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFi val orchestrator = new OrchestratorImpl() val runner = new ConcurrentJobRunnerSpy() + val repartitioner = new JobRepartitionerNull val state = new PipelineStateSpy() - orchestrator.runJobs(Seq(job1, job2, job3))(conf, state, appContext, runner, spark) + orchestrator.runJobs(Seq(job1, job2, job3))(conf, state, appContext, runner, repartitioner, spark) assert(state.completedStatuses.length == 3) @@ -80,8 +83,9 @@ class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFi val orchestrator = new OrchestratorImpl() val runner = new ConcurrentJobRunnerSpy(includeFails = true) val state = new PipelineStateSpy() + val repartitioner = new JobRepartitionerNull - orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, spark) + orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, repartitioner, spark) assert(state.completedStatuses.length == 4) @@ -114,8 +118,9 @@ class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFi val orchestrator = new OrchestratorImpl() val runner = new ConcurrentJobRunnerSpy() val state = new PipelineStateSpy() + val repartitioner = new JobRepartitionerNull - orchestrator.runJobs(Seq(job1, job4))(conf, state, appContext, runner, spark) + orchestrator.runJobs(Seq(job1, job4))(conf, state, appContext, runner, repartitioner, spark) assert(state.completedStatuses.length == 2) @@ -141,8 +146,9 @@ class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFi val orchestrator = new OrchestratorImpl() val runner = new ConcurrentJobRunnerSpy(includeFails = true) val state = new PipelineStateSpy() + val repartitioner = new JobRepartitionerNull - orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, spark) + orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, repartitioner, spark) assert(state.completedStatuses.length == 4) @@ -172,8 +178,9 @@ class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFi val orchestrator = new OrchestratorImpl() val runner = new ConcurrentJobRunnerSpy(includeFails = true) val state = new PipelineStateSpy() + val repartitioner = new JobRepartitionerNull - orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, spark) + orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, repartitioner, spark) assert(state.completedStatuses.length == 4) @@ -203,8 +210,9 @@ class OrchestratorSuite extends AnyWordSpec with SparkTestBase with AppContextFi val orchestrator = new OrchestratorImpl() val runner = new ConcurrentJobRunnerSpy(includeFails = true) val state = new PipelineStateSpy() + val repartitioner = new JobRepartitionerNull - orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, spark) + orchestrator.runJobs(Seq(job1, job2, job3, job4))(conf, state, appContext, runner, repartitioner, spark) assert(state.completedStatuses.length == 4)