Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Use enum to represent CAST eval_mode in expr.proto #415

Merged
merged 6 commits into from
Jun 3, 2024
Merged
8 changes: 4 additions & 4 deletions core/src/execution/datafusion/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -346,10 +346,10 @@ impl PhysicalPlanner {
let child = self.create_expr(expr.child.as_ref().unwrap(), input_schema)?;
let datatype = to_arrow_datatype(expr.datatype.as_ref().unwrap());
let timezone = expr.timezone.clone();
let eval_mode = match expr.eval_mode.as_str() {
"ANSI" => EvalMode::Ansi,
"TRY" => EvalMode::Try,
"LEGACY" => EvalMode::Legacy,
let eval_mode = match expr.eval_mode {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can make use of the generated protobuf enum here rather than hard-code numberic values. This would also allow us to remove the "other" error handling case.

                let eval_mode = match spark_expression::EvalMode::try_from(expr.eval_mode)? {
                    spark_expression::EvalMode::Legacy => EvalMode::Legacy,
                    spark_expression::EvalMode::Try => EvalMode::Try,
                    spark_expression::EvalMode::Ansi => EvalMode::Ansi,
                };

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the suggestion, Andy.
I have made changes to planner.rs as per suggestion.

0 => EvalMode::Legacy,
1 => EvalMode::Try,
2 => EvalMode::Ansi,
other => {
return Err(ExecutionError::GeneralError(format!(
"Invalid Cast EvalMode: \"{other}\""
Expand Down
12 changes: 10 additions & 2 deletions core/src/execution/proto/expr.proto
Original file line number Diff line number Diff line change
Expand Up @@ -233,12 +233,20 @@ message Remainder {
DataType return_type = 4;
}

enum EvalMode {
LEGACY = 0;
TRY = 1;
ANSI = 2;
}

message Cast {
Expr child = 1;
DataType datatype = 2;
string timezone = 3;
// LEGACY, ANSI, or TRY
string eval_mode = 4;
// Depreciateid: LEGACY, ANSI, or TRY - preserved for backward compatibity
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No need for backwards compatibility here since we have not published a release yet

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed string eval_mode.

string eval_mode_string= 4; // for backward compatibility
EvalMode eval_mode = 5; // New enum field

}

message Equal {
Expand Down
20 changes: 18 additions & 2 deletions spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@

package org.apache.comet.serde

import java.util.Locale

import scala.collection.JavaConverters._

import org.apache.spark.internal.Logging
Expand Down Expand Up @@ -525,6 +527,18 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
* @return
* The protobuf representation of the expression, or None if the expression is not supported
*/

def stringToEvalMode(evalModeStr: String): ExprOuterClass.EvalMode =
evalModeStr.toUpperCase(Locale.ROOT) match {
case "LEGACY" => ExprOuterClass.EvalMode.LEGACY
case "TRY" => ExprOuterClass.EvalMode.TRY
case "ANSI" => ExprOuterClass.EvalMode.ANSI
case _ =>
throw new IllegalArgumentException(
"Invalid eval mode"
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would be good to include the invalid value in the error message

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have updated the code as follows.

   5   def stringToEvalMode(evalModeStr: String): ExprOuterClass.EvalMode =
   4     evalModeStr.toUpperCase(Locale.ROOT) match {
   3       case "LEGACY" => ExprOuterClass.EvalMode.LEGACY
   2       case "TRY" => ExprOuterClass.EvalMode.TRY
   1       case "ANSI" => ExprOuterClass.EvalMode.ANSI
 536       case invalid =>
   1         throw new IllegalArgumentException(
   2           s"Invalid eval mode '$invalid' "
   3         ) // Assuming we want to catch errors strictly
   4     }

) // Assuming we want to catch errors strictly
}

def exprToProto(
expr: Expression,
input: Seq[Attribute],
Expand All @@ -535,12 +549,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
childExpr: Option[Expr],
evalMode: String): Option[Expr] = {
val dataType = serializeDataType(dt)
val evalModeEnum = stringToEvalMode(evalMode) // Convert string to enum

if (childExpr.isDefined && dataType.isDefined) {
val castBuilder = ExprOuterClass.Cast.newBuilder()
castBuilder.setChild(childExpr.get)
castBuilder.setDatatype(dataType.get)
castBuilder.setEvalMode(evalMode)
// castBuilder.setEvalMode(evalMode)
andygrove marked this conversation as resolved.
Show resolved Hide resolved
castBuilder.setEvalMode(evalModeEnum) // Set the enum in protobuf

val timeZone = timeZoneId.getOrElse("UTC")
castBuilder.setTimezone(timeZone)
Expand Down Expand Up @@ -1207,7 +1223,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde {
.newBuilder()
.setChild(e)
.setDatatype(serializeDataType(IntegerType).get)
.setEvalMode("LEGACY") // year is not affected by ANSI mode
.setEvalMode(stringToEvalMode("LEGACY")) // year is not affected by ANSI mode
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can just use the enum directly here rather than convert from string.

Suggested change
.setEvalMode(stringToEvalMode("LEGACY")) // year is not affected by ANSI mode
// year is not affected by ANSI mode so we use LEGACY mode here
.setEvalMode(ExprOuterClass.EvalMode.LEGACY)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for another great suggestion.
Updated accordingly to QueryPlanSerde.scala L1226.

.build())
.build()
})
Expand Down
Loading