Trait polars_core::chunked_array::ops::TakeRandom

pub trait TakeRandom {
    type Item;

    fn get(&self, index: usize) -> Option<Self::Item>;

    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item>
    where
        Self: Sized,
    { ... }
}

Random access

Required Associated Types

type Item

Required Methods

fn get(&self, index: usize) -> Option<Self::Item>

Get a nullable value by index.

Panics

Panics if index >= self.len()

Provided Methods

unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item>where
    Self: Sized,

Get a value by index and ignore the null bit.

Safety

Does not do bound checks.

Examples found in repository

src/chunked_array/ops/take/take_single.rs (line 90)

    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        (*self).get_unchecked(index)
    }

src/chunked_array/ops/compare_inner.rs (line 89)

    unsafe fn eq_element_unchecked(&self, idx_a: usize, idx_b: usize) -> bool {
        // no nulls so we can do unchecked
        self.get_unchecked(idx_a) == self.get_unchecked(idx_b)
    }
}

/// Create a type that implements PartialEqInner
pub(crate) trait IntoPartialEqInner<'a> {
    /// Create a type that implements `TakeRandom`.
    fn into_partial_eq_inner(self) -> Box<dyn PartialEqInner + 'a>;
}

/// We use a trait object because we want to call this from Series and cannot use a typed enum.
impl<'a, T> IntoPartialEqInner<'a> for &'a ChunkedArray<T>
where
    T: PolarsNumericType,
{
    fn into_partial_eq_inner(self) -> Box<dyn PartialEqInner + 'a> {
        let mut chunks = self.downcast_iter();

        if self.chunks.len() == 1 {
            let arr = chunks.next().unwrap();

            if !self.has_validity() {
                let t = NumTakeRandomCont {
                    slice: arr.values(),
                };
                Box::new(t)
            } else {
                let t = NumTakeRandomSingleChunk::<'_, T::Native>::new(arr);
                Box::new(t)
            }
        } else {
            let t = NumTakeRandomChunked::<'_, T::Native> {
                chunks: chunks.collect(),
                chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
            };
            Box::new(t)
        }
    }
}

impl<'a> IntoPartialEqInner<'a> for &'a Utf8Chunked {
    fn into_partial_eq_inner(self) -> Box<dyn PartialEqInner + 'a> {
        match self.chunks.len() {
            1 => {
                let arr = self.downcast_iter().next().unwrap();
                let t = Utf8TakeRandomSingleChunk { arr };
                Box::new(t)
            }
            _ => {
                let chunks = self.downcast_chunks();
                let t = Utf8TakeRandom {
                    chunks,
                    chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
                };
                Box::new(t)
            }
        }
    }
}

#[cfg(feature = "dtype-binary")]
impl<'a> IntoPartialEqInner<'a> for &'a BinaryChunked {
    fn into_partial_eq_inner(self) -> Box<dyn PartialEqInner + 'a> {
        match self.chunks.len() {
            1 => {
                let arr = self.downcast_iter().next().unwrap();
                let t = BinaryTakeRandomSingleChunk { arr };
                Box::new(t)
            }
            _ => {
                let chunks = self.downcast_chunks();
                let t = BinaryTakeRandom {
                    chunks,
                    chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
                };
                Box::new(t)
            }
        }
    }
}

impl<'a> IntoPartialEqInner<'a> for &'a BooleanChunked {
    fn into_partial_eq_inner(self) -> Box<dyn PartialEqInner + 'a> {
        match self.chunks.len() {
            1 => {
                let arr = self.downcast_iter().next().unwrap();
                let t = BoolTakeRandomSingleChunk { arr };
                Box::new(t)
            }
            _ => {
                let chunks = self.downcast_chunks();
                let t = BoolTakeRandom {
                    chunks,
                    chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
                };
                Box::new(t)
            }
        }
    }
}

// Partial ordering implementations

fn fallback<T: PartialEq>(a: T) -> Ordering {
    // nan != nan
    // this is a simple way to check if it is nan
    // without convincing the compiler we deal with floats
    #[allow(clippy::eq_op)]
    if a != a {
        Ordering::Less
    } else {
        Ordering::Greater
    }
}

impl<T> PartialOrdInner for NumTakeRandomCont<'_, T>
where
    T: Copy + PartialOrd + Sync,
{
    unsafe fn cmp_element_unchecked(&self, idx_a: usize, idx_b: usize) -> Ordering {
        // no nulls so we can do unchecked
        let a = self.get_unchecked(idx_a);
        let b = self.get_unchecked(idx_b);
        a.partial_cmp(&b).unwrap_or_else(|| fallback(a))
    }

src/chunked_array/logical/categorical/mod.rs (line 149)

    unsafe fn get_any_value_unchecked(&self, i: usize) -> AnyValue<'_> {
        match self.logical.0.get_unchecked(i) {
            Some(i) => AnyValue::Categorical(i, self.get_rev_map()),
            None => AnyValue::Null,
        }
    }

src/chunked_array/ops/take/take_random.rs (line 99)

    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        match self {
            Self::SingleNoNull(s) => s.get_unchecked(index),
            Self::Single(s) => s.get_unchecked(index),
            Self::Multi(m) => m.get_unchecked(index),
        }
    }
}

pub enum TakeRandBranch2<S, M> {
    Single(S),
    Multi(M),
}

impl<S, M, I> TakeRandom for TakeRandBranch2<S, M>
where
    S: TakeRandom<Item = I>,
    M: TakeRandom<Item = I>,
{
    type Item = I;

    fn get(&self, index: usize) -> Option<Self::Item> {
        match self {
            Self::Single(s) => s.get(index),
            Self::Multi(m) => m.get(index),
        }
    }

    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        match self {
            Self::Single(s) => s.get_unchecked(index),
            Self::Multi(m) => m.get_unchecked(index),
        }
    }
}

#[allow(clippy::type_complexity)]
impl<'a, T> IntoTakeRandom<'a> for &'a ChunkedArray<T>
where
    T: PolarsNumericType,
{
    type Item = T::Native;
    type TakeRandom = TakeRandBranch3<
        NumTakeRandomCont<'a, T::Native>,
        NumTakeRandomSingleChunk<'a, T::Native>,
        NumTakeRandomChunked<'a, T::Native>,
    >;

    #[inline]
    fn take_rand(&self) -> Self::TakeRandom {
        let mut chunks = self.downcast_iter();

        if self.chunks.len() == 1 {
            let arr = chunks.next().unwrap();

            if !self.has_validity() {
                let t = NumTakeRandomCont {
                    slice: arr.values(),
                };
                TakeRandBranch3::SingleNoNull(t)
            } else {
                let t = NumTakeRandomSingleChunk::new(arr);
                TakeRandBranch3::Single(t)
            }
        } else {
            let t = NumTakeRandomChunked {
                chunks: chunks.collect(),
                chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
            };
            TakeRandBranch3::Multi(t)
        }
    }
}

pub struct Utf8TakeRandom<'a> {
    pub(crate) chunks: Chunks<'a, Utf8Array<i64>>,
    pub(crate) chunk_lens: Vec<IdxSize>,
}

impl<'a> TakeRandom for Utf8TakeRandom<'a> {
    type Item = &'a str;

    #[inline]
    fn get(&self, index: usize) -> Option<Self::Item> {
        take_random_get!(self, index)
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        take_random_get_unchecked!(self, index)
    }
}

pub struct Utf8TakeRandomSingleChunk<'a> {
    pub(crate) arr: &'a Utf8Array<i64>,
}

impl<'a> TakeRandom for Utf8TakeRandomSingleChunk<'a> {
    type Item = &'a str;

    #[inline]
    fn get(&self, index: usize) -> Option<Self::Item> {
        take_random_get_single!(self, index)
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        if self.arr.is_valid_unchecked(index) {
            Some(self.arr.value_unchecked(index))
        } else {
            None
        }
    }
}

impl<'a> IntoTakeRandom<'a> for &'a Utf8Chunked {
    type Item = &'a str;
    type TakeRandom = TakeRandBranch2<Utf8TakeRandomSingleChunk<'a>, Utf8TakeRandom<'a>>;

    fn take_rand(&self) -> Self::TakeRandom {
        match self.chunks.len() {
            1 => {
                let arr = self.downcast_iter().next().unwrap();
                let t = Utf8TakeRandomSingleChunk { arr };
                TakeRandBranch2::Single(t)
            }
            _ => {
                let chunks = self.downcast_chunks();
                let t = Utf8TakeRandom {
                    chunks,
                    chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
                };
                TakeRandBranch2::Multi(t)
            }
        }
    }
}

#[cfg(feature = "dtype-binary")]
pub struct BinaryTakeRandom<'a> {
    pub(crate) chunks: Chunks<'a, BinaryArray<i64>>,
    pub(crate) chunk_lens: Vec<IdxSize>,
}

#[cfg(feature = "dtype-binary")]
impl<'a> TakeRandom for BinaryTakeRandom<'a> {
    type Item = &'a [u8];

    #[inline]
    fn get(&self, index: usize) -> Option<Self::Item> {
        take_random_get!(self, index)
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        take_random_get_unchecked!(self, index)
    }
}

#[cfg(feature = "dtype-binary")]
pub struct BinaryTakeRandomSingleChunk<'a> {
    pub(crate) arr: &'a BinaryArray<i64>,
}

#[cfg(feature = "dtype-binary")]
impl<'a> TakeRandom for BinaryTakeRandomSingleChunk<'a> {
    type Item = &'a [u8];

    #[inline]
    fn get(&self, index: usize) -> Option<Self::Item> {
        take_random_get_single!(self, index)
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        if self.arr.is_valid_unchecked(index) {
            Some(self.arr.value_unchecked(index))
        } else {
            None
        }
    }
}

#[cfg(feature = "dtype-binary")]
impl<'a> IntoTakeRandom<'a> for &'a BinaryChunked {
    type Item = &'a [u8];
    type TakeRandom = TakeRandBranch2<BinaryTakeRandomSingleChunk<'a>, BinaryTakeRandom<'a>>;

    fn take_rand(&self) -> Self::TakeRandom {
        match self.chunks.len() {
            1 => {
                let arr = self.downcast_iter().next().unwrap();
                let t = BinaryTakeRandomSingleChunk { arr };
                TakeRandBranch2::Single(t)
            }
            _ => {
                let chunks = self.downcast_chunks();
                let t = BinaryTakeRandom {
                    chunks,
                    chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
                };
                TakeRandBranch2::Multi(t)
            }
        }
    }
}

impl<'a> IntoTakeRandom<'a> for &'a BooleanChunked {
    type Item = bool;
    type TakeRandom = TakeRandBranch2<BoolTakeRandomSingleChunk<'a>, BoolTakeRandom<'a>>;

    fn take_rand(&self) -> Self::TakeRandom {
        match self.chunks.len() {
            1 => {
                let arr = self.downcast_iter().next().unwrap();
                let t = BoolTakeRandomSingleChunk { arr };
                TakeRandBranch2::Single(t)
            }
            _ => {
                let chunks = self.downcast_chunks();
                let t = BoolTakeRandom {
                    chunks,
                    chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
                };
                TakeRandBranch2::Multi(t)
            }
        }
    }
}

impl<'a> IntoTakeRandom<'a> for &'a ListChunked {
    type Item = Series;
    type TakeRandom = TakeRandBranch2<ListTakeRandomSingleChunk<'a>, ListTakeRandom<'a>>;

    fn take_rand(&self) -> Self::TakeRandom {
        let mut chunks = self.downcast_iter();
        if self.chunks.len() == 1 {
            let t = ListTakeRandomSingleChunk {
                arr: chunks.next().unwrap(),
                name: self.name(),
            };
            TakeRandBranch2::Single(t)
        } else {
            let t = ListTakeRandom {
                ca: self,
                chunks: chunks.collect(),
                chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
            };
            TakeRandBranch2::Multi(t)
        }
    }
}

pub struct NumTakeRandomChunked<'a, T>
where
    T: NumericNative,
{
    pub(crate) chunks: Vec<&'a PrimitiveArray<T>>,
    pub(crate) chunk_lens: Vec<IdxSize>,
}

impl<'a, T> TakeRandom for NumTakeRandomChunked<'a, T>
where
    T: NumericNative,
{
    type Item = T;

    #[inline]
    fn get(&self, index: usize) -> Option<Self::Item> {
        take_random_get!(self, index)
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        take_random_get_unchecked!(self, index)
    }
}

pub struct NumTakeRandomCont<'a, T> {
    pub(crate) slice: &'a [T],
}

impl<'a, T> TakeRandom for NumTakeRandomCont<'a, T>
where
    T: Copy,
{
    type Item = T;

    #[inline]
    fn get(&self, index: usize) -> Option<Self::Item> {
        self.slice.get(index).copied()
    }

    #[inline]
    unsafe fn get_unchecked(&self, index: usize) -> Option<Self::Item> {
        Some(*self.slice.get_unchecked(index))
    }
}

pub struct TakeRandomBitmap<'a> {
    bytes: &'a [u8],
    offset: usize,
}

impl<'a> TakeRandomBitmap<'a> {
    pub(crate) fn new(bitmap: &'a Bitmap) -> Self {
        let (bytes, offset, _) = bitmap.as_slice();
        Self { bytes, offset }
    }

    unsafe fn get_unchecked(&self, index: usize) -> bool {
        get_bit_unchecked(self.bytes, self.offset + index)
    }
}

pub struct NumTakeRandomSingleChunk<'a, T>
where
    T: NumericNative,
{
    pub(crate) vals: &'a [T],
    pub(crate) validity: TakeRandomBitmap<'a>,
}

impl<'a, T: NumericNative> NumTakeRandomSingleChunk<'a, T> {
    pub(crate) fn new(arr: &'a PrimitiveArray<T>) -> Self {
        let validity = TakeRandomBitmap::new(arr.validity().unwrap());
        let vals = arr.values();
        NumTakeRandomSingleChunk { vals, validity }
    }
}

impl<'a, T> TakeRandom for NumTakeRandomSingleChunk<'a, T>
where
    T: NumericNative,
{
    type Item = T;

    #[inline]
    fn get(&self, index: usize) -> Option<Self::Item> {
        if index < self.vals.len() {
            unsafe { self.get_unchecked(index) }
        } else {
            None
        }
    }

src/chunked_array/logical/categorical/ops/take_random.rs (line 39)

    unsafe fn cmp_element_unchecked(&self, idx_a: usize, idx_b: usize) -> Ordering {
        let a = self
            .cats
            .get_unchecked(idx_a)
            .map(|cat| self.rev_map.value_unchecked(cat as usize));
        let b = self
            .cats
            .get_unchecked(idx_b)
            .map(|cat| self.rev_map.value_unchecked(cat as usize));
        a.partial_cmp(&b).unwrap()
    }
}

pub(crate) struct CategoricalTakeRandomGlobal<'a> {
    rev_map_part_1: &'a PlHashMap<u32, u32>,
    rev_map_part_2: &'a Utf8Array<i64>,
    cats: TakeCats<'a>,
}
impl<'a> CategoricalTakeRandomGlobal<'a> {
    pub(crate) fn new(ca: &'a CategoricalChunked) -> Self {
        // should be rechunked upstream
        assert_eq!(ca.logical.chunks.len(), 1, "implementation error");
        if let RevMapping::Global(rev_map_part_1, rev_map_part_2, _) = &**ca.get_rev_map() {
            let cats = ca.logical().take_rand();
            Self {
                rev_map_part_1,
                rev_map_part_2,
                cats,
            }
        } else {
            unreachable!()
        }
    }
}

impl PartialOrdInner for CategoricalTakeRandomGlobal<'_> {
    unsafe fn cmp_element_unchecked(&self, idx_a: usize, idx_b: usize) -> Ordering {
        let a = self.cats.get_unchecked(idx_a).map(|cat| {
            let idx = self.rev_map_part_1.get(&cat).unwrap();
            self.rev_map_part_2.value_unchecked(*idx as usize)
        });
        let b = self.cats.get_unchecked(idx_b).map(|cat| {
            let idx = self.rev_map_part_1.get(&cat).unwrap();
            self.rev_map_part_2.value_unchecked(*idx as usize)
        });
        a.partial_cmp(&b).unwrap()
    }

src/frame/hash_join/mod.rs (line 217)

    fn zip_outer_join_column(
        &self,
        right_column: &Series,
        opt_join_tuples: &[(Option<IdxSize>, Option<IdxSize>)],
    ) -> Series {
        let right_ca = self.unpack_series_matching_type(right_column).unwrap();

        let left_rand_access = self.take_rand();
        let right_rand_access = right_ca.take_rand();

        opt_join_tuples
            .iter()
            .map(|(opt_left_idx, opt_right_idx)| {
                if let Some(left_idx) = opt_left_idx {
                    unsafe { left_rand_access.get_unchecked(*left_idx as usize) }
                } else {
                    unsafe {
                        let right_idx = opt_right_idx.unwrap_unchecked();
                        right_rand_access.get_unchecked(right_idx as usize)
                    }
                }
            })
            .collect_trusted::<ChunkedArray<T>>()
            .into_series()
    }

Additional examples can be found in:

• src/chunked_array/ops/aggregate.rs
• src/chunked_array/ops/take/mod.rs

Implementors

impl TakeRandom for BooleanChunked

type Item = bool

impl TakeRandom for ListChunked

type Item = Series

impl<'a> TakeRandom for &'a Utf8Chunked

type Item = &'a str

impl<'a> TakeRandom for BoolTakeRandom<'a>

type Item = bool

impl<'a> TakeRandom for BoolTakeRandomSingleChunk<'a>

type Item = bool

impl<'a> TakeRandom for ListTakeRandom<'a>

type Item = Series

impl<'a> TakeRandom for ListTakeRandomSingleChunk<'a>

type Item = Series

impl<'a> TakeRandom for Utf8TakeRandom<'a>

type Item = &'a str

impl<'a> TakeRandom for Utf8TakeRandomSingleChunk<'a>

type Item = &'a str

impl<'a, T> TakeRandom for &'a ChunkedArray<T>where
    T: PolarsNumericType,

type Item = <T as PolarsNumericType>::Native

impl<'a, T> TakeRandom for NumTakeRandomChunked<'a, T>where
    T: NumericNative,

type Item = T

impl<'a, T> TakeRandom for NumTakeRandomCont<'a, T>where
    T: Copy,

type Item = T

impl<'a, T> TakeRandom for NumTakeRandomSingleChunk<'a, T>where
    T: NumericNative,

type Item = T

impl<'a, T: PolarsObject> TakeRandom for &'a ObjectChunked<T>

type Item = &'a T

impl<'a, T: PolarsObject> TakeRandom for ObjectTakeRandom<'a, T>

type Item = &'a T

impl<'a, T: PolarsObject> TakeRandom for ObjectTakeRandomSingleChunk<'a, T>

type Item = &'a T

impl<N, S, M, I> TakeRandom for TakeRandBranch3<N, S, M>where
    N: TakeRandom<Item = I>,
    S: TakeRandom<Item = I>,
    M: TakeRandom<Item = I>,

type Item = I

impl<S, M, I> TakeRandom for TakeRandBranch2<S, M>where
    S: TakeRandom<Item = I>,
    M: TakeRandom<Item = I>,

type Item = I

impl<T> TakeRandom for ChunkedArray<T>where
    T: PolarsNumericType,

type Item = <T as PolarsNumericType>::Native