Skip to content

Commit d75d8a6

Browse files
committed
Use size_hint in PyArray::from_iter.
Even though the size hint is not guaranteed to be correct, it is a better than a fixed guess.
1 parent 021676e commit d75d8a6

File tree

1 file changed

+7
-5
lines changed

1 file changed

+7
-5
lines changed

src/array.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -743,8 +743,8 @@ impl<T: TypeNum> PyArray<T, Ix1> {
743743
/// Construct one-dimension PyArray from a type which implements
744744
/// [`IntoIterator`](https://doc.rust-lang.org/std/iter/trait.IntoIterator.html).
745745
///
746-
/// This method can allocate memory multiple times and not fast.
747-
/// When you can use [from_exact_iter](method.from_exact_iter.html), we recommend to use it.
746+
/// If no reliable [`size_hint`](https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.size_hint) is available,
747+
/// this method can allocate memory multiple time, which can hurt performance.
748748
///
749749
/// # Example
750750
/// ```
@@ -758,13 +758,15 @@ impl<T: TypeNum> PyArray<T, Ix1> {
758758
/// # }
759759
/// ```
760760
pub fn from_iter(py: Python<'_>, iter: impl IntoIterator<Item = T>) -> &Self {
761-
let mut capacity = 512 / mem::size_of::<T>();
761+
let iter = iter.into_iter();
762+
let (min_len, max_len) = iter.size_hint();
763+
let mut capacity = max_len.unwrap_or(min_len.max(512 / mem::size_of::<T>()));
762764
let array = Self::new(py, [capacity], false);
763765
let mut length = 0;
764766
unsafe {
765-
for (i, item) in iter.into_iter().enumerate() {
767+
for (i, item) in iter.enumerate() {
766768
length += 1;
767-
if length >= capacity {
769+
if length > capacity {
768770
capacity *= 2;
769771
array
770772
.resize(capacity)

0 commit comments

Comments
 (0)