address comments

apache · Mar 11, 2024 · 3c8c34f · 3c8c34f
1 parent 9ae1ccb
commit 3c8c34f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/SparkPlanExecApiImpl.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/SparkPlanExecApiImpl.scala
@@ -229,14 +229,16 @@ class SparkPlanExecApiImpl extends SparkPlanExecApi {
         val projectTransformer = ProjectExecTransformer(projectList, newChild)
         val sortOrder = SortOrder(projectTransformer.output.head, Ascending)
         val sortByHashCode = SortExecTransformer(Seq(sortOrder), global = false, projectTransformer)
-        val projectValidationResult = projectTransformer.doValidate()
-        val sortValidationResult = sortByHashCode.doValidate()
-        if (projectValidationResult.isValid && sortValidationResult.isValid) {
-          ColumnarShuffleExchangeExec(shuffle, sortByHashCode, sortByHashCode.output.drop(1))
+        val dropSortColumnTransformer = ProjectExecTransformer(projectList.drop(1), sortByHashCode)
+        if (dropSortColumnTransformer.doValidate().isValid) {
+          ColumnarShuffleExchangeExec(
+            shuffle,
+            dropSortColumnTransformer,
+            dropSortColumnTransformer.output)
         } else {
           TransformHints.tagNotTransformable(
             shuffle,
-            if (projectValidationResult.isValid) sortValidationResult else projectValidationResult)
+            dropSortColumnTransformer.doValidate().reason.get)
           shuffle.withNewChildren(newChild :: Nil)
         }
       case _ =>

diff --git a/backends-velox/src/test/scala/io/glutenproject/execution/TestOperator.scala b/backends-velox/src/test/scala/io/glutenproject/execution/TestOperator.scala
@@ -1186,9 +1186,10 @@ class TestOperator extends VeloxWholeStageTransformerSuite {
     runQueryAndCompare("SELECT /*+ REPARTITION(3) */ l_orderkey, l_partkey FROM lineitem") {
       /*
         ColumnarExchange RoundRobinPartitioning(3), REPARTITION_BY_NUM, [l_orderkey#16L, l_partkey#17L)
-      +- ^(2) SortExecTransformer [hash_partition_key#302 ASC NULLS FIRST], false, 0
-         +- ^(2) ProjectExecTransformer [hash(l_orderkey#16L, l_partkey#17L) AS hash_partition_key#302, l_orderkey#16L, l_partkey#17L]
-            +- ^(2) BatchScanExecTransformer[l_orderkey#16L, l_partkey#17L] ParquetScan DataFilters: [], Format: parquet, Location: InMemoryFileIndex(1 paths)[..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<l_orderkey:bigint,l_partkey:bigint>, PushedFilters: [] RuntimeFilters: []
+        +- ^(2) ProjectExecTransformer [l_orderkey#16L, l_partkey#17L]
+          +- ^(2) SortExecTransformer [hash_partition_key#302 ASC NULLS FIRST], false, 0
+            +- ^(2) ProjectExecTransformer [hash(l_orderkey#16L, l_partkey#17L) AS hash_partition_key#302, l_orderkey#16L, l_partkey#17L]
+                +- ^(2) BatchScanExecTransformer[l_orderkey#16L, l_partkey#17L] ParquetScan DataFilters: [], Format: parquet, Location: InMemoryFileIndex(1 paths)[..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<l_orderkey:bigint,l_partkey:bigint>, PushedFilters: [] RuntimeFilters: []
        */
       checkOperatorMatch[SortExecTransformer]
     }