fedor/Notes
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Vendor dependencies

Browse Source
This commit is contained in:
Fedor Lyanguzov
2026-08-01 16:11:49 +03:00
parent 7f139a0241
commit 6b5e7f0f8b
29706 changed files with 9575646 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notes-service/vendor/arrow-arith/src/aggregate.rs
Vendored
+1753
View File
File diff suppressed because it is too large Load Diff
notes-service/vendor/arrow-arith/src/arithmetic.rs
Vendored
+347
View File
@@ -0,0 +1,347 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Defines basic arithmetic kernels for `PrimitiveArrays`.
//!
//! These kernels can leverage SIMD if available on your system. Currently no runtime
//! detection is provided, you should enable the specific SIMD intrinsics using
//! `RUSTFLAGS="-C target-feature=+avx2"` for example. See the documentation
//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
use crate::arity::*;
use arrow_array::types::*;
use arrow_array::*;
use arrow_buffer::ArrowNativeType;
use arrow_buffer::i256;
use arrow_schema::*;
use std::cmp::min;
use std::sync::Arc;
/// Returns the precision and scale of the result of a multiplication of two decimal types,
/// and the divisor for fixed point multiplication.
fn get_fixed_point_info(
left: (u8, i8),
right: (u8, i8),
required_scale: i8,
) -> Result<(u8, i8, i256), ArrowError> {
let product_scale = left.1 + right.1;
let precision = min(left.0 + right.0 + 1, DECIMAL128_MAX_PRECISION);
if required_scale > product_scale {
return Err(ArrowError::ComputeError(format!(
"Required scale {required_scale} is greater than product scale {product_scale}",
)));
}
let divisor = i256::from_i128(10).pow_wrapping((product_scale - required_scale) as u32);
Ok((precision, product_scale, divisor))
}
/// Perform `left * right` operation on two decimal arrays. If either left or right value is
/// null then the result is also null.
///
/// This performs decimal multiplication which allows precision loss if an exact representation
/// is not possible for the result, according to the required scale. In the case, the result
/// will be rounded to the required scale.
///
/// If the required scale is greater than the product scale, an error is returned.
///
/// This doesn't detect overflow. Once overflowing, the result will wrap around.
///
/// It is implemented for compatibility with precision loss `multiply` function provided by
/// other data processing engines. For multiplication with precision loss detection, use
/// `multiply_dyn` or `multiply_dyn_checked` instead.
pub fn multiply_fixed_point_dyn(
left: &dyn Array,
right: &dyn Array,
required_scale: i8,
) -> Result<ArrayRef, ArrowError> {
match (left.data_type(), right.data_type()) {
(DataType::Decimal128(_, _), DataType::Decimal128(_, _)) => {
let left = left.as_any().downcast_ref::<Decimal128Array>().unwrap();
let right = right.as_any().downcast_ref::<Decimal128Array>().unwrap();
multiply_fixed_point(left, right, required_scale).map(|a| Arc::new(a) as ArrayRef)
}
(_, _) => Err(ArrowError::CastError(format!(
"Unsupported data type {}, {}",
left.data_type(),
right.data_type()
))),
}
}
/// Perform `left * right` operation on two decimal arrays. If either left or right value is
/// null then the result is also null.
///
/// This performs decimal multiplication which allows precision loss if an exact representation
/// is not possible for the result, according to the required scale. In the case, the result
/// will be rounded to the required scale.
///
/// If the required scale is greater than the product scale, an error is returned.
///
/// It is implemented for compatibility with precision loss `multiply` function provided by
/// other data processing engines. For multiplication with precision loss detection, use
/// `multiply` or `multiply_checked` instead.
pub fn multiply_fixed_point_checked(
left: &PrimitiveArray<Decimal128Type>,
right: &PrimitiveArray<Decimal128Type>,
required_scale: i8,
) -> Result<PrimitiveArray<Decimal128Type>, ArrowError> {
let (precision, product_scale, divisor) = get_fixed_point_info(
(left.precision(), left.scale()),
(right.precision(), right.scale()),
required_scale,
)?;
if required_scale == product_scale {
return try_binary::<_, _, _, Decimal128Type>(left, right, |a, b| a.mul_checked(b))?
.with_precision_and_scale(precision, required_scale);
}
try_binary::<_, _, _, Decimal128Type>(left, right, |a, b| {
let a = i256::from_i128(a);
let b = i256::from_i128(b);
let mut mul = a.wrapping_mul(b);
mul = divide_and_round::<Decimal256Type>(mul, divisor);
mul.to_i128().ok_or_else(|| {
ArrowError::ArithmeticOverflow(format!("Overflow happened on: {a:?} * {b:?}"))
})
})
.and_then(|a| a.with_precision_and_scale(precision, required_scale))
}
/// Perform `left * right` operation on two decimal arrays. If either left or right value is
/// null then the result is also null.
///
/// This performs decimal multiplication which allows precision loss if an exact representation
/// is not possible for the result, according to the required scale. In the case, the result
/// will be rounded to the required scale.
///
/// If the required scale is greater than the product scale, an error is returned.
///
/// This doesn't detect overflow. Once overflowing, the result will wrap around.
/// For an overflow-checking variant, use `multiply_fixed_point_checked` instead.
///
/// It is implemented for compatibility with precision loss `multiply` function provided by
/// other data processing engines. For multiplication with precision loss detection, use
/// `multiply` or `multiply_checked` instead.
pub fn multiply_fixed_point(
left: &PrimitiveArray<Decimal128Type>,
right: &PrimitiveArray<Decimal128Type>,
required_scale: i8,
) -> Result<PrimitiveArray<Decimal128Type>, ArrowError> {
let (precision, product_scale, divisor) = get_fixed_point_info(
(left.precision(), left.scale()),
(right.precision(), right.scale()),
required_scale,
)?;
if required_scale == product_scale {
return binary(left, right, |a, b| a.mul_wrapping(b))?
.with_precision_and_scale(precision, required_scale);
}
binary::<_, _, _, Decimal128Type>(left, right, |a, b| {
let a = i256::from_i128(a);
let b = i256::from_i128(b);
let mut mul = a.wrapping_mul(b);
mul = divide_and_round::<Decimal256Type>(mul, divisor);
mul.as_i128()
})
.and_then(|a| a.with_precision_and_scale(precision, required_scale))
}
/// Divide a decimal native value by given divisor and round the result.
fn divide_and_round<I>(input: I::Native, div: I::Native) -> I::Native
where
I: DecimalType,
I::Native: ArrowNativeTypeOp,
{
let d = input.div_wrapping(div);
let r = input.mod_wrapping(div);
let half = div.div_wrapping(I::Native::from_usize(2).unwrap());
let half_neg = half.neg_wrapping();
// Round result
match input >= I::Native::ZERO {
true if r >= half => d.add_wrapping(I::Native::ONE),
false if r <= half_neg => d.sub_wrapping(I::Native::ONE),
_ => d,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::numeric::mul;
#[test]
fn test_decimal_multiply_allow_precision_loss() {
// Overflow happening as i128 cannot hold multiplying result.
// [123456789]
let a = Decimal128Array::from(vec![123456789000000000000000000])
.with_precision_and_scale(38, 18)
.unwrap();
// [10]
let b = Decimal128Array::from(vec![10000000000000000000])
.with_precision_and_scale(38, 18)
.unwrap();
let err = mul(&a, &b).unwrap_err();
assert!(
err.to_string().contains(
"Overflow happened on: 123456789000000000000000000 * 10000000000000000000"
)
);
// Allow precision loss.
let result = multiply_fixed_point_checked(&a, &b, 28).unwrap();
// [1234567890]
let expected = Decimal128Array::from(vec![12345678900000000000000000000000000000])
.with_precision_and_scale(38, 28)
.unwrap();
assert_eq!(&expected, &result);
assert_eq!(
result.value_as_string(0),
"1234567890.0000000000000000000000000000"
);
// Rounding case
// [0.000000000000000001, 123456789.555555555555555555, 1.555555555555555555]
let a = Decimal128Array::from(vec![1, 123456789555555555555555555, 1555555555555555555])
.with_precision_and_scale(38, 18)
.unwrap();
// [1.555555555555555555, 11.222222222222222222, 0.000000000000000001]
let b = Decimal128Array::from(vec![1555555555555555555, 11222222222222222222, 1])
.with_precision_and_scale(38, 18)
.unwrap();
let result = multiply_fixed_point_checked(&a, &b, 28).unwrap();
// [
// 0.0000000000000000015555555556,
// 1385459527.2345679012071330528765432099,
// 0.0000000000000000015555555556
// ]
let expected = Decimal128Array::from(vec![
15555555556,
13854595272345679012071330528765432099,
15555555556,
])
.with_precision_and_scale(38, 28)
.unwrap();
assert_eq!(&expected, &result);
// Rounded the value "1385459527.234567901207133052876543209876543210".
assert_eq!(
result.value_as_string(1),
"1385459527.2345679012071330528765432099"
);
assert_eq!(result.value_as_string(0), "0.0000000000000000015555555556");
assert_eq!(result.value_as_string(2), "0.0000000000000000015555555556");
let a = Decimal128Array::from(vec![1230])
.with_precision_and_scale(4, 2)
.unwrap();
let b = Decimal128Array::from(vec![1000])
.with_precision_and_scale(4, 2)
.unwrap();
// Required scale is same as the product of the input scales. Behavior is same as multiply.
let result = multiply_fixed_point_checked(&a, &b, 4).unwrap();
assert_eq!(result.precision(), 9);
assert_eq!(result.scale(), 4);
let expected = mul(&a, &b).unwrap();
assert_eq!(expected.as_ref(), &result);
// Required scale cannot be larger than the product of the input scales.
let result = multiply_fixed_point_checked(&a, &b, 5).unwrap_err();
assert!(
result
.to_string()
.contains("Required scale 5 is greater than product scale 4")
);
}
#[test]
fn test_decimal_multiply_allow_precision_loss_overflow() {
// [99999999999123456789]
let a = Decimal128Array::from(vec![99999999999123456789000000000000000000])
.with_precision_and_scale(38, 18)
.unwrap();
// [9999999999910]
let b = Decimal128Array::from(vec![9999999999910000000000000000000])
.with_precision_and_scale(38, 18)
.unwrap();
let err = multiply_fixed_point_checked(&a, &b, 28).unwrap_err();
assert!(err.to_string().contains(
"Overflow happened on: 99999999999123456789000000000000000000 * 9999999999910000000000000000000"
));
let result = multiply_fixed_point(&a, &b, 28).unwrap();
let expected = Decimal128Array::from(vec![62946009661555981610246871926660136960])
.with_precision_and_scale(38, 28)
.unwrap();
assert_eq!(&expected, &result);
}
#[test]
fn test_decimal_multiply_fixed_point() {
// [123456789]
let a = Decimal128Array::from(vec![123456789000000000000000000])
.with_precision_and_scale(38, 18)
.unwrap();
// [10]
let b = Decimal128Array::from(vec![10000000000000000000])
.with_precision_and_scale(38, 18)
.unwrap();
// `multiply` overflows on this case.
let err = mul(&a, &b).unwrap_err();
assert_eq!(
err.to_string(),
"Arithmetic overflow: Overflow happened on: 123456789000000000000000000 * 10000000000000000000"
);
// Avoid overflow by reducing the scale.
let result = multiply_fixed_point(&a, &b, 28).unwrap();
// [1234567890]
let expected = Decimal128Array::from(vec![12345678900000000000000000000000000000])
.with_precision_and_scale(38, 28)
.unwrap();
assert_eq!(&expected, &result);
assert_eq!(
result.value_as_string(0),
"1234567890.0000000000000000000000000000"
);
}
}
notes-service/vendor/arrow-arith/src/arity.rs
Vendored
+528
View File
@@ -0,0 +1,528 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Kernels for operating on [`PrimitiveArray`]s
use arrow_array::builder::BufferBuilder;
use arrow_array::*;
use arrow_buffer::ArrowNativeType;
use arrow_buffer::MutableBuffer;
use arrow_buffer::buffer::NullBuffer;
use arrow_data::ArrayData;
use arrow_schema::ArrowError;
/// See [`PrimitiveArray::unary`]
pub fn unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> PrimitiveArray<O>
where
I: ArrowPrimitiveType,
O: ArrowPrimitiveType,
F: Fn(I::Native) -> O::Native,
{
array.unary(op)
}
/// See [`PrimitiveArray::unary_mut`]
pub fn unary_mut<I, F>(
array: PrimitiveArray<I>,
op: F,
) -> Result<PrimitiveArray<I>, PrimitiveArray<I>>
where
I: ArrowPrimitiveType,
F: Fn(I::Native) -> I::Native,
{
array.unary_mut(op)
}
/// See [`PrimitiveArray::try_unary`]
pub fn try_unary<I, F, O>(array: &PrimitiveArray<I>, op: F) -> Result<PrimitiveArray<O>, ArrowError>
where
I: ArrowPrimitiveType,
O: ArrowPrimitiveType,
F: Fn(I::Native) -> Result<O::Native, ArrowError>,
{
array.try_unary(op)
}
/// See [`PrimitiveArray::try_unary_mut`]
pub fn try_unary_mut<I, F>(
array: PrimitiveArray<I>,
op: F,
) -> Result<Result<PrimitiveArray<I>, ArrowError>, PrimitiveArray<I>>
where
I: ArrowPrimitiveType,
F: Fn(I::Native) -> Result<I::Native, ArrowError>,
{
array.try_unary_mut(op)
}
/// Allies a binary infallable function to two [`PrimitiveArray`]s,
/// producing a new [`PrimitiveArray`]
///
/// # Details
///
/// Given two arrays of length `len`, calls `op(a[i], b[i])` for `i` in `0..len`, collecting
/// the results in a [`PrimitiveArray`].
///
/// If any index is null in either `a` or `b`, the
/// corresponding index in the result will also be null
///
/// Like [`unary`], the `op` is evaluated for every element in the two arrays,
/// including those elements which are NULL. This is beneficial as the cost of
/// the operation is low compared to the cost of branching, and especially when
/// the operation can be vectorised, however, requires `op` to be infallible for
/// all possible values of its inputs
///
/// # Errors
///
/// * if the arrays have different lengths.
///
/// # Example
/// ```
/// # use arrow_arith::arity::binary;
/// # use arrow_array::{Float32Array, Int32Array};
/// # use arrow_array::types::Int32Type;
/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8), Some(7.2)]);
/// let b = Int32Array::from(vec![1, 2, 4, 9]);
/// // compute int(a) + b for each element
/// let c = binary(&a, &b, |a, b| a as i32 + b).unwrap();
/// assert_eq!(c, Int32Array::from(vec![Some(6), None, Some(10), Some(16)]));
/// ```
pub fn binary<A, B, F, O>(
a: &PrimitiveArray<A>,
b: &PrimitiveArray<B>,
op: F,
) -> Result<PrimitiveArray<O>, ArrowError>
where
A: ArrowPrimitiveType,
B: ArrowPrimitiveType,
O: ArrowPrimitiveType,
F: Fn(A::Native, B::Native) -> O::Native,
{
if a.len() != b.len() {
return Err(ArrowError::ComputeError(
"Cannot perform binary operation on arrays of different length".to_string(),
));
}
if a.is_empty() {
return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
}
let nulls = NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref());
let values = a
.values()
.into_iter()
.zip(b.values())
.map(|(l, r)| op(*l, *r));
let buffer: Vec<_> = values.collect();
Ok(PrimitiveArray::new(buffer.into(), nulls))
}
/// Applies a binary and infallible function to values in two arrays, replacing
/// the values in the first array in place.
///
/// # Details
///
/// Given two arrays of length `len`, calls `op(a[i], b[i])` for `i` in
/// `0..len`, modifying the [`PrimitiveArray`] `a` in place, if possible.
///
/// If any index is null in either `a` or `b`, the corresponding index in the
/// result will also be null.
///
/// # Buffer Reuse
///
/// If the underlying buffers in `a` are not shared with other arrays, mutates
/// the underlying buffer in place, without allocating.
///
/// If the underlying buffer in `a` are shared, returns Err(self)
///
/// Like [`unary`] the provided function is evaluated for every index, ignoring validity. This
/// is beneficial when the cost of the operation is low compared to the cost of branching, and
/// especially when the operation can be vectorised, however, requires `op` to be infallible
/// for all possible values of its inputs
///
/// # Errors
///
/// * If the arrays have different lengths
/// * If the array is not mutable (see "Buffer Reuse")
///
/// # See Also
///
/// * Documentation on [`PrimitiveArray::unary_mut`] for operating on [`ArrayRef`].
///
/// # Example
/// ```
/// # use arrow_arith::arity::binary_mut;
/// # use arrow_array::{Float32Array, Int32Array};
/// # use arrow_array::types::Int32Type;
/// // compute a + b for each element
/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8)]);
/// let b = Int32Array::from(vec![Some(1), None, Some(2)]);
/// // compute a + b, updating the value in a in place if possible
/// let a = binary_mut(a, &b, |a, b| a + b as f32).unwrap().unwrap();
/// // a is updated in place
/// assert_eq!(a, Float32Array::from(vec![Some(6.1), None, Some(8.8)]));
/// ```
///
/// # Example with shared buffers
/// ```
/// # use arrow_arith::arity::binary_mut;
/// # use arrow_array::Float32Array;
/// # use arrow_array::types::Int32Type;
/// let a = Float32Array::from(vec![Some(5.1f32), None, Some(6.8)]);
/// let b = Float32Array::from(vec![Some(1.0f32), None, Some(2.0)]);
/// // a_clone shares the buffer with a
/// let a_cloned = a.clone();
/// // try to update a in place, but it is shared. Returns Err(a)
/// let a = binary_mut(a, &b, |a, b| a + b).unwrap_err();
/// assert_eq!(a_cloned, a);
/// // drop shared reference
/// drop(a_cloned);
/// // now a is not shared, so we can update it in place
/// let a = binary_mut(a, &b, |a, b| a + b).unwrap().unwrap();
/// assert_eq!(a, Float32Array::from(vec![Some(6.1), None, Some(8.8)]));
/// ```
pub fn binary_mut<T, U, F>(
a: PrimitiveArray<T>,
b: &PrimitiveArray<U>,
op: F,
) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
where
T: ArrowPrimitiveType,
U: ArrowPrimitiveType,
F: Fn(T::Native, U::Native) -> T::Native,
{
if a.len() != b.len() {
return Ok(Err(ArrowError::ComputeError(
"Cannot perform binary operation on arrays of different length".to_string(),
)));
}
if a.is_empty() {
return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty(
&T::DATA_TYPE,
))));
}
let mut builder = a.into_builder()?;
builder
.values_slice_mut()
.iter_mut()
.zip(b.values())
.for_each(|(l, r)| *l = op(*l, *r));
let array = builder.finish();
// The builder has the null buffer from `a`, it is not changed.
let nulls = NullBuffer::union(array.logical_nulls().as_ref(), b.logical_nulls().as_ref());
let array_builder = array.into_data().into_builder().nulls(nulls);
let array_data = unsafe { array_builder.build_unchecked() };
Ok(Ok(PrimitiveArray::<T>::from(array_data)))
}
/// Applies the provided fallible binary operation across `a` and `b`.
///
/// This will return any error encountered, or collect the results into
/// a [`PrimitiveArray`]. If any index is null in either `a`
/// or `b`, the corresponding index in the result will also be null
///
/// Like [`try_unary`] the function is only evaluated for non-null indices
///
/// # Error
///
/// Return an error if the arrays have different lengths or
/// the operation is under erroneous
pub fn try_binary<A: ArrayAccessor, B: ArrayAccessor, F, O>(
a: A,
b: B,
op: F,
) -> Result<PrimitiveArray<O>, ArrowError>
where
O: ArrowPrimitiveType,
F: Fn(A::Item, B::Item) -> Result<O::Native, ArrowError>,
{
if a.len() != b.len() {
return Err(ArrowError::ComputeError(
"Cannot perform a binary operation on arrays of different length".to_string(),
));
}
if a.is_empty() {
return Ok(PrimitiveArray::from(ArrayData::new_empty(&O::DATA_TYPE)));
}
let len = a.len();
if a.null_count() == 0 && b.null_count() == 0 {
try_binary_no_nulls(len, a, b, op)
} else {
let nulls =
NullBuffer::union(a.logical_nulls().as_ref(), b.logical_nulls().as_ref()).unwrap();
let mut buffer = BufferBuilder::<O::Native>::new(len);
buffer.append_n_zeroed(len);
let slice = buffer.as_slice_mut();
nulls.try_for_each_valid_idx(|idx| {
unsafe {
*slice.get_unchecked_mut(idx) = op(a.value_unchecked(idx), b.value_unchecked(idx))?
};
Ok::<_, ArrowError>(())
})?;
let values = buffer.finish().into();
Ok(PrimitiveArray::new(values, Some(nulls)))
}
}
/// Applies the provided fallible binary operation across `a` and `b` by mutating the mutable
/// [`PrimitiveArray`] `a` with the results.
///
/// Returns any error encountered, or collects the results into a [`PrimitiveArray`] as return
/// value. If any index is null in either `a` or `b`, the corresponding index in the result will
/// also be null.
///
/// Like [`try_unary`] the function is only evaluated for non-null indices.
///
/// See [`binary_mut`] for errors and buffer reuse information.
pub fn try_binary_mut<T, F>(
a: PrimitiveArray<T>,
b: &PrimitiveArray<T>,
op: F,
) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
where
T: ArrowPrimitiveType,
F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
{
if a.len() != b.len() {
return Ok(Err(ArrowError::ComputeError(
"Cannot perform binary operation on arrays of different length".to_string(),
)));
}
let len = a.len();
if a.is_empty() {
return Ok(Ok(PrimitiveArray::from(ArrayData::new_empty(
&T::DATA_TYPE,
))));
}
if a.null_count() == 0 && b.null_count() == 0 {
try_binary_no_nulls_mut(len, a, b, op)
} else {
let nulls =
create_union_null_buffer(a.logical_nulls().as_ref(), b.logical_nulls().as_ref())
.unwrap();
let mut builder = a.into_builder()?;
let slice = builder.values_slice_mut();
let r = nulls.try_for_each_valid_idx(|idx| {
unsafe {
*slice.get_unchecked_mut(idx) =
op(*slice.get_unchecked(idx), b.value_unchecked(idx))?
};
Ok::<_, ArrowError>(())
});
if let Err(err) = r {
return Ok(Err(err));
}
let array_builder = builder.finish().into_data().into_builder();
let array_data = unsafe { array_builder.nulls(Some(nulls)).build_unchecked() };
Ok(Ok(PrimitiveArray::<T>::from(array_data)))
}
}
/// Computes the union of the nulls in two optional [`NullBuffer`] which
/// is not shared with the input buffers.
///
/// The union of the nulls is the same as `NullBuffer::union(lhs, rhs)` but
/// it does not increase the reference count of the null buffer.
fn create_union_null_buffer(
lhs: Option<&NullBuffer>,
rhs: Option<&NullBuffer>,
) -> Option<NullBuffer> {
match (lhs, rhs) {
(Some(lhs), Some(rhs)) => Some(NullBuffer::new(lhs.inner() & rhs.inner())),
(Some(n), None) | (None, Some(n)) => Some(NullBuffer::new(n.inner() & n.inner())),
(None, None) => None,
}
}
/// This intentional inline(never) attribute helps LLVM optimize the loop.
#[inline(never)]
fn try_binary_no_nulls<A: ArrayAccessor, B: ArrayAccessor, F, O>(
len: usize,
a: A,
b: B,
op: F,
) -> Result<PrimitiveArray<O>, ArrowError>
where
O: ArrowPrimitiveType,
F: Fn(A::Item, B::Item) -> Result<O::Native, ArrowError>,
{
let mut buffer = MutableBuffer::new(len * O::Native::get_byte_width());
for idx in 0..len {
unsafe {
buffer.push_unchecked(op(a.value_unchecked(idx), b.value_unchecked(idx))?);
};
}
Ok(PrimitiveArray::new(buffer.into(), None))
}
/// This intentional inline(never) attribute helps LLVM optimize the loop.
#[inline(never)]
fn try_binary_no_nulls_mut<T, F>(
len: usize,
a: PrimitiveArray<T>,
b: &PrimitiveArray<T>,
op: F,
) -> Result<Result<PrimitiveArray<T>, ArrowError>, PrimitiveArray<T>>
where
T: ArrowPrimitiveType,
F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
{
let mut builder = a.into_builder()?;
let slice = builder.values_slice_mut();
for idx in 0..len {
unsafe {
match op(*slice.get_unchecked(idx), b.value_unchecked(idx)) {
Ok(value) => *slice.get_unchecked_mut(idx) = value,
Err(err) => return Ok(Err(err)),
};
};
}
Ok(Ok(builder.finish()))
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::types::*;
use std::sync::Arc;
#[test]
#[allow(deprecated)]
fn test_unary_f64_slice() {
let input = Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
let input_slice = input.slice(1, 4);
let result = unary(&input_slice, |n| n.round());
assert_eq!(
result,
Float64Array::from(vec![None, Some(7.0), None, Some(7.0)])
);
}
#[test]
fn test_binary_mut() {
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
let c = binary_mut(a, &b, |l, r| l + r).unwrap().unwrap();
let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
assert_eq!(c, expected);
}
#[test]
fn test_binary_mut_null_buffer() {
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
let b = Int32Array::from(vec![Some(10), Some(11), Some(12), Some(13), Some(14)]);
let r1 = binary_mut(a, &b, |a, b| a + b).unwrap();
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
let b = Int32Array::new(
vec![10, 11, 12, 13, 14].into(),
Some(vec![true, true, true, true, true].into()),
);
// unwrap here means that no copying occured
let r2 = binary_mut(a, &b, |a, b| a + b).unwrap();
assert_eq!(r1.unwrap(), r2.unwrap());
}
#[test]
fn test_try_binary_mut() {
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().unwrap();
let expected = Int32Array::from(vec![Some(16), None, Some(12), None, Some(6)]);
assert_eq!(c, expected);
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
let b = Int32Array::from(vec![1, 2, 3, 4, 5]);
let c = try_binary_mut(a, &b, |l, r| Ok(l + r)).unwrap().unwrap();
let expected = Int32Array::from(vec![16, 16, 12, 12, 6]);
assert_eq!(c, expected);
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
let b = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]);
let _ = try_binary_mut(a, &b, |l, r| {
if l == 1 {
Err(ArrowError::InvalidArgumentError(
"got error".parse().unwrap(),
))
} else {
Ok(l + r)
}
})
.unwrap()
.expect_err("should got error");
}
#[test]
fn test_try_binary_mut_null_buffer() {
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
let b = Int32Array::from(vec![Some(10), Some(11), Some(12), Some(13), Some(14)]);
let r1 = try_binary_mut(a, &b, |a, b| Ok(a + b)).unwrap();
let a = Int32Array::from(vec![Some(3), Some(4), Some(5), Some(6), None]);
let b = Int32Array::new(
vec![10, 11, 12, 13, 14].into(),
Some(vec![true, true, true, true, true].into()),
);
// unwrap here means that no copying occured
let r2 = try_binary_mut(a, &b, |a, b| Ok(a + b)).unwrap();
assert_eq!(r1.unwrap(), r2.unwrap());
}
#[test]
fn test_unary_dict_mut() {
let values = Int32Array::from(vec![Some(10), Some(20), None]);
let keys = Int8Array::from_iter_values([0, 0, 1, 2]);
let dictionary = DictionaryArray::new(keys, Arc::new(values));
let updated = dictionary.unary_mut::<_, Int32Type>(|x| x + 1).unwrap();
let typed = updated.downcast_dict::<Int32Array>().unwrap();
assert_eq!(typed.value(0), 11);
assert_eq!(typed.value(1), 11);
assert_eq!(typed.value(2), 21);
let values = updated.values();
assert!(values.is_null(2));
}
}
notes-service/vendor/arrow-arith/src/bitwise.rs
Vendored
+392
View File
@@ -0,0 +1,392 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Module contains bitwise operations on arrays
use crate::arity::{binary, unary};
use arrow_array::*;
use arrow_buffer::ArrowNativeType;
use arrow_schema::ArrowError;
use num_traits::{WrappingShl, WrappingShr};
use std::ops::{BitAnd, BitOr, BitXor, Not};
/// The helper function for bitwise operation with two array
fn bitwise_op<T, F>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
op: F,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
F: Fn(T::Native, T::Native) -> T::Native,
{
binary(left, right, op)
}
/// Perform `left & right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_and<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitAnd<Output = T::Native>,
{
bitwise_op(left, right, |a, b| a & b)
}
/// Perform `left | right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_or<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitOr<Output = T::Native>,
{
bitwise_op(left, right, |a, b| a | b)
}
/// Perform `left ^ right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_xor<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitXor<Output = T::Native>,
{
bitwise_op(left, right, |a, b| a ^ b)
}
/// Perform bitwise `left << right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_shift_left<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: WrappingShl<Output = T::Native>,
{
bitwise_op(left, right, |a, b| {
let b = b.as_usize();
a.wrapping_shl(b as u32)
})
}
/// Perform bitwise `left >> right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_shift_right<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: WrappingShr<Output = T::Native>,
{
bitwise_op(left, right, |a, b| {
let b = b.as_usize();
a.wrapping_shr(b as u32)
})
}
/// Perform `!array` operation on array. If array value is null
/// then the result is also null.
pub fn bitwise_not<T>(array: &PrimitiveArray<T>) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: Not<Output = T::Native>,
{
Ok(unary(array, |value| !value))
}
/// Perform `left & !right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn bitwise_and_not<T>(
left: &PrimitiveArray<T>,
right: &PrimitiveArray<T>,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitAnd<Output = T::Native>,
T::Native: Not<Output = T::Native>,
{
bitwise_op(left, right, |a, b| a & !b)
}
/// Perform bitwise `and` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_and_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitAnd<Output = T::Native>,
{
Ok(unary(array, |value| value & scalar))
}
/// Perform bitwise `or` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_or_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitOr<Output = T::Native>,
{
Ok(unary(array, |value| value | scalar))
}
/// Perform bitwise `xor` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_xor_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: BitXor<Output = T::Native>,
{
Ok(unary(array, |value| value ^ scalar))
}
/// Perform bitwise `left << right` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_shift_left_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: WrappingShl<Output = T::Native>,
{
Ok(unary(array, |value| {
let scalar = scalar.as_usize();
value.wrapping_shl(scalar as u32)
}))
}
/// Perform bitwise `left >> right` every value in an array with the scalar. If any value in the array is null then the
/// result is also null.
pub fn bitwise_shift_right_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>, ArrowError>
where
T: ArrowNumericType,
T::Native: WrappingShr<Output = T::Native>,
{
Ok(unary(array, |value| {
let scalar = scalar.as_usize();
value.wrapping_shr(scalar as u32)
}))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_bitwise_and_array() -> Result<(), ArrowError> {
// unsigned value
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12)]);
let expected = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let result = bitwise_and(&left, &right)?;
assert_eq!(expected, result);
// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = Int32Array::from(vec![Some(5), Some(-10), Some(8), Some(12)]);
let expected = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let result = bitwise_and(&left, &right)?;
assert_eq!(expected, result);
Ok(())
}
#[test]
fn test_bitwise_shift_left() {
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(8)]);
let right = UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12), Some(u64::MAX)]);
let expected = UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(0)]);
let result = bitwise_shift_left(&left, &right).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_shift_left_scalar() {
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(8)]);
let scalar = 2;
let expected = UInt64Array::from(vec![Some(4), Some(8), None, Some(16), Some(32)]);
let result = bitwise_shift_left_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_shift_right() {
let left = UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(3)]);
let right = UInt64Array::from(vec![Some(5), Some(10), Some(8), Some(12), Some(65)]);
let expected = UInt64Array::from(vec![Some(1), Some(2), None, Some(4), Some(1)]);
let result = bitwise_shift_right(&left, &right).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_shift_right_scalar() {
let left = UInt64Array::from(vec![Some(32), Some(2048), None, Some(16384), Some(3)]);
let scalar = 2;
let expected = UInt64Array::from(vec![Some(8), Some(512), None, Some(4096), Some(0)]);
let result = bitwise_shift_right_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_and_array_scalar() {
// unsigned value
let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
let scalar = 7;
let expected = UInt64Array::from(vec![Some(7), Some(2), None, Some(4)]);
let result = bitwise_and_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let scalar = -20;
let expected = Int32Array::from(vec![Some(0), Some(0), None, Some(4)]);
let result = bitwise_and_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_or_array() {
// unsigned value
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
let expected = UInt64Array::from(vec![Some(7), Some(7), None, Some(13)]);
let result = bitwise_or(&left, &right).unwrap();
assert_eq!(expected, result);
// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
let expected = Int32Array::from(vec![Some(-7), Some(-5), None, Some(13)]);
let result = bitwise_or(&left, &right).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_not_array() {
// unsigned value
let array = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let expected = UInt64Array::from(vec![
Some(18446744073709551614),
Some(18446744073709551613),
None,
Some(18446744073709551611),
]);
let result = bitwise_not(&array).unwrap();
assert_eq!(expected, result);
// signed value
let array = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let expected = Int32Array::from(vec![Some(-2), Some(-3), None, Some(-5)]);
let result = bitwise_not(&array).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_and_not_array() {
// unsigned value
let left = UInt64Array::from(vec![Some(8), Some(2), None, Some(4)]);
let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
let expected = UInt64Array::from(vec![Some(8), Some(2), None, Some(0)]);
let result = bitwise_and_not(&left, &right).unwrap();
assert_eq!(expected, result);
assert_eq!(
bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
result
);
// signed value
let left = Int32Array::from(vec![Some(2), Some(1), None, Some(3)]);
let right = Int32Array::from(vec![Some(-7), Some(-5), Some(8), Some(13)]);
let expected = Int32Array::from(vec![Some(2), Some(0), None, Some(2)]);
let result = bitwise_and_not(&left, &right).unwrap();
assert_eq!(expected, result);
assert_eq!(
bitwise_and(&left, &bitwise_not(&right).unwrap()).unwrap(),
result
);
}
#[test]
fn test_bitwise_or_array_scalar() {
// unsigned value
let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
let scalar = 7;
let expected = UInt64Array::from(vec![Some(15), Some(7), None, Some(7)]);
let result = bitwise_or_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let scalar = 20;
let expected = Int32Array::from(vec![Some(21), Some(22), None, Some(20)]);
let result = bitwise_or_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_xor_array() {
// unsigned value
let left = UInt64Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = UInt64Array::from(vec![Some(7), Some(5), Some(8), Some(13)]);
let expected = UInt64Array::from(vec![Some(6), Some(7), None, Some(9)]);
let result = bitwise_xor(&left, &right).unwrap();
assert_eq!(expected, result);
// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let right = Int32Array::from(vec![Some(-7), Some(5), Some(8), Some(-13)]);
let expected = Int32Array::from(vec![Some(-8), Some(7), None, Some(-9)]);
let result = bitwise_xor(&left, &right).unwrap();
assert_eq!(expected, result);
}
#[test]
fn test_bitwise_xor_array_scalar() {
// unsigned value
let left = UInt64Array::from(vec![Some(15), Some(2), None, Some(4)]);
let scalar = 7;
let expected = UInt64Array::from(vec![Some(8), Some(5), None, Some(3)]);
let result = bitwise_xor_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
// signed value
let left = Int32Array::from(vec![Some(1), Some(2), None, Some(4)]);
let scalar = -20;
let expected = Int32Array::from(vec![Some(-19), Some(-18), None, Some(-24)]);
let result = bitwise_xor_scalar(&left, scalar).unwrap();
assert_eq!(expected, result);
}
}
notes-service/vendor/arrow-arith/src/boolean.rs
Vendored
+974
View File
@@ -0,0 +1,974 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Defines boolean kernels on Arrow `BooleanArray`'s, e.g. `AND`, `OR` and `NOT`.
//!
//! These kernels can leverage SIMD if available on your system. Currently no runtime
//! detection is provided, you should enable the specific SIMD intrinsics using
//! `RUSTFLAGS="-C target-feature=+avx2"` for example. See the documentation
//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
use arrow_array::*;
use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper};
use arrow_buffer::{BooleanBuffer, NullBuffer, buffer_bin_and_not};
use arrow_schema::ArrowError;
/// Logical 'and' boolean values with Kleene logic
///
/// # Behavior
///
/// This function behaves as follows with nulls:
///
/// * `true` and `null` = `null`
/// * `null` and `true` = `null`
/// * `false` and `null` = `false`
/// * `null` and `false` = `false`
/// * `null` and `null` = `null`
///
/// In other words, in this context a null value really means \"unknown\",
/// and an unknown value 'and' false is always false.
/// For a different null behavior, see function \"and\".
///
/// # Example
///
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::and_kleene;
/// let a = BooleanArray::from(vec![Some(true), Some(false), None]);
/// let b = BooleanArray::from(vec![None, None, None]);
/// let and_ab = and_kleene(&a, &b).unwrap();
/// assert_eq!(and_ab, BooleanArray::from(vec![None, Some(false), None]));
/// ```
///
/// # Fails
///
/// If the operands have different lengths
pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
let left_values = left.values();
let right_values = right.values();
let buffer = match (left.nulls(), right.nulls()) {
(None, None) => None,
(Some(left_null_buffer), None) => {
// The right side has no null values.
// The final null bit is set only if:
// 1. left null bit is set, or
// 2. right data bit is false (because null AND false = false).
Some(bitwise_bin_op_helper(
left_null_buffer.buffer(),
left_null_buffer.offset(),
right_values.inner(),
right_values.offset(),
left.len(),
|a, b| a | !b,
))
}
(None, Some(right_null_buffer)) => {
// Same as above
Some(bitwise_bin_op_helper(
right_null_buffer.buffer(),
right_null_buffer.offset(),
left_values.inner(),
left_values.offset(),
left.len(),
|a, b| a | !b,
))
}
(Some(left_null_buffer), Some(right_null_buffer)) => {
// Follow the same logic above. Both sides have null values.
// Assume a is left null bits, b is left data bits, c is right null bits,
// d is right data bits.
// The final null bits are:
// (a | (c & !d)) & (c | (a & !b))
Some(bitwise_quaternary_op_helper(
[
left_null_buffer.buffer(),
left_values.inner(),
right_null_buffer.buffer(),
right_values.inner(),
],
[
left_null_buffer.offset(),
left_values.offset(),
right_null_buffer.offset(),
right_values.offset(),
],
left.len(),
|a, b, c, d| (a | (c & !d)) & (c | (a & !b)),
))
}
};
let nulls = buffer.map(|b| NullBuffer::new(BooleanBuffer::new(b, 0, left.len())));
Ok(BooleanArray::new(left_values & right_values, nulls))
}
/// Logical 'or' boolean values with Kleene logic
///
/// # Behavior
///
/// This function behaves as follows with nulls:
///
/// * `true` or `null` = `true`
/// * `null` or `true` = `true`
/// * `false` or `null` = `null`
/// * `null` or `false` = `null`
/// * `null` or `null` = `null`
///
/// In other words, in this context a null value really means \"unknown\",
/// and an unknown value 'or' true is always true.
/// For a different null behavior, see function \"or\".
///
/// # Example
///
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::or_kleene;
/// let a = BooleanArray::from(vec![Some(true), Some(false), None]);
/// let b = BooleanArray::from(vec![None, None, None]);
/// let or_ab = or_kleene(&a, &b).unwrap();
/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), None, None]));
/// ```
///
/// # Fails
///
/// If the operands have different lengths
pub fn or_kleene(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
let left_values = left.values();
let right_values = right.values();
let buffer = match (left.nulls(), right.nulls()) {
(None, None) => None,
(Some(left_nulls), None) => {
// The right side has no null values.
// The final null bit is set only if:
// 1. left null bit is set, or
// 2. right data bit is true (because null OR true = true).
Some(bitwise_bin_op_helper(
left_nulls.buffer(),
left_nulls.offset(),
right_values.inner(),
right_values.offset(),
left.len(),
|a, b| a | b,
))
}
(None, Some(right_nulls)) => {
// Same as above
Some(bitwise_bin_op_helper(
right_nulls.buffer(),
right_nulls.offset(),
left_values.inner(),
left_values.offset(),
left.len(),
|a, b| a | b,
))
}
(Some(left_nulls), Some(right_nulls)) => {
// Follow the same logic above. Both sides have null values.
// Assume a is left null bits, b is left data bits, c is right null bits,
// d is right data bits.
// The final null bits are:
// (a | (c & d)) & (c | (a & b))
Some(bitwise_quaternary_op_helper(
[
left_nulls.buffer(),
left_values.inner(),
right_nulls.buffer(),
right_values.inner(),
],
[
left_nulls.offset(),
left_values.offset(),
right_nulls.offset(),
right_values.offset(),
],
left.len(),
|a, b, c, d| (a | (c & d)) & (c | (a & b)),
))
}
};
let nulls = buffer.map(|b| NullBuffer::new(BooleanBuffer::new(b, 0, left.len())));
Ok(BooleanArray::new(left_values | right_values, nulls))
}
/// Helper function to implement binary kernels
pub(crate) fn binary_boolean_kernel<F>(
left: &BooleanArray,
right: &BooleanArray,
op: F,
) -> Result<BooleanArray, ArrowError>
where
F: Fn(&BooleanBuffer, &BooleanBuffer) -> BooleanBuffer,
{
if left.len() != right.len() {
return Err(ArrowError::ComputeError(
"Cannot perform bitwise operation on arrays of different length".to_string(),
));
}
let nulls = NullBuffer::union(left.nulls(), right.nulls());
let values = op(left.values(), right.values());
Ok(BooleanArray::new(values, nulls))
}
/// Performs `AND` operation on two arrays. If either left or right value is null then the
/// result is also null.
/// # Error
/// This function errors when the arrays have different lengths.
/// # Example
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::and;
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
/// let and_ab = and(&a, &b).unwrap();
/// assert_eq!(and_ab, BooleanArray::from(vec![Some(false), Some(true), None]));
/// ```
pub fn and(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
binary_boolean_kernel(left, right, |a, b| a & b)
}
/// Performs `OR` operation on two arrays. If either left or right value is null then the
/// result is also null.
/// # Error
/// This function errors when the arrays have different lengths.
/// # Example
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::or;
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
/// let or_ab = or(&a, &b).unwrap();
/// assert_eq!(or_ab, BooleanArray::from(vec![Some(true), Some(true), None]));
/// ```
pub fn or(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
binary_boolean_kernel(left, right, |a, b| a | b)
}
/// Performs `AND_NOT` operation on two arrays. If either left or right value is null then the
/// result is also null.
/// # Error
/// This function errors when the arrays have different lengths.
/// # Example
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::{and, not, and_not};
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let b = BooleanArray::from(vec![Some(true), Some(true), Some(false)]);
/// let andn_ab = and_not(&a, &b).unwrap();
/// assert_eq!(andn_ab, BooleanArray::from(vec![Some(false), Some(false), None]));
/// // It's equal to and(left, not(right))
/// assert_eq!(andn_ab, and(&a, &not(&b).unwrap()).unwrap());
pub fn and_not(left: &BooleanArray, right: &BooleanArray) -> Result<BooleanArray, ArrowError> {
binary_boolean_kernel(left, right, |a, b| {
let buffer = buffer_bin_and_not(a.inner(), b.offset(), b.inner(), a.offset(), a.len());
BooleanBuffer::new(buffer, left.offset(), left.len())
})
}
/// Performs unary `NOT` operation on an arrays. If value is null then the result is also
/// null.
/// # Error
/// This function never errors. It returns an error for consistency.
/// # Example
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::not;
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let not_a = not(&a).unwrap();
/// assert_eq!(not_a, BooleanArray::from(vec![Some(true), Some(false), None]));
/// ```
pub fn not(left: &BooleanArray) -> Result<BooleanArray, ArrowError> {
let nulls = left.nulls().cloned();
let values = !left.values();
Ok(BooleanArray::new(values, nulls))
}
/// Returns a non-null [BooleanArray] with whether each value of the array is null.
/// # Error
/// This function never errors.
/// # Example
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::is_null;
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let a_is_null = is_null(&a).unwrap();
/// assert_eq!(a_is_null, BooleanArray::from(vec![false, false, true]));
/// ```
pub fn is_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
let values = match input.logical_nulls() {
None => BooleanBuffer::new_unset(input.len()),
Some(nulls) => !nulls.inner(),
};
Ok(BooleanArray::new(values, None))
}
/// Returns a non-null [BooleanArray] with whether each value of the array is not null.
/// # Error
/// This function never errors.
/// # Example
/// ```rust
/// # use arrow_array::BooleanArray;
/// # use arrow_arith::boolean::is_not_null;
/// let a = BooleanArray::from(vec![Some(false), Some(true), None]);
/// let a_is_not_null = is_not_null(&a).unwrap();
/// assert_eq!(a_is_not_null, BooleanArray::from(vec![true, true, false]));
/// ```
pub fn is_not_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
let values = match input.logical_nulls() {
None => BooleanBuffer::new_set(input.len()),
Some(n) => n.inner().clone(),
};
Ok(BooleanArray::new(values, None))
}
#[cfg(test)]
mod tests {
use arrow_buffer::ScalarBuffer;
use arrow_schema::{DataType, Field, UnionFields};
use super::*;
use std::sync::Arc;
#[test]
fn test_bool_array_and() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![false, true, false, true]);
let c = and(&a, &b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![false, true, false, true]);
let c = or(&a, &b).unwrap();
let expected = BooleanArray::from(vec![false, true, true, true]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_and_not() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![false, true, false, true]);
let c = and_not(&a, &b).unwrap();
let expected = BooleanArray::from(vec![false, false, true, false]);
assert_eq!(c, expected);
assert_eq!(c, and(&a, &not(&b).unwrap()).unwrap());
}
#[test]
fn test_bool_array_or_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = or(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
None,
None,
None,
Some(false),
Some(true),
None,
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_boolean_array_kleene_no_remainder() {
let n = 1024;
let a = BooleanArray::from(vec![true; n]);
let b = BooleanArray::from(vec![None; n]);
let result = or_kleene(&a, &b).unwrap();
assert_eq!(result, a);
}
#[test]
fn test_bool_array_and_kleene_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = and_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
Some(false),
None,
Some(false),
Some(false),
Some(false),
None,
Some(false),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or_kleene_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = or_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
None,
Some(true),
None,
Some(false),
Some(true),
Some(true),
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or_kleene_right_sided_nulls() {
let a = BooleanArray::from(vec![false, false, false, true, true, true]);
// ensure null bitmap of a is absent
assert!(a.nulls().is_none());
let b = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(false),
None,
]);
// ensure null bitmap of b is present
assert!(b.nulls().is_some());
let c = or_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_or_kleene_left_sided_nulls() {
let a = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(false),
None,
]);
// ensure null bitmap of b is absent
assert!(a.nulls().is_some());
let b = BooleanArray::from(vec![false, false, false, true, true, true]);
// ensure null bitmap of a is present
assert!(b.nulls().is_none());
let c = or_kleene(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
Some(true),
Some(false),
None,
Some(true),
Some(true),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_not() {
let a = BooleanArray::from(vec![false, true]);
let c = not(&a).unwrap();
let expected = BooleanArray::from(vec![true, false]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_not_sliced() {
let a = BooleanArray::from(vec![None, Some(true), Some(false), None, Some(true)]);
let a = a.slice(1, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = not(a).unwrap();
let expected = BooleanArray::from(vec![Some(false), Some(true), None, Some(false)]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_and_nulls() {
let a = BooleanArray::from(vec![
None,
None,
None,
Some(false),
Some(false),
Some(false),
Some(true),
Some(true),
Some(true),
]);
let b = BooleanArray::from(vec![
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
None,
Some(false),
Some(true),
]);
let c = and(&a, &b).unwrap();
let expected = BooleanArray::from(vec![
None,
None,
None,
None,
Some(false),
Some(false),
None,
Some(false),
Some(true),
]);
assert_eq!(c, expected);
}
#[test]
fn test_bool_array_and_sliced_same_offset() {
let a = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, false, true, true,
]);
let b = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, true, false, true,
]);
let a = a.slice(8, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let b = b.slice(8, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_sliced_same_offset_mod8() {
let a = BooleanArray::from(vec![
false, false, true, true, false, false, false, false, false, false, false, false,
]);
let b = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, true, false, true,
]);
let a = a.slice(0, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let b = b.slice(8, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_sliced_offset1() {
let a = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, false, true, true,
]);
let b = BooleanArray::from(vec![false, true, false, true]);
let a = a.slice(8, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, &b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_sliced_offset2() {
let a = BooleanArray::from(vec![false, false, true, true]);
let b = BooleanArray::from(vec![
false, false, false, false, false, false, false, false, false, true, false, true,
]);
let b = b.slice(8, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(&a, b).unwrap();
let expected = BooleanArray::from(vec![false, false, false, true]);
assert_eq!(expected, c);
}
#[test]
fn test_bool_array_and_nulls_offset() {
let a = BooleanArray::from(vec![None, Some(false), Some(true), None, Some(true)]);
let a = a.slice(1, 4);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
let b = BooleanArray::from(vec![
None,
None,
Some(true),
Some(false),
Some(true),
Some(true),
]);
let b = b.slice(2, 4);
let b = b.as_any().downcast_ref::<BooleanArray>().unwrap();
let c = and(a, b).unwrap();
let expected = BooleanArray::from(vec![Some(false), Some(false), None, Some(true)]);
assert_eq!(expected, c);
}
#[test]
fn test_nonnull_array_is_null() {
let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
let res = is_null(a.as_ref()).unwrap();
let expected = BooleanArray::from(vec![false, false, false, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nonnull_array_with_offset_is_null() {
let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
let a = a.slice(8, 4);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, false, false, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nonnull_array_is_not_null() {
let a = Int32Array::from(vec![1, 2, 3, 4]);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, true, true, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nonnull_array_with_offset_is_not_null() {
let a = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1]);
let a = a.slice(8, 4);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, true, true, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_is_null() {
let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, true, false, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_with_offset_is_null() {
let a = Int32Array::from(vec![
None,
None,
None,
None,
None,
None,
None,
None,
// offset 8, previous None values are skipped by the slice
Some(1),
None,
Some(2),
None,
Some(3),
Some(4),
None,
None,
]);
let a = a.slice(8, 4);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, true, false, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_is_not_null() {
let a = Int32Array::from(vec![Some(1), None, Some(3), None]);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, false, true, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_nullable_array_with_offset_is_not_null() {
let a = Int32Array::from(vec![
None,
None,
None,
None,
None,
None,
None,
None,
// offset 8, previous None values are skipped by the slice
Some(1),
None,
Some(2),
None,
Some(3),
Some(4),
None,
None,
]);
let a = a.slice(8, 4);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, false, true, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_null_array_is_null() {
let a = NullArray::new(3);
let res = is_null(&a).unwrap();
let expected = BooleanArray::from(vec![true, true, true]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_null_array_is_not_null() {
let a = NullArray::new(3);
let res = is_not_null(&a).unwrap();
let expected = BooleanArray::from(vec![false, false, false]);
assert_eq!(expected, res);
assert!(res.nulls().is_none());
}
#[test]
fn test_dense_union_is_null() {
// union of [{A=1}, {A=}, {B=3.2}, {B=}, {C="a"}, {C=}]
let int_array = Int32Array::from(vec![Some(1), None]);
let float_array = Float64Array::from(vec![Some(3.2), None]);
let str_array = StringArray::from(vec![Some("a"), None]);
let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
let offsets = [0, 1, 0, 1, 0, 1]
.into_iter()
.collect::<ScalarBuffer<i32>>();
let children = vec![
Arc::new(int_array) as Arc<dyn Array>,
Arc::new(float_array),
Arc::new(str_array),
];
let array = UnionArray::try_new(union_fields(), type_ids, Some(offsets), children).unwrap();
let result = is_null(&array).unwrap();
let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
assert_eq!(expected, &result);
}
#[test]
fn test_sparse_union_is_null() {
// union of [{A=1}, {A=}, {B=3.2}, {B=}, {C="a"}, {C=}]
let int_array = Int32Array::from(vec![Some(1), None, None, None, None, None]);
let float_array = Float64Array::from(vec![None, None, Some(3.2), None, None, None]);
let str_array = StringArray::from(vec![None, None, None, None, Some("a"), None]);
let type_ids = [0, 0, 1, 1, 2, 2].into_iter().collect::<ScalarBuffer<i8>>();
let children = vec![
Arc::new(int_array) as Arc<dyn Array>,
Arc::new(float_array),
Arc::new(str_array),
];
let array = UnionArray::try_new(union_fields(), type_ids, None, children).unwrap();
let result = is_null(&array).unwrap();
let expected = &BooleanArray::from(vec![false, true, false, true, false, true]);
assert_eq!(expected, &result);
}
fn union_fields() -> UnionFields {
[
(0, Arc::new(Field::new("A", DataType::Int32, true))),
(1, Arc::new(Field::new("B", DataType::Float64, true))),
(2, Arc::new(Field::new("C", DataType::Utf8, true))),
]
.into_iter()
.collect()
}
}
notes-service/vendor/arrow-arith/src/lib.rs
Vendored
+33
View File
@@ -0,0 +1,33 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Arrow arithmetic and aggregation kernels
#![doc(
html_logo_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_white-bg.svg",
html_favicon_url = "https://arrow.apache.org/img/arrow-logo_chevrons_black-txt_transparent-bg.svg"
)]
#![cfg_attr(docsrs, feature(doc_cfg))]
#![warn(missing_docs)]
pub mod aggregate;
#[doc(hidden)] // Kernels to be removed in a future release
pub mod arithmetic;
pub mod arity;
pub mod bitwise;
pub mod boolean;
pub mod numeric;
pub mod temporal;
notes-service/vendor/arrow-arith/src/numeric.rs
Vendored
+2178
View File
File diff suppressed because it is too large Load Diff
notes-service/vendor/arrow-arith/src/temporal.rs
Vendored
+1972
View File
File diff suppressed because it is too large Load Diff