Skip to content

Commit f572ec1

Browse files
HaoYang670HaoYang670
authored
Update try_binary and checked_ops, and remove math_checked_op (#2717)
* update try_binary delete math_checked_op update the return type of checked ops Signed-off-by: remzi <13716567376yh@gmail.com> * float div not panic on zero Signed-off-by: remzi <13716567376yh@gmail.com> * fix nan test Signed-off-by: remzi <13716567376yh@gmail.com> * add float divide by zero Signed-off-by: remzi <13716567376yh@gmail.com> * add float tests Signed-off-by: remzi <13716567376yh@gmail.com> * fix compile error Signed-off-by: remzi <13716567376yh@gmail.com> Signed-off-by: remzi <13716567376yh@gmail.com>
1 parent 43d912c commit f572ec1

4 files changed

Lines changed: 153 additions & 149 deletions

File tree

‎arrow/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ serde_json = { version = "1.0", default-features = false, features = ["std"], op
5151
indexmap = { version = "1.9", default-features = false, features = ["std"] }
5252
rand = { version = "0.8", default-features = false, features = ["std", "std_rng"], optional = true }
5353
num = { version = "0.4", default-features = false, features = ["std"] }
54-
half = { version = "2.0", default-features = false }
54+
half = { version = "2.0", default-features = false, features = ["num-traits"]}
5555
hashbrown = { version = "0.12", default-features = false }
5656
csv_crate = { version = "1.1", default-features = false, optional = true, package = "csv" }
5757
regex = { version = "1.5.6", default-features = false, features = ["std", "unicode"] }

‎arrow/src/compute/kernels/arithmetic.rs‎

Lines changed: 93 additions & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -78,32 +78,6 @@ where
7878
Ok(binary(left, right, op))
7979
}
8080

81-
/// This is similar to `math_op` as it performs given operation between two input primitive arrays.
82-
/// But the given operation can return `None` if overflow is detected. For the case, this function
83-
/// returns an `Err`.
84-
fn math_checked_op<LT, RT, F>(
85-
left: &PrimitiveArray<LT>,
86-
right: &PrimitiveArray<RT>,
87-
op: F,
88-
) -> Result<PrimitiveArray<LT>>
89-
where
90-
LT: ArrowNumericType,
91-
RT: ArrowNumericType,
92-
F: Fn(LT::Native, RT::Native) -> Option<LT::Native>,
93-
{
94-
if left.len() != right.len() {
95-
return Err(ArrowError::ComputeError(
96-
"Cannot perform math operation on arrays of different length".to_string(),
97-
));
98-
}
99-
100-
try_binary(left, right, |a, b| {
101-
op(a, b).ok_or_else(|| {
102-
ArrowError::ComputeError(format!("Overflow happened on: {:?}, {:?}", a, b))
103-
})
104-
})
105-
}
106-
10781
/// Helper function for operations where a valid `0` on the right array should
10882
/// result in an [ArrowError::DivideByZero], namely the division and modulo operations
10983
///
@@ -121,26 +95,9 @@ where
12195
LT: ArrowNumericType,
12296
RT: ArrowNumericType,
12397
RT::Native: One + Zero,
124-
F: Fn(LT::Native, RT::Native) -> Option<LT::Native>,
98+
F: Fn(LT::Native, RT::Native) -> Result<LT::Native>,
12599
{
126-
if left.len() != right.len() {
127-
return Err(ArrowError::ComputeError(
128-
"Cannot perform math operation on arrays of different length".to_string(),
129-
));
130-
}
131-
132-
try_binary(left, right, |l, r| {
133-
if r.is_zero() {
134-
Err(ArrowError::DivideByZero)
135-
} else {
136-
op(l, r).ok_or_else(|| {
137-
ArrowError::ComputeError(format!(
138-
"Overflow happened on: {:?}, {:?}",
139-
l, r
140-
))
141-
})
142-
}
143-
})
100+
try_binary(left, right, op)
144101
}
145102

146103
/// Helper function for operations where a valid `0` on the right array should
@@ -161,16 +118,12 @@ fn math_checked_divide_op_on_iters<T, F>(
161118
where
162119
T: ArrowNumericType,
163120
T::Native: One + Zero,
164-
F: Fn(T::Native, T::Native) -> T::Native,
121+
F: Fn(T::Native, T::Native) -> Result<T::Native>,
165122
{
166123
let buffer = if null_bit_buffer.is_some() {
167124
let values = left.zip(right).map(|(left, right)| {
168125
if let (Some(l), Some(r)) = (left, right) {
169-
if r.is_zero() {
170-
Err(ArrowError::DivideByZero)
171-
} else {
172-
Ok(op(l, r))
173-
}
126+
op(l, r)
174127
} else {
175128
Ok(T::default_value())
176129
}
@@ -179,15 +132,10 @@ where
179132
unsafe { Buffer::try_from_trusted_len_iter(values) }
180133
} else {
181134
// no value is null
182-
let values = left.map(|l| l.unwrap()).zip(right.map(|r| r.unwrap())).map(
183-
|(left, right)| {
184-
if right.is_zero() {
185-
Err(ArrowError::DivideByZero)
186-
} else {
187-
Ok(op(left, right))
188-
}
189-
},
190-
);
135+
let values = left
136+
.map(|l| l.unwrap())
137+
.zip(right.map(|r| r.unwrap()))
138+
.map(|(left, right)| op(left, right));
191139
// Safety: Iterator comes from a PrimitiveArray which reports its size correctly
192140
unsafe { Buffer::try_from_trusted_len_iter(values) }
193141
}?;
@@ -654,7 +602,7 @@ where
654602
K: ArrowNumericType,
655603
T: ArrowNumericType,
656604
T::Native: One + Zero,
657-
F: Fn(T::Native, T::Native) -> T::Native,
605+
F: Fn(T::Native, T::Native) -> Result<T::Native>,
658606
{
659607
if left.len() != right.len() {
660608
return Err(ArrowError::ComputeError(format!(
@@ -725,7 +673,7 @@ where
725673
T: ArrowNumericType,
726674
T::Native: ArrowNativeTypeOp,
727675
{
728-
math_checked_op(left, right, |a, b| a.add_checked(b))
676+
try_binary(left, right, |a, b| a.add_checked(b))
729677
}
730678

731679
/// Perform `left + right` operation on two arrays. If either left or right value is null
@@ -826,11 +774,7 @@ where
826774
T: ArrowNumericType,
827775
T::Native: ArrowNativeTypeOp,
828776
{
829-
try_unary(array, |value| {
830-
value.add_checked(scalar).ok_or_else(|| {
831-
ArrowError::CastError(format!("Overflow: adding {:?} to {:?}", scalar, value))
832-
})
833-
})
777+
try_unary(array, |value| value.add_checked(scalar))
834778
}
835779

836780
/// Add every value in an array by a scalar. If any value in the array is null then the
@@ -863,12 +807,8 @@ where
863807
T: ArrowNumericType,
864808
T::Native: ArrowNativeTypeOp,
865809
{
866-
try_unary_dyn::<_, T>(array, |value| {
867-
value.add_checked(scalar).ok_or_else(|| {
868-
ArrowError::CastError(format!("Overflow: adding {:?} to {:?}", scalar, value))
869-
})
870-
})
871-
.map(|a| Arc::new(a) as ArrayRef)
810+
try_unary_dyn::<_, T>(array, |value| value.add_checked(scalar))
811+
.map(|a| Arc::new(a) as ArrayRef)
872812
}
873813

874814
/// Perform `left - right` operation on two arrays. If either left or right value is null
@@ -900,7 +840,7 @@ where
900840
T: ArrowNumericType,
901841
T::Native: ArrowNativeTypeOp,
902842
{
903-
math_checked_op(left, right, |a, b| a.sub_checked(b))
843+
try_binary(left, right, |a, b| a.sub_checked(b))
904844
}
905845

906846
/// Perform `left - right` operation on two arrays. If either left or right value is null
@@ -953,14 +893,7 @@ where
953893
T: ArrowNumericType,
954894
T::Native: ArrowNativeTypeOp + Zero,
955895
{
956-
try_unary(array, |value| {
957-
value.sub_checked(scalar).ok_or_else(|| {
958-
ArrowError::CastError(format!(
959-
"Overflow: subtracting {:?} from {:?}",
960-
scalar, value
961-
))
962-
})
963-
})
896+
try_unary(array, |value| value.sub_checked(scalar))
964897
}
965898

966899
/// Subtract every value in an array by a scalar. If any value in the array is null then the
@@ -991,15 +924,8 @@ where
991924
T: ArrowNumericType,
992925
T::Native: ArrowNativeTypeOp,
993926
{
994-
try_unary_dyn::<_, T>(array, |value| {
995-
value.sub_checked(scalar).ok_or_else(|| {
996-
ArrowError::CastError(format!(
997-
"Overflow: subtracting {:?} from {:?}",
998-
scalar, value
999-
))
1000-
})
1001-
})
1002-
.map(|a| Arc::new(a) as ArrayRef)
927+
try_unary_dyn::<_, T>(array, |value| value.sub_checked(scalar))
928+
.map(|a| Arc::new(a) as ArrayRef)
1003929
}
1004930

1005931
/// Perform `-` operation on an array. If value is null then the result is also null.
@@ -1052,7 +978,7 @@ where
1052978
T: ArrowNumericType,
1053979
T::Native: ArrowNativeTypeOp,
1054980
{
1055-
math_checked_op(left, right, |a, b| a.mul_checked(b))
981+
try_binary(left, right, |a, b| a.mul_checked(b))
1056982
}
1057983

1058984
/// Perform `left * right` operation on two arrays. If either left or right value is null
@@ -1105,14 +1031,7 @@ where
11051031
T: ArrowNumericType,
11061032
T::Native: ArrowNativeTypeOp + Zero + One,
11071033
{
1108-
try_unary(array, |value| {
1109-
value.mul_checked(scalar).ok_or_else(|| {
1110-
ArrowError::CastError(format!(
1111-
"Overflow: multiplying {:?} by {:?}",
1112-
value, scalar,
1113-
))
1114-
})
1115-
})
1034+
try_unary(array, |value| value.mul_checked(scalar))
11161035
}
11171036

11181037
/// Multiply every value in an array by a scalar. If any value in the array is null then the
@@ -1143,15 +1062,8 @@ where
11431062
T: ArrowNumericType,
11441063
T::Native: ArrowNativeTypeOp,
11451064
{
1146-
try_unary_dyn::<_, T>(array, |value| {
1147-
value.mul_checked(scalar).ok_or_else(|| {
1148-
ArrowError::CastError(format!(
1149-
"Overflow: multiplying {:?} by {:?}",
1150-
value, scalar
1151-
))
1152-
})
1153-
})
1154-
.map(|a| Arc::new(a) as ArrayRef)
1065+
try_unary_dyn::<_, T>(array, |value| value.mul_checked(scalar))
1066+
.map(|a| Arc::new(a) as ArrayRef)
11551067
}
11561068

11571069
/// Perform `left % right` operation on two arrays. If either left or right value is null
@@ -1170,7 +1082,13 @@ where
11701082
a % b
11711083
});
11721084
#[cfg(not(feature = "simd"))]
1173-
return math_checked_divide_op(left, right, |a, b| Some(a % b));
1085+
return try_binary(left, right, |a, b| {
1086+
if b.is_zero() {
1087+
Err(ArrowError::DivideByZero)
1088+
} else {
1089+
Ok(a % b)
1090+
}
1091+
});
11741092
}
11751093

11761094
/// Perform `left / right` operation on two arrays. If either left or right value is null
@@ -1225,12 +1143,17 @@ where
12251143
pub fn divide_dyn(left: &dyn Array, right: &dyn Array) -> Result<ArrayRef> {
12261144
match left.data_type() {
12271145
DataType::Dictionary(_, _) => {
1228-
typed_dict_math_op!(left, right, |a, b| a / b, math_divide_checked_op_dict)
1146+
typed_dict_math_op!(
1147+
left,
1148+
right,
1149+
|a, b| a.div_checked(b),
1150+
math_divide_checked_op_dict
1151+
)
12291152
}
12301153
_ => {
12311154
downcast_primitive_array!(
12321155
(left, right) => {
1233-
math_checked_divide_op(left, right, |a, b| Some(a / b)).map(|a| Arc::new(a) as ArrayRef)
1156+
math_checked_divide_op(left, right, |a, b| a.div_checked(b)).map(|a| Arc::new(a) as ArrayRef)
12341157
}
12351158
_ => Err(ArrowError::CastError(format!(
12361159
"Unsupported data type {}, {}",
@@ -1331,15 +1254,8 @@ where
13311254
return Err(ArrowError::DivideByZero);
13321255
}
13331256

1334-
try_unary_dyn::<_, T>(array, |value| {
1335-
value.div_checked(divisor).ok_or_else(|| {
1336-
ArrowError::CastError(format!(
1337-
"Overflow: dividing {:?} by {:?}",
1338-
value, divisor
1339-
))
1340-
})
1341-
})
1342-
.map(|a| Arc::new(a) as ArrayRef)
1257+
try_unary_dyn::<_, T>(array, |value| value.div_checked(divisor))
1258+
.map(|a| Arc::new(a) as ArrayRef)
13431259
}
13441260

13451261
#[cfg(test)]
@@ -2134,31 +2050,57 @@ mod tests {
21342050

21352051
#[test]
21362052
#[should_panic(expected = "DivideByZero")]
2137-
fn test_primitive_array_divide_by_zero_with_checked() {
2053+
fn test_int_array_divide_by_zero_with_checked() {
21382054
let a = Int32Array::from(vec![15]);
21392055
let b = Int32Array::from(vec![0]);
21402056
divide_checked(&a, &b).unwrap();
21412057
}
21422058

2059+
#[test]
2060+
#[should_panic(expected = "DivideByZero")]
2061+
fn test_f32_array_divide_by_zero_with_checked() {
2062+
let a = Float32Array::from(vec![15.0]);
2063+
let b = Float32Array::from(vec![0.0]);
2064+
divide_checked(&a, &b).unwrap();
2065+
}
2066+
21432067
#[test]
21442068
#[should_panic(expected = "attempt to divide by zero")]
2145-
fn test_primitive_array_divide_by_zero() {
2069+
fn test_int_array_divide_by_zero() {
21462070
let a = Int32Array::from(vec![15]);
21472071
let b = Int32Array::from(vec![0]);
21482072
divide(&a, &b).unwrap();
21492073
}
21502074

2075+
#[test]
2076+
fn test_f32_array_divide_by_zero() {
2077+
let a = Float32Array::from(vec![1.5, 0.0, -1.5]);
2078+
let b = Float32Array::from(vec![0.0, 0.0, 0.0]);
2079+
let result = divide(&a, &b).unwrap();
2080+
assert_eq!(result.value(0), f32::INFINITY);
2081+
assert!(result.value(1).is_nan());
2082+
assert_eq!(result.value(2), f32::NEG_INFINITY);
2083+
}
2084+
21512085
#[test]
21522086
#[should_panic(expected = "DivideByZero")]
2153-
fn test_primitive_array_divide_dyn_by_zero() {
2087+
fn test_int_array_divide_dyn_by_zero() {
21542088
let a = Int32Array::from(vec![15]);
21552089
let b = Int32Array::from(vec![0]);
21562090
divide_dyn(&a, &b).unwrap();
21572091
}
21582092

21592093
#[test]
21602094
#[should_panic(expected = "DivideByZero")]
2161-
fn test_primitive_array_divide_dyn_by_zero_dict() {
2095+
fn test_f32_array_divide_dyn_by_zero() {
2096+
let a = Float32Array::from(vec![1.5]);
2097+
let b = Float32Array::from(vec![0.0]);
2098+
divide_dyn(&a, &b).unwrap();
2099+
}
2100+
2101+
#[test]
2102+
#[should_panic(expected = "DivideByZero")]
2103+
fn test_int_array_divide_dyn_by_zero_dict() {
21622104
let mut builder =
21632105
PrimitiveDictionaryBuilder::<Int8Type, Int32Type>::with_capacity(1, 1);
21642106
builder.append(15).unwrap();
@@ -2174,14 +2116,38 @@ mod tests {
21742116

21752117
#[test]
21762118
#[should_panic(expected = "DivideByZero")]
2177-
fn test_primitive_array_modulus_by_zero() {
2119+
fn test_f32_dict_array_divide_dyn_by_zero() {
2120+
let mut builder =
2121+
PrimitiveDictionaryBuilder::<Int8Type, Float32Type>::with_capacity(1, 1);
2122+
builder.append(1.5).unwrap();
2123+
let a = builder.finish();
2124+
2125+
let mut builder =
2126+
PrimitiveDictionaryBuilder::<Int8Type, Float32Type>::with_capacity(1, 1);
2127+
builder.append(0.0).unwrap();
2128+
let b = builder.finish();
2129+
2130+
divide_dyn(&a, &b).unwrap();
2131+
}
2132+
2133+
#[test]
2134+
#[should_panic(expected = "DivideByZero")]
2135+
fn test_i32_array_modulus_by_zero() {
21782136
let a = Int32Array::from(vec![15]);
21792137
let b = Int32Array::from(vec![0]);
21802138
modulus(&a, &b).unwrap();
21812139
}
21822140

21832141
#[test]
2184-
fn test_primitive_array_divide_f64() {
2142+
#[should_panic(expected = "DivideByZero")]
2143+
fn test_f32_array_modulus_by_zero() {
2144+
let a = Float32Array::from(vec![1.5]);
2145+
let b = Float32Array::from(vec![0.0]);
2146+
modulus(&a, &b).unwrap();
2147+
}
2148+
2149+
#[test]
2150+
fn test_f64_array_divide() {
21852151
let a = Float64Array::from(vec![15.0, 15.0, 8.0]);
21862152
let b = Float64Array::from(vec![5.0, 6.0, 8.0]);
21872153
let c = divide(&a, &b).unwrap();

0 commit comments

Comments
 (0)