polars_arrow/legacy/kernels/take_agg/
mod.rs

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

polars_arrow/legacy/kernels/take_agg/
mod.rs
