Skip to content

Commit 0f98674

Browse files
[AURON #2468] Support binary sort keys in native range shuffle (#2469)
**Which issue does this PR close?** Closes #2468 **Rationale for this change** Auron currently rejects `BinaryType` sort keys when converting Spark `RangePartitioning` to a native shuffle exchange. Binary values are already supported by Auron's Arrow type conversion and native row encoding, so this restriction causes queries that order or range-partition by a binary column to fall back unnecessarily. **What changes are included in this PR?** Allows a top-level `BinaryType` sort key in native range partitioning. Updates the existing binary range-partitioning test to assert that `NativeShuffleExchangeExec` is used. The general supported-type check is unchanged, so binary support is not enabled for unrelated operators or data source scans. **Are there any user-facing changes?** Queries that range-partition by a binary column can now use Auron's native shuffle path. There are no user-facing API or configuration changes. **How was this patch tested?** UT. **Was this patch authored or co-authored using generative AI tooling?** - [x] Yes - [ ] No Codex was used only to Review and understand the codebase Signed-off-by: weimingdiit <weimingdiit@gmail.com> Co-authored-by: Shilun Fan <slfan1989@apache.org> Signed-off-by: Shilun Fan <slfan1989@apache.org>
1 parent 2412b01 commit 0f98674

2 files changed

Lines changed: 6 additions & 2 deletions

File tree

spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import org.apache.spark.sql.{AuronQueryTest, Row}
2020
import org.apache.spark.sql.auron.join.JoinBuildSides.{JoinBuildLeft, JoinBuildRight}
2121
import org.apache.spark.sql.execution.auron.plan.NativeFilterBase
2222
import org.apache.spark.sql.execution.auron.plan.NativeShuffledHashJoinBase
23+
import org.apache.spark.sql.execution.auron.plan.NativeShuffleExchangeExec
2324
import org.apache.spark.sql.execution.auron.plan.NativeSortMergeJoinBase
2425
import org.apache.spark.sql.execution.joins.auron.plan.NativeBroadcastJoinExec
2526

@@ -112,7 +113,10 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ
112113
withTable("t1", "t2") {
113114
sql("create table t1(c1 binary, c2 int) using parquet")
114115
sql("insert into t1 values (cast('test1' as binary), 1), (cast('test2' as binary), 2)")
115-
checkSparkAnswerAndOperator("select c2 from t1 order by c1")
116+
val df = checkSparkAnswerAndOperator("select c2 from t1 order by c1")
117+
assert(collectFirst(df.queryExecution.executedPlan) { case e: NativeShuffleExchangeExec =>
118+
e
119+
}.isDefined)
116120
}
117121
}
118122

spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -447,7 +447,7 @@ object AuronConverters extends Logging {
447447
outputPartitioning match {
448448
case partitioning: RangePartitioning =>
449449
val unsupportedOrderType = partitioning.ordering
450-
.find(e => !isTypeSupported(e.dataType))
450+
.find(e => !isTypeSupported(e.dataType) && e.dataType != BinaryType)
451451
assert(
452452
unsupportedOrderType.isEmpty,
453453
s"Unsupported order type in range partitioning: ${unsupportedOrderType.get}")

0 commit comments

Comments
 (0)