pyo3-arrow 0.21.0

use std::fmt::{self, Display};
use std::sync::{Arc, Mutex};

use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, StructArray};
use arrow_schema::{ArrowError, Field, Schema, SchemaRef};
use pyo3::exceptions::{PyIOError, PyStopIteration, PyTypeError, PyValueError};
use pyo3::intern;
use pyo3::prelude::*;
use pyo3::types::{PyCapsule, PyIterator, PySequence, PyTuple, PyType};
use pyo3::Borrowed;

use crate::error::PyArrowResult;
use crate::export::{Arro3RecordBatch, Arro3Schema, Arro3Table};
use crate::ffi::from_python::utils::import_stream_pycapsule;
use crate::ffi::to_python::nanoarrow::to_nanoarrow_array_stream;
use crate::ffi::to_python::to_stream_pycapsule;
use crate::ffi::to_schema_pycapsule;
use crate::input::AnyRecordBatch;
use crate::schema::display_schema;
use crate::utils::schema_equals;
use crate::{PyRecordBatch, PySchema, PyTable};
use arrow_array::ArrayIterator;

/// Input to `from_batches`: either a sequence of record batches, which is read
/// eagerly, or any other iterable, which is consumed lazily one batch at a time.
enum RecordBatchInput {
    Sequence(Vec<PyRecordBatch>),
    Iterator(Py<PyIterator>),
}

impl<'a, 'py> FromPyObject<'a, 'py> for RecordBatchInput {
    type Error = PyErr;

    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
        // Only lists, tuples, and objects registered as `collections.abc.Sequence`
        // pass this cast and are read eagerly. Generators, iterators, and other
        // merely-iterable objects are never materialized here.
        //
        // Checking for a sequence before extracting means that a sequence element
        // that is not a record batch raises immediately, instead of the error
        // being deferred to the first read.
        if ob.cast::<PySequence>().is_ok() {
            return Ok(Self::Sequence(ob.extract()?));
        }
        match ob.try_iter() {
            Ok(iter) => Ok(Self::Iterator(iter.unbind())),
            Err(_) => Err(PyTypeError::new_err(format!(
                "Expected a sequence or iterable of record batches, got {}",
                ob.get_type().name()?
            ))),
        }
    }
}

impl RecordBatchInput {
    fn into_record_batch_reader(
        self,
        schema: SchemaRef,
    ) -> PyResult<Box<dyn RecordBatchReader + Send>> {
        match self {
            Self::Sequence(batches) => {
                let batches = batches
                    .into_iter()
                    .map(|batch| batch.into_inner())
                    .collect::<Vec<_>>();
                for (index, batch) in batches.iter().enumerate() {
                    check_batch_schema(index, batch, &schema)?;
                }
                Ok(Box::new(RecordBatchIterator::new(
                    batches.into_iter().map(Ok),
                    schema,
                )))
            }
            Self::Iterator(iter) => Ok(Box::new(PyIteratorRecordBatchReader {
                schema,
                iter,
                index: 0,
            })),
        }
    }
}

/// Check that a batch's schema matches the schema declared for the reader.
///
/// Uses the same comparison as table construction, which matches pyarrow's
/// `RecordBatchReader.from_batches`: names, types, and nullability must match, metadata
/// is ignored. The error message also mirrors pyarrow's.
fn check_batch_schema(batch_index: usize, batch: &RecordBatch, expected: &Schema) -> PyResult<()> {
    let actual = batch.schema();
    if schema_equals(&actual, expected) {
        Ok(())
    } else {
        Err(PyValueError::new_err(format!(
            "Schema at index {batch_index} was different:\n{}vs\n{}",
            SchemaDisplay(actual.as_ref()),
            SchemaDisplay(expected),
        )))
    }
}

/// Formats a [Schema] one field per line, for error messages.
struct SchemaDisplay<'a>(&'a Schema);

impl Display for SchemaDisplay<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        display_schema(self.0, f)
    }
}

/// A [RecordBatchReader] that lazily pulls batches from a Python iterator.
///
/// Each call to `next()` attaches to the Python interpreter and advances the
/// underlying Python iterator by one item.
struct PyIteratorRecordBatchReader {
    schema: SchemaRef,
    iter: Py<PyIterator>,
    /// Number of batches pulled so far, for error messages.
    index: usize,
}

impl Iterator for PyIteratorRecordBatchReader {
    type Item = Result<RecordBatch, ArrowError>;

    fn next(&mut self) -> Option<Self::Item> {
        Python::attach(|py| {
            let item = self.iter.bind(py).clone().next()?;
            let batch = item
                .and_then(|obj| obj.extract::<PyRecordBatch>())
                .map(|batch| batch.into_inner())
                .and_then(|batch| {
                    check_batch_schema(self.index, &batch, &self.schema)?;
                    Ok(batch)
                })
                .map_err(|err| ArrowError::ExternalError(Box::new(err)));
            self.index += 1;
            Some(batch)
        })
    }
}

impl RecordBatchReader for PyIteratorRecordBatchReader {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

/// A Python-facing Arrow record batch reader.
///
/// This is a wrapper around a [RecordBatchReader].
#[pyclass(
    module = "arro3.core._core",
    name = "RecordBatchReader",
    subclass,
    frozen
)]
pub struct PyRecordBatchReader(pub(crate) Mutex<Option<Box<dyn RecordBatchReader + Send>>>);

impl PyRecordBatchReader {
    /// Construct a new PyRecordBatchReader from an existing [RecordBatchReader].
    pub fn new(reader: Box<dyn RecordBatchReader + Send>) -> Self {
        Self(Mutex::new(Some(reader)))
    }

    /// Construct from a raw Arrow C Stream capsule
    pub fn from_arrow_pycapsule(capsule: &Bound<PyCapsule>) -> PyResult<Self> {
        let stream = import_stream_pycapsule(capsule)?;
        let stream_reader = arrow_array::ffi_stream::RecordBatchStreamReader::try_new(stream)
            .map_err(|err| PyValueError::new_err(err.to_string()))?;

        Ok(Self::new(Box::new(stream_reader)))
    }

    /// Consume this reader and convert into a [RecordBatchReader].
    ///
    /// The reader can only be consumed once. Calling `into_reader`
    pub fn into_reader(self) -> PyResult<Box<dyn RecordBatchReader + Send>> {
        let stream = self
            .0
            .lock()
            .unwrap()
            .take()
            .ok_or(PyIOError::new_err("Cannot read from closed stream."))?;
        Ok(stream)
    }

    /// Consume this reader and create a [PyTable] object
    pub fn into_table(self) -> PyArrowResult<PyTable> {
        let stream = self
            .0
            .lock()
            .unwrap()
            .take()
            .ok_or(PyIOError::new_err("Cannot read from closed stream."))?;
        let schema = stream.schema();
        let mut batches = vec![];
        for batch in stream {
            batches.push(batch?);
        }
        Ok(PyTable::try_new(batches, schema)?)
    }

    /// Access the [SchemaRef] of this RecordBatchReader.
    ///
    /// If the stream has already been consumed, this method will error.
    pub fn schema_ref(&self) -> PyResult<SchemaRef> {
        let inner = self.0.lock().unwrap();
        let stream = inner
            .as_ref()
            .ok_or(PyIOError::new_err("Stream already closed."))?;
        Ok(stream.schema())
    }

    /// Export this to a Python `arro3.core.RecordBatchReader`.
    pub fn to_arro3<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
        arro3_mod
            .getattr(intern!(py, "RecordBatchReader"))?
            .call_method1(
                intern!(py, "from_arrow_pycapsule"),
                PyTuple::new(py, vec![self.__arrow_c_stream__(py, None)?])?,
            )
    }

    /// Export this to a Python `arro3.core.RecordBatchReader`.
    pub fn into_arro3(self, py: Python) -> PyResult<Bound<PyAny>> {
        let arro3_mod = py.import(intern!(py, "arro3.core"))?;
        let reader = self
            .0
            .lock()
            .unwrap()
            .take()
            .ok_or(PyIOError::new_err("Cannot read from closed stream"))?;
        let capsule = Self::to_stream_pycapsule(py, reader, None)?;
        arro3_mod
            .getattr(intern!(py, "RecordBatchReader"))?
            .call_method1(
                intern!(py, "from_arrow_pycapsule"),
                PyTuple::new(py, vec![capsule])?,
            )
    }

    /// Export this to a Python `nanoarrow.ArrayStream`.
    pub fn to_nanoarrow<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        to_nanoarrow_array_stream(py, &self.__arrow_c_stream__(py, None)?)
    }

    /// Export to a pyarrow.RecordBatchReader
    ///
    /// Requires pyarrow >=15
    pub fn into_pyarrow(self, py: Python) -> PyResult<Bound<PyAny>> {
        let pyarrow_mod = py.import(intern!(py, "pyarrow"))?;
        let record_batch_reader_class = pyarrow_mod.getattr(intern!(py, "RecordBatchReader"))?;
        record_batch_reader_class.call_method1(
            intern!(py, "from_stream"),
            PyTuple::new(py, vec![self.into_pyobject(py)?])?,
        )
    }

    pub(crate) fn to_stream_pycapsule<'py>(
        py: Python<'py>,
        reader: Box<dyn RecordBatchReader + Send>,
        requested_schema: Option<Bound<'py, PyCapsule>>,
    ) -> PyArrowResult<Bound<'py, PyCapsule>> {
        let schema = reader.schema().clone();
        let array_reader = reader.into_iter().map(|maybe_batch| {
            let arr: ArrayRef = Arc::new(StructArray::from(maybe_batch?));
            Ok(arr)
        });
        let array_reader = Box::new(ArrayIterator::new(
            array_reader,
            Field::new_struct("", schema.fields().clone(), false)
                .with_metadata(schema.metadata.clone())
                .into(),
        ));
        to_stream_pycapsule(py, array_reader, requested_schema)
    }
}

impl From<Box<dyn RecordBatchReader + Send>> for PyRecordBatchReader {
    fn from(value: Box<dyn RecordBatchReader + Send>) -> Self {
        Self::new(value)
    }
}

impl Display for PyRecordBatchReader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "arro3.core.RecordBatchReader")?;
        writeln!(f, "-----------------------")?;
        if let Ok(schema) = self.schema_ref() {
            display_schema(&schema, f)
        } else {
            writeln!(f, "Closed stream")
        }
    }
}

#[pymethods]
impl PyRecordBatchReader {
    fn __arrow_c_schema__<'py>(&'py self, py: Python<'py>) -> PyArrowResult<Bound<'py, PyCapsule>> {
        to_schema_pycapsule(py, self.schema_ref()?.as_ref())
    }

    #[pyo3(signature = (requested_schema=None))]
    fn __arrow_c_stream__<'py>(
        &'py self,
        py: Python<'py>,
        requested_schema: Option<Bound<'py, PyCapsule>>,
    ) -> PyArrowResult<Bound<'py, PyCapsule>> {
        let reader = self
            .0
            .lock()
            .unwrap()
            .take()
            .ok_or(PyIOError::new_err("Cannot read from closed stream"))?;
        Self::to_stream_pycapsule(py, reader, requested_schema)
    }

    // Return self
    // https://stackoverflow.com/a/52056290
    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
        slf
    }

    fn __next__(&self) -> PyArrowResult<Arro3RecordBatch> {
        self.read_next_batch()
    }

    fn __repr__(&self) -> String {
        self.to_string()
    }

    #[classmethod]
    fn from_arrow(_cls: &Bound<PyType>, input: AnyRecordBatch) -> PyArrowResult<Self> {
        let reader = input.into_reader()?;
        Ok(Self::new(reader))
    }

    #[classmethod]
    #[pyo3(name = "from_arrow_pycapsule")]
    fn from_arrow_pycapsule_py(_cls: &Bound<PyType>, capsule: &Bound<PyCapsule>) -> PyResult<Self> {
        Self::from_arrow_pycapsule(capsule)
    }

    #[classmethod]
    fn from_batches(
        _cls: &Bound<PyType>,
        schema: PySchema,
        batches: RecordBatchInput,
    ) -> PyResult<Self> {
        Ok(Self::new(
            batches.into_record_batch_reader(schema.into_inner())?,
        ))
    }

    #[classmethod]
    fn from_stream(_cls: &Bound<PyType>, data: &Bound<PyAny>) -> PyResult<Self> {
        data.extract()
    }

    #[getter]
    fn closed(&self) -> bool {
        self.0.lock().unwrap().is_none()
    }

    fn read_all(&self) -> PyArrowResult<Arro3Table> {
        let stream = self
            .0
            .lock()
            .unwrap()
            .take()
            .ok_or(PyIOError::new_err("Cannot read from closed stream."))?;
        let schema = stream.schema();
        let mut batches = vec![];
        for batch in stream {
            batches.push(batch?);
        }
        Ok(PyTable::try_new(batches, schema)?.into())
    }

    fn read_next_batch(&self) -> PyArrowResult<Arro3RecordBatch> {
        let mut stream = self
            .0
            .lock()
            .unwrap()
            .take()
            .ok_or(PyIOError::new_err("Cannot read from closed stream."))?;
        let next_batch = stream.next();
        // Put the stream back so the reader can be used again.
        self.0.lock().unwrap().replace(stream);

        match next_batch {
            Some(batch) => Ok(batch?.into()),
            None => Err(PyStopIteration::new_err("").into()),
        }
    }

    #[getter]
    fn schema(&self) -> PyResult<Arro3Schema> {
        Ok(PySchema::new(self.schema_ref()?.clone()).into())
    }
}