rluvaton commented on code in PR #8785:
URL: https://github.com/apache/arrow-rs/pull/8785#discussion_r2492243366
##########
arrow-array/src/iterator.rs:
##########
@@ -102,6 +102,35 @@ impl<T: ArrayAccessor> Iterator for ArrayIter<T> {
Some(self.current_end - self.current),
)
}
+
+ fn nth(&mut self, n: usize) -> Option<Self::Item> {
+ // Check if we can advance to the desired offset
+ match self.current.checked_add(n) {
+ // Yes, and still within bounds
+ Some(new_current) if new_current < self.current_end => {
+ self.current = new_current;
+ }
+
+ // Either overflow or would exceed current_end
+ _ => {
+ self.current = self.current_end;
+ return None;
+ }
+ }
+
+ self.next()
+ }
Review Comment:
The default implementation does it in `O(n)` while we can do it in constant
time.
we also can't implement `advance_by` which is used by `nth` in the default
implementation as it is unstable
This is the default implementation:
```rust
#[inline]
#[unstable(feature = "iter_advance_by", reason = "recently added", issue =
"77404")]
fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
/// Helper trait to specialize `advance_by` via `try_fold` for `Sized`
iterators.
trait SpecAdvanceBy {
fn spec_advance_by(&mut self, n: usize) -> Result<(),
NonZero<usize>>;
}
impl<I: Iterator + ?Sized> SpecAdvanceBy for I {
default fn spec_advance_by(&mut self, n: usize) -> Result<(),
NonZero<usize>> {
for i in 0..n {
if self.next().is_none() {
// SAFETY: `i` is always less than `n`.
return Err(unsafe { NonZero::new_unchecked(n - i) });
}
}
Ok(())
}
}
impl<I: Iterator> SpecAdvanceBy for I {
fn spec_advance_by(&mut self, n: usize) -> Result<(),
NonZero<usize>> {
let Some(n) = NonZero::new(n) else {
return Ok(());
};
let res = self.try_fold(n, |n, _| NonZero::new(n.get() - 1));
match res {
None => Ok(()),
Some(n) => Err(n),
}
}
}
self.spec_advance_by(n)
}
#[inline]
#[stable(feature = "rust1", since = "1.0.0")]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
self.advance_by(n).ok()?;
self.next()
}
```
From [Rust source
code](https://github.com/rust-lang/rust/blob/f15a7f38580ddbdc1a23909dd05cf6cc6d9f3919/library/core/src/iter/traits/iterator.rs#L262-L377)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
