From f078351992c63d73d149218e35ec9cedee41689d Mon Sep 17 00:00:00 2001 From: 1aguna Date: Sat, 30 Oct 2021 22:25:58 -0700 Subject: [PATCH] Added bitwise operations (#553) --- Cargo.toml | 4 +++ benches/bitwise.rs | 69 +++++++++++++++++++++++++++++++++++++ src/compute/bitwise.rs | 49 ++++++++++++++++++++++++++ src/compute/mod.rs | 1 + tests/it/compute/bitwise.rs | 41 ++++++++++++++++++++++ tests/it/compute/mod.rs | 1 + 6 files changed, 165 insertions(+) create mode 100644 benches/bitwise.rs create mode 100644 src/compute/bitwise.rs create mode 100644 tests/it/compute/bitwise.rs diff --git a/Cargo.toml b/Cargo.toml index 783300ca679..38ff035869b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -249,3 +249,7 @@ harness = false [[bench]] name = "avro_read" harness = false + +[[bench]] +name = "bitwise" +harness = false \ No newline at end of file diff --git a/benches/bitwise.rs b/benches/bitwise.rs new file mode 100644 index 00000000000..c91b5aad394 --- /dev/null +++ b/benches/bitwise.rs @@ -0,0 +1,69 @@ +use std::ops::{BitOr, BitXor, Not}; + +use criterion::{criterion_group, criterion_main, Criterion}; +use num_traits::NumCast; + +use arrow2::{ + array::PrimitiveArray, + compute::bitwise::*, + datatypes::DataType, + types::NativeType, + util::bench_util::{ + create_boolean_array, create_primitive_array, create_primitive_array_with_seed, + }, +}; + +fn bench_or(lhs: &PrimitiveArray, rhs: &PrimitiveArray) +where + T: NativeType + BitOr + NumCast, +{ + criterion::black_box(or(lhs, rhs)).unwrap(); +} + +fn bench_xor(lhs: &PrimitiveArray, rhs: &PrimitiveArray) +where + T: NativeType + BitXor + NumCast, +{ + criterion::black_box(xor(lhs, rhs)).unwrap(); +} + +fn bench_and(lhs: &PrimitiveArray, rhs: &PrimitiveArray) +where + T: NativeType + BitAnd + NumCast, +{ + criterion::black_box(and(lhs, rhs)).unwrap(); +} + +fn bench_not(arr: &PrimitiveArray) +where + T: NativeType + Not + NumCast, +{ + criterion::black_box(not(arr)); +} + +fn add_benchmark(c: &mut Criterion) { + (10..=20).step_by(2).for_each(|log2_size| { + let size = 2usize.pow(log2_size); + let arr_a = create_primitive_array_with_seed::(size, DataType::UInt64, 0.0, 43); + let arr_b = create_primitive_array_with_seed::(size, DataType::UInt64, 0.0, 42); + + c.bench_function(&format!("or 2^{}", log2_size), |b| { + b.iter(|| bench_or(&arr_a, &arr_b)) + }); + + c.bench_function(&format!("xor 2^{}", log2_size), |b| { + b.iter(|| bench_xor(&arr_a, &arr_b)) + }); + + c.bench_function(&format!("and 2^{}", log2_size), |b| { + b.iter(|| bench_and(&arr_a, &arr_b)) + }); + + c.bench_function(&format!("not 2^{}", log2_size), |b| { + b.iter(|| bench_not(&arr_a)) + }); + }); +} + +criterion_group!(benches, add_benchmark); +criterion_main!(benches); diff --git a/src/compute/bitwise.rs b/src/compute/bitwise.rs new file mode 100644 index 00000000000..127c3cc9344 --- /dev/null +++ b/src/compute/bitwise.rs @@ -0,0 +1,49 @@ +use std::ops::{BitAnd, BitOr, BitXor, Not}; + +use crate::array::{Array, PrimitiveArray}; +use crate::compute::arithmetics::basic::check_same_type; +use crate::compute::arity::{binary, unary}; +use crate::error::Result; +use crate::types::NativeType; + +/// Performs `OR` operation on two arrays. +/// # Error +/// This function errors when the arrays have different lengths or are different types. +pub fn or(lhs: &PrimitiveArray, rhs: &PrimitiveArray) -> Result> +where + T: NativeType + BitOr, +{ + check_same_type(lhs, rhs)?; + binary(lhs, rhs, lhs.data_type().clone(), |a, b| a | b) +} + +/// Performs `XOR` operation on two arrays. +/// # Error +/// This function errors when the arrays have different lengths or are different types. +pub fn xor(lhs: &PrimitiveArray, rhs: &PrimitiveArray) -> Result> +where + T: NativeType + BitXor, +{ + check_same_type(lhs, rhs)?; + binary(lhs, rhs, lhs.data_type().clone(), |a, b| a ^ b) +} + +/// Performs `AND` operation on two arrays. +/// # Error +/// This function errors when the arrays have different lengths or are different types. +pub fn and(lhs: &PrimitiveArray, rhs: &PrimitiveArray) -> Result> +where + T: NativeType + BitAnd, +{ + check_same_type(lhs, rhs)?; + binary(lhs, rhs, lhs.data_type().clone(), |a, b| a & b) +} + +/// Performs `OR` operation on one array. +pub fn not(arr: &PrimitiveArray) -> PrimitiveArray +where + T: NativeType + Not, +{ + let op = move |a: T| !a; + unary(arr, op, arr.data_type().clone()) +} diff --git a/src/compute/mod.rs b/src/compute/mod.rs index 51fcca52ef5..a48b87d18e4 100644 --- a/src/compute/mod.rs +++ b/src/compute/mod.rs @@ -14,6 +14,7 @@ pub mod aggregate; pub mod arithmetics; pub mod arity; +pub mod bitwise; pub mod boolean; pub mod boolean_kleene; pub mod cast; diff --git a/tests/it/compute/bitwise.rs b/tests/it/compute/bitwise.rs new file mode 100644 index 00000000000..92b8b47b4ab --- /dev/null +++ b/tests/it/compute/bitwise.rs @@ -0,0 +1,41 @@ +use arrow2::array::*; +use arrow2::compute::bitwise::*; + +#[test] +fn test_xor() { + let a = Int32Array::from(&[Some(2), Some(4), Some(6), Some(7)]); + let b = Int32Array::from(&[None, Some(6), Some(9), Some(7)]); + let result = xor(&a, &b).unwrap(); + let expected = Int32Array::from(&[None, Some(2), Some(15), Some(0)]); + + assert_eq!(result, expected); +} + +#[test] +fn test_and() { + let a = Int32Array::from(&[Some(1), Some(2), Some(15)]); + let b = Int32Array::from(&[None, Some(2), Some(6)]); + let result = and(&a, &b).unwrap(); + let expected = Int32Array::from(&[None, Some(2), Some(6)]); + + assert_eq!(result, expected); +} + +#[test] +fn test_or() { + let a = Int32Array::from(&[Some(1), Some(2), Some(0)]); + let b = Int32Array::from(&[None, Some(2), Some(0)]); + let result = or(&a, &b).unwrap(); + let expected = Int32Array::from(&[None, Some(2), Some(0)]); + + assert_eq!(result, expected); +} + +#[test] +fn test_not() { + let a = Int8Array::from(&[None, Some(1i8), Some(-100i8)]); + let result = not(&a); + let expected = Int8Array::from(&[None, Some(-2), Some(99)]); + + assert_eq!(result, expected); +} diff --git a/tests/it/compute/mod.rs b/tests/it/compute/mod.rs index 78eced0b882..732f1319303 100644 --- a/tests/it/compute/mod.rs +++ b/tests/it/compute/mod.rs @@ -1,5 +1,6 @@ mod aggregate; mod arithmetics; +mod bitwise; mod boolean; mod boolean_kleene; mod cast;