include!(concat!(env!("OUT_DIR"), "/simd_lanes.rs"));
use std::marker::PhantomData;
use minarrow::{
Bitmask, BooleanAVT, BooleanArray, FloatArray, Integer, IntegerArray, Length, MaskedArray,
Offset, StringArray, Vec64,
aliases::{FloatAVT, IntegerAVT},
enums::error::KernelError,
vec64,
};
use num_traits::{Float, Num, NumCast, One, Zero};
use minarrow::StringAVT;
use minarrow::utils::confirm_mask_capacity;
use crate::kernels::sort::total_cmp_f;
#[inline(always)]
fn new_null_mask(len: usize) -> Bitmask {
Bitmask::new_set_all(len, false)
}
#[inline(always)]
fn prealloc_vec<T: Copy>(len: usize) -> Vec64<T> {
let mut v = Vec64::<T>::with_capacity(len);
unsafe { v.set_len(len) };
v
}
#[inline(always)]
fn rolling_push_pop_to<T, FAdd, FRem>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
mut add: FAdd,
mut remove: FRem,
zero: T,
out: &mut [T],
out_mask: &mut Bitmask,
) where
T: Copy,
FAdd: FnMut(T, T) -> T,
FRem: FnMut(T, T) -> T,
{
let n = data.len();
assert_eq!(
n,
out.len(),
"rolling_push_pop_to: input/output length mismatch"
);
if subwindow == 0 {
for slot in out.iter_mut() {
*slot = zero;
}
return;
}
let mut agg = zero;
let mut invalids = 0usize;
for i in 0..n {
if mask.map_or(true, |m| unsafe { m.get_unchecked(i) }) {
agg = add(agg, data[i]);
} else {
invalids += 1;
}
if i + 1 > subwindow {
let j = i + 1 - subwindow - 1;
if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
agg = remove(agg, data[j]);
} else {
invalids -= 1;
}
}
if i + 1 < subwindow {
unsafe { out_mask.set_unchecked(i, false) };
out[i] = zero;
} else {
let ok = invalids == 0;
unsafe { out_mask.set_unchecked(i, ok) };
out[i] = agg;
}
}
}
#[inline(always)]
pub fn rolling_push_pop<T, FAdd, FRem>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
add: FAdd,
remove: FRem,
zero: T,
) -> (Vec64<T>, Bitmask)
where
T: Copy,
FAdd: FnMut(T, T) -> T,
FRem: FnMut(T, T) -> T,
{
let n = data.len();
let mut out = prealloc_vec::<T>(n);
let mut out_mask = new_null_mask(n);
rolling_push_pop_to(
data,
mask,
subwindow,
add,
remove,
zero,
&mut out,
&mut out_mask,
);
(out, out_mask)
}
#[inline(always)]
pub fn rolling_extreme_to<T, F>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
mut better: F,
zero: T,
out: &mut [T],
out_mask: &mut Bitmask,
) where
T: Copy,
F: FnMut(&T, &T) -> bool,
{
let n = data.len();
assert_eq!(
n,
out.len(),
"rolling_extreme_to: input/output length mismatch"
);
if subwindow == 0 {
return;
}
for i in 0..n {
if i + 1 < subwindow {
unsafe { out_mask.set_unchecked(i, false) };
out[i] = zero;
continue;
}
let start = i + 1 - subwindow;
let mut found = false;
let mut extreme = zero;
for j in start..=i {
if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
if !found {
extreme = data[j];
found = true;
} else if better(&data[j], &extreme) {
extreme = data[j];
}
} else {
found = false;
break;
}
}
unsafe { out_mask.set_unchecked(i, found) };
out[i] = if found { extreme } else { zero };
}
}
#[inline(always)]
pub fn rolling_extreme<T, F>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
better: F,
zero: T,
) -> (Vec64<T>, Bitmask)
where
T: Copy,
F: FnMut(&T, &T) -> bool,
{
let n = data.len();
let mut out = prealloc_vec::<T>(n);
let mut out_mask = new_null_mask(n);
rolling_extreme_to(data, mask, subwindow, better, zero, &mut out, &mut out_mask);
(out, out_mask)
}
#[inline]
pub fn rolling_sum_int_to<T: Num + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
rolling_push_pop_to(
data,
mask,
subwindow,
|a, b| a + b,
|a, b| a - b,
T::zero(),
out,
out_mask,
);
if mask.is_some() && subwindow > 0 && subwindow - 1 < out.len() {
unsafe { out_mask.set_unchecked(subwindow - 1, false) };
out[subwindow - 1] = T::zero();
}
}
#[inline]
pub fn rolling_sum_int<T: Num + Copy + Zero>(
window: IntegerAVT<'_, T>,
subwindow: usize,
) -> IntegerArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_sum_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_sum_float_to<T: Float + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
rolling_push_pop_to(
data,
mask,
subwindow,
|a, b| a + b,
|a, b| a - b,
T::zero(),
out,
out_mask,
);
if subwindow > 0 && subwindow - 1 < out.len() {
out_mask.set(subwindow - 1, false);
out[subwindow - 1] = T::zero();
}
}
#[inline]
pub fn rolling_sum_float<T: Float + Copy + Zero>(
window: FloatAVT<'_, T>,
subwindow: usize,
) -> FloatArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_sum_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
FloatArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_sum_bool_to(
data: &[i32],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [i32],
out_mask: &mut Bitmask,
) {
rolling_push_pop_to(
data,
mask,
subwindow,
|a, b| a + b,
|a, b| a - b,
0,
out,
out_mask,
);
if subwindow > 0 && subwindow - 1 < out.len() {
out_mask.set(subwindow - 1, false);
out[subwindow - 1] = 0;
}
}
#[inline]
pub fn rolling_sum_bool(window: BooleanAVT<'_, ()>, subwindow: usize) -> IntegerArray<i32> {
let (arr, offset, len) = window;
let bools: Vec<i32> = arr.iter_range(offset, len).map(|b| b as i32).collect();
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<i32>(len);
let mut out_mask = new_null_mask(len);
rolling_sum_bool_to(&bools, mask.as_ref(), subwindow, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_product_int_to<T: Num + Copy + One + Zero + PartialEq>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
let n = data.len();
assert_eq!(
n,
out.len(),
"rolling_product_int_to: input/output length mismatch"
);
if subwindow == 0 {
for slot in out.iter_mut() {
*slot = T::zero();
}
return;
}
let mut nz_product = T::one();
let mut zero_count = 0usize;
let mut invalids = 0usize;
for i in 0..n {
if mask.map_or(true, |m| unsafe { m.get_unchecked(i) }) {
if data[i] == T::zero() {
zero_count += 1;
} else {
nz_product = nz_product * data[i];
}
} else {
invalids += 1;
}
if i + 1 > subwindow {
let j = i + 1 - subwindow - 1;
if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
if data[j] == T::zero() {
zero_count -= 1;
} else {
nz_product = nz_product / data[j];
}
} else {
invalids -= 1;
}
}
if i + 1 < subwindow {
unsafe { out_mask.set_unchecked(i, false) };
out[i] = T::zero();
} else {
let ok = invalids == 0;
unsafe { out_mask.set_unchecked(i, ok) };
out[i] = if zero_count > 0 {
T::zero()
} else {
nz_product
};
}
}
}
#[inline]
pub fn rolling_product_int<T: Num + Copy + One + Zero>(
window: IntegerAVT<'_, T>,
subwindow: usize,
) -> IntegerArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_product_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_product_float_to<T: Float + Copy + One + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
rolling_push_pop_to(
data,
mask,
subwindow,
|a, b| a * b,
|a, b| a / b,
T::one(),
out,
out_mask,
);
}
#[inline]
pub fn rolling_product_float<T: Float + Copy + One + Zero>(
window: FloatAVT<'_, T>,
subwindow: usize,
) -> FloatArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_product_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
FloatArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_product_bool_to(
data: &[i32],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut Bitmask,
out_mask: &mut Bitmask,
) {
let n = data.len();
for i in 0..n {
let start = if i + 1 >= subwindow {
i + 1 - subwindow
} else {
0
};
let mut acc = true;
let mut valid = subwindow > 0 && i + 1 >= subwindow;
for j in start..=i {
let is_valid = mask.map_or(true, |m| unsafe { m.get_unchecked(j) });
if is_valid {
acc &= data[j] != 0;
} else {
valid = false;
break;
}
}
unsafe { out_mask.set_unchecked(i, valid) };
out.set(i, valid && acc);
}
}
#[inline]
pub fn rolling_product_bool(window: BooleanAVT<'_, ()>, subwindow: usize) -> BooleanArray<()> {
let (arr, offset, len) = window;
let n = len;
let mut out_mask = new_null_mask(n);
let mut out = Bitmask::new_set_all(n, false);
for i in 0..n {
let start = if i + 1 >= subwindow {
i + 1 - subwindow
} else {
0
};
let mut acc = true;
let mut valid = subwindow > 0 && i + 1 >= subwindow;
for j in start..=i {
match unsafe { arr.get_unchecked(offset + j) } {
Some(val) => acc &= val,
None => {
valid = false;
break;
}
}
}
unsafe { out_mask.set_unchecked(i, valid) };
out.set(i, valid && acc);
}
BooleanArray {
data: out.into(),
null_mask: Some(out_mask),
len: n,
_phantom: PhantomData,
}
}
#[inline]
pub fn rolling_mean_int_to<T: NumCast + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [f64],
out_mask: &mut Bitmask,
) {
let n = data.len();
if subwindow == 0 {
return;
}
for i in 0..n {
if i + 1 < subwindow {
unsafe { out_mask.set_unchecked(i, false) };
out[i] = 0.0;
continue;
}
let start = i + 1 - subwindow;
let mut sum = 0.0;
let mut valid = true;
for j in start..=i {
if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
sum += num_traits::cast(data[j]).unwrap_or(0.0);
} else {
valid = false;
break;
}
}
unsafe { out_mask.set_unchecked(i, valid) };
out[i] = if valid { sum / subwindow as f64 } else { 0.0 };
}
if subwindow > 0 && subwindow - 1 < out.len() {
unsafe { out_mask.set_unchecked(subwindow - 1, false) };
out[subwindow - 1] = 0.0;
}
}
#[inline]
pub fn rolling_mean_int<T: NumCast + Copy + Zero>(
window: IntegerAVT<'_, T>,
subwindow: usize,
) -> FloatArray<f64> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<f64>(len);
let mut out_mask = new_null_mask(len);
rolling_mean_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
FloatArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_mean_float_to<T: Float + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
let n = data.len();
if subwindow == 0 {
return;
}
for i in 0..n {
if i + 1 < subwindow {
unsafe { out_mask.set_unchecked(i, false) };
out[i] = T::zero();
continue;
}
let start = i + 1 - subwindow;
let mut sum = T::zero();
let mut valid = true;
for j in start..=i {
if mask.map_or(true, |m| unsafe { m.get_unchecked(j) }) {
sum = sum + data[j];
} else {
valid = false;
break;
}
}
unsafe { out_mask.set_unchecked(i, valid) };
out[i] = if valid {
sum / T::from(subwindow as u32).unwrap()
} else {
T::zero()
};
}
if subwindow > 0 && subwindow - 1 < out.len() {
unsafe { out_mask.set_unchecked(subwindow - 1, false) };
out[subwindow - 1] = T::zero();
}
}
#[inline]
pub fn rolling_mean_float<T: Float + Copy + Zero>(
window: FloatAVT<'_, T>,
subwindow: usize,
) -> FloatArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_mean_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
FloatArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_min_int_to<T: Ord + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
rolling_extreme_to(
data,
mask,
subwindow,
|a, b| a < b,
T::zero(),
out,
out_mask,
);
if subwindow > 0 && subwindow - 1 < out.len() {
out_mask.set(subwindow - 1, false);
out[subwindow - 1] = T::zero();
}
}
#[inline]
pub fn rolling_min_int<T: Ord + Copy + Zero>(
window: IntegerAVT<'_, T>,
subwindow: usize,
) -> IntegerArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_min_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_max_int_to<T: Ord + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
rolling_extreme_to(
data,
mask,
subwindow,
|a, b| a > b,
T::zero(),
out,
out_mask,
);
}
#[inline]
pub fn rolling_max_int<T: Ord + Copy + Zero>(
window: IntegerAVT<'_, T>,
subwindow: usize,
) -> IntegerArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_max_int_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_min_float_to<T: Float + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
rolling_extreme_to(
data,
mask,
subwindow,
|a, b| a < b,
T::zero(),
out,
out_mask,
);
if subwindow > 0 && subwindow - 1 < out.len() {
out_mask.set(subwindow - 1, false);
out[subwindow - 1] = T::zero();
}
}
#[inline]
pub fn rolling_min_float<T: Float + Copy + Zero>(
window: FloatAVT<'_, T>,
subwindow: usize,
) -> FloatArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_min_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
FloatArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn rolling_max_float_to<T: Float + Copy + Zero>(
data: &[T],
mask: Option<&Bitmask>,
subwindow: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
rolling_extreme_to(
data,
mask,
subwindow,
|a, b| a > b,
T::zero(),
out,
out_mask,
);
}
#[inline]
pub fn rolling_max_float<T: Float + Copy + Zero>(
window: FloatAVT<'_, T>,
subwindow: usize,
) -> FloatArray<T> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = new_null_mask(len);
rolling_max_float_to(data, mask.as_ref(), subwindow, &mut out, &mut out_mask);
FloatArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
#[inline]
pub fn rolling_count_to(len: usize, subwindow: usize, out: &mut [i32], out_mask: &mut Bitmask) {
for i in 0..len {
let start = if i + 1 >= subwindow {
i + 1 - subwindow
} else {
0
};
let count = (i - start + 1) as i32;
let valid_row = subwindow > 0 && i + 1 >= subwindow;
unsafe { out_mask.set_unchecked(i, valid_row) };
out[i] = if valid_row { count } else { 0 };
}
}
pub fn rolling_count(window: (Offset, Length), subwindow: usize) -> IntegerArray<i32> {
let (_offset, len) = window;
let mut out = prealloc_vec::<i32>(len);
let mut out_mask = new_null_mask(len);
for i in 0..len {
let start = if i + 1 >= subwindow {
i + 1 - subwindow
} else {
0
};
let count = (i - start + 1) as i32;
let valid_row = subwindow > 0 && i + 1 >= subwindow;
unsafe { out_mask.set_unchecked(i, valid_row) };
out[i] = if valid_row { count } else { 0 };
}
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline(always)]
fn rank_numeric_to<T, F>(
data: &[T],
mask: Option<&Bitmask>,
mut cmp: F,
out: &mut [i32],
out_mask: &mut Bitmask,
) where
T: Copy,
F: FnMut(&T, &T) -> std::cmp::Ordering,
{
let n = data.len();
let mut indices: Vec<usize> = (0..n).collect();
indices.sort_by(|&i, &j| cmp(&data[i], &data[j]));
for (rank, &i) in indices.iter().enumerate() {
if mask.map_or(true, |m| unsafe { m.get_unchecked(i) }) {
out[i] = (rank + 1) as i32;
unsafe { out_mask.set_unchecked(i, true) };
}
}
}
#[inline(always)]
fn rank_numeric<T, F>(data: &[T], mask: Option<&Bitmask>, cmp: F) -> IntegerArray<i32>
where
T: Copy,
F: FnMut(&T, &T) -> std::cmp::Ordering,
{
let n = data.len();
let mut out = vec64![0i32; n];
let mut out_mask = Bitmask::new_set_all(n, false);
rank_numeric_to(data, mask, cmp, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline(always)]
pub fn rank_int_to<T: Ord + Copy>(
data: &[T],
mask: Option<&Bitmask>,
out: &mut [i32],
out_mask: &mut Bitmask,
) {
rank_numeric_to(data, mask, T::cmp, out, out_mask);
}
#[inline(always)]
pub fn rank_float_to<T: Float + Copy>(
data: &[T],
mask: Option<&Bitmask>,
out: &mut [i32],
out_mask: &mut Bitmask,
) {
rank_numeric_to(data, mask, total_cmp_f, out, out_mask);
}
#[inline(always)]
pub fn rank_int<T: Ord + Copy>(window: IntegerAVT<T>) -> IntegerArray<i32> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let null_mask = if len != arr.data.len() {
&arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
} else {
&arr.null_mask
};
rank_numeric(data, null_mask.as_ref(), |a, b| a.cmp(b))
}
#[inline(always)]
pub fn rank_float<T: Float + Copy>(window: FloatAVT<T>) -> IntegerArray<i32> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let null_mask = if len != arr.data.len() {
&arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
} else {
&arr.null_mask
};
rank_numeric(data, null_mask.as_ref(), total_cmp_f)
}
#[inline(always)]
pub fn rank_str<T: Integer>(arr: StringAVT<T>) -> Result<IntegerArray<i32>, KernelError> {
let (array, offset, len) = arr;
let mask = array.null_mask.as_ref();
let _ = confirm_mask_capacity(array.len(), mask)?;
let mut tuples: Vec<(usize, &str)> = (0..len)
.filter(|&i| mask.map_or(true, |m| unsafe { m.get_unchecked(offset + i) }))
.map(|i| (i, unsafe { array.get_unchecked(offset + i) }.unwrap_or("")))
.collect();
tuples.sort_by(|a, b| a.1.cmp(&b.1));
let mut out = vec64![0i32; len];
let mut out_mask = new_null_mask(len);
for (rank, (i, _)) in tuples.iter().enumerate() {
out[*i] = (rank + 1) as i32;
unsafe { out_mask.set_unchecked(*i, true) };
}
Ok(IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
})
}
#[inline(always)]
fn dense_rank_numeric_to<T, F, G>(
data: &[T],
mask: Option<&Bitmask>,
mut sort: F,
mut eq: G,
out: &mut [i32],
out_mask: &mut Bitmask,
) -> Result<(), KernelError>
where
T: Copy,
F: FnMut(&T, &T) -> std::cmp::Ordering,
G: FnMut(&T, &T) -> bool,
{
let n = data.len();
let _ = confirm_mask_capacity(n, mask)?;
let mut uniqs: Vec<T> = (0..n)
.filter(|&i| mask.map_or(true, |m| unsafe { m.get_unchecked(i) }))
.map(|i| data[i])
.collect();
uniqs.sort_by(&mut sort);
uniqs.dedup_by(|a, b| eq(&*a, &*b));
for i in 0..n {
if mask.map_or(true, |m| unsafe { m.get_unchecked(i) }) {
let rank = uniqs.binary_search_by(|x| sort(x, &data[i])).unwrap() + 1;
out[i] = rank as i32;
unsafe { out_mask.set_unchecked(i, true) };
} else {
out[i] = 0;
}
}
Ok(())
}
#[inline(always)]
fn dense_rank_numeric<T, F, G>(
data: &[T],
mask: Option<&Bitmask>,
sort: F,
eq: G,
) -> Result<IntegerArray<i32>, KernelError>
where
T: Copy,
F: FnMut(&T, &T) -> std::cmp::Ordering,
G: FnMut(&T, &T) -> bool,
{
let n = data.len();
let mut out = prealloc_vec::<i32>(n);
let mut out_mask = Bitmask::new_set_all(n, false);
dense_rank_numeric_to(data, mask, sort, eq, &mut out, &mut out_mask)?;
Ok(IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
})
}
#[inline(always)]
#[inline(always)]
pub fn dense_rank_int_to<T: Ord + Copy>(
data: &[T],
mask: Option<&Bitmask>,
out: &mut [i32],
out_mask: &mut Bitmask,
) -> Result<(), KernelError> {
dense_rank_numeric_to(data, mask, T::cmp, |a, b| a == b, out, out_mask)
}
#[inline(always)]
pub fn dense_rank_float_to<T: Float + Copy>(
data: &[T],
mask: Option<&Bitmask>,
out: &mut [i32],
out_mask: &mut Bitmask,
) -> Result<(), KernelError> {
dense_rank_numeric_to(data, mask, total_cmp_f, |a, b| a == b, out, out_mask)
}
pub fn dense_rank_int<T: Ord + Copy>(
window: IntegerAVT<T>,
) -> Result<IntegerArray<i32>, KernelError> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let null_mask = if len != arr.data.len() {
&arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
} else {
&arr.null_mask
};
dense_rank_numeric(data, null_mask.as_ref(), |a, b| a.cmp(b), |a, b| a == b)
}
#[inline(always)]
pub fn dense_rank_float<T: Float + Copy>(
window: FloatAVT<T>,
) -> Result<IntegerArray<i32>, KernelError> {
let (arr, offset, len) = window;
let data = &arr.data[offset..offset + len];
let null_mask = if len != arr.data.len() {
&arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
} else {
&arr.null_mask
};
dense_rank_numeric(data, null_mask.as_ref(), total_cmp_f, |a, b| a == b)
}
#[inline(always)]
pub fn dense_rank_str<T: Integer>(arr: StringAVT<T>) -> Result<IntegerArray<i32>, KernelError> {
let (array, offset, len) = arr;
let mask = array.null_mask.as_ref();
let _ = confirm_mask_capacity(array.len(), mask)?;
let mut vals: Vec<&str> = (0..len)
.filter(|&i| mask.map_or(true, |m| unsafe { m.get_unchecked(offset + i) }))
.map(|i| unsafe { array.get_unchecked(offset + i) }.unwrap_or(""))
.collect();
vals.sort();
vals.dedup();
let mut out = prealloc_vec::<i32>(len);
let mut out_mask = Bitmask::new_set_all(len, false);
for i in 0..len {
let valid = mask.map_or(true, |m| unsafe { m.get_unchecked(offset + i) });
if valid {
let rank = vals
.binary_search(&unsafe { array.get_unchecked(offset + i) }.unwrap_or(""))
.unwrap()
+ 1;
out[i] = rank as i32;
unsafe { out_mask.set_unchecked(i, true) };
} else {
out[i] = 0;
}
}
Ok(IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
})
}
#[inline(always)]
fn shift_with_bounds_to<T: Copy>(
data: &[T],
mask: Option<&Bitmask>,
len: usize,
offset_fn: impl Fn(usize) -> Option<usize>,
default: T,
out: &mut [T],
out_mask: &mut Bitmask,
) {
assert_eq!(
len,
out.len(),
"shift_with_bounds_to: input/output length mismatch"
);
for i in 0..len {
if let Some(j) = offset_fn(i) {
out[i] = data[j];
let is_valid = mask.map_or(true, |m| unsafe { m.get_unchecked(j) });
unsafe { out_mask.set_unchecked(i, is_valid) };
} else {
out[i] = default;
}
}
}
#[inline(always)]
fn shift_with_bounds<T: Copy>(
data: &[T],
mask: Option<&Bitmask>,
len: usize,
offset_fn: impl Fn(usize) -> Option<usize>,
default: T,
) -> (Vec64<T>, Bitmask) {
let mut out = prealloc_vec::<T>(len);
let mut out_mask = Bitmask::new_set_all(len, false);
shift_with_bounds_to(data, mask, len, offset_fn, default, &mut out, &mut out_mask);
(out, out_mask)
}
#[inline(always)]
fn shift_str_with_bounds<T: Integer>(
arr: StringAVT<T>,
offset_fn: impl Fn(usize) -> Option<usize>,
) -> Result<StringArray<T>, KernelError> {
let (array, offset, len) = arr;
let src_mask = array.null_mask.as_ref();
let _ = confirm_mask_capacity(array.len(), src_mask)?;
let mut offsets = Vec64::<T>::with_capacity(len + 1);
unsafe {
offsets.set_len(len + 1);
}
offsets[0] = T::zero();
let mut total_bytes = 0;
let mut string_lengths = vec![0usize; len];
for i in 0..len {
let byte_len = if let Some(j) = offset_fn(i) {
let src_idx = offset + j;
let valid = src_mask.map_or(true, |m| unsafe { m.get_unchecked(src_idx) });
if valid {
unsafe { array.get_unchecked(src_idx).unwrap_or("").len() }
} else {
0
}
} else {
0
};
total_bytes += byte_len;
string_lengths[i] = byte_len;
offsets[i + 1] = T::from_usize(total_bytes);
}
let mut data = Vec64::<u8>::with_capacity(total_bytes);
let mut out_mask = Bitmask::new_set_all(len, false);
for i in 0..len {
if let Some(j) = offset_fn(i) {
let src_idx = offset + j;
let valid = src_mask.map_or(true, |m| unsafe { m.get_unchecked(src_idx) });
if valid {
let s = unsafe { array.get_unchecked(src_idx).unwrap_or("") };
data.extend_from_slice(s.as_bytes());
unsafe { out_mask.set_unchecked(i, true) };
continue;
}
}
}
Ok(StringArray {
offsets: offsets.into(),
data: data.into(),
null_mask: Some(out_mask),
})
}
#[inline]
pub fn lag_int_to<T: Copy + Default>(
data: &[T],
mask: Option<&Bitmask>,
n: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
let len = data.len();
shift_with_bounds_to(
data,
mask,
len,
|i| if i >= n { Some(i - n) } else { None },
T::default(),
out,
out_mask,
);
}
#[inline]
pub fn lag_int<T: Copy + Default>(window: IntegerAVT<T>, n: usize) -> IntegerArray<T> {
let (arr, offset, len) = window;
let data_window = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = Bitmask::new_set_all(len, false);
lag_int_to(data_window, mask.as_ref(), n, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn lead_int_to<T: Copy + Default>(
data: &[T],
mask: Option<&Bitmask>,
n: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
let len = data.len();
shift_with_bounds_to(
data,
mask,
len,
|i| if i + n < len { Some(i + n) } else { None },
T::default(),
out,
out_mask,
);
}
#[inline]
pub fn lead_int<T: Copy + Default>(window: IntegerAVT<T>, n: usize) -> IntegerArray<T> {
let (arr, offset, len) = window;
let data_window = &arr.data[offset..offset + len];
let mask = arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len));
let mut out = prealloc_vec::<T>(len);
let mut out_mask = Bitmask::new_set_all(len, false);
lead_int_to(data_window, mask.as_ref(), n, &mut out, &mut out_mask);
IntegerArray {
data: out.into(),
null_mask: Some(out_mask),
}
}
#[inline]
pub fn lag_float_to<T: Copy + num_traits::Zero>(
data: &[T],
mask: Option<&Bitmask>,
n: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
let len = data.len();
shift_with_bounds_to(
data,
mask,
len,
|i| if i >= n { Some(i - n) } else { None },
T::zero(),
out,
out_mask,
);
}
#[inline]
pub fn lead_float_to<T: Copy + num_traits::Zero>(
data: &[T],
mask: Option<&Bitmask>,
n: usize,
out: &mut [T],
out_mask: &mut Bitmask,
) {
let len = data.len();
shift_with_bounds_to(
data,
mask,
len,
|i| if i + n < len { Some(i + n) } else { None },
T::zero(),
out,
out_mask,
);
}
#[inline]
pub fn lag_float<T: Copy + num_traits::Zero>(window: FloatAVT<T>, n: usize) -> FloatArray<T> {
let (arr, offset, len) = window;
let data_window = &arr.data[offset..offset + len];
let mask_window = if len != arr.data.len() {
arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
} else {
arr.null_mask.clone()
};
let (data, null_mask) = shift_with_bounds(
data_window,
mask_window.as_ref(),
len,
|i| if i >= n { Some(i - n) } else { None },
T::zero(),
);
FloatArray {
data: data.into(),
null_mask: Some(null_mask),
}
}
#[inline]
pub fn lead_float<T: Copy + num_traits::Zero>(window: FloatAVT<T>, n: usize) -> FloatArray<T> {
let (arr, offset, len) = window;
let data_window = &arr.data[offset..offset + len];
let mask_window = if len != arr.data.len() {
arr.null_mask.as_ref().map(|m| m.slice_clone(offset, len))
} else {
arr.null_mask.clone()
};
let (data, null_mask) = shift_with_bounds(
data_window,
mask_window.as_ref(),
len,
|i| if i + n < len { Some(i + n) } else { None },
T::zero(),
);
FloatArray {
data: data.into(),
null_mask: Some(null_mask),
}
}
#[inline]
pub fn lag_str<T: Integer>(arr: StringAVT<T>, n: usize) -> Result<StringArray<T>, KernelError> {
shift_str_with_bounds(arr, |i| if i >= n { Some(i - n) } else { None })
}
#[inline]
pub fn lead_str<T: Integer>(arr: StringAVT<T>, n: usize) -> Result<StringArray<T>, KernelError> {
let (_array, _offset, len) = arr;
shift_str_with_bounds(arr, |i| if i + n < len { Some(i + n) } else { None })
}
#[inline(always)]
pub fn shift_int<T: Copy + Default>(window: IntegerAVT<T>, offset: isize) -> IntegerArray<T> {
let (arr, win_offset, win_len) = window;
if offset == 0 {
return IntegerArray {
data: Vec64::from_slice(&arr.data[win_offset..win_offset + win_len]).into(),
null_mask: if win_len != arr.data.len() {
arr.null_mask
.as_ref()
.map(|m| m.slice_clone(win_offset, win_len))
} else {
arr.null_mask.clone()
},
};
} else if offset > 0 {
lead_int((arr, win_offset, win_len), offset as usize)
} else {
lag_int((arr, win_offset, win_len), offset.unsigned_abs())
}
}
#[inline(always)]
pub fn shift_float<T: Copy + num_traits::Zero>(
window: FloatAVT<T>,
offset: isize,
) -> FloatArray<T> {
let (arr, win_offset, win_len) = window;
if offset == 0 {
return FloatArray {
data: Vec64::from_slice(&arr.data[win_offset..win_offset + win_len]).into(),
null_mask: if win_len != arr.data.len() {
arr.null_mask
.as_ref()
.map(|m| m.slice_clone(win_offset, win_len))
} else {
arr.null_mask.clone()
},
};
} else if offset > 0 {
lead_float((arr, win_offset, win_len), offset as usize)
} else {
lag_float((arr, win_offset, win_len), offset.unsigned_abs())
}
}
#[inline(always)]
pub fn shift_str<T: Integer>(
arr: StringAVT<T>,
shift_offset: isize,
) -> Result<StringArray<T>, KernelError> {
if shift_offset == 0 {
let (array, off, len) = arr;
Ok(array.slice_clone(off, len))
} else if shift_offset > 0 {
lead_str(arr, shift_offset as usize)
} else {
lag_str(arr, shift_offset.unsigned_abs())
}
}
#[cfg(test)]
mod tests {
use minarrow::structs::variants::float::FloatArray;
use minarrow::structs::variants::integer::IntegerArray;
use minarrow::structs::variants::string::StringArray;
use minarrow::{Bitmask, BooleanArray};
use super::*;
fn bm(bits: &[bool]) -> Bitmask {
let mut m = Bitmask::new_set_all(bits.len(), false);
for (i, &b) in bits.iter().enumerate() {
m.set(i, b);
}
m
}
fn expect_int<T: PartialEq + std::fmt::Debug + Clone>(
arr: &IntegerArray<T>,
values: &[T],
valid: &[bool],
) {
assert_eq!(arr.data.as_slice(), values);
let mask = arr.null_mask.as_ref().expect("mask missing");
for (i, &v) in valid.iter().enumerate() {
assert_eq!(mask.get(i), v, "mask bit {}", i);
}
}
fn expect_float<T: num_traits::Float + std::fmt::Debug>(
arr: &FloatArray<T>,
values: &[T],
valid: &[bool],
eps: T,
) {
let data = arr.data.as_slice();
assert_eq!(data.len(), values.len());
for (a, b) in data.iter().zip(values.iter()) {
assert!((*a - *b).abs() <= eps, "value mismatch {:?} vs {:?}", a, b);
}
let mask = arr.null_mask.as_ref().expect("mask missing");
for (i, &v) in valid.iter().enumerate() {
assert_eq!(mask.get(i), v);
}
}
#[test]
fn test_rolling_sum_int_basic() {
let arr = IntegerArray::<i32>::from_slice(&[1, 2, 3, 4, 5]);
let res = rolling_sum_int((&arr, 0, arr.len()), 3);
expect_int(&res, &[0, 0, 6, 9, 12], &[false, false, true, true, true]);
}
#[test]
fn test_rolling_sum_int_masked() {
let mut arr = IntegerArray::<i32>::from_slice(&[1, 2, 3, 4]);
arr.null_mask = Some(bm(&[true, false, true, true]));
let res = rolling_sum_int((&arr, 0, arr.len()), 2);
expect_int(
&res,
&[0, 0, 3, 7],
&[false, false, false, true], );
}
#[test]
fn test_rolling_sum_float() {
let arr = FloatArray::<f32>::from_slice(&[1.0, 2.0, 3.0]);
let res = rolling_sum_float((&arr, 0, arr.len()), 2);
expect_float(&res, &[0.0, 0.0, 5.0], &[false, false, true], 1e-6f32);
}
#[test]
fn test_rolling_sum_bool() {
let bools = BooleanArray::from_slice(&[true, true, false, true]);
let res = rolling_sum_bool((&bools, 0, bools.len()), 2); expect_int(&res, &[0, 0, 1, 1], &[false, false, true, true]);
}
#[test]
fn test_rolling_min_max_mean_count() {
let arr = IntegerArray::<i32>::from_slice(&[3, 1, 4, 1, 5]);
let rmin = rolling_min_int((&arr, 0, arr.len()), 2);
expect_int(&rmin, &[0, 0, 1, 1, 1], &[false, false, true, true, true]);
let rmax = rolling_max_int((&arr, 0, arr.len()), 3);
expect_int(&rmax, &[0, 0, 4, 4, 5], &[false, false, true, true, true]);
let rmean = rolling_mean_int((&arr, 0, arr.len()), 2);
expect_float(
&rmean,
&[0.0, 0.0, 2.5, 2.5, 3.0],
&[false, false, true, true, true],
1e-12,
);
let cnt = rolling_count((0, 5), 3);
expect_int(&cnt, &[0, 0, 3, 3, 3], &[false, false, true, true, true]);
}
#[test]
fn test_rank_int_basic() {
let arr = IntegerArray::<i32>::from_slice(&[30, 10, 20]);
let res = rank_int((&arr, 0, arr.len()));
expect_int(&res, &[3, 1, 2], &[true, true, true]);
}
#[test]
fn test_rank_float_with_nulls() {
let mut arr = FloatArray::<f64>::from_slice(&[2.0, 1.0, 3.0]);
arr.null_mask = Some(bm(&[true, false, true]));
let res = rank_float((&arr, 0, arr.len()));
expect_int(&res, &[2, 0, 3], &[true, false, true]);
}
#[test]
fn test_dense_rank_str_duplicates() {
let arr = StringArray::<u32>::from_slice(&["b", "a", "b", "c"]);
let res = dense_rank_str((&arr, 0, arr.len())).unwrap();
expect_int(&res, &[2, 1, 2, 3], &[true, true, true, true]);
}
#[test]
fn test_dense_rank_str_duplicates_chunk() {
let arr = StringArray::<u32>::from_slice(&["x", "b", "a", "b", "c", "y"]);
let res = dense_rank_str((&arr, 1, 4)).unwrap(); expect_int(&res, &[2, 1, 2, 3], &[true, true, true, true]);
}
#[test]
fn test_lag_lead_int() {
let arr = IntegerArray::<i32>::from_slice(&[10, 20, 30, 40]);
let lag1 = lag_int((&arr, 0, arr.len()), 1);
expect_int(&lag1, &[0, 10, 20, 30], &[false, true, true, true]);
let lead2 = lead_int((&arr, 0, arr.len()), 2);
expect_int(&lead2, &[30, 40, 0, 0], &[true, true, false, false]);
}
#[test]
fn test_shift_float_positive_negative_zero() {
let arr = FloatArray::<f32>::from_slice(&[1.0, 2.0, 3.0]);
let s0 = shift_float((&arr, 0, arr.len()), 0);
assert_eq!(s0.data, arr.data);
let s1 = shift_float((&arr, 0, arr.len()), 1);
expect_float(&s1, &[2.0, 3.0, 0.0], &[true, true, false], 1e-6f32);
let s_neg = shift_float((&arr, 0, arr.len()), -1);
expect_float(&s_neg, &[0.0, 1.0, 2.0], &[false, true, true], 1e-6f32);
}
#[test]
fn test_lag_lead_str() {
let arr = StringArray::<u32>::from_slice(&["a", "b", "c"]);
let l1 = lag_str((&arr, 0, arr.len()), 1).unwrap();
assert_eq!(l1.get(0), None);
assert_eq!(l1.get(1), Some("a"));
assert_eq!(l1.get(2), Some("b"));
let d1 = lead_str((&arr, 0, arr.len()), 1).unwrap();
assert_eq!(d1.get(0), Some("b"));
assert_eq!(d1.get(1), Some("c"));
assert_eq!(d1.get(2), None);
}
#[test]
fn test_lag_lead_str_chunk() {
let arr = StringArray::<u32>::from_slice(&["x", "a", "b", "c", "y"]);
let l1 = lag_str((&arr, 1, 3), 1).unwrap();
assert_eq!(l1.get(0), None);
assert_eq!(l1.get(1), Some("a"));
assert_eq!(l1.get(2), Some("b"));
let d1 = lead_str((&arr, 1, 3), 1).unwrap();
assert_eq!(d1.get(0), Some("b"));
assert_eq!(d1.get(1), Some("c"));
assert_eq!(d1.get(2), None);
}
#[test]
fn test_rolling_sum_int_edge_windows() {
let arr = IntegerArray::<i32>::from_slice(&[1, 2, 3, 4, 5]);
let r0 = rolling_sum_int((&arr, 0, arr.len()), 0);
assert_eq!(r0.data.as_slice(), &[0, 0, 0, 0, 0]);
assert_eq!(r0.null_mask.as_ref().unwrap().all_unset(), true);
let r1 = rolling_sum_int((&arr, 0, arr.len()), 1);
assert_eq!(r1.data.as_slice(), &[1, 2, 3, 4, 5]);
assert!(r1.null_mask.as_ref().unwrap().all_set());
let r_large = rolling_sum_int((&arr, 0, arr.len()), 10);
assert_eq!(r_large.data.as_slice(), &[0, 0, 0, 0, 0]);
assert_eq!(r_large.null_mask.as_ref().unwrap().all_unset(), true);
}
#[test]
fn test_rolling_sum_float_masked_nulls_propagate() {
let mut arr = FloatArray::<f32>::from_slice(&[1.0, 2.0, 3.0, 4.0]);
arr.null_mask = Some(bm(&[true, true, false, true]));
let r = rolling_sum_float((&arr, 0, arr.len()), 2);
expect_float(
&r,
&[0.0, 0.0, 2.0, 4.0],
&[false, false, false, false],
1e-6,
);
}
#[test]
fn test_rolling_sum_bool_with_nulls() {
let mut b = BooleanArray::from_slice(&[true, false, true, true]);
b.null_mask = Some(bm(&[true, false, true, true]));
let r = rolling_sum_bool((&b, 0, b.len()), 2);
expect_int(&r, &[0, 0, 1, 2], &[false, false, false, true]);
}
#[test]
fn test_lag_str_null_propagation() {
let mut arr = StringArray::<u32>::from_slice(&["x", "y", "z"]);
arr.null_mask = Some(bm(&[true, false, true])); let lag1 = lag_str((&arr, 0, arr.len()), 1).unwrap();
assert_eq!(lag1.get(0), None); assert_eq!(lag1.get(1), Some("x"));
assert_eq!(lag1.get(2), None); let m = lag1.null_mask.unwrap();
assert_eq!(m.get(0), false);
assert_eq!(m.get(1), true);
assert_eq!(m.get(2), false);
}
#[test]
fn test_lag_str_null_propagation_chunk() {
let mut arr = StringArray::<u32>::from_slice(&["w", "x", "y", "z", "q"]);
arr.null_mask = Some(bm(&[true, true, false, true, true]));
let lag1 = lag_str((&arr, 1, 3), 1).unwrap();
assert_eq!(lag1.get(0), None); assert_eq!(lag1.get(1), Some("x")); assert_eq!(lag1.get(2), None); let m = lag1.null_mask.unwrap();
assert_eq!(m.get(0), false);
assert_eq!(m.get(1), true);
assert_eq!(m.get(2), false);
}
#[test]
fn test_shift_str_large_offset() {
let arr = StringArray::<u32>::from_slice(&["a", "b", "c"]);
let shifted = shift_str((&arr, 0, arr.len()), 10).unwrap(); assert_eq!(shifted.len(), 3);
for i in 0..3 {
assert_eq!(shifted.get(i), None);
assert_eq!(shifted.null_mask.as_ref().unwrap().get(i), false);
}
}
#[test]
fn test_shift_str_large_offset_chunk() {
let arr = StringArray::<u32>::from_slice(&["w", "a", "b", "c", "x"]);
let shifted = shift_str((&arr, 1, 3), 10).unwrap(); assert_eq!(shifted.len(), 3);
for i in 0..3 {
assert_eq!(shifted.get(i), None);
assert_eq!(shifted.null_mask.as_ref().unwrap().get(i), false);
}
}
#[test]
fn test_rank_str_ties_and_nulls() {
let mut arr = StringArray::<u32>::from_slice(&["a", "b", "a", "c"]);
arr.null_mask = Some(bm(&[true, true, false, true]));
let r = rank_str((&arr, 0, arr.len())).unwrap();
expect_int(&r, &[1, 2, 0, 3], &[true, true, false, true]);
}
#[test]
fn test_rank_str_ties_and_nulls_chunk() {
let mut arr = StringArray::<u32>::from_slice(&["w", "a", "b", "a", "c"]);
arr.null_mask = Some(bm(&[true, true, true, false, true]));
let r = rank_str((&arr, 1, 4)).unwrap(); expect_int(&r, &[1, 2, 0, 3], &[true, true, false, true]);
}
}