vortex_array/compute/
numeric.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5
6use vortex_error::VortexResult;
7
8use crate::Array;
9use crate::ArrayRef;
10use crate::arrays::ConstantArray;
11use crate::arrow::Datum;
12use crate::arrow::from_arrow_array_with_len;
13use crate::builtins::ArrayBuiltins;
14use crate::compute::Options;
15use crate::scalar::NumericOperator;
16use crate::scalar::Scalar;
17use crate::scalar_fn::fns::operators::Operator;
18
19/// Point-wise add two numeric arrays.
20///
21/// Errs at runtime if the sum would overflow or underflow.
22///
23/// The result is null at any index that either input is null.
24#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
25pub fn add(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
26    lhs.to_array().binary(rhs.to_array(), Operator::Add)
27}
28
29/// Point-wise add a scalar value to this array on the right-hand-side.
30#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
31pub fn add_scalar(lhs: &dyn Array, rhs: Scalar) -> VortexResult<ArrayRef> {
32    lhs.to_array()
33        .binary(ConstantArray::new(rhs, lhs.len()).to_array(), Operator::Add)
34}
35
36/// Point-wise subtract two numeric arrays.
37#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
38pub fn sub(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
39    lhs.to_array().binary(rhs.to_array(), Operator::Sub)
40}
41
42/// Point-wise subtract a scalar value from this array on the right-hand-side.
43#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
44pub fn sub_scalar(lhs: &dyn Array, rhs: Scalar) -> VortexResult<ArrayRef> {
45    lhs.to_array()
46        .binary(ConstantArray::new(rhs, lhs.len()).to_array(), Operator::Sub)
47}
48
49/// Point-wise multiply two numeric arrays.
50#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
51pub fn mul(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
52    lhs.to_array().binary(rhs.to_array(), Operator::Mul)
53}
54
55/// Point-wise multiply a scalar value into this array on the right-hand-side.
56#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
57pub fn mul_scalar(lhs: &dyn Array, rhs: Scalar) -> VortexResult<ArrayRef> {
58    lhs.to_array()
59        .binary(ConstantArray::new(rhs, lhs.len()).to_array(), Operator::Mul)
60}
61
62/// Point-wise divide two numeric arrays.
63#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
64pub fn div(lhs: &dyn Array, rhs: &dyn Array) -> VortexResult<ArrayRef> {
65    lhs.to_array().binary(rhs.to_array(), Operator::Div)
66}
67
68/// Point-wise divide a scalar value into this array on the right-hand-side.
69#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
70pub fn div_scalar(lhs: &dyn Array, rhs: Scalar) -> VortexResult<ArrayRef> {
71    lhs.to_array()
72        .binary(ConstantArray::new(rhs, lhs.len()).to_array(), Operator::Div)
73}
74
75/// Point-wise numeric operation between two arrays of the same type and length.
76#[deprecated(note = "Use `ArrayBuiltins::binary` instead")]
77pub fn numeric(lhs: &dyn Array, rhs: &dyn Array, op: NumericOperator) -> VortexResult<ArrayRef> {
78    arrow_numeric(lhs, rhs, op)
79}
80
81impl Options for NumericOperator {
82    fn as_any(&self) -> &dyn Any {
83        self
84    }
85}
86
87/// Implementation of numeric operations using the Arrow crate.
88pub(crate) fn arrow_numeric(
89    lhs: &dyn Array,
90    rhs: &dyn Array,
91    operator: NumericOperator,
92) -> VortexResult<ArrayRef> {
93    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
94    let len = lhs.len();
95
96    let left = Datum::try_new(lhs)?;
97    let right = Datum::try_new_with_target_datatype(rhs, left.data_type())?;
98
99    let array = match operator {
100        NumericOperator::Add => arrow_arith::numeric::add(&left, &right)?,
101        NumericOperator::Sub => arrow_arith::numeric::sub(&left, &right)?,
102        NumericOperator::Mul => arrow_arith::numeric::mul(&left, &right)?,
103        NumericOperator::Div => arrow_arith::numeric::div(&left, &right)?,
104    };
105
106    from_arrow_array_with_len(array.as_ref(), len, nullable)
107}
108
109#[cfg(test)]
110#[allow(deprecated)]
111mod test {
112    use vortex_buffer::buffer;
113
114    use crate::IntoArray;
115    use crate::arrays::PrimitiveArray;
116    use crate::assert_arrays_eq;
117    use crate::compute::sub_scalar;
118
119    #[test]
120    fn test_scalar_subtract_unsigned() {
121        let values = buffer![1u16, 2, 3].into_array();
122        let result = sub_scalar(&values, 1u16.into()).unwrap();
123        assert_arrays_eq!(result, PrimitiveArray::from_iter([0u16, 1, 2]));
124    }
125
126    #[test]
127    fn test_scalar_subtract_signed() {
128        let values = buffer![1i64, 2, 3].into_array();
129        let result = sub_scalar(&values, (-1i64).into()).unwrap();
130        assert_arrays_eq!(result, PrimitiveArray::from_iter([2i64, 3, 4]));
131    }
132
133    #[test]
134    fn test_scalar_subtract_nullable() {
135        let values = PrimitiveArray::from_option_iter([Some(1u16), Some(2), None, Some(3)]);
136        let result = sub_scalar(values.as_ref(), Some(1u16).into()).unwrap();
137        assert_arrays_eq!(
138            result,
139            PrimitiveArray::from_option_iter([Some(0u16), Some(1), None, Some(2)])
140        );
141    }
142
143    #[test]
144    fn test_scalar_subtract_float() {
145        let values = buffer![1.0f64, 2.0, 3.0].into_array();
146        let to_subtract = -1f64;
147        let result = sub_scalar(&values, to_subtract.into()).unwrap();
148        assert_arrays_eq!(result, PrimitiveArray::from_iter([2.0f64, 3.0, 4.0]));
149    }
150
151    #[test]
152    fn test_scalar_subtract_float_underflow_is_ok() {
153        let values = buffer![f32::MIN, 2.0, 3.0].into_array();
154        let _results = sub_scalar(&values, 1.0f32.into()).unwrap();
155        let _results = sub_scalar(&values, f32::MAX.into()).unwrap();
156    }
157}
vortex_array/compute/numeric.rs

vortex_array/compute/
numeric.rs
