polars_ops/series/ops/
index.rs

Help
1use num_traits::{Signed, Zero};
2use polars_core::error::{PolarsResult, polars_ensure};
3use polars_core::prelude::arity::unary_elementwise_values;
4use polars_core::prelude::{
5    ChunkedArray, Column, DataType, IDX_DTYPE, IdxCa, PolarsIntegerType, Series,
6};
7use polars_utils::index::ToIdx;
8
9fn convert<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
10where
11    T: PolarsIntegerType,
12    T::Native: ToIdx,
13{
14    let target_len = target_len as u64;
15    Ok(unary_elementwise_values(ca, |v| v.to_idx(target_len)))
16}
17
18pub fn convert_to_unsigned_index(s: &Series, target_len: usize) -> PolarsResult<IdxCa> {
19    let dtype = s.dtype();
20    polars_ensure!(dtype.is_integer(), InvalidOperation: "expected integers as index");
21    if dtype.is_unsigned_integer() {
22        let nulls_before_cast = s.null_count();
23        let out = s.cast(&IDX_DTYPE).unwrap();
24        polars_ensure!(out.null_count() == nulls_before_cast, OutOfBounds: "some integers did not fit polars' index size");
25        return Ok(out.idx().unwrap().clone());
26    }
27    match dtype {
28        DataType::Int64 => {
29            let ca = s.i64().unwrap();
30            convert(ca, target_len)
31        },
32        DataType::Int32 => {
33            let ca = s.i32().unwrap();
34            convert(ca, target_len)
35        },
36        #[cfg(feature = "dtype-i16")]
37        DataType::Int16 => {
38            let ca = s.i16().unwrap();
39            convert(ca, target_len)
40        },
41        #[cfg(feature = "dtype-i8")]
42        DataType::Int8 => {
43            let ca = s.i8().unwrap();
44            convert(ca, target_len)
45        },
46        _ => unreachable!(),
47    }
48}
49
50/// May give false negatives because it ignores the null values.
51fn is_positive_idx_uncertain_impl<T>(ca: &ChunkedArray<T>) -> bool
52where
53    T: PolarsIntegerType,
54    T::Native: Signed,
55{
56    ca.downcast_iter().all(|v| {
57        let values = v.values();
58        let mut all_positive = true;
59
60        // process chunks to autovec but still have early return
61        for chunk in values.chunks(1024) {
62            for v in chunk.iter() {
63                all_positive &= v.is_positive() | v.is_zero()
64            }
65            if !all_positive {
66                return all_positive;
67            }
68        }
69        all_positive
70    })
71}
72
73/// May give false negatives because it ignores the null values.
74pub fn is_positive_idx_uncertain(s: &Series) -> bool {
75    let dtype = s.dtype();
76    debug_assert!(dtype.is_integer(), "expected integers as index");
77    if dtype.is_unsigned_integer() {
78        return true;
79    }
80    match dtype {
81        DataType::Int64 => {
82            let ca = s.i64().unwrap();
83            is_positive_idx_uncertain_impl(ca)
84        },
85        DataType::Int32 => {
86            let ca = s.i32().unwrap();
87            is_positive_idx_uncertain_impl(ca)
88        },
89        #[cfg(feature = "dtype-i16")]
90        DataType::Int16 => {
91            let ca = s.i16().unwrap();
92            is_positive_idx_uncertain_impl(ca)
93        },
94        #[cfg(feature = "dtype-i8")]
95        DataType::Int8 => {
96            let ca = s.i8().unwrap();
97            is_positive_idx_uncertain_impl(ca)
98        },
99        _ => unreachable!(),
100    }
101}
102
103/// May give false negatives because it ignores the null values.
104pub fn is_positive_idx_uncertain_col(c: &Column) -> bool {
105    // @scalar-opt
106    // @partition-opt
107    is_positive_idx_uncertain(c.as_materialized_series())
108}
polars_ops/series/ops/index.rs

polars_ops/series/ops/
index.rs
