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<()> {
let array0 = PrimitiveArray::<i64>::from(&[Some(1), Some(2), Some(3)]);
let array1 = PrimitiveArray::<i64>::from(&[Some(4), None, Some(6)]);
let added = arithmetic_primitive(&array0, Operator::Add, &array1)?;
assert_eq!(
added,
PrimitiveArray::<i64>::from(&[Some(5), None, Some(9)])
);
let subtracted = arithmetic_primitive(&array0, Operator::Subtract, &array1)?;
assert_eq!(
subtracted,
PrimitiveArray::<i64>::from(&[Some(-3), None, Some(-3)])
);
let plus10 = arithmetic_primitive_scalar(&array0, Operator::Add, &10)?;
assert_eq!(
plus10,
PrimitiveArray::<i64>::from(&[Some(11), Some(12), Some(13)])
);
let array0 = &array0 as &dyn Array;
let array1 = &array1 as &dyn Array;
assert!(can_arithmetic(
array0.data_type(),
Operator::Add,
array1.data_type()
));
let added = arithmetic(array0, Operator::Add, array1).unwrap();
assert_eq!(
PrimitiveArray::<i64>::from(&[Some(5), None, Some(9)]),
added.as_ref(),
);
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)])
);
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);
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(())
}