Skip to content

Commit 1ae5c71

Browse files
author
Bhargava Vadlamani
committed
init
1 parent 945f85c commit 1ae5c71

11 files changed

Lines changed: 98 additions & 38 deletions

File tree

‎native/spark-expr/src/comet_scalar_funcs.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,8 @@ pub fn create_comet_physical_fun_with_eval_mode(
115115
make_comet_scalar_udf!("rpad", func, without data_type)
116116
}
117117
"round" => {
118-
make_comet_scalar_udf!("round", spark_round, data_type)
118+
let fail_on_error_s = fail_on_error.unwrap_or(false);
119+
make_comet_scalar_udf!("round", spark_round, data_type, fail_on_error_s)
119120
}
120121
"unscaled_value" => {
121122
let func = Arc::new(spark_unscaled_value);

‎native/spark-expr/src/math_funcs/round.rs‎

Lines changed: 55 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,41 +22,72 @@ use arrow::datatypes::DataType;
2222
use datafusion::common::{exec_err, internal_err, DataFusionError, ScalarValue};
2323
use datafusion::{functions::math::round::round, physical_plan::ColumnarValue};
2424
use std::{cmp::min, sync::Arc};
25+
use crate::SparkError;
2526

2627
macro_rules! integer_round {
27-
($X:expr, $DIV:expr, $HALF:expr) => {{
28+
($X:expr, $DIV:expr, $HALF:expr, $FAIL_ON_ERROR:expr) => {{
2829
let rem = $X % $DIV;
2930
if rem <= -$HALF {
30-
($X - rem).sub_wrapping($DIV)
31+
if($FAIL_ON_ERROR){
32+
match ($X - rem).sub_checked($DIV){
33+
Ok(v) => v,
34+
Err(_) => {
35+
return Err(SparkError::ArithmeticOverflow {from_type: String::from("integer"), }.into());
36+
}
37+
}
38+
}
39+
else{
40+
($X - rem).sub_wrapping($DIV)
41+
}
3142
} else if rem >= $HALF {
32-
($X - rem).add_wrapping($DIV)
43+
if($FAIL_ON_ERROR){
44+
match ($X - rem).add_checked($DIV){
45+
Ok(v) => v,
46+
Err(_) => {
47+
return Err(SparkError::ArithmeticOverflow {from_type: String::from("integer")}.into());
48+
}
49+
}
50+
}
51+
else{
52+
($X - rem).add_wrapping($DIV)
53+
}
3354
} else {
34-
$X - rem
55+
if($FAIL_ON_ERROR){
56+
match ($X - rem).sub_checked(/* i64 */){
57+
Ok(v) => v,
58+
Err(_) => {
59+
return Err(SparkError::ArithmeticOverflow {from_type: String::from("integer")}.into());
60+
}
61+
}
62+
}
63+
else{
64+
$X - rem
65+
}
3566
}
3667
}};
3768
}
3869

3970
macro_rules! round_integer_array {
40-
($ARRAY:expr, $POINT:expr, $TYPE:ty, $NATIVE:ty) => {{
71+
($ARRAY:expr, $POINT:expr, $TYPE:ty, $NATIVE:ty, $FAIL_ON_ERROR:expr) => {{
4172
let array = $ARRAY.as_any().downcast_ref::<$TYPE>().unwrap();
4273
let ten: $NATIVE = 10;
43-
let result: $TYPE = if let Some(div) = ten.checked_pow((-(*$POINT)) as u32) {
74+
let result = if let Some(div) = ten.checked_pow((-(*$POINT)) as u32) {
4475
let half = div / 2;
45-
arrow::compute::kernels::arity::unary(array, |x| integer_round!(x, div, half))
76+
arrow::compute::kernels::arity::try_unary(array, |x| integer_round!(x, div, half, $FAIL_ON_ERROR))
4677
} else {
47-
arrow::compute::kernels::arity::unary(array, |_| 0)
48-
};
78+
arrow::compute::kernels::arity::try_unary(array, |_| 0)
79+
}.map_err(|e| DataFusionError::ArrowError(e))?;
4980
Ok(ColumnarValue::Array(Arc::new(result)))
5081
}};
5182
}
5283

5384
macro_rules! round_integer_scalar {
54-
($SCALAR:expr, $POINT:expr, $TYPE:expr, $NATIVE:ty) => {{
85+
($SCALAR:expr, $POINT:expr, $TYPE:expr, $NATIVE:ty, $FAIL_ON_ERROR:expr) => {{
5586
let ten: $NATIVE = 10;
5687
if let Some(div) = ten.checked_pow((-(*$POINT)) as u32) {
5788
let half = div / 2;
5889
Ok(ColumnarValue::Scalar($TYPE(
59-
$SCALAR.map(|x| integer_round!(x, div, half)),
90+
$SCALAR.map(|x| integer_round!(x, div, half, $FAIL_ON_ERROR)),
6091
)))
6192
} else {
6293
Ok(ColumnarValue::Scalar($TYPE(Some(0))))
@@ -68,6 +99,7 @@ macro_rules! round_integer_scalar {
6899
pub fn spark_round(
69100
args: &[ColumnarValue],
70101
data_type: &DataType,
102+
fail_on_error: bool
71103
) -> Result<ColumnarValue, DataFusionError> {
72104
let value = &args[0];
73105
let point = &args[1];
@@ -76,10 +108,10 @@ pub fn spark_round(
76108
};
77109
match value {
78110
ColumnarValue::Array(array) => match array.data_type() {
79-
DataType::Int64 if *point < 0 => round_integer_array!(array, point, Int64Array, i64),
80-
DataType::Int32 if *point < 0 => round_integer_array!(array, point, Int32Array, i32),
81-
DataType::Int16 if *point < 0 => round_integer_array!(array, point, Int16Array, i16),
82-
DataType::Int8 if *point < 0 => round_integer_array!(array, point, Int8Array, i8),
111+
DataType::Int64 if *point < 0 => round_integer_array!(array, point, Int64Array, i64, fail_on_error),
112+
DataType::Int32 if *point < 0 => round_integer_array!(array, point, Int32Array, i32, fail_on_error),
113+
DataType::Int16 if *point < 0 => round_integer_array!(array, point, Int16Array, i16, fail_on_error),
114+
DataType::Int8 if *point < 0 => round_integer_array!(array, point, Int8Array, i8, fail_on_error),
83115
DataType::Decimal128(_, scale) if *scale >= 0 => {
84116
let f = decimal_round_f(scale, point);
85117
let (precision, scale) = get_precision_scale(data_type);
@@ -93,16 +125,16 @@ pub fn spark_round(
93125
},
94126
ColumnarValue::Scalar(a) => match a {
95127
ScalarValue::Int64(a) if *point < 0 => {
96-
round_integer_scalar!(a, point, ScalarValue::Int64, i64)
128+
round_integer_scalar!(a, point, ScalarValue::Int64, i64, fail_on_error)
97129
}
98130
ScalarValue::Int32(a) if *point < 0 => {
99-
round_integer_scalar!(a, point, ScalarValue::Int32, i32)
131+
round_integer_scalar!(a, point, ScalarValue::Int32, i32, fail_on_error)
100132
}
101133
ScalarValue::Int16(a) if *point < 0 => {
102-
round_integer_scalar!(a, point, ScalarValue::Int16, i16)
134+
round_integer_scalar!(a, point, ScalarValue::Int16, i16, fail_on_error)
103135
}
104136
ScalarValue::Int8(a) if *point < 0 => {
105-
round_integer_scalar!(a, point, ScalarValue::Int8, i8)
137+
round_integer_scalar!(a, point, ScalarValue::Int8, i8, fail_on_error)
106138
}
107139
ScalarValue::Decimal128(a, _, scale) if *scale >= 0 => {
108140
let f = decimal_round_f(scale, point);
@@ -158,7 +190,7 @@ mod test {
158190
]))),
159191
ColumnarValue::Scalar(ScalarValue::Int64(Some(2))),
160192
];
161-
let ColumnarValue::Array(result) = spark_round(&args, &DataType::Float32)? else {
193+
let ColumnarValue::Array(result) = spark_round(&args, &DataType::Float32,false)? else {
162194
unreachable!()
163195
};
164196
let floats = as_float32_array(&result)?;
@@ -176,7 +208,7 @@ mod test {
176208
]))),
177209
ColumnarValue::Scalar(ScalarValue::Int64(Some(2))),
178210
];
179-
let ColumnarValue::Array(result) = spark_round(&args, &DataType::Float64)? else {
211+
let ColumnarValue::Array(result) = spark_round(&args, &DataType::Float64, false)? else {
180212
unreachable!()
181213
};
182214
let floats = as_float64_array(&result)?;
@@ -193,7 +225,7 @@ mod test {
193225
ColumnarValue::Scalar(ScalarValue::Int64(Some(2))),
194226
];
195227
let ColumnarValue::Scalar(ScalarValue::Float32(Some(result))) =
196-
spark_round(&args, &DataType::Float32)?
228+
spark_round(&args, &DataType::Float32, false)?
197229
else {
198230
unreachable!()
199231
};
@@ -209,7 +241,7 @@ mod test {
209241
ColumnarValue::Scalar(ScalarValue::Int64(Some(2))),
210242
];
211243
let ColumnarValue::Scalar(ScalarValue::Float64(Some(result))) =
212-
spark_round(&args, &DataType::Float64)?
244+
spark_round(&args, &DataType::Float64, false)?
213245
else {
214246
unreachable!()
215247
};

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -881,14 +881,15 @@ object QueryPlanSerde extends Logging with CometExprShim {
881881
case UnscaledValue(child) =>
882882
val childExpr = exprToProtoInternal(child, inputs, binding)
883883
val optExpr =
884-
scalarFunctionExprToProtoWithReturnType("unscaled_value", LongType, childExpr)
884+
scalarFunctionExprToProtoWithReturnType("unscaled_value", LongType, false, childExpr)
885885
optExprWithInfo(optExpr, expr, child)
886886

887887
case MakeDecimal(child, precision, scale, true) =>
888888
val childExpr = exprToProtoInternal(child, inputs, binding)
889889
val optExpr = scalarFunctionExprToProtoWithReturnType(
890890
"make_decimal",
891891
DecimalType(precision, scale),
892+
false,
892893
childExpr)
893894
optExprWithInfo(optExpr, expr, child)
894895

@@ -1000,9 +1001,11 @@ object QueryPlanSerde extends Logging with CometExprShim {
10001001
def scalarFunctionExprToProtoWithReturnType(
10011002
funcName: String,
10021003
returnType: DataType,
1004+
failOnError: Boolean,
10031005
args: Option[Expr]*): Option[Expr] = {
10041006
val builder = ExprOuterClass.ScalarFunc.newBuilder()
10051007
builder.setFunc(funcName)
1008+
builder.setFailOnError(failOnError)
10061009
serializeDataType(returnType).flatMap { t =>
10071010
builder.setReturnType(t)
10081011
scalarFunctionExprToProto0(builder, args: _*)
@@ -1012,6 +1015,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
10121015
def scalarFunctionExprToProto(funcName: String, args: Option[Expr]*): Option[Expr] = {
10131016
val builder = ExprOuterClass.ScalarFunc.newBuilder()
10141017
builder.setFunc(funcName)
1018+
builder.setFailOnError(false)
10151019
scalarFunctionExprToProto0(builder, args: _*)
10161020
}
10171021

‎spark/src/main/scala/org/apache/comet/serde/arithmetic.scala‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -325,7 +325,12 @@ object CometRound extends CometExpressionSerde[Round] {
325325
// `scale` must be Int64 type in DataFusion
326326
val scaleExpr = exprToProtoInternal(Literal(_scale.toLong, LongType), inputs, binding)
327327
val optExpr =
328-
scalarFunctionExprToProtoWithReturnType("round", r.dataType, childExpr, scaleExpr)
328+
scalarFunctionExprToProtoWithReturnType(
329+
"round",
330+
r.dataType,
331+
r.ansiEnabled,
332+
childExpr,
333+
scaleExpr)
329334
optExprWithInfo(optExpr, r, r.child)
330335
}
331336

‎spark/src/main/scala/org/apache/comet/serde/arrays.scala‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@ object CometArraysOverlap extends CometExpressionSerde[ArraysOverlap] {
210210
val arraysOverlapScalarExpr = scalarFunctionExprToProtoWithReturnType(
211211
"array_has_any",
212212
BooleanType,
213+
false,
213214
leftArrayExprProto,
214215
rightArrayExprProto)
215216
optExprWithInfo(arraysOverlapScalarExpr, expr, expr.children: _*)
@@ -250,6 +251,7 @@ object CometArrayCompact extends CometExpressionSerde[Expression] {
250251
val arrayCompactScalarExpr = scalarFunctionExprToProtoWithReturnType(
251252
"array_remove_all",
252253
ArrayType(elementType = elementType),
254+
false,
253255
arrayExprProto,
254256
nullLiteralProto)
255257
optExprWithInfo(arrayCompactScalarExpr, expr, expr.children: _*)

‎spark/src/main/scala/org/apache/comet/serde/bitwise.scala‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ object CometBitwiseGet extends CometExpressionSerde[BitwiseGet] {
135135
val argProto = exprToProto(expr.left, inputs, binding)
136136
val posProto = exprToProto(expr.right, inputs, binding)
137137
val bitGetScalarExpr =
138-
scalarFunctionExprToProtoWithReturnType("bit_get", ByteType, argProto, posProto)
138+
scalarFunctionExprToProtoWithReturnType("bit_get", ByteType, false, argProto, posProto)
139139
optExprWithInfo(bitGetScalarExpr, expr, expr.children: _*)
140140
}
141141
}
@@ -147,7 +147,7 @@ object CometBitwiseCount extends CometExpressionSerde[BitwiseCount] {
147147
binding: Boolean): Option[ExprOuterClass.Expr] = {
148148
val childProto = exprToProto(expr.child, inputs, binding)
149149
val bitCountScalarExpr =
150-
scalarFunctionExprToProtoWithReturnType("bit_count", IntegerType, childProto)
150+
scalarFunctionExprToProtoWithReturnType("bit_count", IntegerType, false, childProto)
151151
optExprWithInfo(bitCountScalarExpr, expr, expr.children: _*)
152152
}
153153
}

‎spark/src/main/scala/org/apache/comet/serde/datetime.scala‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -263,7 +263,12 @@ object CometTruncDate extends CometExpressionSerde[TruncDate] {
263263
val childExpr = exprToProtoInternal(expr.date, inputs, binding)
264264
val formatExpr = exprToProtoInternal(expr.format, inputs, binding)
265265
val optExpr =
266-
scalarFunctionExprToProtoWithReturnType("date_trunc", DateType, childExpr, formatExpr)
266+
scalarFunctionExprToProtoWithReturnType(
267+
"date_trunc",
268+
DateType,
269+
false,
270+
childExpr,
271+
formatExpr)
267272
optExprWithInfo(optExpr, expr, expr.date, expr.format)
268273
}
269274
}

‎spark/src/main/scala/org/apache/comet/serde/hash.scala‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ object CometXxHash64 extends CometExpressionSerde[XxHash64] {
4040
.setLongVal(expr.seed)
4141
val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build())
4242
// the seed is put at the end of the arguments
43-
scalarFunctionExprToProtoWithReturnType("xxhash64", LongType, exprs :+ seedExpr: _*)
43+
scalarFunctionExprToProtoWithReturnType("xxhash64", LongType, false, exprs :+ seedExpr: _*)
4444
}
4545
}
4646

@@ -59,7 +59,11 @@ object CometMurmur3Hash extends CometExpressionSerde[Murmur3Hash] {
5959
.setIntVal(expr.seed)
6060
val seedExpr = Some(ExprOuterClass.Expr.newBuilder().setLiteral(seedBuilder).build())
6161
// the seed is put at the end of the arguments
62-
scalarFunctionExprToProtoWithReturnType("murmur3_hash", IntegerType, exprs :+ seedExpr: _*)
62+
scalarFunctionExprToProtoWithReturnType(
63+
"murmur3_hash",
64+
IntegerType,
65+
false,
66+
exprs :+ seedExpr: _*)
6367
}
6468
}
6569

@@ -81,7 +85,7 @@ object CometSha2 extends CometExpressionSerde[Sha2] {
8185

8286
val leftExpr = exprToProtoInternal(expr.left, inputs, binding)
8387
val numBitsExpr = exprToProtoInternal(expr.right, inputs, binding)
84-
scalarFunctionExprToProtoWithReturnType("sha2", StringType, leftExpr, numBitsExpr)
88+
scalarFunctionExprToProtoWithReturnType("sha2", StringType, false, leftExpr, numBitsExpr)
8589
}
8690
}
8791

‎spark/src/main/scala/org/apache/comet/serde/maps.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ object CometMapFromArrays extends CometExpressionSerde[MapFromArrays] {
8585
val valueType = expr.right.dataType.asInstanceOf[ArrayType].elementType
8686
val returnType = MapType(keyType = keyType, valueType = valueType)
8787
val mapFromArraysExpr =
88-
scalarFunctionExprToProtoWithReturnType("map", returnType, keysExpr, valuesExpr)
88+
scalarFunctionExprToProtoWithReturnType("map", returnType, false, keysExpr, valuesExpr)
8989
optExprWithInfo(mapFromArraysExpr, expr, expr.children: _*)
9090
}
9191
}

‎spark/src/main/scala/org/apache/comet/serde/math.scala‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,8 @@ object CometCeil extends CometExpressionSerde[Ceil] {
5050
withInfo(expr, s"Decimal type $t has negative scale")
5151
None
5252
case _ =>
53-
val optExpr = scalarFunctionExprToProtoWithReturnType("ceil", expr.dataType, childExpr)
53+
val optExpr =
54+
scalarFunctionExprToProtoWithReturnType("ceil", expr.dataType, false, childExpr)
5455
optExprWithInfo(optExpr, expr, expr.child)
5556
}
5657
}
@@ -69,7 +70,8 @@ object CometFloor extends CometExpressionSerde[Floor] {
6970
withInfo(expr, s"Decimal type $t has negative scale")
7071
None
7172
case _ =>
72-
val optExpr = scalarFunctionExprToProtoWithReturnType("floor", expr.dataType, childExpr)
73+
val optExpr =
74+
scalarFunctionExprToProtoWithReturnType("floor", expr.dataType, false, childExpr)
7375
optExprWithInfo(optExpr, expr, expr.child)
7476
}
7577
}
@@ -118,7 +120,7 @@ object CometHex extends CometExpressionSerde[Hex] with MathExprBase {
118120
inputs: Seq[Attribute],
119121
binding: Boolean): Option[ExprOuterClass.Expr] = {
120122
val childExpr = exprToProtoInternal(expr.child, inputs, binding)
121-
val optExpr = scalarFunctionExprToProtoWithReturnType("hex", expr.dataType, childExpr)
123+
val optExpr = scalarFunctionExprToProtoWithReturnType("hex", expr.dataType, false, childExpr)
122124
optExprWithInfo(optExpr, expr, expr.child)
123125
}
124126
}
@@ -132,7 +134,12 @@ object CometUnhex extends CometExpressionSerde[Unhex] with MathExprBase {
132134
val failOnErrorExpr = exprToProtoInternal(Literal(expr.failOnError), inputs, binding)
133135

134136
val optExpr =
135-
scalarFunctionExprToProtoWithReturnType("unhex", expr.dataType, childExpr, failOnErrorExpr)
137+
scalarFunctionExprToProtoWithReturnType(
138+
"unhex",
139+
expr.dataType,
140+
false,
141+
childExpr,
142+
failOnErrorExpr)
136143
optExprWithInfo(optExpr, expr, expr.child)
137144
}
138145
}

0 commit comments

Comments
 (0)