diff --git a/examples/arithmetics.rs b/examples/arithmetics.rs new file mode 100644 index 00000000000..3123b951521 --- /dev/null +++ b/examples/arithmetics.rs @@ -0,0 +1,84 @@ +use arrow2::array::{Array, PrimitiveArray}; +use arrow2::compute::arithmetics::*; +use arrow2::compute::arity::{binary, unary}; +use arrow2::datatypes::DataType; +use arrow2::error::Result; + +fn main() -> Result<()> { + // say we have two arrays + let array0 = PrimitiveArray::::from(&[Some(1), Some(2), Some(3)]); + let array1 = PrimitiveArray::::from(&[Some(4), None, Some(6)]); + + // we can add them as follows: + let added = arithmetic_primitive(&array0, Operator::Add, &array1)?; + assert_eq!( + added, + PrimitiveArray::::from(&[Some(5), None, Some(9)]) + ); + + // subtract: + let subtracted = arithmetic_primitive(&array0, Operator::Subtract, &array1)?; + assert_eq!( + subtracted, + PrimitiveArray::::from(&[Some(-3), None, Some(-3)]) + ); + + // add a scalar: + let plus10 = arithmetic_primitive_scalar(&array0, Operator::Add, &10)?; + assert_eq!( + plus10, + PrimitiveArray::::from(&[Some(11), Some(12), Some(13)]) + ); + + // when the array is a trait object, there is a similar API + let array0 = &array0 as &dyn Array; + let array1 = &array1 as &dyn Array; + + // check whether the logical types support addition (they could be any `Array`). + assert!(can_arithmetic( + array0.data_type(), + Operator::Add, + array1.data_type() + )); + + // add them + let added = arithmetic(array0, Operator::Add, array1).unwrap(); + assert_eq!( + PrimitiveArray::::from(&[Some(5), None, Some(9)]), + added.as_ref(), + ); + + // a more exotic implementation: arbitrary binary operations + // this is compiled to SIMD when intrinsics exist. + let array0 = PrimitiveArray::::from(&[Some(1), Some(2), Some(3)]); + let array1 = PrimitiveArray::::from(&[Some(4), None, Some(6)]); + + let op = |x: i64, y: i64| x.pow(2) + y.pow(2); + let r = binary(&array0, &array1, DataType::Int64, op)?; + assert_eq!( + r, + PrimitiveArray::::from(&[Some(1 + 16), None, Some(9 + 36)]) + ); + + // arbitrary unary operations + // this is compiled to SIMD when intrinsics exist. + let array0 = PrimitiveArray::::from(&[Some(4.0), None, Some(6.0)]); + let r = unary( + &array0, + |x| x.cos().powi(2) + x.sin().powi(2), + DataType::Float64, + ); + assert!((r.values()[0] - 1.0).abs() < 0.0001); + assert!(r.is_null(1)); + assert!((r.values()[2] - 1.0).abs() < 0.0001); + + // finally, a transformation that changes types: + let array0 = PrimitiveArray::::from(&[Some(4.4), None, Some(4.6)]); + let rounded = unary(&array0, |x| x.round() as i64, DataType::Int64); + assert_eq!( + rounded, + PrimitiveArray::::from(&[Some(4), None, Some(5)]) + ); + + Ok(()) +} diff --git a/guide/src/SUMMARY.md b/guide/src/SUMMARY.md index a73ef112ecd..cdf5aec37ee 100644 --- a/guide/src/SUMMARY.md +++ b/guide/src/SUMMARY.md @@ -4,6 +4,7 @@ - [The arrow format](./arrow.md) - [Low-level API](./low_level.md) - [High-level API](./high_level.md) +- [Compute](./compute.md) - [Metadata](./metadata.md) - [Foreign interfaces](./ffi.md) - [IO](./io/README.md) diff --git a/guide/src/compute.md b/guide/src/compute.md new file mode 100644 index 00000000000..68948149de0 --- /dev/null +++ b/guide/src/compute.md @@ -0,0 +1,23 @@ +# Compute API + +When compiled with the feature `compute`, this crate offers a wide range of functions to perform both vertical (e.g. add two arrays) and horizontal (compute the sum of an array) operations. + +```rust +{{#include ../../examples/arithmetics.rs}} +``` + +An overview of the implemented functionality. + +* arithmetics, checked, saturating, etc. +* `sum`, `min` and `max` +* `unary`, `binary`, etc. +* `comparison` +* `cast` +* `take`, `filter`, `concat` +* `sort`, `hash`, `merge-sort` +* `if-then-else` +* `nullif` +* `lenght` (of string) +* `hour`, `year` (of temporal logical types) +* `regex` +* (list) `contains` diff --git a/src/compute/arithmetics/mod.rs b/src/compute/arithmetics/mod.rs index ce21ed3c140..496f02d72f3 100644 --- a/src/compute/arithmetics/mod.rs +++ b/src/compute/arithmetics/mod.rs @@ -228,7 +228,8 @@ pub enum Operator { } /// Perform arithmetic operations on two primitive arrays based on the Operator enum -fn arithmetic_primitive( +// +pub fn arithmetic_primitive( lhs: &PrimitiveArray, op: Operator, rhs: &PrimitiveArray, diff --git a/src/compute/arity.rs b/src/compute/arity.rs index 0e99c488e16..e69d8162a21 100644 --- a/src/compute/arity.rs +++ b/src/compute/arity.rs @@ -128,7 +128,8 @@ where /// Applies a binary operations to two primitive arrays. This is the fastest /// way to perform an operation on two primitive array when the benefits of a /// vectorized operation outweighs the cost of branching nulls and non-nulls. -/// +/// # Errors +/// This function errors iff the arrays have a different length. /// # Implementation /// This will apply the function for all values, including those on null slots. /// This implies that the operation must be infallible for any value of the