From 3464c74877a8b5f793ee1a44c6a72beae5ee20b2 Mon Sep 17 00:00:00 2001 From: Peter Bower <37089506+pbower@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:58:26 +0100 Subject: [PATCH] Improve minarrow-py conversions --- minarrow-py/src/convert.rs | 132 +++++++++++++++++++++++++++++++------ 1 file changed, 112 insertions(+), 20 deletions(-) diff --git a/minarrow-py/src/convert.rs b/minarrow-py/src/convert.rs index 5862f2f..0fdbef1 100644 --- a/minarrow-py/src/convert.rs +++ b/minarrow-py/src/convert.rs @@ -19,17 +19,58 @@ use minarrow::ffi::arrow_dtype::{ArrowType, CategoricalIndexType}; use minarrow::arr_str64_opt; #[cfg(feature = "extended_numeric_types")] use minarrow::{arr_i8_opt, arr_i16_opt, arr_u8_opt, arr_u16_opt}; +use minarrow::enums::array::extract_option_values64; +use minarrow::enums::time_units::TimeUnit; use minarrow::{ arr_bool_opt, arr_f32_opt, arr_f64_opt, arr_i32_opt, arr_i64_opt, arr_str32_opt, arr_u32_opt, - arr_u64_opt, Array, Bitmask, CategoricalArray, Scalar, Vec64, + arr_u64_opt, Array, ArrayV, Bitmask, CategoricalArray, DatetimeArray, Scalar, Vec64, }; +use crate::array::PyArray; use crate::arrow_type::PyArrowType; use pyo3::exceptions::{PyIndexError, PyNotImplementedError, PyTypeError, PyValueError}; +use pyo3::ffi; use pyo3::prelude::*; -use pyo3::types::PyBool; +use pyo3::types::{PyBool, PyString}; use pyo3::IntoPyObjectExt; +/// Reads a Python sequence into a single 64-byte aligned `Vec64` buffer. +/// +/// This is the extraction step for the constructors that accept a Python +/// sequence. Extracting through pyo3's `Vec` would allocate through the +/// global allocator and then need a second allocation and copy to reach +/// minarrow's 64-byte buffer alignment, so each element is instead extracted +/// into the aligned buffer as it is read. +/// +/// Acceptance matches pyo3's `Vec` extraction. The input must satisfy the +/// Python sequence protocol and must not be a `str`, otherwise the function +/// raises `TypeError`. +fn read_sequence<'py, T>(data: &Bound<'py, PyAny>) -> PyResult> +where + T: for<'a> FromPyObject<'a, 'py>, + for<'a> PyErr: From<>::Error>, +{ + if data.is_instance_of::() { + return Err(PyTypeError::new_err( + "expected a sequence of values, not a str", + )); + } + // pyo3 gates `extract::>` on the sequence protocol, so acceptance + // here matches it and an object passing the gate supports the length and + // iteration below. + if unsafe { ffi::PySequence_Check(data.as_ptr()) } == 0 { + return Err(PyTypeError::new_err(format!( + "expected a sequence of values, got {}", + data.get_type().name()? + ))); + } + let mut values = Vec64::with_capacity(data.len().unwrap_or(0)); + for item in data.try_iter()? { + values.push(item?.extract()?); + } + Ok(values) +} + /// Build a minarrow `Array` from a Python sequence, inferring the element type /// from its values. `None` becomes null. Integers promote to float when a /// sequence also contains floats. An empty or all-null sequence builds a float @@ -39,21 +80,27 @@ use pyo3::IntoPyObjectExt; /// there is no per-element type check. `bool` is tried before `int` because a /// Python `bool` is an `int` subclass. pub fn build_array(data: &Bound<'_, PyAny>) -> PyResult { - if let Ok(values) = data.extract::>>() { + // Reuse an existing Array as-is rather than re-inferring the dtype from + // its values, which mistypes temporal and categorical columns. The ArrayV + // conversion materialises a windowed array into a standalone copy. + if let Ok(array) = data.extract::>() { + return Ok(ArrayV::from(&array.0).to_array()); + } + if let Ok(values) = read_sequence::>(data) { if !values.iter().all(Option::is_none) { - return Ok(arr_bool_opt!(values.into_iter().collect::>())); + return Ok(arr_bool_opt!(values)); } } - if let Ok(values) = data.extract::>>() { + if let Ok(values) = read_sequence::>(data) { if !values.iter().all(Option::is_none) { - return Ok(arr_i64_opt!(values.into_iter().collect::>())); + return Ok(arr_i64_opt!(values)); } } - if let Ok(values) = data.extract::>>() { - return Ok(arr_f64_opt!(values.into_iter().collect::>())); + if let Ok(values) = read_sequence::>(data) { + return Ok(arr_f64_opt!(values)); } - if let Ok(values) = data.extract::>>() { - return Ok(arr_str32_opt!(values.into_iter().collect::>())); + if let Ok(values) = read_sequence::>(data) { + return Ok(arr_str32_opt!(values)); } Err(PyTypeError::new_err( "Array elements must be bool, int, float, str, or None", @@ -62,14 +109,30 @@ pub fn build_array(data: &Bound<'_, PyAny>) -> PyResult { /// Build a minarrow `Array` from a Python sequence coerced to `dtype`. `None` /// becomes null. A value that does not fit the target type raises the -/// extraction's own `TypeError` or `OverflowError`. Types that need more than a -/// flat sequence, such as categorical and temporal, are not built here. +/// extraction's own `TypeError` or `OverflowError`. +/// +/// Temporal dtypes accept integer values in the unit the dtype declares, so +/// `ArrowType::Timestamp(TimeUnit::Milliseconds, None)` over +/// `[1700000000000, ...]` builds a millisecond-precision timestamp column. pub fn build_array_typed(data: &Bound<'_, PyAny>, dtype: &ArrowType) -> PyResult { macro_rules! build { ($t:ty, $make:ident) => {{ - let values: Vec64> = - data.extract::>>()?.into_iter().collect(); - Ok($make!(values)) + Ok($make!(read_sequence::>(data)?)) + }}; + } + // Temporal dtypes build through DatetimeArray rather than the arr_*_opt + // macros because the array stores the dtype's TimeUnit alongside its + // values. Null handling matches the other types, with None entries + // recorded in the null mask by extract_option_values64. + macro_rules! build_temporal { + ($t:ty, $make:ident, $unit:expr) => {{ + let (values, null_mask) = + extract_option_values64(read_sequence::>(data)?); + Ok(Array::$make(DatetimeArray::<$t>::from_vec64( + values, + null_mask, + Some($unit), + ))) }}; } match dtype { @@ -92,6 +155,15 @@ pub fn build_array_typed(data: &Bound<'_, PyAny>, dtype: &ArrowType) -> PyResult #[cfg(feature = "large_string")] ArrowType::LargeString => build!(String, arr_str64_opt), ArrowType::Dictionary(index) => categorical_from_values(data, index), + ArrowType::Date32 => build_temporal!(i32, from_datetime_i32, TimeUnit::Days), + ArrowType::Date64 => build_temporal!(i64, from_datetime_i64, TimeUnit::Milliseconds), + ArrowType::Time32(unit) | ArrowType::Duration32(unit) => { + build_temporal!(i32, from_datetime_i32, *unit) + } + ArrowType::Time64(unit) | ArrowType::Duration64(unit) => { + build_temporal!(i64, from_datetime_i64, *unit) + } + ArrowType::Timestamp(unit, _) => build_temporal!(i64, from_datetime_i64, *unit), other => Err(PyValueError::new_err(format!( "dtype {} cannot be built from a Python sequence; use from_arrow instead", other @@ -102,6 +174,12 @@ pub fn build_array_typed(data: &Bound<'_, PyAny>, dtype: &ArrowType) -> PyResult /// Parse a dtype string such as `"int32"`, `"f64"`, `"string"`, or `"categorical"`. /// Categorical granularities beyond `UInt32` are accepted only when the matching /// feature is compiled into the build. +/// +/// Temporal dtypes include their unit in the string, so `"timestamp[ms]"` +/// parses to `Timestamp(TimeUnit::Milliseconds, None)` and `"timestamp"` +/// without a unit raises `ValueError`. A timezone cannot be written in a +/// string, so timezone-aware timestamps are constructed as an `ArrowType` +/// and passed to `dtype=` in place of a string. pub fn parse_dtype(name: &str) -> PyResult { Ok(match name.trim().to_ascii_lowercase().as_str() { #[cfg(feature = "extended_numeric_types")] @@ -129,6 +207,23 @@ pub fn parse_dtype(name: &str) -> PyResult { "string" | "str" | "utf8" | "str32" => ArrowType::String, "large_string" | "largestring" | "str64" => ArrowType::LargeString, "bool" | "boolean" => ArrowType::Boolean, + "date32" => ArrowType::Date32, + "date64" => ArrowType::Date64, + "timestamp[s]" | "datetime[s]" => ArrowType::Timestamp(TimeUnit::Seconds, None), + "timestamp[ms]" | "datetime[ms]" => ArrowType::Timestamp(TimeUnit::Milliseconds, None), + "timestamp[us]" | "datetime[us]" => ArrowType::Timestamp(TimeUnit::Microseconds, None), + "timestamp[ns]" | "datetime[ns]" => ArrowType::Timestamp(TimeUnit::Nanoseconds, None), + "timestamp" | "datetime" => { + return Err(PyValueError::new_err( + "a timestamp dtype names its unit, as 'timestamp[ms]' or 'datetime[ms]'. \ + Pass an ArrowType for a timezone-aware timestamp", + )); + } + "date" => { + return Err(PyValueError::new_err( + "a date dtype names its width, as 'date32' (days) or 'date64' (milliseconds)", + )); + } "categorical" | "category" | "cat" => { #[cfg(feature = "default_categorical_8")] { @@ -199,11 +294,8 @@ fn categorical_from_values( data: &Bound<'_, PyAny>, index: &CategoricalIndexType, ) -> PyResult { - let values: Vec64> = data - .extract::>>() - .map_err(|_| PyTypeError::new_err("categorical values must be a list of strings or None"))? - .into_iter() - .collect(); + let values: Vec64> = read_sequence::>(data) + .map_err(|_| PyTypeError::new_err("categorical values must be a list of strings or None"))?; let mut mask = Bitmask::new_set_all(values.len(), true); let mut strings: Vec64<&str> = Vec64::with_capacity(values.len()); for (row, value) in values.iter().enumerate() {