Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Added example. (#242)
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao authored Jul 31, 2021
1 parent 8e146a7 commit eea4d33
Show file tree
Hide file tree
Showing 5 changed files with 112 additions and 2 deletions.
84 changes: 84 additions & 0 deletions examples/arithmetics.rs
Original file line number Diff line number Diff line change
@@ -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::<i64>::from(&[Some(1), Some(2), Some(3)]);
let array1 = PrimitiveArray::<i64>::from(&[Some(4), None, Some(6)]);

// we can add them as follows:
let added = arithmetic_primitive(&array0, Operator::Add, &array1)?;
assert_eq!(
added,
PrimitiveArray::<i64>::from(&[Some(5), None, Some(9)])
);

// subtract:
let subtracted = arithmetic_primitive(&array0, Operator::Subtract, &array1)?;
assert_eq!(
subtracted,
PrimitiveArray::<i64>::from(&[Some(-3), None, Some(-3)])
);

// add a scalar:
let plus10 = arithmetic_primitive_scalar(&array0, Operator::Add, &10)?;
assert_eq!(
plus10,
PrimitiveArray::<i64>::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::<i64>::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::<i64>::from(&[Some(1), Some(2), Some(3)]);
let array1 = PrimitiveArray::<i64>::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::<i64>::from(&[Some(1 + 16), None, Some(9 + 36)])
);

// arbitrary unary operations
// this is compiled to SIMD when intrinsics exist.
let array0 = PrimitiveArray::<f64>::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::<f64>::from(&[Some(4.4), None, Some(4.6)]);
let rounded = unary(&array0, |x| x.round() as i64, DataType::Int64);
assert_eq!(
rounded,
PrimitiveArray::<i64>::from(&[Some(4), None, Some(5)])
);

Ok(())
}
1 change: 1 addition & 0 deletions guide/src/SUMMARY.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
23 changes: 23 additions & 0 deletions guide/src/compute.md
Original file line number Diff line number Diff line change
@@ -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`
3 changes: 2 additions & 1 deletion src/compute/arithmetics/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,8 @@ pub enum Operator {
}

/// Perform arithmetic operations on two primitive arrays based on the Operator enum
fn arithmetic_primitive<T>(
//
pub fn arithmetic_primitive<T>(
lhs: &PrimitiveArray<T>,
op: Operator,
rhs: &PrimitiveArray<T>,
Expand Down
3 changes: 2 additions & 1 deletion src/compute/arity.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit eea4d33

Please sign in to comment.