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

Added support for bitwise operations #553

Merged
merged 15 commits into from
Oct 31, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -249,3 +249,7 @@ harness = false
[[bench]]
name = "avro_read"
harness = false

[[bench]]
name = "bitwise"
harness = false
69 changes: 69 additions & 0 deletions benches/bitwise.rs
Original file line number Diff line number Diff line change
@@ -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<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>)
where
T: NativeType + BitOr<Output = T> + NumCast,
{
criterion::black_box(or(lhs, rhs)).unwrap();
}

fn bench_xor<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>)
where
T: NativeType + BitXor<Output = T> + NumCast,
{
criterion::black_box(xor(lhs, rhs)).unwrap();
}

fn bench_and<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>)
where
T: NativeType + BitAnd<Output = T> + NumCast,
{
criterion::black_box(and(lhs, rhs)).unwrap();
}

fn bench_not<T>(arr: &PrimitiveArray<T>)
where
T: NativeType + Not<Output = T> + 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::<u64>(size, DataType::UInt64, 0.0, 43);
let arr_b = create_primitive_array_with_seed::<u64>(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);
49 changes: 49 additions & 0 deletions src/compute/bitwise.rs
Original file line number Diff line number Diff line change
@@ -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<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
where
T: NativeType + BitOr<Output = T>,
{
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<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
where
T: NativeType + BitXor<Output = T>,
{
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<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>>
where
T: NativeType + BitAnd<Output = T>,
{
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<T>(arr: &PrimitiveArray<T>) -> PrimitiveArray<T>
where
T: NativeType + Not<Output = T>,
{
let op = move |a: T| !a;
unary(arr, op, arr.data_type().clone())
}
1 change: 1 addition & 0 deletions src/compute/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
41 changes: 41 additions & 0 deletions tests/it/compute/bitwise.rs
Original file line number Diff line number Diff line change
@@ -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);
}
1 change: 1 addition & 0 deletions tests/it/compute/mod.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
mod aggregate;
mod arithmetics;
mod bitwise;
mod boolean;
mod boolean_kleene;
mod cast;
Expand Down