This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 7772d4c test: Add test for RoundRobinPartitioning (#54)
7772d4c is described below
commit 7772d4c67653053be327eaae74dd0a3e808846fd
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Tue Feb 20 00:02:15 2024 -0800
test: Add test for RoundRobinPartitioning (#54)
---
.../org/apache/comet/exec/CometShuffleSuite.scala | 24 ++++++++++++++++++++++
1 file changed, 24 insertions(+)
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometShuffleSuite.scala
b/spark/src/test/scala/org/apache/comet/exec/CometShuffleSuite.scala
index dc47482..431a08f 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometShuffleSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometShuffleSuite.scala
@@ -63,6 +63,30 @@ abstract class CometShuffleSuiteBase extends CometTestBase
with AdaptiveSparkPla
import testImplicits._
+ test("RoundRobinPartitioning is supported by columnar shuffle") {
+ withSQLConf(
+ // AQE has `ShuffleStage` which is a leaf node which blocks
+ // collecting `CometShuffleExchangeExec` node.
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
+ CometConf.COMET_COLUMNAR_SHUFFLE_ENABLED.key -> "true") {
+ withParquetTable((0 until 5).map(i => (i, (i + 1).toLong)), "tbl") {
+ val df = sql("SELECT * FROM tbl")
+ val shuffled = df
+ .select($"_1" + 1 as ("a"))
+ .filter($"a" > 4)
+ .repartition(10)
+ .limit(2)
+
+ checkAnswer(shuffled, Row(5) :: Nil)
+ val cometShuffleExecs = checkCometExchange(shuffled, 1, false)
+
+ cometShuffleExecs(0).outputPartitioning.getClass.getName
+ .contains("RoundRobinPartitioning")
+ }
+ }
+ }
+
test("columnar shuffle on array") {
Seq(10, 201).foreach { numPartitions =>
Seq("1.0", "10.0").foreach { ratio =>
