hudi-agent commented on code in PR #19767:
URL: https://github.com/apache/hudi/pull/19767#discussion_r3870813442
##########
hudi-client/hudi-spark-client/src/main/scala/org/apache/spark/sql/BucketPartitionUtils.scala:
##########
@@ -49,10 +57,35 @@ object BucketPartitionUtils extends SparkAdapterSupport {
val getPartitionKey = getPartitionKeyExtractor()
// use internalRow to avoid extra convert.
- val reRdd = df.queryExecution.toRdd
- .keyBy(row => getPartitionKey(row))
- .repartitionAndSortWithinPartitions(partitioner)
- .values
+ val internalRows = df.queryExecution.toRdd
+ val reRdd = if (sortByRecordKey) {
+ val utf8StringFactory = sparkAdapter.getUTF8StringFactory
+ // Use (bucket route, record key) as the shuffle key. Tuple ordering
groups rows by
+ // (partition path, bucket id) first, then orders each bucket by the
full record key.
+ // Scala derives the record-key ordering from HoodieUTF8String's
Comparable implementation,
+ // which uses binary UTF-8 ordering for both Spark 3 and Spark 4.
+ val keyedRows = internalRows.keyBy(row => {
+ // InternalRow may reuse its mutable backing buffer, so keep an
independent key for shuffle.
+ val recordKey = utf8StringFactory.wrapUTF8String(
+ row.getUTF8String(HoodieRecord.RECORD_KEY_META_FIELD_ORD).copy())
+ (getPartitionKey(row), recordKey)
+ })
+ // The record key participates only in sorting; bucket routing remains
unchanged.
+ val bucketRoutePartitioner = new Partitioner {
+ override def numPartitions: Int = partitioner.numPartitions
+
+ override def getPartition(key: Any): Int = {
+ val bucketRoute = key.asInstanceOf[((String, Int),
HoodieUTF8String)]._1
Review Comment:
🤖 nit: have you considered a `case` match instead of `asInstanceOf` here?
Something like `case ((route, _), _: HoodieUTF8String) =>
partitioner.getPartition(route)` would make the expected tuple shape
self-documenting and give a clear `MatchError` rather than a
`ClassCastException` deep in the shuffle path if the key type ever changes.
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag
quality.</i></sub>
##########
hudi-client/hudi-spark-client/src/test/java/org/apache/hudi/client/clustering/run/strategy/TestSingleSparkJobConsistentHashingExecutionStrategy.java:
##########
@@ -0,0 +1,89 @@
+/*
+ * 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.hudi.client.clustering.run.strategy;
+
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.util.collection.ClosableIterator;
+import org.apache.hudi.common.util.collection.LoserTreeMergeIterator;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+class TestSingleSparkJobConsistentHashingExecutionStrategy {
Review Comment:
🤖 nit: could you either move this test into `TestLoserTreeMergeIterator`
(since it only exercises `LoserTreeMergeIterator` directly, which already has
its own test class) or rename this class to something like
`TestSingleSparkJobConsistentHashingLsmMerge`? As-is the class name points
future readers to the wrong abstraction.
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag
quality.</i></sub>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
