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/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new d1d57c68 minor: refactor UnaryExpr handling to make code more concise
(#1065)
d1d57c68 is described below
commit d1d57c6826b37d0c60e312204ea71ae17d44e20a
Author: Andy Grove <[email protected]>
AuthorDate: Fri Nov 8 07:22:29 2024 -0700
minor: refactor UnaryExpr handling to make code more concise (#1065)
* minor: refactor UnaryExpr handling to make more concise
* add documentation
---
.../org/apache/comet/serde/QueryPlanSerde.scala | 71 +++++++++++-----------
1 file changed, 36 insertions(+), 35 deletions(-)
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
index 8bdc886d..c47c2b6a 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -1282,12 +1282,10 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
}
case StringSpace(child) =>
- createUnaryExpr(child, inputs).map { expr =>
- ExprOuterClass.Expr
- .newBuilder()
- .setStringSpace(expr)
- .build()
- }
+ createUnaryExpr(
+ child,
+ inputs,
+ (builder, unaryExpr) => builder.setStringSpace(unaryExpr))
case Hour(child, timeZoneId) =>
val childExpr = exprToProtoInternal(child, inputs)
@@ -1422,20 +1420,10 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
optExprWithInfo(optExpr, expr, child)
case IsNull(child) =>
- createUnaryExpr(child, inputs).map { expr =>
- ExprOuterClass.Expr
- .newBuilder()
- .setIsNull(expr)
- .build()
- }
+ createUnaryExpr(child, inputs, (builder, unaryExpr) =>
builder.setIsNull(unaryExpr))
case IsNotNull(child) =>
- createUnaryExpr(child, inputs).map { expr =>
- ExprOuterClass.Expr
- .newBuilder()
- .setIsNotNull(expr)
- .build()
- }
+ createUnaryExpr(child, inputs, (builder, unaryExpr) =>
builder.setIsNotNull(unaryExpr))
case IsNaN(child) =>
val childExpr = exprToProtoInternal(child, inputs)
@@ -1930,12 +1918,7 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
}
case BitwiseNot(child) =>
- createUnaryExpr(child, inputs).map { expr =>
- ExprOuterClass.Expr
- .newBuilder()
- .setBitwiseNot(expr)
- .build()
- }
+ createUnaryExpr(child, inputs, (builder, unaryExpr) =>
builder.setBitwiseNot(unaryExpr))
case BitwiseOr(left, right) =>
createBinaryExpr(left, right, inputs).map { builder =>
@@ -2001,12 +1984,7 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
in(expr, value, list, inputs, true)
case Not(child) =>
- createUnaryExpr(child, inputs).map { expr =>
- ExprOuterClass.Expr
- .newBuilder()
- .setNot(expr)
- .build()
- }
+ createUnaryExpr(child, inputs, (builder, unaryExpr) =>
builder.setNot(unaryExpr))
case UnaryMinus(child, failOnError) =>
val childExpr = exprToProtoInternal(child, inputs)
@@ -2287,16 +2265,39 @@ object QueryPlanSerde extends Logging with
ShimQueryPlanSerde with CometExprShim
}
}
+ /**
+ * Creates a UnaryExpr by calling exprToProtoInternal for the provided
child expression and
+ * then invokes the supplied function to wrap this UnaryExpr in a
top-level Expr.
+ *
+ * @param child
+ * Spark expression
+ * @param inputs
+ * Inputs to the expression
+ * @param f
+ * Function that accepts an Expr.Builder and a UnaryExpr and builds the
specific top-level
+ * Expr
+ * @return
+ * Some(Expr) or None if not supported
+ */
def createUnaryExpr(
child: Expression,
- inputs: Seq[Attribute]): Option[ExprOuterClass.UnaryExpr] = {
+ inputs: Seq[Attribute],
+ f: (ExprOuterClass.Expr.Builder, ExprOuterClass.UnaryExpr) =>
ExprOuterClass.Expr.Builder)
+ : Option[ExprOuterClass.Expr] = {
val childExpr = exprToProtoInternal(child, inputs)
if (childExpr.isDefined) {
+ // create the generic UnaryExpr message
+ val inner = ExprOuterClass.UnaryExpr
+ .newBuilder()
+ .setChild(childExpr.get)
+ .build()
+ // call the user-supplied function to wrap UnaryExpr in a top-level
Expr
+ // such as Expr.IsNull or Expr.IsNotNull
Some(
- ExprOuterClass.UnaryExpr
- .newBuilder()
- .setChild(childExpr.get)
- .build())
+ f(
+ ExprOuterClass.Expr
+ .newBuilder(),
+ inner).build())
} else {
withInfo(expr, child)
None
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
