Github user HyukjinKwon commented on a diff in the pull request:
https://github.com/apache/spark/pull/12855#discussion_r61863848
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/WriterContainer.scala
---
@@ -363,84 +365,87 @@ private[sql] class DynamicPartitionWriterContainer(
}
def writeRows(taskContext: TaskContext, iterator:
Iterator[InternalRow]): Unit = {
- executorSideSetup(taskContext)
-
- // We should first sort by partition columns, then bucket id, and
finally sorting columns.
- val sortingExpressions: Seq[Expression] = partitionColumns ++
bucketIdExpression ++ sortColumns
- val getSortingKey = UnsafeProjection.create(sortingExpressions,
inputSchema)
-
- val sortingKeySchema = StructType(sortingExpressions.map {
- case a: Attribute => StructField(a.name, a.dataType, a.nullable)
- // The sorting expressions are all `Attribute` except bucket id.
- case _ => StructField("bucketId", IntegerType, nullable = false)
- })
-
- // Returns the data columns to be written given an input row
- val getOutputRow = UnsafeProjection.create(dataColumns, inputSchema)
-
- // Returns the partition path given a partition key.
- val getPartitionString =
- UnsafeProjection.create(Concat(partitionStringExpression) :: Nil,
partitionColumns)
-
- // Sorts the data before write, so that we only need one writer at the
same time.
- // TODO: inject a local sort operator in planning.
- val sorter = new UnsafeKVExternalSorter(
- sortingKeySchema,
- StructType.fromAttributes(dataColumns),
- SparkEnv.get.blockManager,
- SparkEnv.get.serializerManager,
- TaskContext.get().taskMemoryManager().pageSizeBytes)
-
- while (iterator.hasNext) {
- val currentRow = iterator.next()
- sorter.insertKV(getSortingKey(currentRow), getOutputRow(currentRow))
- }
- logInfo(s"Sorting complete. Writing out partition files one at a
time.")
-
- val getBucketingKey: InternalRow => InternalRow = if
(sortColumns.isEmpty) {
- identity
- } else {
-
UnsafeProjection.create(sortingExpressions.dropRight(sortColumns.length).zipWithIndex.map
{
- case (expr, ordinal) => BoundReference(ordinal, expr.dataType,
expr.nullable)
+ if (iterator.hasNext) {
--- End diff --
Here as well. Simply added `iterator.hasNext` check.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
