Implement scalar accessor for chunked array (#162)

arro3-core/python/arro3/core/_core.pyi

@@ -1,4 +1,5 @@
-from typing import Any, Literal, Sequence, overload
+from typing import Any, Iterable, Literal, Sequence, overload
+
 import numpy as np
 from numpy.typing import NDArray
 
@@ -56,6 +57,9 @@ class Array:
         """
     def __eq__(self, other) -> bool: ...
     def __getitem__(self, i: int) -> Scalar: ...
+    # Note: we don't actually implement this, but it's inferred by having a __getitem__
+    # key
+    def __iter__(self) -> Iterable[Scalar]: ...
     def __len__(self) -> int: ...
     def __repr__(self) -> str: ...
     @classmethod
@@ -229,6 +233,10 @@ class ChunkedArray:
         pyarrow array, without copying memory.
         """
     def __eq__(self, other) -> bool: ...
+    def __getitem__(self, i: int) -> Scalar: ...
+    # Note: we don't actually implement this, but it's inferred by having a __getitem__
+    # key
+    def __iter__(self) -> Iterable[Scalar]: ...
     def __len__(self) -> int: ...
     def __repr__(self) -> str: ...
     @classmethod

pyo3-arrow/src/chunked.rs

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use arrow::compute::concat;
 use arrow_array::{Array, ArrayRef};
 use arrow_schema::{ArrowError, DataType, Field, FieldRef};
-use pyo3::exceptions::{PyTypeError, PyValueError};
+use pyo3::exceptions::{PyIndexError, PyTypeError, PyValueError};
 use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyTuple, PyType};
@@ -18,7 +18,7 @@ use crate::ffi::to_python::to_stream_pycapsule;
 use crate::ffi::to_schema_pycapsule;
 use crate::input::AnyArray;
 use crate::interop::numpy::to_numpy::chunked_to_numpy;
-use crate::{PyArray, PyDataType, PyField};
+use crate::{PyArray, PyDataType, PyField, PyScalar};
 
 /// A Python-facing Arrow chunked array.
 ///
@@ -293,6 +293,20 @@ impl PyChunkedArray {
         self.field == other.field && self.chunks == other.chunks
     }
 
+    fn __getitem__(&self, mut i: usize) -> PyArrowResult<PyScalar> {
+        if i >= self.len() {
+            return Err(PyIndexError::new_err("Index out of range").into());
+        }
+        // for chunk in self.ch
+        for chunk in self.chunks() {
+            if i < chunk.len() {
+                return PyScalar::try_new(chunk.slice(i, 1), self.field.clone());
+            }
+            i -= chunk.len();
+        }
+        unreachable!("index in range but past end of last chunk")
+    }
+
     fn __len__(&self) -> usize {
         self.chunks.iter().fold(0, |acc, x| acc + x.len())
     }