| // 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.kudu.backup |
| |
| import org.apache.kudu.backup.Backup.TableMetadataPB |
| import org.apache.kudu.client.AlterTableOptions |
| import org.apache.kudu.client.KuduPartitioner |
| import org.apache.kudu.client.Partition |
| import org.apache.kudu.client.RangePartitionBound |
| import org.apache.kudu.client.RangePartitionWithCustomHashSchema |
| import org.apache.kudu.client.SessionConfiguration.FlushMode |
| import org.apache.kudu.spark.kudu.KuduContext |
| import org.apache.kudu.spark.kudu.RowConverter |
| import org.apache.spark.sql.DataFrame |
| import org.apache.spark.sql.SparkSession |
| import org.apache.yetus.audience.InterfaceAudience |
| import org.apache.yetus.audience.InterfaceStability |
| import org.slf4j.Logger |
| import org.slf4j.LoggerFactory |
| |
| import scala.collection.JavaConverters._ |
| import scala.collection.parallel.ForkJoinTaskSupport |
| import scala.concurrent.forkjoin.ForkJoinPool |
| import scala.util.Failure |
| import scala.util.Success |
| import scala.util.Try |
| |
| /** |
| * The main class for a Kudu restore spark job. |
| */ |
| @InterfaceAudience.Private |
| @InterfaceStability.Unstable |
| object KuduRestore { |
| val log: Logger = LoggerFactory.getLogger(getClass) |
| val ImpalaPrefix = "impala::" |
| |
| /** |
| * Returns the table name in which the data will be restored considering the flags removeImpalaPrefix, |
| * newDatabaseName and tableSuffix |
| */ |
| def getRestoreTableName(fullTableName: String, options: RestoreOptions): String = { |
| // Break the table down into prefix::databaseName.tableName |
| var prefix = "" |
| var databaseName = "" |
| var tableName = fullTableName |
| val hasImpalaPrefix = tableName.startsWith(ImpalaPrefix) |
| if (hasImpalaPrefix) { |
| prefix = ImpalaPrefix |
| tableName = tableName.substring(ImpalaPrefix.length) |
| } |
| val hasDatabase = tableName.contains(".") |
| if (hasDatabase) { |
| databaseName = tableName.substring(0, tableName.indexOf(".") + 1) |
| tableName = tableName.substring(tableName.indexOf(".") + 1) |
| } |
| |
| // If the user does not want the Impala prefix, drop it |
| if (options.removeImpalaPrefix) { |
| prefix = "" |
| } |
| |
| // If there is a databaseName specified by the user, use that |
| if (options.newDatabaseName.nonEmpty) { |
| databaseName = options.newDatabaseName.concat(".") |
| } |
| |
| s"${prefix}${databaseName}${tableName}${options.tableSuffix}" |
| } |
| |
| private def doRestore( |
| tableName: String, |
| context: KuduContext, |
| session: SparkSession, |
| io: BackupIO, |
| options: RestoreOptions, |
| backupMap: Map[String, BackupGraph]): Unit = { |
| if (!backupMap.contains(tableName)) { |
| throw new RuntimeException(s"No valid backups found for table: $tableName") |
| } |
| val graph = backupMap(tableName) |
| val restorePath = graph.restorePath |
| val lastMetadata = restorePath.backups.last.metadata |
| val restoreName = getRestoreTableName(lastMetadata.getTableName, options) |
| val numJobs = restorePath.backups.size |
| var currentJob = 1 |
| restorePath.backups.foreach { backup => |
| session.sparkContext.setJobDescription(s"Kudu Restore($currentJob/$numJobs): $tableName") |
| |
| log.info(s"Restoring table $tableName from path: ${backup.path}") |
| val metadata = backup.metadata |
| val isFullRestore = metadata.getFromMs == 0 |
| // TODO (KUDU-2788): Store the full metadata to compare/validate for each applied partial. |
| |
| // On the full restore we may need to create the table. |
| if (isFullRestore) { |
| if (options.createTables) { |
| log.info(s"Creating restore table $restoreName") |
| // We use the last schema in the restore path when creating the table to |
| // ensure the table is created in its final state. |
| createTableRangePartitionByRangePartition( |
| restoreName, |
| lastMetadata, |
| options.restoreOwner, |
| context) |
| } |
| } |
| val backupSchema = BackupUtils.dataSchema(TableMetadata.getKuduSchema(metadata)) |
| val rowActionCol = backupSchema.fields.last.name |
| |
| var data = session.sqlContext.read |
| .format(metadata.getDataFormat) |
| .schema(backupSchema) |
| .load(backup.path.toString) |
| // Default the the row action column with a value of "UPSERT" so that the |
| // rows from a full backup, which don't have a row action, are upserted. |
| // TODO(ghenke): Consider using INSERT_IGNORE for full backups. |
| .na |
| .fill(RowAction.UPSERT.getValue, Seq(rowActionCol)) |
| |
| // Adjust for dropped and renamed columns. |
| data = adjustSchema(data, metadata, lastMetadata, rowActionCol) |
| val restoreSchema = data.schema |
| |
| // Write the data to Kudu. |
| data.queryExecution.toRdd.foreachPartition { internalRows => |
| val table = context.syncClient.openTable(restoreName) |
| val converter = new RowConverter(table.getSchema, restoreSchema, false) |
| val partitioner = createPartitionFilter(metadata, lastMetadata) |
| val session = context.syncClient.newSession |
| session.setFlushMode(FlushMode.AUTO_FLUSH_BACKGROUND) |
| // In the case of task retries we need to ignore NotFound errors for deleted rows. |
| // This can't occur if DELETE_IGNORE is used, but still needs to be set in the case |
| // DELETE is used for backwards compatibility. |
| session.setIgnoreAllNotFoundRows(true) |
| try for (internalRow <- internalRows) { |
| // Convert the InternalRows to Rows. |
| // This avoids any corruption as reported in SPARK-26880. |
| val row = converter.toRow(internalRow) |
| // Get the operation type based on the row action column. |
| // This will always be the last column in the row. |
| val rowActionValue = row.getByte(row.length - 1) |
| val rowAction = RowAction.fromValue(rowActionValue) |
| // Generate an operation based on the row action. |
| val operation = rowAction match { |
| case RowAction.UPSERT => table.newUpsert() |
| case RowAction.DELETE => { |
| if (context.supportsIgnoreOperations) { |
| table.newDeleteIgnore() |
| } else { |
| table.newDelete() |
| } |
| } |
| case _ => throw new IllegalStateException(s"Unsupported RowAction: $rowAction") |
| } |
| // Convert the Spark row to a partial row and set it on the operation. |
| val partialRow = converter.toPartialRow(row) |
| // Drop rows that are not covered by the partitioner. This is how we |
| // detect a partition which was dropped between backups and filter |
| // out the rows from that dropped partition. |
| if (partitioner.isCovered(partialRow)) { |
| operation.setRow(partialRow) |
| session.apply(operation) |
| } |
| } finally { |
| session.close() |
| } |
| // Fail the task if there are any errors. |
| // It is important to capture all of the errors via getRowErrors and then check |
| // the length because each call to session.getPendingErrors clears the ErrorCollector. |
| val pendingErrors = session.getPendingErrors |
| if (pendingErrors.getRowErrors.nonEmpty) { |
| val errors = pendingErrors.getRowErrors |
| val sample = errors.take(5).map(_.getErrorStatus).mkString |
| if (pendingErrors.isOverflowed) { |
| throw new RuntimeException( |
| s"PendingErrors overflowed. Failed to write at least ${errors.length} rows " + |
| s"to Kudu; Sample errors: $sample") |
| } else { |
| throw new RuntimeException( |
| s"Failed to write ${errors.length} rows to Kudu; Sample errors: $sample") |
| } |
| } |
| } |
| currentJob += 1 |
| } |
| } |
| |
| def run(options: RestoreOptions, session: SparkSession): Boolean = { |
| // Set the job group for all the spark restore jobs. |
| // Note: The job description will be overridden by each Kudu table job. |
| session.sparkContext.setJobGroup(s"Kudu Restore @ ${options.timestampMs}", "Kudu Restore") |
| |
| log.info(s"Restoring from path: ${options.rootPath}") |
| val context = |
| new KuduContext( |
| options.kuduMasterAddresses, |
| session.sparkContext |
| ) |
| val io = new BackupIO(session.sparkContext.hadoopConfiguration, options.rootPath) |
| |
| // Read the required backup metadata. |
| val backupGraphs = io.readBackupGraphsByTableName(options.tables, options.timestampMs) |
| |
| // Key the backupMap by the last table name. |
| val backupMap = backupGraphs |
| .groupBy(_.restorePath.tableName) |
| .mapValues(_.maxBy(_.restorePath.toMs)) |
| |
| // Parallelize the processing. Managing resources of parallel restore jobs is very complex, so |
| // only the simplest possible thing is attempted. Kudu trusts Spark to manage resources. |
| // TODO (KUDU-2832): If the job fails to restore a table it may still create the table, which |
| // will cause subsequent restores to fail unless the table is deleted or the restore suffix is |
| // changed. We ought to try to clean up the mess when a failure happens. |
| val parallelTables = options.tables.par |
| parallelTables.tasksupport = new ForkJoinTaskSupport( |
| new ForkJoinPool(options.numParallelRestores)) |
| val restoreResults = parallelTables.map { tableName => |
| val restoreResult = |
| Try(doRestore(tableName, context, session, io, options, backupMap)) |
| restoreResult match { |
| case Success(()) => |
| log.info(s"Successfully restored table $tableName") |
| case Failure(ex) => |
| if (options.numParallelRestores == 1 && options.failOnFirstError) |
| throw ex |
| else |
| log.error(s"Failed to restore table $tableName", ex) |
| } |
| (tableName, restoreResult) |
| } |
| |
| restoreResults.filter(_._2.isFailure).foreach { |
| case (tableName, ex) => |
| log.error( |
| s"Failed to restore table $tableName: Look back in the logs for the full exception. Error: ${ex.toString}") |
| } |
| !restoreResults.exists(_._2.isFailure) |
| } |
| |
| // Kudu isn't good at creating a lot of tablets at once, and by default tables may only be created |
| // with at most 60 tablets. Additional tablets can be added later by adding range partitions. So, |
| // to restore tables with more tablets than that, we need to create the table piece-by-piece. This |
| // does so in the simplest way: creating the table with the first range partition, if there is |
| // one, and then altering it to add the rest of the partitions, one partition at a time. |
| private def createTableRangePartitionByRangePartition( |
| restoreName: String, |
| metadata: TableMetadataPB, |
| restoreOwner: Boolean, |
| context: KuduContext): Unit = { |
| // Create the table with the first range partition (or none if there are none). |
| val schema = TableMetadata.getKuduSchema(metadata) |
| val options = TableMetadata.getCreateTableOptionsWithoutRangePartitions(metadata, restoreOwner) |
| // Returns the range bounds of the ranges that contain the table wide hash schema. |
| val boundsWithoutHashSchema = TableMetadata.getRangeBoundPartialRows(metadata) |
| // Returns the range bounds and hash schema of the ranges that contain a custom hash schema. |
| val boundsWithCustomHashSchema = |
| TableMetadata.getRangeBoundsPartialRowsWithHashSchemas(metadata) |
| if (boundsWithoutHashSchema.nonEmpty) { |
| // Adds the first range partition with table wide hash schema through create. |
| boundsWithoutHashSchema.headOption.foreach(bound => { |
| val (lower, upper) = bound |
| options.addRangePartition(lower, upper) |
| }) |
| context.createTable(restoreName, schema, options) |
| |
| // Add the rest of the range partitions with table wide hash schema through alters. |
| boundsWithoutHashSchema.tail.foreach(bound => { |
| val (lower, upper) = bound |
| val options = new AlterTableOptions() |
| options.addRangePartition(lower, upper) |
| context.syncClient.alterTable(restoreName, options) |
| }) |
| |
| // Adds range partitions with custom hash schema through alters. |
| boundsWithCustomHashSchema.foreach(bound => { |
| val rangePartition = new RangePartitionWithCustomHashSchema( |
| bound.lowerBound, |
| bound.upperBound, |
| RangePartitionBound.INCLUSIVE_BOUND, |
| RangePartitionBound.EXCLUSIVE_BOUND) |
| bound.hashSchemas.asScala.foreach { hp => |
| val columnNames = hp.getColumnIds.asScala.map { id => |
| schema.getColumnByIndex(id).getName |
| } |
| rangePartition.addHashPartitions(columnNames.asJava, hp.getNumBuckets, hp.getSeed) |
| } |
| val options = new AlterTableOptions() |
| options.addRangePartition(rangePartition) |
| context.syncClient.alterTable(restoreName, options) |
| }) |
| } else if (boundsWithCustomHashSchema.nonEmpty) { |
| // Adds first range partition with custom hash schema through create. |
| boundsWithCustomHashSchema.headOption.foreach(bound => { |
| val rangePartition = new RangePartitionWithCustomHashSchema( |
| bound.lowerBound, |
| bound.upperBound, |
| RangePartitionBound.INCLUSIVE_BOUND, |
| RangePartitionBound.EXCLUSIVE_BOUND) |
| bound.hashSchemas.asScala.foreach { hp => |
| val columnNames = hp.getColumnIds.asScala.map { id => |
| schema.getColumnByIndex(id).getName |
| } |
| rangePartition.addHashPartitions(columnNames.asJava, hp.getNumBuckets, hp.getSeed) |
| } |
| options.addRangePartition(rangePartition) |
| }) |
| context.createTable(restoreName, schema, options) |
| // Adds rest of range partitions with custom hash schema through alters. |
| boundsWithCustomHashSchema.tail.foreach(bound => { |
| val rangePartition = new RangePartitionWithCustomHashSchema( |
| bound.lowerBound, |
| bound.upperBound, |
| RangePartitionBound.INCLUSIVE_BOUND, |
| RangePartitionBound.EXCLUSIVE_BOUND) |
| bound.hashSchemas.asScala.foreach { hp => |
| val columnNames = hp.getColumnIds.asScala.map { id => |
| schema.getColumnByIndex(id).getName |
| } |
| rangePartition.addHashPartitions(columnNames.asJava, hp.getNumBuckets, hp.getSeed) |
| } |
| val options = new AlterTableOptions() |
| options.addRangePartition(rangePartition) |
| context.syncClient.alterTable(restoreName, options) |
| }) |
| } |
| } |
| |
| /** |
| * Returns a modified DataFrame with columns adjusted to match the lastMetadata. |
| */ |
| private def adjustSchema( |
| df: DataFrame, |
| currentMetadata: TableMetadataPB, |
| lastMetadata: TableMetadataPB, |
| rowActionCol: String): DataFrame = { |
| log.info("Adjusting columns to handle alterations") |
| val idToName = lastMetadata.getColumnIdsMap.asScala.map(_.swap) |
| // Ignore the rowActionCol, which isn't a real column. |
| val currentColumns = currentMetadata.getColumnIdsMap.asScala.filter(_._1 != rowActionCol) |
| var result = df |
| // First drop all the columns that no longer exist. |
| // This is required to be sure a rename doesn't collide with an old column. |
| currentColumns.foreach { |
| case (colName, id) => |
| if (!idToName.contains(id)) { |
| // If the last metadata doesn't contain the id, the column is dropped. |
| log.info(s"Dropping the column $colName from backup data") |
| result = result.drop(colName) |
| } |
| } |
| // Then rename all the columns that were renamed in the last metadata. |
| currentColumns.foreach { |
| case (colName, id) => |
| if (idToName.contains(id) && idToName(id) != colName) { |
| // If the final name doesn't match the current name, the column is renamed. |
| log.info(s"Renamed the column $colName to ${idToName(id)} in backup data") |
| result = result.withColumnRenamed(colName, idToName(id)) |
| } |
| } |
| result |
| } |
| |
| /** |
| * Creates a KuduPartitioner that can be used to filter out rows for the current |
| * backup data which no longer apply to partitions in the last metadata. |
| * |
| * In order to do this, tablet metadata are compared in the current metadata to the |
| * last metadata. Tablet IDs that are not in the final metadata are filtered out and |
| * the remaining tablet metadata is used to create a KuduPartitioner. The resulting |
| * KuduPartitioner can then be used to filter out rows that are no longer valid |
| * because those rows will fall into a non-covered range. |
| */ |
| private def createPartitionFilter( |
| currentMetadata: TableMetadataPB, |
| lastMetadata: TableMetadataPB): KuduPartitioner = { |
| val lastTablets = lastMetadata.getTabletsMap |
| val validTablets = |
| currentMetadata.getTabletsMap.asScala.flatMap { |
| case (id, pm) => |
| if (lastTablets.containsKey(id)) { |
| // Create the partition object needed for the KuduPartitioner. |
| val partition = new Partition( |
| pm.getPartitionKeyStart.toByteArray, |
| pm.getPartitionKeyEnd.toByteArray, |
| pm.getHashBucketsList) |
| Some((id, partition)) |
| } else { |
| // Ignore tablets that are no longer valid |
| None |
| } |
| } |
| val partitionSchema = TableMetadata.getPartitionSchema(currentMetadata) |
| new KuduPartitioner(partitionSchema, validTablets.asJava) |
| } |
| |
| def main(args: Array[String]): Unit = { |
| val options = RestoreOptions |
| .parse(args) |
| .getOrElse(throw new IllegalArgumentException("could not parse the arguments")) |
| |
| val session = SparkSession |
| .builder() |
| .appName("Kudu Table Restore") |
| .getOrCreate() |
| |
| val isRunSuccessful: Boolean = run(options, session) |
| if (!isRunSuccessful) { |
| throw new RuntimeException("Kudu Table Restore application failed!") |
| } |
| |
| session.stop() |
| } |
| } |
