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::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    I: IntoIterator<Item = usize>,
21>(
22    arr: &PrimitiveArray<T>,
23    indices: I,
24) -> impl Iterator<Item = T> {
25    debug_assert!(arr.null_count() == 0);
26    let array_values = arr.values().as_slice();
27
28    indices
29        .into_iter()
30        .map(|idx| *array_values.get_unchecked(idx))
31}
32
33/// Take kernel for single chunk and an iterator as index.
34/// # Safety
35/// caller must ensure iterators indexes are in bounds
36#[inline]
37pub unsafe fn take_agg_primitive_iter_unchecked<T: NativeType, I: IntoIterator<Item = usize>>(
38    arr: &PrimitiveArray<T>,
39    indices: I,
40) -> impl Iterator<Item = T> {
41    let array_values = arr.values().as_slice();
42    let validity = arr.validity().unwrap();
43
44    indices
45        .into_iter()
46        .filter(|&idx| validity.get_bit_unchecked(idx))
47        .map(|idx| *array_values.get_unchecked(idx))
48}
49
50/// Take kernel for single chunk and an iterator as index.
51/// # Safety
52/// caller must ensure iterators indexes are in bounds
53#[inline]
54pub unsafe fn take_agg_primitive_iter_unchecked_count_nulls<
55    T: NativeType + ToPrimitive,
56    I: IntoIterator<Item = usize>,
57    TOut,
58    F: Fn(TOut, T) -> TOut,
59>(
60    arr: &PrimitiveArray<T>,
61    indices: I,
62    init: TOut,
63    f: F,
64    len: IdxSize,
65) -> Option<(TOut, IdxSize)> {
66    let array_values = arr.values().as_slice();
67    let validity = arr.validity().expect("null buffer should be there");
68
69    let mut null_count = 0 as IdxSize;
70    let out = indices.into_iter().fold(init, |acc, idx| {
71        if validity.get_bit_unchecked(idx) {
72            f(acc, *array_values.get_unchecked(idx))
73        } else {
74            null_count += 1;
75            acc
76        }
77    });
78    if null_count == len {
79        None
80    } else {
81        Some((out, null_count))
82    }
83}
84
85/// Take kernel for single chunk and an iterator as index.
86/// # Safety
87/// caller must ensure iterators indexes are in bounds
88#[inline]
89pub unsafe fn take_agg_bin_iter_unchecked<
90    'a,
91    I: IntoIterator<Item = usize>,
92    F: Fn(&'a [u8], &'a [u8]) -> &'a [u8],
93>(
94    arr: &'a BinaryViewArray,
95    indices: I,
96    f: F,
97    len: IdxSize,
98) -> Option<&'a [u8]> {
99    let mut null_count = 0 as IdxSize;
100    let validity = arr.validity().unwrap();
101
102    let out = indices
103        .into_iter()
104        .map(|idx| {
105            if validity.get_bit_unchecked(idx) {
106                Some(arr.value_unchecked(idx))
107            } else {
108                None
109            }
110        })
111        .reduce(|acc, opt_val| match (acc, opt_val) {
112            (Some(acc), Some(str_val)) => Some(f(acc, str_val)),
113            (_, None) => {
114                null_count += 1;
115                acc
116            },
117            (None, Some(str_val)) => Some(str_val),
118        });
119    if null_count == len {
120        None
121    } else {
122        out.flatten()
123    }
124}
125
126/// # Safety
127/// caller must ensure iterators indexes are in bounds
128#[inline]
129pub unsafe fn take_agg_bin_iter_unchecked_arg<
130    'a,
131    I: IntoIterator<Item = usize>,
132    F: Fn((IdxSize, &'a [u8]), (IdxSize, &'a [u8])) -> (IdxSize, &'a [u8]),
133>(
134    arr: &'a BinaryViewArray,
135    indices: I,
136    f: F,
137) -> Option<IdxSize> {
138    let validity = arr.validity().unwrap();
139
140    indices
141        .into_iter()
142        .enumerate()
143        .filter_map(|(pos, idx)| {
144            if validity.get_bit_unchecked(idx) {
145                Some((pos as IdxSize, arr.value_unchecked(idx)))
146            } else {
147                None
148            }
149        })
150        .reduce(f)
151        .map(|(pos, _)| pos)
152}
153
154/// Take kernel for single chunk and an iterator as index.
155/// # Safety
156/// caller must ensure iterators indexes are in bounds
157#[inline]
158pub unsafe fn take_agg_bin_iter_unchecked_no_null<
159    'a,
160    I: IntoIterator<Item = usize>,
161    F: Fn(&'a [u8], &'a [u8]) -> &'a [u8],
162>(
163    arr: &'a BinaryViewArray,
164    indices: I,
165    f: F,
166) -> Option<&'a [u8]> {
167    indices
168        .into_iter()
169        .map(|idx| arr.value_unchecked(idx))
170        .reduce(|acc, str_val| f(acc, str_val))
171}
172
173/// # Safety
174/// caller must ensure iterators indexes are in bounds
175#[inline]
176pub unsafe fn take_agg_bin_iter_unchecked_no_null_arg<
177    'a,
178    I: IntoIterator<Item = usize>,
179    F: Fn((IdxSize, &'a [u8]), (IdxSize, &'a [u8])) -> (IdxSize, &'a [u8]),
180>(
181    arr: &'a BinaryViewArray,
182    indices: I,
183    f: F,
184) -> Option<IdxSize> {
185    indices
186        .into_iter()
187        .enumerate()
188        .map(|(pos, idx)| (pos as IdxSize, arr.value_unchecked(idx)))
189        .reduce(f)
190        .map(|(pos, _)| pos)
191}
polars_arrow/legacy/kernels/take_agg/mod.rs

polars_arrow/legacy/kernels/take_agg/
mod.rs
