polars_ops/series/ops/
index.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
use num_traits::{Signed, Zero};
use polars_core::error::{polars_ensure, PolarsResult};
use polars_core::prelude::arity::unary_elementwise_values;
use polars_core::prelude::{
    ChunkedArray, Column, DataType, IdxCa, PolarsIntegerType, Series, IDX_DTYPE,
};
use polars_utils::index::ToIdx;

fn convert<T>(ca: &ChunkedArray<T>, target_len: usize) -> PolarsResult<IdxCa>
where
    T: PolarsIntegerType,
    T::Native: ToIdx,
{
    let target_len = target_len as u64;
    Ok(unary_elementwise_values(ca, |v| v.to_idx(target_len)))
}

pub fn convert_to_unsigned_index(s: &Series, target_len: usize) -> PolarsResult<IdxCa> {
    let dtype = s.dtype();
    polars_ensure!(dtype.is_integer(), InvalidOperation: "expected integers as index");
    if dtype.is_unsigned_integer() {
        let nulls_before_cast = s.null_count();
        let out = s.cast(&IDX_DTYPE).unwrap();
        polars_ensure!(out.null_count() == nulls_before_cast, OutOfBounds: "some integers did not fit polars' index size");
        return Ok(out.idx().unwrap().clone());
    }
    match dtype {
        DataType::Int64 => {
            let ca = s.i64().unwrap();
            convert(ca, target_len)
        },
        DataType::Int32 => {
            let ca = s.i32().unwrap();
            convert(ca, target_len)
        },
        #[cfg(feature = "dtype-i16")]
        DataType::Int16 => {
            let ca = s.i16().unwrap();
            convert(ca, target_len)
        },
        #[cfg(feature = "dtype-i8")]
        DataType::Int8 => {
            let ca = s.i8().unwrap();
            convert(ca, target_len)
        },
        _ => unreachable!(),
    }
}

/// May give false negatives because it ignores the null values.
fn is_positive_idx_uncertain_impl<T>(ca: &ChunkedArray<T>) -> bool
where
    T: PolarsIntegerType,
    T::Native: Signed,
{
    ca.downcast_iter().all(|v| {
        let values = v.values();
        let mut all_positive = true;

        // process chunks to autovec but still have early return
        for chunk in values.chunks(1024) {
            for v in chunk.iter() {
                all_positive &= v.is_positive() | v.is_zero()
            }
            if !all_positive {
                return all_positive;
            }
        }
        all_positive
    })
}

/// May give false negatives because it ignores the null values.
pub fn is_positive_idx_uncertain(s: &Series) -> bool {
    let dtype = s.dtype();
    debug_assert!(dtype.is_integer(), "expected integers as index");
    if dtype.is_unsigned_integer() {
        return true;
    }
    match dtype {
        DataType::Int64 => {
            let ca = s.i64().unwrap();
            is_positive_idx_uncertain_impl(ca)
        },
        DataType::Int32 => {
            let ca = s.i32().unwrap();
            is_positive_idx_uncertain_impl(ca)
        },
        #[cfg(feature = "dtype-i16")]
        DataType::Int16 => {
            let ca = s.i16().unwrap();
            is_positive_idx_uncertain_impl(ca)
        },
        #[cfg(feature = "dtype-i8")]
        DataType::Int8 => {
            let ca = s.i8().unwrap();
            is_positive_idx_uncertain_impl(ca)
        },
        _ => unreachable!(),
    }
}

/// May give false negatives because it ignores the null values.
pub fn is_positive_idx_uncertain_col(c: &Column) -> bool {
    // @scalar-opt
    // @partition-opt
    is_positive_idx_uncertain(c.as_materialized_series())
}