polars_arrow/legacy/kernels/take_agg/
mod.rs

1//! kernels that combine take and aggregations.
2mod boolean;
3mod var;
4
5pub use boolean::*;
6use num_traits::{NumCast, ToPrimitive};
7use polars_utils::IdxSize;
8pub use var::*;
9
10use crate::array::{Array, BinaryViewArray, BooleanArray, PrimitiveArray};
11use crate::types::NativeType;
12
13/// Take kernel for single chunk without nulls and an iterator as index.
14/// # Safety
15/// caller must ensure iterators indexes are in bounds
16#[inline]
17pub unsafe fn take_agg_no_null_primitive_iter_unchecked<
18    T: NativeType + ToPrimitive,
19    TOut: NumCast + NativeType,
20    I: IntoIterator<Item = usize>,
21    F: Fn(TOut, TOut) -> TOut,
22>(
23    arr: &PrimitiveArray<T>,
24    indices: I,
25    f: F,
26) -> Option<TOut> {
27    debug_assert!(arr.null_count() == 0);
28    let array_values = arr.values().as_slice();
29
30    indices
31        .into_iter()
32        .map(|idx| TOut::from(*array_values.get_unchecked(idx)).unwrap_unchecked())
33        .reduce(f)
34}
35
36/// Take kernel for single chunk and an iterator as index.
37/// # Safety
38/// caller must ensure iterators indexes are in bounds
39#[inline]
40pub unsafe fn take_agg_primitive_iter_unchecked<
41    T: NativeType,
42    I: IntoIterator<Item = usize>,
43    F: Fn(T, T) -> T,
44>(
45    arr: &PrimitiveArray<T>,
46    indices: I,
47    f: F,
48) -> Option<T> {
49    let array_values = arr.values().as_slice();
50    let validity = arr.validity().unwrap();
51
52    indices
53        .into_iter()
54        .filter(|&idx| validity.get_bit_unchecked(idx))
55        .map(|idx| *array_values.get_unchecked(idx))
56        .reduce(f)
57}
58
59/// Take kernel for single chunk and an iterator as index.
60/// # Safety
61/// caller must enure iterators indexes are in bounds
62#[inline]
63pub unsafe fn take_agg_primitive_iter_unchecked_count_nulls<
64    T: NativeType + ToPrimitive,
65    TOut: NumCast + NativeType,
66    I: IntoIterator<Item = usize>,
67    F: Fn(TOut, TOut) -> TOut,
68>(
69    arr: &PrimitiveArray<T>,
70    indices: I,
71    f: F,
72    init: TOut,
73    len: IdxSize,
74) -> Option<(TOut, IdxSize)> {
75    let array_values = arr.values().as_slice();
76    let validity = arr.validity().expect("null buffer should be there");
77
78    let mut null_count = 0 as IdxSize;
79    let out = indices.into_iter().fold(init, |acc, idx| {
80        if validity.get_bit_unchecked(idx) {
81            f(
82                acc,
83                NumCast::from(*array_values.get_unchecked(idx)).unwrap_unchecked(),
84            )
85        } else {
86            null_count += 1;
87            acc
88        }
89    });
90    if null_count == len {
91        None
92    } else {
93        Some((out, null_count))
94    }
95}
96
97/// Take kernel for single chunk and an iterator as index.
98/// # Safety
99/// caller must ensure iterators indexes are in bounds
100#[inline]
101pub unsafe fn take_agg_bin_iter_unchecked<
102    'a,
103    I: IntoIterator<Item = usize>,
104    F: Fn(&'a [u8], &'a [u8]) -> &'a [u8],
105>(
106    arr: &'a BinaryViewArray,
107    indices: I,
108    f: F,
109    len: IdxSize,
110) -> Option<&'a [u8]> {
111    let mut null_count = 0 as IdxSize;
112    let validity = arr.validity().unwrap();
113
114    let out = indices
115        .into_iter()
116        .map(|idx| {
117            if validity.get_bit_unchecked(idx) {
118                Some(arr.value_unchecked(idx))
119            } else {
120                None
121            }
122        })
123        .reduce(|acc, opt_val| match (acc, opt_val) {
124            (Some(acc), Some(str_val)) => Some(f(acc, str_val)),
125            (_, None) => {
126                null_count += 1;
127                acc
128            },
129            (None, Some(str_val)) => Some(str_val),
130        });
131    if null_count == len {
132        None
133    } else {
134        out.flatten()
135    }
136}
137
138/// Take kernel for single chunk and an iterator as index.
139/// # Safety
140/// caller must ensure iterators indexes are in bounds
141#[inline]
142pub unsafe fn take_agg_bin_iter_unchecked_no_null<
143    'a,
144    I: IntoIterator<Item = usize>,
145    F: Fn(&'a [u8], &'a [u8]) -> &'a [u8],
146>(
147    arr: &'a BinaryViewArray,
148    indices: I,
149    f: F,
150) -> Option<&'a [u8]> {
151    indices
152        .into_iter()
153        .map(|idx| arr.value_unchecked(idx))
154        .reduce(|acc, str_val| f(acc, str_val))
155}
polars_arrow/legacy/kernels/take_agg/mod.rs

polars_arrow/legacy/kernels/take_agg/
mod.rs
