Skip to content

Commit 0589ddf

Browse files
yew1ebyew1eb
committed
[AURON #1746] Introduce NativeTakeOrderedAndProjectExec to fuse TakeOrdered + Project
1 parent a553d4a commit 0589ddf

7 files changed

Lines changed: 84 additions & 31 deletions

File tree

‎spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,8 @@ import org.apache.spark.sql.execution.auron.plan.NativeShuffleExchangeBase
8787
import org.apache.spark.sql.execution.auron.plan.NativeShuffleExchangeExec
8888
import org.apache.spark.sql.execution.auron.plan.NativeSortBase
8989
import org.apache.spark.sql.execution.auron.plan.NativeSortExec
90-
import org.apache.spark.sql.execution.auron.plan.NativeTakeOrderedBase
91-
import org.apache.spark.sql.execution.auron.plan.NativeTakeOrderedExec
90+
import org.apache.spark.sql.execution.auron.plan.NativeTakeOrderedAndProjectBase
91+
import org.apache.spark.sql.execution.auron.plan.NativeTakeOrderedAndProjectExec
9292
import org.apache.spark.sql.execution.auron.plan.NativeUnionBase
9393
import org.apache.spark.sql.execution.auron.plan.NativeUnionExec
9494
import org.apache.spark.sql.execution.auron.plan.NativeWindowBase
@@ -328,17 +328,18 @@ class ShimsImpl extends Shims with Logging {
328328
child: SparkPlan): NativeSortBase =
329329
NativeSortExec(sortOrder, global, child)
330330

331-
override def createNativeTakeOrderedExec(
331+
override def createNativeTakeOrderedAndProjectExec(
332332
limit: Long,
333333
sortOrder: Seq[SortOrder],
334-
child: SparkPlan): NativeTakeOrderedBase =
335-
NativeTakeOrderedExec(limit, sortOrder, child)
334+
projectList: Seq[NamedExpression],
335+
child: SparkPlan): NativeTakeOrderedAndProjectBase =
336+
NativeTakeOrderedAndProjectExec(limit, sortOrder, projectList, child)
336337

337338
override def createNativePartialTakeOrderedExec(
338339
limit: Long,
339340
sortOrder: Seq[SortOrder],
340341
child: SparkPlan,
341-
metrics: Map[String, SQLMetric]): NativePartialTakeOrderedBase =
342+
metrics: Map[String, SQLMetric]): NativePartialTakeOrderedAndProjectBase =
342343
NativePartialTakeOrderedExec(limit, sortOrder, child, metrics)
343344

344345
override def createNativeUnionExec(

‎spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativePartialTakeOrderedExec.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ case class NativePartialTakeOrderedExec(
2727
sortOrder: Seq[SortOrder],
2828
override val child: SparkPlan,
2929
override val metrics: Map[String, SQLMetric])
30-
extends NativePartialTakeOrderedBase(limit, sortOrder, child, metrics) {
30+
extends NativePartialTakeOrderedAndProjectBase(limit, sortOrder, child, metrics) {
3131

3232
@sparkver("3.2 / 3.3 / 3.4 / 3.5")
3333
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =

spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedExec.scala renamed to spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedAndProjectExec.scala

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,17 @@
1616
*/
1717
package org.apache.spark.sql.execution.auron.plan
1818

19-
import org.apache.spark.sql.catalyst.expressions.SortOrder
19+
import org.apache.spark.sql.catalyst.expressions.{NamedExpression, SortOrder}
2020
import org.apache.spark.sql.execution.SparkPlan
2121

2222
import org.apache.auron.sparkver
2323

24-
case class NativeTakeOrderedExec(
24+
case class NativeTakeOrderedAndProjectExec(
2525
limit: Long,
2626
sortOrder: Seq[SortOrder],
27+
projectList: Seq[NamedExpression],
2728
override val child: SparkPlan)
28-
extends NativeTakeOrderedBase(limit, sortOrder, child) {
29+
extends NativeTakeOrderedAndProjectBase(limit, sortOrder, projectList, child) {
2930

3031
@sparkver("3.2 / 3.3 / 3.4 / 3.5")
3132
override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.auron
18+
19+
import org.apache.spark.sql.AuronQueryTest
20+
import org.apache.spark.sql.execution.auron.plan.{NativeFilterExec, NativeTakeOrderedAndProjectExec}
21+
22+
class AuronExecSuite extends AuronQueryTest with BaseAuronSQLSuite {
23+
24+
test("TakeOrderedAndProject") {
25+
withTempView("t1") {
26+
sql("create table t1(id INT, name STRING) using parquet")
27+
sql("insert into t1 values(1, 'a'),(2, 'b'),(3, 'c'),(3, 'c'),(4, 'd'),(5, 'e')")
28+
29+
// executeCollect (collect rows directly to driver)
30+
var df = checkSparkAnswerAndOperator(
31+
"SELECT id + 42, name, length(name) FROM t1 order by id limit 4")
32+
val firstNode = collect(stripAQEPlan(df.queryExecution.executedPlan)) { case exec =>
33+
exec
34+
}.head
35+
assert(firstNode.isInstanceOf[NativeTakeOrderedAndProjectExec])
36+
37+
// doExecuteNative
38+
df = checkSparkAnswerAndOperator(
39+
"select * from (SELECT id, id + 42, length(name) FROM t1 order by id limit 4) where id > 1")
40+
assert(collectFirst(df.queryExecution.executedPlan) {
41+
case f: NativeFilterExec if f.child.isInstanceOf[NativeTakeOrderedAndProjectExec] => true
42+
}.isDefined)
43+
}
44+
}
45+
}

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -783,17 +783,11 @@ object AuronConverters extends Logging {
783783

784784
def convertTakeOrderedAndProjectExec(exec: TakeOrderedAndProjectExec): SparkPlan = {
785785
logDebugPlanConversion(exec)
786-
val nativeTakeOrdered = Shims.get.createNativeTakeOrderedExec(
786+
Shims.get.createNativeTakeOrderedAndProjectExec(
787787
exec.limit,
788788
exec.sortOrder,
789+
exec.projectList,
789790
addRenameColumnsExec(convertToNative(exec.child)))
790-
791-
if (exec.projectList != exec.child.output) {
792-
val project = ProjectExec(exec.projectList, nativeTakeOrdered)
793-
tryConvert(project, convertProjectExec)
794-
} else {
795-
nativeTakeOrdered
796-
}
797791
}
798792

799793
def convertHashAggregateExec(exec: HashAggregateExec): SparkPlan = {

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -149,16 +149,17 @@ abstract class Shims {
149149
global: Boolean,
150150
child: SparkPlan): NativeSortBase
151151

152-
def createNativeTakeOrderedExec(
152+
def createNativeTakeOrderedAndProjectExec(
153153
limit: Long,
154154
sortOrder: Seq[SortOrder],
155-
child: SparkPlan): NativeTakeOrderedBase
155+
projectList: Seq[NamedExpression],
156+
child: SparkPlan): NativeTakeOrderedAndProjectBase
156157

157158
def createNativePartialTakeOrderedExec(
158159
limit: Long,
159160
sortOrder: Seq[SortOrder],
160161
child: SparkPlan,
161-
metrics: Map[String, SQLMetric]): NativePartialTakeOrderedBase
162+
metrics: Map[String, SQLMetric]): NativePartialTakeOrderedAndProjectBase
162163

163164
def createNativeUnionExec(children: Seq[SparkPlan], output: Seq[Attribute]): NativeUnionBase
164165

spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedBase.scala renamed to spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeTakeOrderedAndProjectBase.scala

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,16 +27,14 @@ import org.apache.spark.sql.auron.NativeRDD
2727
import org.apache.spark.sql.auron.NativeSupports
2828
import org.apache.spark.sql.auron.Shims
2929
import org.apache.spark.sql.catalyst.InternalRow
30-
import org.apache.spark.sql.catalyst.expressions.Ascending
31-
import org.apache.spark.sql.catalyst.expressions.Attribute
32-
import org.apache.spark.sql.catalyst.expressions.NullsFirst
33-
import org.apache.spark.sql.catalyst.expressions.SortOrder
30+
import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, NamedExpression, NullsFirst, SortOrder}
3431
import org.apache.spark.sql.catalyst.expressions.codegen.LazilyGeneratedOrdering
3532
import org.apache.spark.sql.catalyst.plans.physical.Partitioning
3633
import org.apache.spark.sql.catalyst.plans.physical.SinglePartition
3734
import org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning
3835
import org.apache.spark.sql.execution.SparkPlan
3936
import org.apache.spark.sql.execution.UnaryExecNode
37+
import org.apache.spark.sql.execution.auron.plan.NativeProjectBase.getNativeProjectBuilder
4038
import org.apache.spark.sql.execution.metric.SQLMetric
4139

4240
import org.apache.auron.metric.SparkMetricNode
@@ -46,9 +44,10 @@ import org.apache.auron.protobuf.PhysicalPlanNode
4644
import org.apache.auron.protobuf.PhysicalSortExprNode
4745
import org.apache.auron.protobuf.SortExecNode
4846

49-
abstract class NativeTakeOrderedBase(
47+
abstract class NativeTakeOrderedAndProjectBase(
5048
limit: Long,
5149
sortOrder: Seq[SortOrder],
50+
projectList: Seq[NamedExpression],
5251
override val child: SparkPlan)
5352
extends UnaryExecNode
5453
with NativeSupports {
@@ -76,6 +75,8 @@ abstract class NativeTakeOrderedBase(
7675
.build()
7776
}
7877

78+
private def nativeProject = getNativeProjectBuilder(projectList).buildPartial()
79+
7980
override def executeCollect(): Array[InternalRow] = {
8081
val partial = Shims.get.createNativePartialTakeOrderedExec(limit, sortOrder, child, metrics)
8182
val ord = new LazilyGeneratedOrdering(sortOrder, output)
@@ -114,6 +115,7 @@ abstract class NativeTakeOrderedBase(
114115

115116
// check whether native converting is supported
116117
nativeSortExprs
118+
nativeProject
117119

118120
override def doExecuteNative(): NativeRDD = {
119121
val partial = Shims.get.createNativePartialTakeOrderedExec(limit, sortOrder, child, metrics)
@@ -126,6 +128,7 @@ abstract class NativeTakeOrderedBase(
126128
val shuffled = Shims.get.createNativeShuffleExchangeExec(SinglePartition, partial)
127129
val shuffledRDD = NativeHelper.executeNative(shuffled)
128130
val nativeSortExprs = this.nativeSortExprs
131+
val nativeProject = this.nativeProject
129132

130133
// take top-K from the final partition
131134
new NativeRDD(
@@ -137,19 +140,27 @@ abstract class NativeTakeOrderedBase(
137140
rddShuffleReadFull = false,
138141
(_, taskContext) => {
139142
val inputPartition = shuffledRDD.partitions(0)
140-
val nativeTakeOrderedExec = SortExecNode
143+
val sortExec = SortExecNode
141144
.newBuilder()
142145
.setInput(shuffledRDD.nativePlan(inputPartition, taskContext))
143146
.addAllExpr(nativeSortExprs.asJava)
144147
.setFetchLimit(FetchLimit.newBuilder().setLimit(limit))
145148
.build()
146-
PhysicalPlanNode.newBuilder().setSort(nativeTakeOrderedExec).build()
149+
val nativeTakeOrderedExec = PhysicalPlanNode.newBuilder().setSort(sortExec).build()
150+
151+
if (projectList != child.output) {
152+
val nativeTakeOrderedAndProjectExec =
153+
nativeProject.toBuilder.setInput(nativeTakeOrderedExec).build()
154+
PhysicalPlanNode.newBuilder().setProjection(nativeTakeOrderedAndProjectExec).build()
155+
} else {
156+
nativeTakeOrderedExec
157+
}
147158
},
148-
friendlyName = "NativeRDD.FinalTakeOrdered")
159+
friendlyName = "NativeRDD.FinalTakeOrderedAndProject")
149160
}
150161
}
151162

152-
abstract class NativePartialTakeOrderedBase(
163+
abstract class NativePartialTakeOrderedAndProjectBase(
153164
limit: Long,
154165
sortOrder: Seq[SortOrder],
155166
override val child: SparkPlan,
@@ -197,6 +208,6 @@ abstract class NativePartialTakeOrderedBase(
197208
.build()
198209
PhysicalPlanNode.newBuilder().setSort(nativeTakeOrderedExec).build()
199210
},
200-
friendlyName = "NativeRDD.PartialTakeOrdered")
211+
friendlyName = "NativeRDD.PartialTakeOrderedAndProject")
201212
}
202213
}

0 commit comments

Comments
 (0)