// SPDX-FileCopyrightText: 2016-2025 Copyright The Apache Software Foundation

// SPDX-FileCopyrightText: 2025 Copyright the Vortex contributors

// SPDX-License-Identifier: Apache-2.0

// SPDX-FileComment: Derived from upstream file arrow-pyarrow/src/main at commit 549709fb at https://github.com/apache/arrow-rs

// SPDX-FileNotice: https://github.com/apache/arrow-rs/blob/549709fbdf91cd1f6c263a7e4540c542b6fecf6b/NOTICE.txt

#![allow(clippy::same_name_method)]

use std::convert::From;

use std::convert::TryFrom;

use std::ffi::CStr;

use std::ptr::addr_of;

use std::sync::Arc;

use arrow_array::RecordBatch;

use arrow_array::RecordBatchIterator;

use arrow_array::RecordBatchOptions;

use arrow_array::RecordBatchReader;

use arrow_array::StructArray;

use arrow_array::ffi;

use arrow_array::ffi::FFI_ArrowArray;

use arrow_array::ffi::FFI_ArrowSchema;

use arrow_array::ffi_stream::ArrowArrayStreamReader;

use arrow_array::ffi_stream::FFI_ArrowArrayStream;

use arrow_data::ArrayData;

use arrow_schema::ArrowError;

use arrow_schema::DataType;

use arrow_schema::Field;

use arrow_schema::Schema;

use pyo3::exceptions::PyTypeError;

use pyo3::exceptions::PyValueError;

use pyo3::ffi::Py_uintptr_t;

use pyo3::ffi::c_str;

use pyo3::import_exception;

use pyo3::intern;

use pyo3::prelude::*;

use pyo3::types::PyCapsule;

use pyo3::types::PyTuple;

use crate::classes::array_class;

use crate::classes::data_type_class;

use crate::classes::field_class;

use crate::classes::record_batch_reader_class;

use crate::classes::schema_class;

const SCHEMA_NAME: &CStr = c_str!("arrow_schema");

const ARRAY_NAME: &CStr = c_str!("arrow_array");

const ARRAY_STREAM_NAME: &CStr = c_str!("arrow_array_stream");

import_exception!(pyarrow, ArrowException);

/// Represents an exception raised by PyArrow.

pub type PyArrowException = ArrowException;

fn to_py_err(err: ArrowError) -> PyErr {

PyArrowException::new_err(err.to_string())

}

/// Trait for converting Python objects to arrow-rs types.

pub trait FromPyArrow<'a, 'py>: Sized {

/// Convert a Python object to an arrow-rs type.

///

/// Takes a GIL-bound value from Python and returns a result with the arrow-rs type.

fn from_pyarrow(value: &Borrowed<'a, 'py, PyAny>) -> PyResult<Self>;

}

/// Create a new PyArrow object from a arrow-rs type.

pub trait ToPyArrow {

/// Convert the implemented type into a Python object without consuming it.

fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>>;

}

/// Convert an arrow-rs type into a PyArrow object.

pub trait IntoPyArrow {

/// Convert the implemented type into a Python object while consuming it.

fn into_pyarrow(self, py: Python) -> PyResult<Py<PyAny>>;

}

impl<'py> FromPyArrow<'_, 'py> for DataType {

fn from_pyarrow(value: &Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {

let py = value.py();

if !value.hasattr(intern!(py, "__arrow_c_schema__"))? {

return Err(PyValueError::new_err(

"Expected __arrow_c_schema__ attribute to be set.",

));

}

let capsule = value.getattr(intern!(py, "__arrow_c_schema__"))?.call0()?;

let capsule = capsule.cast::<PyCapsule>()?;

let schema_ptr = unsafe {

capsule

.pointer_checked(Some(SCHEMA_NAME))?

.cast::<FFI_ArrowSchema>()

.as_ref()

};

DataType::try_from(schema_ptr).map_err(to_py_err)

}

}

impl ToPyArrow for DataType {

fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {

let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?;

let dtype = data_type_class(py)?.call_method1(

intern!(py, "_import_from_c"),

(&raw const c_schema as Py_uintptr_t,),

)?;

Ok(dtype.into())

}

}

impl<'py> FromPyArrow<'_, 'py> for Field {

fn from_pyarrow(value: &Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {

let py = value.py();

if !value.hasattr(intern!(py, "__arrow_c_schema__"))? {

return Err(PyValueError::new_err(

"Expected __arrow_c_schema__ attribute to be set.",

));

}

let capsule = value.getattr(intern!(py, "__arrow_c_schema__"))?.call0()?;

let capsule = capsule.cast::<PyCapsule>()?;

let schema_ptr = unsafe {

capsule

.pointer_checked(Some(SCHEMA_NAME))?

.cast::<FFI_ArrowSchema>()

.as_ref()

};

let field = Field::try_from(schema_ptr).map_err(to_py_err)?;

Ok(field)

}

}

impl ToPyArrow for Field {

fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {

let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?;

let dtype = field_class(py)?.call_method1(

intern!(py, "_import_from_c"),

(&raw const c_schema as Py_uintptr_t,),

)?;

Ok(dtype.into())

}

}

impl<'py> FromPyArrow<'_, 'py> for Schema {

fn from_pyarrow(value: &Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {

let py = value.py();

if !value.hasattr(intern!(py, "__arrow_c_schema__"))? {

return Err(PyValueError::new_err(

"Expected __arrow_c_schema__ attribute to be set.",

));

}

let capsule = value.getattr(intern!(py, "__arrow_c_schema__"))?.call0()?;

let capsule = capsule.cast::<PyCapsule>()?;

let schema_ptr = unsafe {

capsule

.pointer_checked(Some(SCHEMA_NAME))?

.cast::<FFI_ArrowSchema>()

.as_ref()

};

let schema = Schema::try_from(schema_ptr).map_err(to_py_err)?;

Ok(schema)

}

}

impl ToPyArrow for Schema {

fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {

let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?;

let schema = schema_class(py)?.call_method1(

intern!(py, "_import_from_c"),

(&raw const c_schema as Py_uintptr_t,),

)?;

Ok(schema.into())

}

}

impl<'py> FromPyArrow<'_, 'py> for ArrayData {

fn from_pyarrow(value: &Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {

let py = value.py();

if !value.hasattr(intern!(py, "__arrow_c_array__"))? {

return Err(PyValueError::new_err(

"Expected __arrow_c_array__ attribute to be set.",

));

}

let tuple = value.getattr(intern!(py, "__arrow_c_array__"))?.call0()?;

if !tuple.is_instance_of::<PyTuple>() {

return Err(PyTypeError::new_err(

"Expected __arrow_c_array__ to return a tuple.",

));

}

let schema_capsule = tuple.get_item(0)?;

let schema_capsule = schema_capsule.cast::<PyCapsule>()?;

let array_capsule = tuple.get_item(1)?;

let array_capsule = array_capsule.cast::<PyCapsule>()?;

let schema_ptr = unsafe {

schema_capsule

.pointer_checked(Some(SCHEMA_NAME))?

.cast::<FFI_ArrowSchema>()

.as_ref()

};

let array_ptr = array_capsule

.pointer_checked(Some(ARRAY_NAME))?

.cast::<FFI_ArrowArray>()

.as_ptr();

let array = unsafe { FFI_ArrowArray::from_raw(array_ptr) };

unsafe { ffi::from_ffi(array, schema_ptr) }.map_err(to_py_err)

}

}

impl ToPyArrow for ArrayData {

fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {

let array = FFI_ArrowArray::new(self);

let schema = FFI_ArrowSchema::try_from(self.data_type()).map_err(to_py_err)?;

let array = array_class(py)?.call_method1(

intern!(py, "_import_from_c"),

(

addr_of!(array) as Py_uintptr_t,

addr_of!(schema) as Py_uintptr_t,

),

)?;

Ok(array.unbind())

}

}

impl<'py> FromPyArrow<'_, 'py> for RecordBatch {

fn from_pyarrow(value: &Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {

let py = value.py();

if !value.hasattr(intern!(py, "__arrow_c_array__"))? {

return Err(PyValueError::new_err(

"Expected __arrow_c_array__ attribute to be set.",

));

}

let tuple = value.getattr(intern!(py, "__arrow_c_array__"))?.call0()?;

if !tuple.is_instance_of::<PyTuple>() {

return Err(PyTypeError::new_err(

"Expected __arrow_c_array__ to return a tuple.",

));

}

let schema_capsule = tuple.get_item(0)?;

let schema_capsule = schema_capsule.cast::<PyCapsule>()?;

let array_capsule = tuple.get_item(1)?;

let array_capsule = array_capsule.cast::<PyCapsule>()?;

let schema_ptr = unsafe {

schema_capsule

.pointer_checked(Some(SCHEMA_NAME))?

.cast::<FFI_ArrowSchema>()

.as_ref()

};

let array_ptr = array_capsule

.pointer_checked(Some(ARRAY_NAME))?

.cast::<FFI_ArrowArray>()

.as_ptr();

let ffi_array = unsafe { FFI_ArrowArray::from_raw(array_ptr) };

let mut array_data = unsafe { ffi::from_ffi(ffi_array, schema_ptr) }.map_err(to_py_err)?;

if !matches!(array_data.data_type(), DataType::Struct(_)) {

return Err(PyTypeError::new_err(

"Expected Struct type from __arrow_c_array.",

));

}

let options = RecordBatchOptions::default().with_row_count(Some(array_data.len()));

// Ensure data is aligned (by potentially copying the buffers).

// This is needed because some python code (for example the

// python flight client) produces unaligned buffers

// See https://github.com/apache/arrow/issues/43552 for details

array_data.align_buffers();

let array = StructArray::from(array_data);

// StructArray does not embed metadata from schema. We need to override

// the output schema with the schema from the capsule.

let schema = Arc::new(Schema::try_from(schema_ptr).map_err(to_py_err)?);

let (_fields, columns, nulls) = array.into_parts();

assert_eq!(

nulls.map(|n| n.null_count()).unwrap_or_default(),

0,

"Cannot convert nullable StructArray to RecordBatch, see StructArray documentation"

);

RecordBatch::try_new_with_options(schema, columns, &options).map_err(to_py_err)

}

}

impl ToPyArrow for RecordBatch {

fn to_pyarrow(&self, py: Python) -> PyResult<Py<PyAny>> {

// Workaround apache/arrow#37669 by returning RecordBatchIterator

let reader = RecordBatchIterator::new(vec![Ok(self.clone())], self.schema());

let reader: Box<dyn RecordBatchReader + Send> = Box::new(reader);

let py_reader = reader.into_pyarrow(py)?;

py_reader.call_method0(py, intern!(py, "read_next_batch"))

}

}

/// Supports conversion from `pyarrow.RecordBatchReader` to [ArrowArrayStreamReader].

impl<'py> FromPyArrow<'_, 'py> for ArrowArrayStreamReader {

fn from_pyarrow(value: &Borrowed<'_, 'py, PyAny>) -> PyResult<Self> {

let py = value.py();

if !value.hasattr(intern!(py, "__arrow_c_stream__"))? {

return Err(PyValueError::new_err(

"Expected __arrow_c_stream__ attribute to be set.",

));

}

let capsule = value.getattr(intern!(py, "__arrow_c_stream__"))?.call0()?;

let capsule = capsule.cast::<PyCapsule>()?;

let array_ptr = capsule

.pointer_checked(Some(ARRAY_STREAM_NAME))?

.cast::<FFI_ArrowArrayStream>()

.as_ptr();

let stream = unsafe { FFI_ArrowArrayStream::from_raw(array_ptr) };

let stream_reader = ArrowArrayStreamReader::try_new(stream)

.map_err(|err| PyValueError::new_err(err.to_string()))?;

Ok(stream_reader)

}

}

/// Convert a [`RecordBatchReader`] into a `pyarrow.RecordBatchReader`.

impl IntoPyArrow for Box<dyn RecordBatchReader + Send> {

// We can't implement `ToPyArrow` for `T: RecordBatchReader + Send` because

// there is already a blanket implementation for `T: ToPyArrow`.

fn into_pyarrow(self, py: Python) -> PyResult<Py<PyAny>> {

let mut stream = FFI_ArrowArrayStream::new(self);

let args = PyTuple::new(py, [&raw mut stream as Py_uintptr_t])?;

let reader =

record_batch_reader_class(py)?.call_method1(intern!(py, "_import_from_c"), args)?;

Ok(Py::from(reader))

}

}