[SPARK-19563][SQL] avoid unnecessary sort in FileFormatWriter #16898
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -38,10 +38,9 @@ import org.apache.spark.sql.catalyst.expressions._ | |
| import org.apache.spark.sql.catalyst.plans.physical.HashPartitioning | ||
| import org.apache.spark.sql.catalyst.InternalRow | ||
| import org.apache.spark.sql.catalyst.util.DateTimeUtils | ||
| import org.apache.spark.sql.execution.{QueryExecution, SortExec, SQLExecution} | ||
| import org.apache.spark.sql.types.{StringType, StructType} | ||
| import org.apache.spark.util.{SerializableConfiguration, Utils} | ||
|
|
||
|
|
||
| /** A helper object for writing FileFormat data out to a location. */ | ||
|
|
@@ -64,9 +63,9 @@ object FileFormatWriter extends Logging { | |
| val serializableHadoopConf: SerializableConfiguration, | ||
| val outputWriterFactory: OutputWriterFactory, | ||
| val allColumns: Seq[Attribute], | ||
| val dataColumns: Seq[Attribute], | ||
| val partitionColumns: Seq[Attribute], | ||
| val bucketIdExpression: Option[Expression], | ||
| val path: String, | ||
| val customPartitionLocations: Map[TablePartitionSpec, String], | ||
| val maxRecordsPerFile: Long) | ||
|
|
@@ -108,9 +107,21 @@ object FileFormatWriter extends Logging { | |
| job.setOutputValueClass(classOf[InternalRow]) | ||
| FileOutputFormat.setOutputPath(job, new Path(outputSpec.outputPath)) | ||
|
|
||
| val allColumns = queryExecution.logical.output | ||
| val partitionSet = AttributeSet(partitionColumns) | ||
| val dataColumns = queryExecution.logical.output.filterNot(partitionSet.contains) | ||
|
|
||
| val bucketIdExpression = bucketSpec.map { spec => | ||
| val bucketColumns = spec.bucketColumnNames.map(c => dataColumns.find(_.name == c).get) | ||
| // Use `HashPartitioning.partitionIdExpression` as our bucket id expression, so that we can | ||
| // guarantee the data distribution is same between shuffle and bucketed data source, which | ||
| // enables us to only shuffle one side when join a bucketed table and a normal one. | ||
| HashPartitioning(bucketColumns, spec.numBuckets).partitionIdExpression | ||
| } | ||
| val sortColumns = bucketSpec.toSeq.flatMap { | ||
| spec => spec.sortColumnNames.map(c => dataColumns.find(_.name == c).get) | ||
| } | ||
|
|
||
| // Note: prepareWrite has side effect. It sets "job". | ||
| val outputWriterFactory = | ||
| fileFormat.prepareWrite(sparkSession, job, options, dataColumns.toStructType) | ||
|
|
@@ -119,23 +130,45 @@ object FileFormatWriter extends Logging { | |
| uuid = UUID.randomUUID().toString, | ||
| serializableHadoopConf = new SerializableConfiguration(job.getConfiguration), | ||
| outputWriterFactory = outputWriterFactory, | ||
| allColumns = allColumns, | ||
| dataColumns = dataColumns, | ||
| partitionColumns = partitionColumns, | ||
| bucketIdExpression = bucketIdExpression, | ||
| path = outputSpec.outputPath, | ||
| customPartitionLocations = outputSpec.customPartitionLocations, | ||
| maxRecordsPerFile = options.get("maxRecordsPerFile").map(_.toLong) | ||
| .getOrElse(sparkSession.sessionState.conf.maxRecordsPerFile) | ||
| ) | ||
|
|
||
| // We should first sort by partition columns, then bucket id, and finally sorting columns. | ||
| val requiredOrdering = partitionColumns ++ bucketIdExpression ++ sortColumns | ||
| // the sort order doesn't matter | ||
| val actualOrdering = queryExecution.executedPlan.outputOrdering.map(_.child) | ||
|
There was a problem hiding this comment. @cloud-fan would it be possible to use the logical plan rather than the executedPlan? If the optimizer decides the data is already sorted according according to the logical plan the executedPlan won't include the fields.
Contributor
Author
There was a problem hiding this comment. That would be great, but may need some refactoring. |
||
| val orderingMatched = if (requiredOrdering.length > actualOrdering.length) { | ||
| false | ||
| } else { | ||
| requiredOrdering.zip(actualOrdering).forall { | ||
| case (requiredOrder, childOutputOrder) => | ||
| requiredOrder.semanticEquals(childOutputOrder) | ||
|
Member
There was a problem hiding this comment. Because
Contributor
Author
There was a problem hiding this comment. it's |
||
| } | ||
| } | ||
|
|
||
| SQLExecution.withNewExecutionId(sparkSession, queryExecution) { | ||
| // This call shouldn't be put into the `try` block below because it only initializes and | ||
| // prepares the job, any exception thrown from here shouldn't cause abortJob() to be called. | ||
| committer.setupJob(job) | ||
|
|
||
| try { | ||
| val rdd = if (orderingMatched) { | ||
| queryExecution.toRdd | ||
| } else { | ||
| SortExec( | ||
| requiredOrdering.map(SortOrder(_, Ascending)), | ||
| global = false, | ||
| child = queryExecution.executedPlan).execute() | ||
| } | ||
|
|
||
| val ret = sparkSession.sparkContext.runJob(rdd, | ||
| (taskContext: TaskContext, iter: Iterator[InternalRow]) => { | ||
| executeTask( | ||
| description = description, | ||
|
|
@@ -189,7 +222,7 @@ object FileFormatWriter extends Logging { | |
| committer.setupTask(taskAttemptContext) | ||
|
|
||
| val writeTask = | ||
| if (description.partitionColumns.isEmpty && description.bucketIdExpression.isEmpty) { | ||
| new SingleDirectoryWriteTask(description, taskAttemptContext, committer) | ||
| } else { | ||
| new DynamicPartitionWriteTask(description, taskAttemptContext, committer) | ||
|
|
@@ -287,31 +320,16 @@ object FileFormatWriter extends Logging { | |
| * multiple directories (partitions) or files (bucketing). | ||
| */ | ||
| private class DynamicPartitionWriteTask( | ||
| desc: WriteJobDescription, | ||
|
Member
There was a problem hiding this comment.
Contributor
Author
There was a problem hiding this comment. I'd like to change both to make it consistent. |
||
| taskAttemptContext: TaskAttemptContext, | ||
| committer: FileCommitProtocol) extends ExecuteWriteTask { | ||
|
|
||
| // currentWriter is initialized whenever we see a new key | ||
| private var currentWriter: OutputWriter = _ | ||
|
|
||
| /** Expressions that given partition columns build a path string like: col1=val/col2=val/... */ | ||
| private def partitionPathExpression: Seq[Expression] = { | ||
| desc.partitionColumns.zipWithIndex.flatMap { case (c, i) => | ||
| // TODO: use correct timezone for partition values. | ||
| val escaped = ScalaUDF( | ||
| ExternalCatalogUtils.escapePathName _, | ||
|
|
@@ -325,35 +343,46 @@ object FileFormatWriter extends Logging { | |
| } | ||
|
|
||
| /** | ||
| * Opens a new OutputWriter given a partition key and optional bucket id. | ||
| * If bucket id is specified, we will append it to the end of the file name, but before the | ||
|
||
| * file extension, e.g. part-r-00009-ea518ad4-455a-4431-b471-d24e03814677-00002.gz.parquet | ||
| * | ||
| * @param partColsAndBucketId a row consisting of partition columns and a bucket id for the | ||
| * current row. | ||
| * @param getPartitionPath a function that projects the partition values into a path string. | ||
| * @param fileCounter the number of files that have been written in the past for this specific | ||
| * partition. This is used to limit the max number of records written for a | ||
| * single file. The value should start from 0. | ||
| * @param updatedPartitions the set of updated partition paths, we should add the new partition | ||
| * path of this writer to it. | ||
| */ | ||
| private def newOutputWriter( | ||
| partColsAndBucketId: InternalRow, | ||
| getPartitionPath: UnsafeProjection, | ||
| fileCounter: Int, | ||
| updatedPartitions: mutable.Set[String]): Unit = { | ||
| val partDir = if (desc.partitionColumns.isEmpty) { | ||
| None | ||
| } else { | ||
| Option(getPartitionPath(partColsAndBucketId).getString(0)) | ||
| } | ||
| partDir.foreach(updatedPartitions.add) | ||
|
|
||
| // If the bucketId expression is defined, the bucketId column is right after the partition | ||
| // columns. | ||
| val bucketId = if (desc.bucketIdExpression.isDefined) { | ||
| BucketingUtils.bucketIdToString(partColsAndBucketId.getInt(desc.partitionColumns.length)) | ||
| } else { | ||
| "" | ||
| } | ||
|
|
||
| // This must be in a form that matches our bucketing format. See BucketingUtils. | ||
| val ext = f"$bucketId.c$fileCounter%03d" + | ||
| desc.outputWriterFactory.getFileExtension(taskAttemptContext) | ||
|
|
||
| val customPath = partDir match { | ||
| case Some(dir) => | ||
| desc.customPartitionLocations.get(PartitioningUtils.parsePathFragment(dir)) | ||
| case _ => | ||
| None | ||
| } | ||
|
|
@@ -363,80 +392,42 @@ object FileFormatWriter extends Logging { | |
| committer.newTaskTempFile(taskAttemptContext, partDir, ext) | ||
| } | ||
|
|
||
| currentWriter = desc.outputWriterFactory.newInstance( | ||
| path = path, | ||
| dataSchema = desc.dataColumns.toStructType, | ||
| context = taskAttemptContext) | ||
| } | ||
|
|
||
| override def execute(iter: Iterator[InternalRow]): Set[String] = { | ||
| val getPartitionColsAndBucketId = UnsafeProjection.create( | ||
| desc.partitionColumns ++ desc.bucketIdExpression, desc.allColumns) | ||
|
|
||
| // Generates the partition path given the row generated by `getPartitionColsAndBucketId`. | ||
| val getPartPath = UnsafeProjection.create( | ||
| Seq(Concat(partitionPathExpression)), desc.partitionColumns) | ||
|
|
||
| // Returns the data columns to be written given an input row | ||
| val getOutputRow = UnsafeProjection.create(desc.dataColumns, desc.allColumns) | ||
|
|
||
| // If anything below fails, we should abort the task. | ||
| var recordsInFile: Long = 0L | ||
| var fileCounter = 0 | ||
| var currentPartColsAndBucketId: UnsafeRow = null | ||
| val updatedPartitions = mutable.Set[String]() | ||
| for (row <- iter) { | ||
| val nextPartColsAndBucketId = getPartitionColsAndBucketId(row) | ||
|
||
| if (currentPartColsAndBucketId != nextPartColsAndBucketId) { | ||
| // See a new partition or bucket - write to a new partition dir (or a new bucket file). | ||
| currentPartColsAndBucketId = nextPartColsAndBucketId.copy() | ||
| logDebug(s"Writing partition: $currentPartColsAndBucketId") | ||
|
|
||
| recordsInFile = 0 | ||
| fileCounter = 0 | ||
|
|
||
| releaseResources() | ||
| newOutputWriter(currentPartColsAndBucketId, getPartPath, fileCounter, updatedPartitions) | ||
| } else if (desc.maxRecordsPerFile > 0 && | ||
| recordsInFile >= desc.maxRecordsPerFile) { | ||
| // Exceeded the threshold in terms of the number of records per file. | ||
| // Create a new file by increasing the file counter. | ||
| recordsInFile = 0 | ||
|
|
@@ -445,10 +436,10 @@ object FileFormatWriter extends Logging { | |
| s"File counter $fileCounter is beyond max value $MAX_FILE_COUNTER") | ||
|
|
||
| releaseResources() | ||
| newOutputWriter(currentPartColsAndBucketId, getPartPath, fileCounter, updatedPartitions) | ||
| } | ||
|
|
||
| currentWriter.write(getOutputRow(row)) | ||
| recordsInFile += 1 | ||
| } | ||
| releaseResources() | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If we rewrite it to
val dataColumns = allColumns.filterNot(partitionColumns.contains), we do not needpartitionSetThere was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it's so minor, I'll fix it in my next PR