Skip to content

Commit 56ddd51

Browse files
chore: refactor GreaterThan out of QueryPlanSerde
- Create comparisons.scala following the pattern from math/array expressions. - Implements CometGreaterThan as proof of concept for issue #2019.
1 parent 76cf5af commit 56ddd51

File tree

2 files changed

+118
-10
lines changed

2 files changed

+118
-10
lines changed

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,8 @@ object QueryPlanSerde extends Logging with CometExprShim {
120120
classOf[MapKeys] -> CometMapKeys,
121121
classOf[MapValues] -> CometMapValues,
122122
classOf[MapFromArrays] -> CometMapFromArrays,
123-
classOf[GetMapValue] -> CometMapExtract)
123+
classOf[GetMapValue] -> CometMapExtract,
124+
classOf[GreaterThan] -> CometGreaterThan)
124125

125126
def emitWarning(reason: String): Unit = {
126127
logWarning(s"Comet native execution is disabled due to: $reason")
@@ -801,15 +802,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
801802
binding,
802803
(builder, binaryExpr) => builder.setNeqNullSafe(binaryExpr))
803804

804-
case GreaterThan(left, right) =>
805-
createBinaryExpr(
806-
expr,
807-
left,
808-
right,
809-
inputs,
810-
binding,
811-
(builder, binaryExpr) => builder.setGt(binaryExpr))
812-
813805
case GreaterThanOrEqual(left, right) =>
814806
createBinaryExpr(
815807
expr,
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.comet.serde
21+
import org.apache.comet.CometSparkSessionExtensions.withInfo
22+
import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
23+
import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GreaterThan}
24+
25+
import scala.collection.JavaConverters._
26+
27+
object CometGreaterThan extends CometExpressionSerde with ComparisonBase {
28+
override def convert(
29+
expr: Expression,
30+
inputs: Seq[Attribute],
31+
binding: Boolean): Option[ExprOuterClass.Expr] = {
32+
val greaterThan = expr.asInstanceOf[GreaterThan]
33+
34+
createBinaryExpr(
35+
expr,
36+
greaterThan.left,
37+
greaterThan.right,
38+
inputs,
39+
binding,
40+
(builder, binaryExpr) => builder.setGt(binaryExpr))
41+
}
42+
}
43+
44+
sealed trait ComparisonBase {
45+
46+
/**
47+
* Creates a UnaryExpr by calling exprToProtoInternal for the provided child expression and then
48+
* invokes the supplied function to wrap this UnaryExpr in a top-level Expr.
49+
*
50+
* @param child
51+
* Spark expression
52+
* @param inputs
53+
* Inputs to the expression
54+
* @param f
55+
* Function that accepts an Expr.Builder and a UnaryExpr and builds the specific top-level
56+
* Expr
57+
* @return
58+
* Some(Expr) or None if not supported
59+
*/
60+
def createUnaryExpr(
61+
expr: Expression,
62+
child: Expression,
63+
inputs: Seq[Attribute],
64+
binding: Boolean,
65+
f: (ExprOuterClass.Expr.Builder, ExprOuterClass.UnaryExpr) => ExprOuterClass.Expr.Builder)
66+
: Option[ExprOuterClass.Expr] = {
67+
val childExpr = exprToProtoInternal(child, inputs, binding) // TODO review
68+
if (childExpr.isDefined) {
69+
// create the generic UnaryExpr message
70+
val inner = ExprOuterClass.UnaryExpr
71+
.newBuilder()
72+
.setChild(childExpr.get)
73+
.build()
74+
// call the user-supplied function to wrap UnaryExpr in a top-level Expr
75+
// such as Expr.IsNull or Expr.IsNotNull
76+
Some(
77+
f(
78+
ExprOuterClass.Expr
79+
.newBuilder(),
80+
inner).build())
81+
} else {
82+
withInfo(expr, child)
83+
None
84+
}
85+
}
86+
87+
def createBinaryExpr(
88+
expr: Expression,
89+
left: Expression,
90+
right: Expression,
91+
inputs: Seq[Attribute],
92+
binding: Boolean,
93+
f: (ExprOuterClass.Expr.Builder, ExprOuterClass.BinaryExpr) => ExprOuterClass.Expr.Builder)
94+
: Option[ExprOuterClass.Expr] = {
95+
val leftExpr = exprToProtoInternal(left, inputs, binding)
96+
val rightExpr = exprToProtoInternal(right, inputs, binding)
97+
if (leftExpr.isDefined && rightExpr.isDefined) {
98+
// create the generic BinaryExpr message
99+
val inner = ExprOuterClass.BinaryExpr
100+
.newBuilder()
101+
.setLeft(leftExpr.get)
102+
.setRight(rightExpr.get)
103+
.build()
104+
// call the user-supplied function to wrap BinaryExpr in a top-level Expr
105+
// such as Expr.And or Expr.Or
106+
Some(
107+
f(
108+
ExprOuterClass.Expr
109+
.newBuilder(),
110+
inner).build())
111+
} else {
112+
withInfo(expr, left, right)
113+
None
114+
}
115+
}
116+
}

0 commit comments

Comments
 (0)