Github user davies commented on a diff in the pull request:
https://github.com/apache/spark/pull/10638#discussion_r49158190
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/WriterContainer.scala
---
@@ -461,54 +396,55 @@ private[sql] class DynamicPartitionWriterContainer(
// If anything below fails, we should abort the task.
try {
- // If there is no sorting columns, we set sorter to null and try the
hash-based writing first,
- // and fill the sorter if there are too many writers and we need to
fall back on sorting.
- // If there are sorting columns, then we have to sort the data
anyway, and no need to try the
- // hash-based writing first.
- var sorter: UnsafeKVExternalSorter = if (sortColumns.nonEmpty) {
- new UnsafeKVExternalSorter(
- sortingKeySchema,
- StructType.fromAttributes(dataColumns),
- SparkEnv.get.blockManager,
- TaskContext.get().taskMemoryManager().pageSizeBytes)
+ // 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,
+ 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 needNewWriter: (UnsafeRow, UnsafeRow) => Boolean = if
(sortColumns.isEmpty) {
+ (key1, key2) => key1 != key2
} else {
- null
+ // Given 2 sorting keys, we should ignore the sorting columns to
check if these 2 keys
+ // belong to same bucket.
+ val bucketingKey =
sortingExpressions.dropRight(sortColumns.length).zipWithIndex.map {
+ case (expr, ordinal) => BoundReference(ordinal, expr.dataType,
expr.nullable)
+ }
+ // Use 2 projection here as `UnsafeProjection` will share one
result row during execution.
+ val getBucketingKey1 = UnsafeProjection.create(bucketingKey)
+ val getBucketingKey2 = UnsafeProjection.create(bucketingKey)
+
+ (key1, key2) => key1 == null || getBucketingKey1(key1) !=
getBucketingKey2(key2)
--- End diff --
You should save the generated key:
```
currentKey = getBucketingKey(sortedIterator.getKey).copy()
```
compare with:
```
if (currentKey != getBucketingKey(sortedIterator.getKey)) {
}
```
---
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
