Skip to content

Commit

Permalink
fix large idx take chunked (#4282)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ritchie46
ritchie46 committed Aug 5, 2022
1 parent 1e3f703 commit 5487704
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 21 deletions.
14 changes: 7 additions & 7 deletions polars/polars-core/src/chunked_array/ops/compare_inner.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ where
} else {
let t = NumTakeRandomChunked::<'_, T::Native> {
chunks: chunks.collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
Box::new(t)
}
Expand All @@ -131,7 +131,7 @@ impl<'a> IntoPartialEqInner<'a> for &'a Utf8Chunked {
let chunks = self.downcast_chunks();
let t = Utf8TakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
Box::new(t)
}
Expand All @@ -151,7 +151,7 @@ impl<'a> IntoPartialEqInner<'a> for &'a BooleanChunked {
let chunks = self.downcast_chunks();
let t = BoolTakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
Box::new(t)
}
Expand Down Expand Up @@ -212,7 +212,7 @@ where
} else {
let t = NumTakeRandomChunked::<'_, T::Native> {
chunks: chunks.collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
Box::new(t)
}
Expand All @@ -231,7 +231,7 @@ impl<'a> IntoPartialOrdInner<'a> for &'a Utf8Chunked {
let chunks = self.downcast_chunks();
let t = Utf8TakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
Box::new(t)
}
Expand All @@ -251,7 +251,7 @@ impl<'a> IntoPartialOrdInner<'a> for &'a BooleanChunked {
let chunks = self.downcast_chunks();
let t = BoolTakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
Box::new(t)
}
Expand Down Expand Up @@ -304,7 +304,7 @@ impl<'a, T: PolarsObject> IntoPartialEqInner<'a> for &'a ObjectChunked<T> {
let chunks = self.downcast_chunks();
let t = ObjectTakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
Box::new(t)
}
Expand Down
32 changes: 18 additions & 14 deletions polars/polars-core/src/chunked_array/ops/take/take_random.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,10 @@ use std::convert::TryFrom;

macro_rules! take_random_get {
($self:ident, $index:ident) => {{
let (chunk_idx, arr_idx) =
crate::utils::index_to_chunked_index($self.chunk_lens.iter().copied(), $index as u32);
let (chunk_idx, arr_idx) = crate::utils::index_to_chunked_index(
$self.chunk_lens.iter().copied(),
$index as IdxSize,
);

// Safety:
// bounds are checked above
Expand All @@ -29,8 +31,10 @@ macro_rules! take_random_get {

macro_rules! take_random_get_unchecked {
($self:ident, $index:ident) => {{
let (chunk_idx, arr_idx) =
crate::utils::index_to_chunked_index($self.chunk_lens.iter().copied(), $index as u32);
let (chunk_idx, arr_idx) = crate::utils::index_to_chunked_index(
$self.chunk_lens.iter().copied(),
$index as IdxSize,
);

// Safety:
// bounds are checked above
Expand Down Expand Up @@ -155,7 +159,7 @@ where
} else {
let t = NumTakeRandomChunked {
chunks: chunks.collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
TakeRandBranch3::Multi(t)
}
Expand All @@ -164,7 +168,7 @@ where

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

impl<'a> TakeRandom for Utf8TakeRandom<'a> {
Expand Down Expand Up @@ -218,7 +222,7 @@ impl<'a> IntoTakeRandom<'a> for &'a Utf8Chunked {
let chunks = self.downcast_chunks();
let t = Utf8TakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
TakeRandBranch2::Multi(t)
}
Expand All @@ -241,7 +245,7 @@ impl<'a> IntoTakeRandom<'a> for &'a BooleanChunked {
let chunks = self.downcast_chunks();
let t = BoolTakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
TakeRandBranch2::Multi(t)
}
Expand All @@ -265,7 +269,7 @@ impl<'a> IntoTakeRandom<'a> for &'a ListChunked {
let t = ListTakeRandom {
ca: self,
chunks: chunks.collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
TakeRandBranch2::Multi(t)
}
Expand All @@ -277,7 +281,7 @@ where
T: NumericNative,
{
pub(crate) chunks: Vec<&'a PrimitiveArray<T>>,
pub(crate) chunk_lens: Vec<u32>,
pub(crate) chunk_lens: Vec<IdxSize>,
}

impl<'a, T> TakeRandom for NumTakeRandomChunked<'a, T>
Expand Down Expand Up @@ -377,7 +381,7 @@ where

pub struct BoolTakeRandom<'a> {
pub(crate) chunks: Chunks<'a, BooleanArray>,
pub(crate) chunk_lens: Vec<u32>,
pub(crate) chunk_lens: Vec<IdxSize>,
}

impl<'a> TakeRandom for BoolTakeRandom<'a> {
Expand Down Expand Up @@ -419,7 +423,7 @@ impl<'a> TakeRandom for BoolTakeRandomSingleChunk<'a> {
pub struct ListTakeRandom<'a> {
ca: &'a ListChunked,
chunks: Vec<&'a ListArray<i64>>,
chunk_lens: Vec<u32>,
chunk_lens: Vec<IdxSize>,
}

impl<'a> TakeRandom for ListTakeRandom<'a> {
Expand Down Expand Up @@ -476,7 +480,7 @@ impl<'a> TakeRandom for ListTakeRandomSingleChunk<'a> {
#[cfg(feature = "object")]
pub struct ObjectTakeRandom<'a, T: PolarsObject> {
pub(crate) chunks: Chunks<'a, ObjectArray<T>>,
pub(crate) chunk_lens: Vec<u32>,
pub(crate) chunk_lens: Vec<IdxSize>,
}

#[cfg(feature = "object")]
Expand Down Expand Up @@ -533,7 +537,7 @@ impl<'a, T: PolarsObject> IntoTakeRandom<'a> for &'a ObjectChunked<T> {
} else {
let t = ObjectTakeRandom {
chunks,
chunk_lens: self.chunks.iter().map(|a| a.len() as u32).collect(),
chunk_lens: self.chunks.iter().map(|a| a.len() as IdxSize).collect(),
};
TakeRandBranch2::Multi(t)
}
Expand Down

0 comments on commit 5487704

Please sign in to comment.